ExternalStream

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

外部配置的 CUDA 串流的封裝器。

此類別用於封裝在其他函式庫中配置的串流，以促進資料交換和多函式庫互動。

注意

此類別不管理串流生命週期，使用者有責任在使用此類別時保持參考的串流處於活動狀態。

參數

• stream_ptr (int) – cudaStream_t 值的整數表示。在外部配置。

• device (torch.device 或 int, optional) – 串流最初配置的裝置。如果裝置指定不正確，則後續使用此串流的啟動可能會失敗。

query()

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

返回值

一個布林值，指示此串流中的所有核心是否已完成。

返回類型

bool

record_event(event=None)

記錄一個事件。

參數

event (torch.cuda.Event, 可選) – 要記錄的事件。 如果未提供，將分配一個新的事件。

返回值

已記錄的事件。

synchronize()

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

注意

這是 cudaStreamSynchronize() 的一個封裝器：更多資訊請參閱 CUDA 串流文檔。

wait_event(event)

使所有未來提交到串流的工作等待一個事件。

參數

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

注意

這是 cudaStreamWaitEvent() 的一個封裝器：更多資訊請參閱 CUDA 串流文檔。

此函數返回時不會等待 event：只會影響未來的操作。

wait_stream(stream)

與另一個串流同步。

所有未來提交到此串流的工作將等待，直到在調用時提交到給定串流的所有核心完成。

參數

stream (Stream) – 要同步的串流。

注意

此函數返回時不會等待 stream 中目前排隊的核心：只會影響未來的操作。