Repository: cupy/cupy

Issue:
Detection of device synchronization
Device synchronization is a source of performance degradation, but currently even some of CuPy functions implicitly (whether intentionally or not) do that (#2797).

It would be nice to have a mechanism to detect device synchronization.

### Proposal

Add a context manager within which any operation that synchronizes the device causes an error. The context manager sets a thread-local configuration variable.

```py
with cupyx.allow_synchronize(False):
    ... # operations that are not supposed to synchronize
```


In low-level CuPy code that cause synchronization like `ndarray.get`, insert `_declare_synchronize()` call as follows:

```py
cpdef _declare_synchronize():
    ... # Raise an error if within allow_synchronize(False).

cdef class ndarray:
    cpdef get(self, ...):
        _declare_synchronize()
       ...
```

### Application in unit tests

Given the above implemented, we could apply it in `testing.numpy_cupy_equal` for example. If the test code has expected device synchronization, the test implementation can enclose the code with `cupyx.allow_synchronize(True)` explicitly.

### Notes
* `ndarray.get` may copy asynchronously if certain conditions are met (if using a non-default stream and given an explicit `out` with pinned memory). If's difficult for the code to predict whether it's actually synchronous or not. Maybe it's safe to separate asynchronous behavior as another method like `ndarray.get_async`.


Discussion:
I think this is great and could be pretty useful to improve performance of user codes.

One minor detail that just comes to my mind is that instead of explicitly declare every method that synchronizes, could we make the cuda runtime call used to synchronize devices to check the flag set by the context manager?

This will solve the second issue, but it will not detect calls that “could” synchronize but they are not doing it due to certain runtime conditions. (I don’t know if your proposal is just to error on anything that may synchronize which is also legit).
+1 for the runtime module to check the flag if possible. This seems more robust, and also less painful for implementing new functions. At least we don’t need to chase after which function syncs.
Even `runtime` API wrapper doesn't know if the call synchronizes (because it depends on the pointer attributes), so it won't solve the problem.

It doesn't synchronize only if the memory is pinned memory in memcpy calls? or are there other cases where behavior is unclear.

I guess that being able to know if the host memory is page locked or not should be a sufficient condition for this, but I am afraid this is not easy 😢
I'm not sure about other cases, but checking pointer attributes might not be roboust, and also incurs some overhead as well.



Resolving patch:
diff --git a/cupy/core/syncdetect.py b/cupy/core/syncdetect.py
new file mode 100644
--- /dev/null
+++ b/cupy/core/syncdetect.py
@@ -0,0 +1,54 @@
+import contextlib
+import threading
+
+
+_thread_local = threading.local()
+
+
+class DeviceSynchronized(RuntimeError):
+    """Raised when device synchronization is detected while disallowed.
+
+    .. seealso:: :func:`cupyx.allow_synchronize`
+
+    """
+
+    def __init__(self, message=None):
+        if message is None:
+            message = 'Device synchronization was detected while disallowed.'
+        super().__init__(message)
+
+
+def _is_allowed():
+    # Returns whether device synchronization is allowed in the current thread.
+    try:
+        return _thread_local.allowed
+    except AttributeError:
+        _thread_local.allowed = True
+        return True
+
+
+def _declare_synchronize():
+    # Raises DeviceSynchronized if device synchronization is disallowed in
+    # the current thread.
+    if not _is_allowed():
+        raise DeviceSynchronized()
+
+
+@contextlib.contextmanager
+def allow_synchronize(allow):
+    """Allows or disallows device synchronization temporarily in the current \
+thread.
+
+    If device synchronization is detected, :class:`cupyx.DeviceSynchronized`
+    will be raised.
+
+    Note that there can be false negatives and positives.
+    Device synchronization outside CuPy will not be detected.
+    """
+
+    old = _is_allowed()
+    _thread_local.allowed = allow
+    try:
+        yield
+    finally:
+        _thread_local.allowed = old
diff --git a/cupyx/__init__.py b/cupyx/__init__.py
--- a/cupyx/__init__.py
+++ b/cupyx/__init__.py
@@ -12,3 +12,7 @@
 from cupyx._ufunc_config import errstate  # NOQA
 from cupyx._ufunc_config import geterr  # NOQA
 from cupyx._ufunc_config import seterr  # NOQA
+
+
+from cupy.core.syncdetect import allow_synchronize  # NOQA
+from cupy.core.syncdetect import DeviceSynchronized  # NOQA


Test patch:
diff --git a/tests/cupy_tests/core_tests/test_syncdetect.py b/tests/cupy_tests/core_tests/test_syncdetect.py
new file mode 100644
--- /dev/null
+++ b/tests/cupy_tests/core_tests/test_syncdetect.py
@@ -0,0 +1,33 @@
+import unittest
+
+import pytest
+
+import cupy
+import cupyx
+
+
+class TestSyncDetect(unittest.TestCase):
+
+    def test_disallowed(self):
+        a = cupy.array([2, 3])
+        with cupyx.allow_synchronize(False):
+            with pytest.raises(cupyx.DeviceSynchronized):
+                a.get()
+
+    def test_allowed(self):
+        a = cupy.array([2, 3])
+        with cupyx.allow_synchronize(True):
+            a.get()
+
+    def test_nested_disallowed(self):
+        a = cupy.array([2, 3])
+        with cupyx.allow_synchronize(True):
+            with cupyx.allow_synchronize(False):
+                with pytest.raises(cupyx.DeviceSynchronized):
+                    a.get()
+
+    def test_nested_allowed(self):
+        a = cupy.array([2, 3])
+        with cupyx.allow_synchronize(False):
+            with cupyx.allow_synchronize(True):
+                a.get()