[Feature] Import PyTorch's CUDA stream management

[yaox12]

Description

Resolve #4467.

Add stream-related utilities:

1. Add RecordStream() for NDArray and HeteroGraph on the C++ side,
2. Add record_stream() for dgl.heterograph on the Python side.
3. Replace all runtime::CUDAThreadEntry::ThreadLocal()->stream with runtime::getCurrentCUDAStream() to ensuring we're using PyTorch's current stream. With this change, we don't need dgl.set_stream() method or dgl.cuda.stream() context, and just use torch's stream management instead. This will simplify the user experience.
4. Add unit tests.

Checklist

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented
• To the best of my knowledge, examples are either not affected by this change,
  or have been fixed to be compatible with this change
• Related issue is referred in this PR

[yaox12]

Paste @chang-l 's comments in #4480 (comment).

I (Xin) created a method runtime::getCurrentCUDAStream() that would query the current stream from pytorch, as a replacement of runtime::CUDAThreadEntry::ThreadLocal()->stream.

(Chang): I see... in this case, I wonder (1) are we going to keep SetStream/CreateStream/... methods in DGL core library? Or, they are going to in fact set/create pytorch stream instead of cuda stream? (2) if pytorch current stream is changing outside, would that introduces any sync issue to DGL core library? (e.g., when users trying to overlap some operations between dgl and pytorch and they want to call torch.cuda.set_stream when dgl library is running? or it should almost never happen?)

[chang-l]

@yaox12 Thanks for pulling my comment here :)
In general, I agree on using the current stream from pytorch to launch every dgl kernels. There are two things I suggest to fix in the meanwhile:
(1) Update SetStream/CreateStream methods as well, to set/create from pytorch stream pool.
(2) Update doc to remind users that all dgl functions use the same current stream as backend (pytorch), so if they changed to another stream, they should always restore the stream before calling next dgl function, for example better using with instead of set_stream.

[jermainewang]

Reassigned the reviewers. For CUDA related PRs, please reach out to @isratnisa for review. Changes to TensorAdapter needs @BarclayII 's approval.

[nv-dlasalle]

I think we also need to delete the CUDAThreadEntry::stream member, to ensure it is not used after this change.

How do we handle the fact that PyTorch's streams are not created with the cudaStreamNonBlocking flag, so to get real parallelism, users will need to do their training on the non-default stream, and dataloading on other streams as well.

[yaox12]

I think we also need to delete the CUDAThreadEntry::stream member, to ensure it is not used after this change.

What if TensorAdaptor is not available? Should runtime::getCurrentCUDAStream() just return the default stream?

How do we handle the fact that PyTorch's streams are not created with the cudaStreamNonBlocking flag, so to get real parallelism, users will need to do their training on the non-default stream, and dataloading on other streams as well.

PyTorch's streams are created in
https://github.com/pytorch/pytorch/blob/d00cabae7bc894b951d4b2c9c24c7d95bebd86e1/c10/cuda/CUDAStream.cpp#L165-L168

    C10_CUDA_CHECK(cudaStreamCreateWithPriority(
        &lowpri_stream, kDefaultFlags, kLowPriority));
    C10_CUDA_CHECK(cudaStreamCreateWithPriority(
        &hipri_stream, kDefaultFlags, kHighPriority));

with the flag
https://github.com/pytorch/pytorch/blob/d00cabae7bc894b951d4b2c9c24c7d95bebd86e1/c10/cuda/CUDAStream.cpp#L25

static constexpr unsigned int kDefaultFlags = cudaStreamNonBlocking;

[nv-dlasalle]

I think we also need to delete the CUDAThreadEntry::stream member, to ensure it is not used after this change.

What if TensorAdaptor is not available? Should runtime::getCurrentCUDAStream() just return the default stream?

If tensoradapter is not available, I don't think we can safely pass streams between DGL and a backend, so I think returning the default stream is the correct behavior. I think this also makes #3998 more of a priority, since if tensoradapter fails to load we lose more than just faster memory management.

How do we handle the fact that PyTorch's streams are not created with the cudaStreamNonBlocking flag, so to get real parallelism, users will need to do their training on the non-default stream, and dataloading on other streams as well.

PyTorch's streams are created in https://github.com/pytorch/pytorch/blob/d00cabae7bc894b951d4b2c9c24c7d95bebd86e1/c10/cuda/CUDAStream.cpp#L165-L168

    C10_CUDA_CHECK(cudaStreamCreateWithPriority(
        &lowpri_stream, kDefaultFlags, kLowPriority));
    C10_CUDA_CHECK(cudaStreamCreateWithPriority(
        &hipri_stream, kDefaultFlags, kHighPriority));

with the flag https://github.com/pytorch/pytorch/blob/d00cabae7bc894b951d4b2c9c24c7d95bebd86e1/c10/cuda/CUDAStream.cpp#L25

static constexpr unsigned int kDefaultFlags = cudaStreamNonBlocking;

Sorry, you're right, I read the pytorch source wrong.

[dgl-bot]

Commit ID: ef1d3a3

Build ID: 16

Status: ✅ CI test succeeded

Report path: link

Full logs path: link