ExternalStream#

classtorch.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 orderto facilitate data exchange and multi-library interactions.

Note

This class doesn’t manage the stream life-cycle, it is the userresponsibility to keep the referenced stream alive while this class isbeing used.

Parameters

• stream_ptr (int) – Integer representation of thecudaStream_t value.allocated externally.

• device (torch.device orint,optional) – the device where the streamwas originally allocated. If device is specified incorrectly,subsequent launches using this stream may fail.

query()[source]#

Check if all the work submitted has been completed.

Returns

A boolean indicating if all kernels in this stream are completed.

Return type

bool

record_event(event=None)[source]#

Record an event.

Parameters

event (torch.cuda.Event,optional) – event to record. If not given, a new onewill be allocated.

Returns

Recorded event.

synchronize()[source]#

Wait for all the kernels in this stream to complete.

Note

This is a wrapper aroundcudaStreamSynchronize(): seeCUDA Stream documentation for more info.

wait_event(event)[source]#

Make 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 aroundcudaStreamWaitEvent(): seeCUDA Stream documentation for more info.

This function returns without waiting forevent: only futureoperations are affected.

wait_stream(stream)[source]#

Synchronize with another stream.

All future work submitted to this stream will wait until all kernelssubmitted 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 enqueuedkernels instream: only future operations are affected.