使用 with 语句的 Cupy 流同步

Question

这两个代码是否等价？

代码 1：

with cp.cuda.Stream(non_blocking=False) as stream:
    # do stuff
    stream.synchronize()

代码 2：

with cp.cuda.Stream(non_blocking=False) as stream:
    # do stuff
stream.synchronize()

Answer 1

是的，这两个是等价的。 with cp.cuda.Stream(non_blocking=False) as stream: 切换 CuPy 的 current stream，即块内的所有代码将在 stream.

上执行

在 CuPy 中，流在退出上下文管理器时不会自动同步，因此用户可以完全控制同步时间。这在处理多个流时很方便，例如：

s1 = cp.cuda.Stream(non_blocking=False)
s2 = cp.cuda.Stream(non_blocking=False)

for i in range(10):
    with s1:
        # do some work
    with s2:
        # do some work

s1.synchronize()
s2.synchronize()

使用 with 语句的 Cupy 流同步

Cupy streams synchronization using with statement

cuda

stream

cupy