Repository: numba/numba

Issue:
CUDA: assignment to device array elements is asynchronous, and shouldn't be
Running the following example on Windows with a Quadro RTX 8000:

```python
from numba import cuda
import numpy as np

ary = cuda.mapped_array(2, dtype=np.int32)
ary[:] = 0
print(ary)

ary_v = ary.view('u1')
ary_v[1] = 1
ary_v[5] = 1
# cuda.synchronize()
print(ary)
assert sum(ary) == 512
```

produces the following output for me:

```
[0 0]
[256   0]
Traceback (most recent call last):
  File ".\repro.py", line 13, in <module>
    assert sum(ary) == 512
```

Uncommenting the call to `cuda.synchronize()` yields the following output instead:

```
[0 0]
[256 256]
```

This is because `__setitem__` for a device array uses a kernel to set the element: https://github.com/numba/numba/blob/master/numba/cuda/cudadrv/devicearray.py#L629


Discussion:
From https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#mapped-memory:

> Since mapped page-locked memory is shared between host and device however, the application must synchronize memory accesses using streams or events (see Asynchronous Concurrent Execution) to avoid any potential read-after-write, write-after-read, or write-after-write hazards. 

I think we can/should do something internally in Numba to satisfy this requirement.

Resolving patch:
diff --git a/numba/cuda/cudadrv/devicearray.py b/numba/cuda/cudadrv/devicearray.py
--- a/numba/cuda/cudadrv/devicearray.py
+++ b/numba/cuda/cudadrv/devicearray.py
@@ -567,12 +567,20 @@ def __setitem__(self, key, value):
     def setitem(self, key, value, stream=0):
         """Do `__setitem__(key, value)` with CUDA stream
         """
-        return self._so_getitem(key, value, stream)
+        return self._do_setitem(key, value, stream=stream)
 
     def _do_setitem(self, key, value, stream=0):
 
         stream = self._default_stream(stream)
 
+        # If the array didn't have a default stream, and the user didn't provide
+        # a stream, then we will use the default stream for the assignment
+        # kernel and synchronize on it.
+        synchronous = not stream
+        if synchronous:
+            ctx = devices.get_context()
+            stream = ctx.get_default_stream()
+
         # (1) prepare LHS
 
         arr = self._dummy.__getitem__(key)
@@ -613,6 +621,8 @@ def _do_setitem(self, key, value, stream=0):
 
         n_elements = functools.reduce(operator.mul, lhs.shape, 1)
         _assign_kernel(lhs.ndim).forall(n_elements, stream=stream)(lhs, rhs)
+        if synchronous:
+            stream.synchronize()
 
 
 class IpcArrayHandle(object):


Test patch:
diff --git a/numba/cuda/tests/cudadrv/test_cuda_array_slicing.py b/numba/cuda/tests/cudadrv/test_cuda_array_slicing.py
--- a/numba/cuda/tests/cudadrv/test_cuda_array_slicing.py
+++ b/numba/cuda/tests/cudadrv/test_cuda_array_slicing.py
@@ -3,7 +3,8 @@
 import numpy as np
 
 from numba import cuda
-from numba.cuda.testing import unittest, CUDATestCase
+from numba.cuda.testing import unittest, CUDATestCase, skip_on_cudasim
+from unittest.mock import patch
 
 
 class CudaArrayIndexing(CUDATestCase):
@@ -308,6 +309,67 @@ def test_incompatible_shape(self):
                 "dimension 5",  # simulator
             ])
 
+    @skip_on_cudasim('cudasim does not use streams and operates synchronously')
+    def test_sync(self):
+        # There should be a synchronization when no stream is supplied
+        darr = cuda.to_device(np.arange(5))
+
+        with patch.object(cuda.cudadrv.driver.Stream, 'synchronize',
+                          return_value=None) as mock_sync:
+            darr[0] = 10
+
+        mock_sync.assert_called_once()
+
+    @skip_on_cudasim('cudasim does not use streams and operates synchronously')
+    def test_no_sync_default_stream(self):
+        # There should not be a synchronization when the array has a default
+        # stream, whether it is the default stream, the legacy default stream,
+        # the per-thread default stream, or another stream.
+        streams = (cuda.stream(), cuda.default_stream(),
+                   cuda.legacy_default_stream(),
+                   cuda.per_thread_default_stream())
+
+        for stream in streams:
+            darr = cuda.to_device(np.arange(5), stream=stream)
+
+            with patch.object(cuda.cudadrv.driver.Stream, 'synchronize',
+                              return_value=None) as mock_sync:
+                darr[0] = 10
+
+            mock_sync.assert_not_called()
+
+    @skip_on_cudasim('cudasim does not use streams and operates synchronously')
+    def test_no_sync_supplied_stream(self):
+        # There should not be a synchronization when a stream is supplied for
+        # the setitem call, whether it is the default stream, the legacy default
+        # stream, the per-thread default stream, or another stream.
+        streams = (cuda.stream(), cuda.default_stream(),
+                   cuda.legacy_default_stream(),
+                   cuda.per_thread_default_stream())
+
+        for stream in streams:
+            darr = cuda.to_device(np.arange(5))
+
+            with patch.object(cuda.cudadrv.driver.Stream, 'synchronize',
+                              return_value=None) as mock_sync:
+                darr.setitem(0, 10, stream=stream)
+
+            mock_sync.assert_not_called()
+
+    @unittest.skip('Requires PR #6367')
+    def test_issue_6505(self):
+        # On Windows, the writes to ary_v would not be visible prior to the
+        # assertion, due to the assignment being done with a kernel launch that
+        # returns asynchronously - there should now be a sync after the kernel
+        # launch to ensure that the writes are always visible.
+        ary = cuda.mapped_array(2, dtype=np.int32)
+        ary[:] = 0
+
+        ary_v = ary.view('u1')
+        ary_v[1] = 1
+        ary_v[5] = 1
+        self.assertEqual(sum(ary), 512)
+
 
 if __name__ == '__main__':
     unittest.main()