Cupy streams synchronization using with statement

Question

Are the two codes equivalent?

Code 1:

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

Code 2:

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

Answer 1

Yes, these two are equivalent. with cp.cuda.Stream(non_blocking=False) as stream: switches the current stream of CuPy, i.e., all the code within the block will be performed on stream.

In CuPy, streams are not synchronized automatically on exiting the context manager so that users can fully control the synchronization timing. This is handy when dealing with multiple streams, e.g.:

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()