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 中當前排隊的核心完成的情況下返回：隻影響未來的操作。