流¶

class torch.mtia.Stream(device, *, priority)¶

一个按序执行相关任务的异步队列，遵循先入先出（FIFO）原则。它可以控制或同步其他 Stream 的执行，或阻塞当前主机线程，以确保任务的正确排序。它支持使用 with 语句作为上下文管理器，确保 with 块内的算子在该流上运行。

有关适用于所有设备的精确语义的详细描述，请参阅对 CUDA 语义行为的深入说明。

参数

device (torch.device, 可选) – Stream 所需的设备。如果未给出，将使用当前的加速器类型。
priority (int, 可选) – 流的优先级，应为 0 或负数，负数表示优先级更高。默认情况下，流的优先级为 0。

返回

一个 torch.Stream 对象。

返回类型

流

示例

>>> with torch.Stream(device='cuda') as s_cuda:
>>>     a = torch.randn(10, 5, device='cuda')
>>>     b = torch.randn(5, 10, device='cuda')
>>>     c = torch.mm(a, b)

query() → bool¶

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

返回: 一个布尔值，指示此流中的所有内核是否已完成。
返回类型: bool

示例

>>> s_cuda = torch.Stream(device='cuda')
>>> s_cuda.query()
True

record_event(event) → Event¶

记录一个事件。将其加入 Stream 队列，以允许从 FIFO 队列的当前点进行进一步同步。

参数: event (torch.Event, 可选) – 要记录的事件。如果未给出，将分配一个新的事件。
返回: 记录的事件。
返回类型: Event

示例

>>> s_cuda = torch.Stream(device='cuda')
>>> e_cuda = s_cuda.record_event()

synchronize() → None¶

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

示例

>>> s_cuda = torch.Stream(device='cuda')
>>> s_cuda.synchronize()

wait_event(event) → None¶

使提交到此流的所有后续工作等待某个事件。

参数: event (torch.Event) – 要等待的事件。

示例

>>> s1_cuda = torch.Stream(device='cuda')
>>> s2_cuda = torch.Stream(device='cuda')
>>> e_cuda = s1_cuda.record_event()
>>> s2_cuda.wait_event(e_cuda)

wait_stream(stream) → None¶

与另一个流同步。所有提交到此流的后续工作将等待给定的流中所有已提交的内核完成后才开始执行。

参数: stream (torch.Stream) – 要同步的流。

示例

>>> s1_cuda = torch.Stream(device='cuda')
>>> s2_cuda = torch.Stream(device='cuda')
>>> s2_cuda.wait_stream(s1_cuda)

流¶

文档

教程

资源