cupy.cuda.Stream#

cupy.cuda.Stream(null=False, non_blocking=False, ptds=False)[source]#

CUDA 流。

这个类以 RAII 的方式处理 CUDA 流句柄,即当 Stream 实例被 GC 销毁时,其句柄也会被销毁。

注意,如果 nullptds 都为 False,则会创建一个普通的新流。

参数:
  • null (bool) – 如果为 True,则该流是一个空流(即与所有流同步的默认流)。注意,你也可以使用 Stream.null 单例对象,而不是创建一个新的空流对象。

  • ptds (bool) – 如果为 TruenullFalse,则使用每线程默认流。注意,你也可以使用 Stream.ptds 单例对象,而不是创建一个新的每线程默认流对象。

  • non_blocking (bool) – 如果为 Truenullptds 都为 False,则该流不与 NULL 流同步。

变量:
  • ~Stream.ptr (intptr_t) – 原始流句柄。

  • ~Stream.device_id (int) – 创建流的设备的 ID。值 -1 用于单例流对象。

方法

__enter__(self)#
__exit__(self, *args)#
add_callback(self, callback, arg)#

添加一个回调,在所有排队的工作完成后调用。

参数:
  • callback (function) – 回调函数。它必须接受三个参数(Stream 对象、int 错误状态和用户数据对象),且无返回值。

  • arg (object) – 回调函数的参数。

注意

如果可能,请使用 launch_host_func() 方法而不是此方法,因为此方法在 CUDA 未来版本中可能会被弃用和移除。

begin_capture(self, mode=None)#

开始流捕获以构建 CUDA 图。

调用此函数必须与调用 end_capture() 配对以完成捕获。

# create a non-blocking stream for the purpose of capturing
s1 = cp.cuda.Stream(non_blocking=True)
with s1:
    s1.begin_capture()
    # ... perform operations to construct a graph ...
    g = s1.end_capture()

# the returned graph can be launched on any stream (including s1)
g.launch(stream=s1)
s1.synchronize()

s2 = cp.cuda.Stream()
with s2:
    g.launch()
s2.synchronize()
参数:

mode (int) – 流捕获模式。默认为 streamCaptureModeRelaxed

注意

在流捕获期间,不允许同步的设备-主机传输。这对 CuPy API 有特殊的意义,因为一些内部需要同步传输的函数将无法按预期工作并会引发异常。有关 CUDA 流捕获的更多限制,请参阅 CUDA 编程指南。

注意

目前 HIP 不支持此功能。

end_capture(self)#

结束流捕获并检索构建的 CUDA 图。

返回:

封装捕获工作的 CUDA 图对象。

返回类型:

cupy.cuda.Graph

注意

目前 HIP 不支持此功能。

另请参阅

cudaStreamEndCapture()

is_capturing(self)#

检查流是否正在捕获。

返回:

如果成功查询捕获状态,则返回的值表示捕获状态。如果此类查询非法,则可能引发异常,详细信息请参阅 CUDA 编程指南。

返回类型:

bool

launch_host_func(self, callback, arg)#

在所有排队工作完成后,在主机上启动回调。

参数:
  • callback (function) – 回调函数。它必须只接受一个参数(用户数据对象),且无返回值。

  • arg (object) – 回调函数的参数。

注意

如果可能,建议使用此方法而不是 add_callback(),后者在 CUDA 未来版本中可能会被弃用和移除。

另请参阅

cudaLaunchHostFunc()

record(self, event=None)#

在流上记录一个事件。

参数:

event (Nonecupy.cuda.Event) – CUDA 事件。如果为 None,则创建一个新的普通事件并使用。

返回:

记录的事件。

返回类型:

cupy.cuda.Event

synchronize(self)#

等待流完成所有排队的工作。

use(self)#

使此流成为当前流。

如果想临时切换流,请使用 with 语句。

wait_event(self, event)#

使流等待一个事件。

此流上的未来工作将在事件完成后执行。

参数:

event (cupy.cuda.Event) – CUDA 事件。

__eq__(self, other)#
__ne__(value, /)#

返回 self!=value。

__lt__(value, /)#

返回 self<value。

__le__(value, /)#

返回 self<=value。

__gt__(value, /)#

返回 self>value。

__ge__(value, /)#

返回 self>=value。

属性

done#

如果此流上的所有工作已完成,则为 True。

null = <Stream 0 (device -1)>#
ptds = <Stream 2 (device -1)>#