Repository: cupy/cupy

Issue:
Outdated (?) documentation on `ndarray.get()`/`asnumpy()` synchronization behavior    
### Description

Currently, this is what the doc says about the stream argument (in both APIs):
<img width="536" alt="截圖 2023-08-26 上午1 03 46" src="https://github.com/cupy/cupy/assets/5534781/ccaa7d9f-74fb-4a12-a02c-40528a89c74f">
the problematic part is "otherwise, the copy is synchronous" which is not true, because the current stream is used and the copy is synchronous only when the current stream is the null stream:
https://github.com/cupy/cupy/blob/501dc69aa1246607638ced5298c48572e8bb079f/cupy/_core/core.pyx#L1851-L1858



### Idea or request for content

I suggest that we just fix the docs to match the current behavior (which existed for years) since no one came to complain 🙂


Discussion:
Thanks for the catch, sounds good to me :+1:
> I suggest that we just fix the docs to match the current behavior (which existed for years) since no one came to complain 

@kmaehashi @emcastillo @takagi Unfortunately, I feel obliged to revisit this issue and take back the above statement, because we actually had hit bugs. The first one is discovered recently in cuQuantum Python:
https://github.com/NVIDIA/cuQuantum/blob/c202ee7ad23ce9a5072a32b45d36e8310a20eeb5/python/cuquantum/cutensornet/tensor_network.py#L774-L775
and the second one was from an internal study of CuPy (sorry for lack of pointers). 

Basically, any usage pattern like this
```py
with cp.cuda.Stream(non_blocking=True):
    out_d = cp.do_something(...)
    out_h = cp.asnumpy(out_d)  # same for using get()
    ... = np.do_something(out_h)
```
(which is all over the place) is technically a race, because `cp.do_something()` could be a compute-intensive workload on GPU, but running asynchronously so that by the time `np.do_something()` fires, the actual data is still being produced on GPU.

Unfortunately, neither the [data transfer tutorial](https://docs.cupy.dev/en/stable/user_guide/basic.html#move-array-from-a-device-to-the-host) nor the API docs had mentioned about the potential race due to asynchronousity, and most users (including experienced developers like me 😞) would not pay attention to sync after `asnumpy()`/`get()`. 

In other to be backward compatible (no user code change) while ensuring correctness, I propose the following two solutions:
1. Deem the current API docs (see the screenshot in the beginning of this issue) correct, this means
   - any advanced users seeking async D2H copies should pass a `stream` object explicitly (as already hinted)
   - when `stream` is not given, don't do any special treatment for stream 0, we always call the blocking `ptr.copy_to_host()`
2. Remove the `stream` argument in favor of a bool `nonblocking`, following PyTorch's [`Tensor.to()`](https://pytorch.org/docs/stable/generated/torch.Tensor.to.html#torch.Tensor.to)
   - when `nonblocking` is `True` (default), copy is blocking
   - when `nonblocking` is `False`, copy is nonblocking

In either case, the copy shall be done on the given stream (if any) or the current stream, and users are required to set up the stream dependency (as they would do for any async APIs). Please share your thoughts, thank you 🙂
My apologies, I missed your mention. :bow:
Let me confirm if my understanding is correct: when performing D2H copy on non-default stream, `cudaMemcpyAsync` is synchronous with respect to the host thread if `non_blocking=False`, and asynchronous if `non_blocking=True`?

How about adding `async=False` argument to `ndarray.get()`/`asnumpy()`?
Hi @kmaehashi 

> Let me confirm if my understanding is correct: when performing D2H copy on non-default stream, `cudaMemcpyAsync` is synchronous with respect to the host thread if `non_blocking=False`, and asynchronous if `non_blocking=True`?

No, it's simpler than that. Regardless what `non_blocking` is (which refers to `cudaStreamNonBlocking` and is w.r.t the legacy default stream, not the host), if a non-default stream is in use, currently the D2H copy is asynchronous.

> How about adding `async=False` argument to `ndarray.get()`/`asnumpy()`?

Fine with me, I can work on this once we agree to fix it 🙂 
> No, it's simpler than that. Regardless what `non_blocking` is (which refers to `cudaStreamNonBlocking` and is w.r.t the legacy default stream, not the host), if a non-default stream is in use, currently the D2H copy is asynchronous.

Thanks Leo, got it! The current implementation is surprising to me as well...

> > How about adding `async=False` argument to `ndarray.get()`/`asnumpy()`?
> 
> Fine with me, I can work on this once we agree to fix it 🙂

I think `async=False` is what most users expect, and also intuitive for advanced users (rather than asking them to explicitly specify `stream`). This is also consistent with the proposal of adding `async=True` to `cupy.asarray` & `cupy.array`  (https://github.com/cupy/cupy/issues/7821#issuecomment-1695185384).

Resolving patch:
diff --git a/cupy/__init__.py b/cupy/__init__.py
--- a/cupy/__init__.py
+++ b/cupy/__init__.py
@@ -792,29 +792,35 @@ def ndim(a):
 from cupy._core import from_dlpack  # NOQA
 
 
-def asnumpy(a, stream=None, order='C', out=None):
+def asnumpy(a, stream=None, order='C', out=None, *, blocking=True):
     """Returns an array on the host memory from an arbitrary source array.
 
     Args:
         a: Arbitrary object that can be converted to :class:`numpy.ndarray`.
-        stream (cupy.cuda.Stream): CUDA stream object. If it is specified, then
-            the device-to-host copy runs asynchronously. Otherwise, the copy is
-            synchronous. Note that if ``a`` is not a :class:`cupy.ndarray`
+        stream (cupy.cuda.Stream): CUDA stream object. If given, the
+            stream is used to perform the copy. Otherwise, the current
+            stream is used. Note that if ``a`` is not a :class:`cupy.ndarray`
             object, then this argument has no effect.
         order ({'C', 'F', 'A'}): The desired memory layout of the host
-            array. When ``order`` is 'A', it uses 'F' if ``a`` is
-            fortran-contiguous and 'C' otherwise.
+            array. When ``order`` is 'A', it uses 'F' if the array is
+            fortran-contiguous and 'C' otherwise. The ``order`` will be
+            ignored if ``out`` is specified.
         out (numpy.ndarray): The output array to be written to. It must have
             compatible shape and dtype with those of ``a``'s.
+        blocking (bool): If set to ``False``, the copy runs asynchronously
+            on the given (if given) or current stream, and users are
+            responsible for ensuring the stream order. Default is ``True``,
+            so the copy is synchronous (with respect to the host).
 
     Returns:
         numpy.ndarray: Converted array on the host memory.
 
     """
     if isinstance(a, ndarray):
-        return a.get(stream=stream, order=order, out=out)
+        return a.get(stream=stream, order=order, out=out, blocking=blocking)
     elif hasattr(a, "__cuda_array_interface__"):
-        return array(a).get(stream=stream, order=order, out=out)
+        return array(a).get(
+            stream=stream, order=order, out=out, blocking=blocking)
     else:
         temp = _numpy.asarray(a, order=order)
         if out is not None:


Test patch:
diff --git a/tests/cupy_tests/core_tests/test_ndarray_get.py b/tests/cupy_tests/core_tests/test_ndarray_get.py
--- a/tests/cupy_tests/core_tests/test_ndarray_get.py
+++ b/tests/cupy_tests/core_tests/test_ndarray_get.py
@@ -12,9 +12,9 @@ class TestArrayGet(unittest.TestCase):
     def setUp(self):
         self.stream = cuda.Stream.null
 
-    def check_get(self, f, stream, order='C'):
+    def check_get(self, f, stream, order='C', blocking=True):
         a_gpu = f(cupy)
-        a_cpu = a_gpu.get(stream, order=order)
+        a_cpu = a_gpu.get(stream, order=order, blocking=blocking)
         if stream:
             stream.synchronize()
         b_cpu = f(numpy)
@@ -45,6 +45,15 @@ def contiguous_array(xp):
             return testing.shaped_arange((3,), xp, dtype, order)
         self.check_get(contiguous_array, self.stream, order)
 
+    @testing.for_orders('CFA')
+    @testing.for_all_dtypes()
+    def test_contiguous_array_stream_nonblocking(self, dtype, order):
+        # Note: This is just a smoking gun test, the real test is done for
+        # testing cupy.asnumpy(), which under the hood calls .get().
+        def contiguous_array(xp):
+            return testing.shaped_arange((3,), xp, dtype, order)
+        self.check_get(contiguous_array, self.stream, order, False)
+
     @testing.for_orders('CFA')
     @testing.for_all_dtypes()
     def test_non_contiguous_array_stream(self, dtype, order):
diff --git a/tests/cupy_tests/test_numpy_interop.py b/tests/cupy_tests/test_numpy_interop.py
--- a/tests/cupy_tests/test_numpy_interop.py
+++ b/tests/cupy_tests/test_numpy_interop.py
@@ -1,3 +1,4 @@
+import contextlib
 import unittest
 
 import numpy
@@ -148,3 +149,32 @@ def test_asnumpy_out(self):
         testing.assert_array_equal(x, y)
         assert isinstance(y.base, cupy.cuda.PinnedMemoryPointer)
         assert y.base.ptr == y.ctypes.data
+
+    @pytest.mark.parametrize('blocking', (True, False))
+    def test_asnumpy_blocking(self, blocking):
+        prefactor = 4
+        a = cupy.random.random(prefactor*128*1024*1024, dtype=cupy.float64)
+        cupy.cuda.Device().synchronize()
+
+        # Idea: perform D2H copy on a nonblocking stream, during which we try
+        # to "corrupt" the host data via NumPy operation. If the copy is
+        # properly ordered, corruption would not be possible. Here we craft a
+        # problem size and use pinned memory to ensure the failure can be
+        # always triggered. (The CUDART API reference ("API synchronization
+        # behavior") states that copying between device and pageable memory
+        # "might be" synchronous, whereas between device and page-locked
+        # memory "should be" fully asynchronous.)
+        s = cupy.cuda.Stream(non_blocking=True)
+        with s:
+            c = cupyx.empty_pinned(a.shape, dtype=a.dtype)
+            cupy.asnumpy(a, out=c, blocking=blocking)
+            c[c.size//2:] = -1.  # potential data race
+        s.synchronize()
+
+        a[c.size//2:] = -1.
+        if not blocking:
+            ctx = pytest.raises(AssertionError)
+        else:
+            ctx = contextlib.nullcontext()
+        with ctx:
+            assert cupy.allclose(a, c)