外部流

类torch.cuda.ExternalStream(stream_ptr, device=None, **kwargs)[源代码]

CUDA流的外部分配封装器。

此类用于封装在其他库中分配的流，以促进数据交换和多个库之间的交互。

注意

此类不管理流的生命周期。在使用此类期间，确保引用的流保持活跃是用户的责任。

参数

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

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

查询()

检查所有提交的任务是否已经完成。

返回值

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

返回类型

bool

record_event(event=None)

记录事件。

参数

event (torch.cuda.Event, 可选) – 需要记录的事件。如果不指定，则会自动创建一个新事件。

返回值

已记录的事件

同步()

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

注意

这是对 cudaStreamSynchronize() 的封装：更多详情请参阅 CUDA Stream 文档。

等待事件(事件)

让未来提交到流中的所有任务都等待一个事件。

参数

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

注意

这是一个围绕cudaStreamWaitEvent()的封装：更多详情请参阅CUDA Stream文档。

此函数不会等待event就直接返回：只会影响后续的操作。

wait_stream(stream)

与其他流进行同步。

提交到此流的所有 future 工作都将等待在调用时已提交到该流的全部内核完成。

参数

stream (Stream) - 要进行同步的流。

注意

此函数无需等待当前在 stream 中已排队的内核即可返回，仅影响未来的操作。