ExternalStream

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

围绕外部分配的CUDA流进行的封装。

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

注意

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

Parameters

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

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

query()

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

Returns

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

record_event(event=None)

记录一个事件。

Parameters

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

Returns

记录的事件。

synchronize()

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

注意

这是对 cudaStreamSynchronize() 的封装：更多信息请参见 CUDA Stream 文档。

wait_event(event)

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

Parameters

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

注意

这是对 cudaStreamWaitEvent() 的封装：更多信息请参见 CUDA Stream 文档。

此函数在返回时不等待 event：只有未来的操作会受到影响。

wait_stream(stream)

与另一个流同步。

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

Parameters

流 (流) – 要同步的流。

注意

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