ExternalStream¶
-
class
torch.cuda.
ExternalStream
(stream_ptr, device=None, **kwargs)[source]¶ Wrapper around an externally allocated CUDA stream.
This class is used to wrap streams allocated in other libraries in order to facilitate data exchange and multi-library interactions.
Note
This class doesn’t manage the stream life-cycle, it is the user responsibility to keep the referenced stream alive while this class is being used.
- Parameters
stream_ptr (int) – Integer representation of the cudaStream_t value. allocated externally.
device (torch.device or int, optional) – the device where the stream was originally allocated. if device is specified incorrectly, subsequent launches using this stream may fail.
-
query
()¶ Checks if all the work submitted has been completed.
- Returns
A boolean indicating if all kernels in this stream are completed.
-
record_event
(event=None)¶ Records an event.
- Parameters
event (torch.cuda.Event, optional) – event to record. If not given, a new one will be allocated.
- Returns
Recorded event.
-
synchronize
()¶ Wait for all the kernels in this stream to complete.
Note
This is a wrapper around
cudaStreamSynchronize()
: see CUDA Stream documentation for more info.
-
wait_event
(event)¶ Makes all future work submitted to the stream wait for an event.
- Parameters
event (torch.cuda.Event) – an event to wait for.
Note
This is a wrapper around
cudaStreamWaitEvent()
: see CUDA Stream documentation for more info.This function returns without waiting for
event
: only future operations are affected.
-
wait_stream
(stream)¶ Synchronizes with another stream.
All future work submitted to this stream will wait until all kernels submitted to a given stream at the time of call complete.
- Parameters
stream (Stream) – a stream to synchronize.
Note
This function returns without waiting for currently enqueued kernels in
stream
: only future operations are affected.