ExternalStream

class torch.cuda.ExternalStream(stream_ptr, device=None, **kwargs)

包装一个外部分配的 CUDA 流。

此类用于包装在其他库中分配的流，以方便数据交换和多库交互。

注意

此类不管理流的生命周期，用户在使用此类时有责任确保引用的流保持活动状态。

参数

• stream_ptr (int) – 外部分配的 cudaStream_t 值的整数表示。

• device (torch.device 或 int, 可选) – 流最初分配的设备。如果设备指定不正确，则使用此流的后续启动可能会失败。

query()

检查所有提交的工作是否已完成。

返回

一个布尔值，指示此流中的所有内核是否已完成。

返回类型

布尔值

record_event(event=None)

Record an event.

参数

event (torch.cuda.Event, 可选) – 要记录的事件。如果未给出，将分配一个新的事件。

返回

记录的事件。

synchronize()

等待此流中的所有内核完成。

注意

这是 cudaStreamSynchronize() 的包装器：有关更多信息，请参阅 CUDA Stream 文档。

wait_event(event)

使提交到此流的所有未来工作等待一个事件。

参数

event (torch.cuda.Event) – 要等待的事件。

注意

这是 cudaStreamWaitEvent() 的包装器：有关更多信息，请参阅 CUDA Stream 文档。

此函数将在不等待 event 的情况下返回：仅影响未来的操作。

wait_stream(stream)

与另一个流同步。

提交到此流的所有未来工作将等待，直到调用时提交到给定流的所有内核完成。

参数

stream (Stream) – 要同步的流。

注意

此函数将在不等待 stream 中当前排队的内核的情况下返回：仅影响未来的操作。