Repository: numba/numba

Issue:
CUDA simulator atomic operation differs to documentation
According to [docs](https://numba.pydata.org/numba-doc/dev/cuda/intrinsics.html), the CUDA atomic `add` and `max` operations return the old value at the index location, but the fake [implementation](https://github.com/numba/numba/blob/master/numba/cuda/simulator/kernelapi.py#L106) in the CUDA simulator does not return anything.

Is that intentional in the simulator? Does the actual CUDA operation have a return value?
Thank you.
CUDA simulator atomic operation differs to documentation
According to [docs](https://numba.pydata.org/numba-doc/dev/cuda/intrinsics.html), the CUDA atomic `add` and `max` operations return the old value at the index location, but the fake [implementation](https://github.com/numba/numba/blob/master/numba/cuda/simulator/kernelapi.py#L106) in the CUDA simulator does not return anything.

Is that intentional in the simulator? Does the actual CUDA operation have a return value?
Thank you.


Discussion:
@geromidg thanks for asking about this on the Numba issue tracker. @gmarkall , could you perhaps take a look at this one.
I can - sounds like a bug, I'll see if I can reproduce.
@gmarkall awesome, thanks!
Executable reproducer:

```python
from numba import cuda
import numpy as np


@cuda.jit
def atomic_test(x):
    x[1] = cuda.atomic.add(x, 0, 1)


x = np.zeros(2, dtype=np.int32)
INITIAL = 10
x[0] = INITIAL

atomic_test[1, 1](x)
assert x[1] == INITIAL
```

Works fine on hardware, on the simulator falls over with:

```
$ NUMBA_ENABLE_CUDASIM=1 python repro.py 
Traceback (most recent call last):
  File "repro.py", line 14, in <module>
    atomic_test[1, 1](x)
  File "/home/gmarkall/numbadev/numba/numba/cuda/simulator/kernel.py", line 106, in __call__
    bm.run(grid_point, *fake_args)
  File "/home/gmarkall/numbadev/numba/numba/cuda/simulator/kernel.py", line 271, in run
    reraise(*(t.exception))
  File "/home/gmarkall/numbadev/numba/numba/core/utils.py", line 79, in reraise
    raise value.with_traceback(tb)
  File "/home/gmarkall/numbadev/numba/numba/cuda/simulator/kernel.py", line 170, in run
    super(BlockThread, self).run()
  File "/home/gmarkall/miniconda3/envs/numbaenv/lib/python3.7/threading.py", line 870, in run
    self._target(*self._args, **self._kwargs)
  File "/home/gmarkall/numbadev/numba/numba/cuda/simulator/kernel.py", line 249, in target
    self._f(*args)
  File "repro.py", line 7, in atomic_test
    x[1] = cuda.atomic.add(x, 0, 1)
  File "/home/gmarkall/numbadev/numba/numba/cuda/simulator/cudadrv/devicearray.py", line 72, in __setitem__
    self._item.__setitem__(idx, val)
  File "/home/gmarkall/numbadev/numba/numba/cuda/simulator/cudadrv/devicearray.py", line 131, in __setitem__
    return self._ary_access.__setitem__(idx, val)
TypeError: tid=[0, 0, 0] ctaid=[0, 0, 0]: int() argument must be a string, a bytes-like object or a number, not 'NoneType'
```
I've started working on a fix in https://github.com/gmarkall/numba/tree/grm-issue-5458 - I need to fix up some of the tests I've added first (issues around handling NaNs correctly).
@geromidg thanks for asking about this on the Numba issue tracker. @gmarkall , could you perhaps take a look at this one.
I can - sounds like a bug, I'll see if I can reproduce.
@gmarkall awesome, thanks!
Executable reproducer:

```python
from numba import cuda
import numpy as np


@cuda.jit
def atomic_test(x):
    x[1] = cuda.atomic.add(x, 0, 1)


x = np.zeros(2, dtype=np.int32)
INITIAL = 10
x[0] = INITIAL

atomic_test[1, 1](x)
assert x[1] == INITIAL
```

Works fine on hardware, on the simulator falls over with:

```
$ NUMBA_ENABLE_CUDASIM=1 python repro.py 
Traceback (most recent call last):
  File "repro.py", line 14, in <module>
    atomic_test[1, 1](x)
  File "/home/gmarkall/numbadev/numba/numba/cuda/simulator/kernel.py", line 106, in __call__
    bm.run(grid_point, *fake_args)
  File "/home/gmarkall/numbadev/numba/numba/cuda/simulator/kernel.py", line 271, in run
    reraise(*(t.exception))
  File "/home/gmarkall/numbadev/numba/numba/core/utils.py", line 79, in reraise
    raise value.with_traceback(tb)
  File "/home/gmarkall/numbadev/numba/numba/cuda/simulator/kernel.py", line 170, in run
    super(BlockThread, self).run()
  File "/home/gmarkall/miniconda3/envs/numbaenv/lib/python3.7/threading.py", line 870, in run
    self._target(*self._args, **self._kwargs)
  File "/home/gmarkall/numbadev/numba/numba/cuda/simulator/kernel.py", line 249, in target
    self._f(*args)
  File "repro.py", line 7, in atomic_test
    x[1] = cuda.atomic.add(x, 0, 1)
  File "/home/gmarkall/numbadev/numba/numba/cuda/simulator/cudadrv/devicearray.py", line 72, in __setitem__
    self._item.__setitem__(idx, val)
  File "/home/gmarkall/numbadev/numba/numba/cuda/simulator/cudadrv/devicearray.py", line 131, in __setitem__
    return self._ary_access.__setitem__(idx, val)
TypeError: tid=[0, 0, 0] ctaid=[0, 0, 0]: int() argument must be a string, a bytes-like object or a number, not 'NoneType'
```
I've started working on a fix in https://github.com/gmarkall/numba/tree/grm-issue-5458 - I need to fix up some of the tests I've added first (issues around handling NaNs correctly).

Resolving patch:
diff --git a/numba/cuda/cudadrv/nvvm.py b/numba/cuda/cudadrv/nvvm.py
--- a/numba/cuda/cudadrv/nvvm.py
+++ b/numba/cuda/cudadrv/nvvm.py
@@ -388,18 +388,19 @@ def get(self):
 }
 """
 
+
 ir_numba_atomic_max = """
 define internal {T} @___numba_atomic_{T}_max({T}* %ptr, {T} %val) alwaysinline {{
 entry:
     %ptrval = load volatile {T}, {T}* %ptr
-    ; Check if val is a NaN and return *ptr early if so
-    %valnan = fcmp uno {T} %val, %val
+    ; Check if either value is a NaN and return *ptr early if so
+    %valnan = fcmp uno {T} %val, %ptrval
     br i1 %valnan, label %done, label %lt_check
 
 lt_check:
     %dold = phi {T} [ %ptrval, %entry ], [ %dcas, %attempt ]
-    ; Continue attempts if dold < val or dold is NaN (using ult semantics)
-    %lt = fcmp ult {T} %dold, %val
+    ; Continue attempts if dold < val
+    %lt = fcmp nnan olt {T} %dold, %val
     br i1 %lt, label %attempt, label %done
 
 attempt:
@@ -423,14 +424,14 @@ def get(self):
 define internal {T} @___numba_atomic_{T}_min({T}* %ptr, {T} %val) alwaysinline{{
 entry:
     %ptrval = load volatile {T}, {T}* %ptr
-    ; Check if val is a NaN and return *ptr early if so
-    %valnan = fcmp uno {T} %val, %val
+    ; Check if either value is a NaN and return *ptr early if so
+    %valnan = fcmp uno {T} %val, %ptrval
     br i1 %valnan, label %done, label %gt_check
 
 gt_check:
     %dold = phi {T} [ %ptrval, %entry ], [ %dcas, %attempt ]
-    ; Continue attempts if dold > val or dold is NaN (using ugt semantics)
-    %lt = fcmp ugt {T} %dold, %val
+    ; Continue attempts if dold > val
+    %lt = fcmp nnan ogt {T} %dold, %val
     br i1 %lt, label %attempt, label %done
 
 attempt:
diff --git a/numba/cuda/simulator/kernelapi.py b/numba/cuda/simulator/kernelapi.py
--- a/numba/cuda/simulator/kernelapi.py
+++ b/numba/cuda/simulator/kernelapi.py
@@ -105,29 +105,21 @@ def array(self, shape, dtype):
 class FakeCUDAAtomic(object):
     def add(self, array, index, val):
         with addlock:
+            old = array[index]
             array[index] += val
+        return old
 
     def max(self, array, index, val):
         with maxlock:
-            # CUDA Python's semantics for max differ from Numpy's Python's,
-            # so we have special handling here (CUDA Python treats NaN as
-            # missing data).
-            if np.isnan(array[index]):
-                array[index] = val
-            elif np.isnan(val):
-                return
-            array[index] = max(array[index], val)
+            old = array[index]
+            array[index] = max(old, val)
+        return old
 
     def min(self, array, index, val):
         with minlock:
-            # CUDA Python's semantics for min differ from Numpy's Python's,
-            # so we have special handling here (CUDA Python treats NaN as
-            # missing data).
-            if np.isnan(array[index]):
-                array[index] = val
-            elif np.isnan(val):
-                return
-            array[index] = min(array[index], val)
+            old = array[index]
+            array[index] = min(old, val)
+        return old
 
     def compare_and_swap(self, array, old, val):
         with caslock:
diff --git a/numba/cuda/stubs.py b/numba/cuda/stubs.py
--- a/numba/cuda/stubs.py
+++ b/numba/cuda/stubs.py
@@ -450,10 +450,7 @@ class add(Stub):
     class max(Stub):
         """max(ary, idx, val)
 
-        Perform atomic ary[idx] = max(ary[idx], val). NaN is treated as a
-        missing value, so max(NaN, n) == max(n, NaN) == n. Note that this
-        differs from Python and Numpy behaviour, where max(a, b) is always
-        a when either a or b is a NaN.
+        Perform atomic ary[idx] = max(ary[idx], val).
 
         Supported on int32, int64, uint32, uint64, float32, float64 operands only.
 
@@ -464,12 +461,12 @@ class max(Stub):
     class min(Stub):
         """min(ary, idx, val)
 
-        Perform atomic ary[idx] = min(ary[idx], val). NaN is treated as a
-        missing value, so min(NaN, n) == min(n, NaN) == n. Note that this
-        differs from Python and Numpy behaviour, where min(a, b) is always
-        a when either a or b is a NaN.
+        Perform atomic ary[idx] = min(ary[idx], val).
 
         Supported on int32, int64, uint32, uint64, float32, float64 operands only.
+
+        Returns the old value at the index location as if it is loaded
+        atomically.
         """
 
     class compare_and_swap(Stub):


Test patch:
diff --git a/numba/cuda/tests/cudapy/test_atomics.py b/numba/cuda/tests/cudapy/test_atomics.py
--- a/numba/cuda/tests/cudapy/test_atomics.py
+++ b/numba/cuda/tests/cudapy/test_atomics.py
@@ -391,24 +391,48 @@ def test_atomic_max_double_oneindex(self):
         gold = np.max(vals)
         np.testing.assert_equal(res, gold)
 
-    def test_atomic_max_nan_location(self):
+    # Taken together, _test_atomic_minmax_nan_location and
+    # _test_atomic_minmax_nan_val check that NaNs are treated similarly to the
+    # way they are in Python / NumPy - that is, {min,max}(a, b) == a if either
+    # a or b is a NaN. For the atomics, this means that the max is taken as the
+    # value stored in the memory location rather than the value supplied - i.e.
+    # for:
+    #
+    #    cuda.atomic.{min,max}(ary, idx, val)
+    #
+    # the result will be ary[idx] for either of ary[idx] or val being NaN.
+
+    def _test_atomic_minmax_nan_location(self, func):
+
+        cuda_func = cuda.jit('void(float64[:], float64[:,:])')(func)
+
         vals = np.random.randint(0, 128, size=(1,1)).astype(np.float64)
-        gold = vals.copy().reshape(1)
         res = np.zeros(1, np.float64) + np.nan
-        cuda_func = cuda.jit('void(float64[:], float64[:,:])')(atomic_max)
         cuda_func[1, 1](res, vals)
+        np.testing.assert_equal(res, [np.nan])
 
-        np.testing.assert_equal(res, gold)
+    def _test_atomic_minmax_nan_val(self, func):
+        cuda_func = cuda.jit('void(float64[:], float64[:,:])')(func)
 
-    def test_atomic_max_nan_val(self):
         res = np.random.randint(0, 128, size=1).astype(np.float64)
         gold = res.copy()
         vals = np.zeros((1, 1), np.float64) + np.nan
-        cuda_func = cuda.jit('void(float64[:], float64[:,:])')(atomic_max)
         cuda_func[1, 1](res, vals)
 
         np.testing.assert_equal(res, gold)
 
+    def test_atomic_min_nan_location(self):
+        self._test_atomic_minmax_nan_location(atomic_min)
+
+    def test_atomic_max_nan_location(self):
+        self._test_atomic_minmax_nan_location(atomic_max)
+
+    def test_atomic_min_nan_val(self):
+        self._test_atomic_minmax_nan_val(atomic_min)
+
+    def test_atomic_max_nan_val(self):
+        self._test_atomic_minmax_nan_val(atomic_max)
+
     def test_atomic_max_double_shared(self):
         vals = np.random.randint(0, 32, size=32).astype(np.float64)
         res = np.zeros(1, np.float64)
@@ -443,6 +467,69 @@ def test_atomic_compare_and_swap(self):
         np.testing.assert_array_equal(expect_res, res)
         np.testing.assert_array_equal(expect_out, out)
 
+    # Tests that the atomic add, min, and max operations return the old value -
+    # in the simulator, they did not (see Issue #5458). The max and min have
+    # special handling for NaN values, so we explicitly test with a NaN in the
+    # array being modified and the value provided.
+
+    def _test_atomic_returns_old(self, kernel, initial):
+        x = np.zeros(2, dtype=np.float32)
+        x[0] = initial
+        kernel[1, 1](x)
+        if np.isnan(initial):
+            self.assertTrue(np.isnan(x[1]))
+        else:
+            self.assertEqual(x[1], initial)
+
+    def test_atomic_add_returns_old(self):
+        @cuda.jit
+        def kernel(x):
+            x[1] = cuda.atomic.add(x, 0, 1)
+
+        self._test_atomic_returns_old(kernel, 10)
+
+    def test_atomic_max_returns_old(self):
+        @cuda.jit
+        def kernel(x):
+            x[1] = cuda.atomic.max(x, 0, 1)
+
+        self._test_atomic_returns_old(kernel, 10)
+
+    def test_atomic_max_returns_old_nan_in_array(self):
+        @cuda.jit
+        def kernel(x):
+            x[1] = cuda.atomic.max(x, 0, 1)
+
+        self._test_atomic_returns_old(kernel, np.nan)
+
+    def test_atomic_max_returns_old_nan_val(self):
+        @cuda.jit
+        def kernel(x):
+            x[1] = cuda.atomic.max(x, 0, np.nan)
+
+        self._test_atomic_returns_old(kernel, 10)
+
+    def test_atomic_min_returns_old(self):
+        @cuda.jit
+        def kernel(x):
+            x[1] = cuda.atomic.min(x, 0, 11)
+
+        self._test_atomic_returns_old(kernel, 10)
+
+    def test_atomic_min_returns_old_nan_in_array(self):
+        @cuda.jit
+        def kernel(x):
+            x[1] = cuda.atomic.min(x, 0, 11)
+
+        self._test_atomic_returns_old(kernel, np.nan)
+
+    def test_atomic_min_returns_old_nan_val(self):
+        @cuda.jit
+        def kernel(x):
+            x[1] = cuda.atomic.min(x, 0, 11)
+
+        self._test_atomic_returns_old(kernel, np.nan)
+
 
 if __name__ == '__main__':
     unittest.main()