[Feature] Unify the cuda stream used in core library

[chang-l]

Description

To fix #4471, I changed all hard-coded default stream to DGL internal stream.

• Enforced CopyDataFromTo and its associates to use stream thr_entry->stream
• Fix all CUDA_KERNEL_CALL to use thr_entry->stream
• Fix cub submodule to use thr_entry->stream
• Initialized dgl stream to torch.cuda.current_stream for each thread, to ensure backward propag. still uses the right stream.

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

Tests

• Initial test:
  With the script from [Bug] CUDA stream synchronization issue between pytorch and DGL internal functions #4436, run on a side stream and it will not crash anymore.

• GCN example
  Switch to a non-default stream (torch.cuda.set_stream()), then align DGL stream using with dgl._ffi.streams.stream(s)
  Results are consistent and profiler shows 100% kernels launch on the specified side stream.

• Graphsage example
  Switch to a non-default stream (torch.cuda.set_stream()), then align DGL stream using with dgl._ffi.streams.stream(s)
  Results are consistent and profiler shows ~70% kernels launch on the specified side stream, other (non-default) streams are all for non-blocking data-transfer in layer-wise inference stage (issued from pytorch). Cross-checked/confirmed with the reference profiling from default stream results.

• GAT example
  Switch to a non-default stream (torch.cuda.set_stream()), then align DGL stream using with dgl._ffi.streams.stream(s)
  Results are consistent and profiler shows ~99.5% kernels launch on the specified side stream, the rest (<0.5%) are from cudaMemcpy issued on the legacy default stream (auto-sync blocking stream). These cudaMemcpy functions are ALL called from the following function:
  

  dgl/src/array/cuda/macro.cuh

  Lines 16 to 17 in 2dd06b2

  	#define BCAST_IDX_CTX_SWITCH(BCAST, EDGE_MAP, CTX, LHS_OFF, RHS_OFF, ...) do { \
  	const BcastOff &info = (BCAST); \

• Heter-RGCN example
  Switch to a non-default stream in pytorch, then align to the same stream for DGL.
  Results are consistent and profiler shows 100% kernels launch on the specified side stream.

• Multi-GPU runs (graphsage, and rgcn)
  For multi-gpu runs, some un-specified stream launches are from (1) NCCL and (2) pytorch.

[dgl-bot]

To trigger regression tests:

• @dgl-bot run [instance-type] [which tests] [compare-with-branch];
  For example: @dgl-bot run g4dn.4xlarge all dmlc/master or @dgl-bot run c5.9xlarge kernel,api dmlc/master

[chang-l]

Once I resolve the above comments, there are still two remaining issues for this PR, according to my profiling:

(1) This PR depends on Xin's PR: #4472, to completely unify the dgl internal stream.
(2) In pytorch, backward propag. uses a different thread, which makes the dgl stream CUDAThreadEntry->ThreadLocal()->stream out of control within backward propg section, i.e., it will always use the default stream for a different thread.
I am thinking maybe we should avoid declaring cuda stream as a thread-local variable, because cuda stream itself is thread-safe...

Please let me know your thoughts @yaox12 @nv-dlasalle @BarclayII @TristonC

[yaox12]

Yesterday I had a brief sync with @jermainewang and he suggested we always get the current stream from PyTorch (via TensorAdaptor maybe) instead of CUDAThreadEntry->ThreadLocal()->stream. How do you think about it? @chang-l

Besides, I think PyTorch also stores the current CUDA stream in a thread-local variable. How does it handle the backward thread?
https://github.com/pytorch/pytorch/blob/6a3ecda5a25025d48bbc5f0215db8c338745ef79/c10/cuda/CUDAStream.cpp#L137-L138

[chang-l]

always get the current stream from PyTorch (via TensorAdaptor maybe) instead of CUDAThreadEntry->ThreadLocal()->stream.

Are we going to always query the current_stream from pytorch for EVERY kernel calls/mem transfers? I think yes because we have no control on pytorch's stream and users may change torch current stream anywhere in their code, yet the new stream has no way propagate to DGL. This would make code a bit messy and may introduce some overhead.

However, I agree that packing dgl stream and pytorch stream together might be more convenient and safer, from users perspective. Ideally, for now, users should set their desired stream once (for dgl and its backend), and not worry about or not encourage to 'touch' cuda-stream inside the library. How about integrate torch.set_stream in dgl.set_stream, so that (1) users only need to set stream once and only through DGL, (2) implying them not invoking torch.set_stream when coupled with dgl?

In terms of pytorch's implementation of backward propag., I think they use a global stream pool, defined here. And they sync/record the forward streams at the entry point of autograd(backward) engine through function InputBuffer::add. Here are their comment descriptions: https://github.com/pytorch/pytorch/blob/ece0002c4beaebaf083dc75b7bf8ceb19edf7a0b/torch/csrc/autograd/engine.cpp#L164-L168
https://github.com/pytorch/pytorch/blob/6a3ecda5a25025d48bbc5f0215db8c338745ef79/c10/cuda/CUDAStream.h#L13-L53

[chang-l]

Maybe I can pin to torch's current stream (sth like stream = torch.current_stream()) here in the constructor of CUDAThreadEntry:

dgl/src/runtime/cuda/cuda_device_api.cc

Lines 300 to 302 in 2b76674

	CUDAThreadEntry::CUDAThreadEntry()
	: pool(kDLGPU, CUDADeviceAPI::Global()) {
	}

This way whenever it initiates a new CUDAThreadEntry instance at each thread, it can setup the correct stream (through torch of course).

[yaox12]

Maybe I can pin to torch's current stream (sth like stream = torch.current_stream()) here in the constructor of CUDAThreadEntry:

dgl/src/runtime/cuda/cuda_device_api.cc

Lines 300 to 302 in 2b76674

	CUDAThreadEntry::CUDAThreadEntry()
	: pool(kDLGPU, CUDADeviceAPI::Global()) {
	}

This way whenever it initiates a new CUDAThreadEntry instance at each thread, it can setup the correct stream (through torch of course).

Makes sense to me.

[chang-l]

Could this be error-prone of not setting device here? Now, dgl internal cuda stream is stored as a thread-local variable in CUDAThreadEntry. I believe it is not exposed/related to the current device id in use (DGLContext), right?

Now, I guess it relies on whether the device has been appropriately set at beginning by pytorch.

[yaox12]

the device has been appropriately set at beginning by pytorch.

I think it should be fine. For the single GPU case, PyTorch will use device 0 by default. While for multi-GPU DDP, users are required to set the device for each process.

[chang-l]

I tested with several examples and haven't seen any "stream leakage" (pls check the updated PR description).

[yaox12]

the device has been appropriately set at beginning by pytorch.

I think it should be fine. For the single GPU case, PyTorch will use device 0 by default. While for multi-GPU DDP, users are required to set the device for each process.

[chang-l]

Performed tests on hetero-rgcn and multi-gpu cases (updated PR description). No stream leakage found within dgl library calls.

[yaox12]

LGTM.

[dgl-bot]

Commit ID: de0b9d5

Build ID: 14

Status: ✅ CI test succeeded

Report path: link

Full logs path: link

[yaox12]

After merging this PR, I realized that to keep up with PyTorch's CUDA stream, users have to call dgl.cuda.set_stream() after torch.cuda.set_stream() every time, or use dgl.cuda.stream() context together with torch.cuda.stream().

To simplify this, I think

always query the current_stream from pytorch for EVERY kernel calls/mem transfers? I think yes because we have no control on pytorch's stream and users may change torch current stream anywhere in their code, yet the new stream has no way propagate to DGL.

is acceptable, for example, PyTorch will call c10::cuda::getCurrentCUDAStream() everytime when it needs a stream.

To make it, I created a method runtime::getCurrentCUDAStream() that would query the current stream from pytorch, as a replacement of runtime::CUDAThreadEntry::ThreadLocal()->stream in #4503.

[chang-l]

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

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?)