ExternalStream

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

外部分配的 CUDA Stream 的包装类。

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

注意

此类不管理 Stream 的生命周期，用户有责任在使用此类期间保持引用的 Stream 存活。

参数

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

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

query()

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

返回

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

返回类型

bool

record_event(event=None)

记录一个事件。

参数

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

返回

记录的事件。

synchronize()

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

注意

这是 cudaStreamSynchronize() 的包装：请参阅 CUDA Stream 文档了解更多信息。

wait_event(event)

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

参数

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

注意

这是 cudaStreamWaitEvent() 的包装：请参阅 CUDA Stream 文档了解更多信息。

此函数在不等 event 完成的情况下返回：只影响未来的操作。

wait_stream(stream)

与另一个 Stream 同步。

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

参数

stream (Stream) – 要同步的 Stream。

注意

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