ExternalStream — PyTorch 2.7 documentation (original) (raw)

class torch.cuda.ExternalStream(stream_ptr, device=None, **kwargs)[source][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()[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 one will be allocated.

Returns

Recorded event.

synchronize()[source]¶

Wait for all the kernels in this stream to complete.

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

This function returns without waiting for event: only future operations are affected.

wait_stream(stream)[source]¶

Synchronize 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.