Tensor CUDA Stream API#
ACUDA Stream is a linear sequence of execution that belongs to a specific CUDA device.The PyTorch C++ API supports CUDA streams with the CUDAStream class and useful helper functions to make streaming operations easy.You can find them inCUDAStream.h. This note provides more details on how to use Pytorch C++ CUDA Stream APIs.
Acquiring CUDA stream#
Pytorch’s C++ API provides the following ways to acquire CUDA stream:
Acquire a new stream from the CUDA stream pool, streams are preallocated from the pool and returned in a round-robin fashion.
CUDAStreamgetStreamFromPool(constboolisHighPriority=false,DeviceIndexdevice=-1);
Tip
You can request a stream from the high priority pool by setting isHighPriority to true, or a stream for a specific deviceby setting device index (defaulting to the current CUDA stream’s device index).
Acquire the default CUDA stream for the passed CUDA device, or for the current device if no device index is passed.
CUDAStreamgetDefaultCUDAStream(DeviceIndexdevice_index=-1);
Tip
The default stream is where most computation occurs when you aren’t explicitly using streams.
Acquire the current CUDA stream, for the CUDA device with index
device_index, or for the current device if no device index is passed.
CUDAStreamgetCurrentCUDAStream(DeviceIndexdevice_index=-1);
Tip
The current CUDA stream will usually be the default CUDA stream for the device, but it may be different if someonecalledsetCurrentCUDAStream or usedStreamGuard orCUDAStreamGuard.
Set CUDA stream#
Pytorch’s C++ API provides the following ways to set CUDA stream:
Set the current stream on the device of the passed in stream to be the passed in stream.
voidsetCurrentCUDAStream(CUDAStreamstream);
Attention
This function may have nothing to do with the current device. It only changes the current stream on the stream’s device.We recommend usingCUDAStreamGuard, instead, since it switches to the stream’s device and makes it the current stream on that device.CUDAStreamGuard will also restore the current device and stream when it’s destroyed
Use
CUDAStreamGuardto switch to a CUDA stream within a scope, it is defined inCUDAStreamGuard.h
Tip
UseCUDAMultiStreamGuard if you need to set streams on multiple CUDA devices.
CUDA Stream Usage Examples#
Acquiring and setting CUDA stream on the same device
// This example shows how to acquire and set CUDA stream on the same device.// `at::cuda::setCurrentCUDAStream` is used to set current CUDA stream// create a tensor on device 0torch::Tensortensor0=torch::ones({2,2},torch::device(torch::kCUDA));// get a new CUDA stream from CUDA stream pool on device 0at::cuda::CUDAStreammyStream=at::cuda::getStreamFromPool();// set current CUDA stream from default stream to `myStream` on device 0at::cuda::setCurrentCUDAStream(myStream);// sum() on tensor0 uses `myStream` as current CUDA streamtensor0.sum();// get the default CUDA stream on device 0at::cuda::CUDAStreamdefaultStream=at::cuda::getDefaultCUDAStream();// set current CUDA stream back to default CUDA stream on device 0at::cuda::setCurrentCUDAStream(defaultStream);// sum() on tensor0 uses `defaultStream` as current CUDA streamtensor0.sum();
// This example is the same as previous example, but explicitly specify device// index and use CUDA stream guard to set current CUDA stream// create a tensor on device 0torch::Tensortensor0=torch::ones({2,2},torch::device(torch::kCUDA));// get a new stream from CUDA stream pool on device 0at::cuda::CUDAStreammyStream=at::cuda::getStreamFromPool(false,0);// set the current CUDA stream to `myStream` within the scope using CUDA stream guard{at::cuda::CUDAStreamGuardguard(myStream);// current CUDA stream is `myStream` from here till the end of bracket.// sum() on tensor0 uses `myStream` as current CUDA streamtensor0.sum();}// current CUDA stream is reset to default CUDA stream after CUDA stream guard is destroyed// sum() on tensor0 uses default CUDA stream on device 0 as current CUDA streamtensor0.sum();
Attention
Above code is running on the same CUDA device.setCurrentCUDAStream will always set current CUDA stream on current device,but note thatsetCurrentCUDAStream actually set current stream on the device of passed in CUDA stream.
Acquiring and setting CUDA streams on multiple devices.
// This example shows how to acquire and set CUDA stream on two devices.// acquire new CUDA streams from CUDA stream pool on device 0 and device 1at::cuda::CUDAStreammyStream0=at::cuda::getStreamFromPool(false,0);at::cuda::CUDAStreammyStream1=at::cuda::getStreamFromPool(false,1);// set current CUDA stream to `myStream0` on device 0at::cuda::setCurrentCUDAStream(myStream0);// set current CUDA stream to `myStream1` on device 1at::cuda::setCurrentCUDAStream(myStream1);// create a tensor on device 0, no need to specify device index since// current device index is 0torch::Tensortensor0=torch::ones({2,2},torch::device(at::kCUDA));// sum() on tensor0 use `myStream0` as current CUDA stream on device 0tensor0.sum();// change the current device index to 1 by using CUDA device guard within a bracket scope{at::cuda::CUDAGuarddevice_guard{1};// create a tensor on device 1torch::Tensortensor1=torch::ones({2,2},torch::device(at::kCUDA));// sum() on tensor 1 uses `myStream1` as current CUDA stream on device 1tensor1.sum();}// current device is reset to device 0 after device_guard is destroyed// acquire a new CUDA stream on device 1at::cuda::CUDAStreammyStream1_1=at::cuda::getStreamFromPool(false,1);// create a new tensor on device 1torch::Tensortensor1=torch::ones({2,2},torch::device({torch::kCUDA,1}));// change the current device index to 1 and current CUDA stream on device 1// to `myStream1_1` using CUDA stream guard within a scope{at::cuda::CUDAStreamGuardstream_guard(myStream1_1);// sum() on tensor1 use `myStream1_1` as current CUDA stream on device 1tensor1.sum();}// current device is reset to device 0 and current CUDA stream on device 1 is// reset to `myStream1`// sum() on tensor1 uses `myStream1` as current CUDA stream on device 1tensor1.sum();
Working with CUDA multistream guard
// This example shows how to use CUDA multistream guard to set// two streams on two devices at the same time.// create two tensor, one on device 0, one on device 1torch::Tensortensor0=torch::ones({2,2},torch::device({torch::kCUDA,0}));torch::Tensortensor1=torch::ones({2,2},torch::device({torch::kCUDA,1}));// acquire new CUDA streams from CUDA stream pool on device 0 and device 1at::cuda::CUDAStreammyStream0=at::cuda::getStreamFromPool(false,0);at::cuda::CUDAStreammyStream1=at::cuda::getStreamFromPool(false,1);// set current CUDA stream on device 0 to `myStream0` and// set current CUDA stream on device 1 to `myStream1` CUDA using multistream guard{at::cuda::CUDAMultiStreamGuardmulti_guard({myStream0,myStream1});// sum() on tensor0 uses `myStream0` as current CUDA stream on device 0tensor0.sum();// sum() on tensor1 uses `myStream1` as current CUDA stream on device 1tensor1.sum();}// current CUDA stream on device 0 is reset to default CUDA stream on device 0// current CUDA stream on device 1 is reset to default CUDA stream on device 1// sum() on tensor0 uses default CUDA stream as current CUDA stream on device 0tensor0.sum();// sum() on tensor1 uses default CUDA stream as current CUDA stream on device 1tensor1.sum();
Attention
CUDAMultiStreamGuard does not change current device index, it only changes the stream oneach passed in stream’s device. Other than scope controlling, this guard is equivalent tocallingsetCurrentCUDAStream on each passed in stream.
A skeleton example for handling CUDA streams on multiple devices
// This is a skeleton example that shows how to handle CUDA streams on multiple devices// Suppose you want to do work on the non-default stream on two devices simultaneously, and we// already have streams on both devices in two vectors. The following code shows three ways// of acquiring and setting the streams.// Usage 0: acquire CUDA stream and set current CUDA stream with `setCurrentCUDAStream`// Create a CUDA stream vector `streams0` on device 0std::vector<at::cuda::CUDAStream>streams0={at::cuda::getDefaultCUDAStream(),at::cuda::getStreamFromPool()};// set current stream as `streams0[0]` on device 0at::cuda::setCurrentCUDAStream(streams0[0]);// create a CUDA stream vector `streams1` on device using CUDA device guardstd::vector<at::cuda::CUDAStream>streams1;{// device index is set to 1 within this scopeat::cuda::CUDAGuarddevice_guard(1);streams1.push_back(at::cuda::getDefaultCUDAStream());streams1.push_back(at::cuda::getStreamFromPool());}// device index is reset to 0 after device_guard is destroyed// set current stream as `streams1[0]` on device 1at::cuda::setCurrentCUDAStream(streams1[0]);// Usage 1: use CUDA device guard to change the current device index only{at::cuda::CUDAGuarddevice_guard(1);// current device index is changed to 1 within scope// current CUDA stream is still `streams1[0]` on device 1, no change}// current device index is reset to 0 after `device_guard` is destroyed// Usage 2: use CUDA stream guard to change both current device index and current CUDA stream.{at::cuda::CUDAStreamGuardstream_guard(streams1[1]);// current device index and current CUDA stream are set to 1 and `streams1[1]` within scope}// current device index and current CUDA stream are reset to 0 and `streams0[0]` after// stream_guard is destroyed// Usage 3: use CUDA multi-stream guard to change multiple streams on multiple devices{// This is the same as calling `torch::cuda::setCurrentCUDAStream` on both streamsat::cuda::CUDAMultiStreamGuardmulti_guard({streams0[1],streams1[1]});// current device index is not change, still 0// current CUDA stream on device 0 and device 1 are set to `streams0[1]` and `streams1[1]`}// current CUDA stream on device 0 and device 1 are reset to `streams0[0]` and `streams1[0]`// after `multi_guard` is destroyed.
