Merge pull request #1319 from IntelPython/fix/fetch_pfi_returns

Capturing return value from fetch_* atomic functions

Author: Diptorup Deb

numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_atomic_ref_overloads.py

@@ -67,7 +67,7 @@ def _parse_enum_or_int_literal_(literal_int) -> int:
 def _intrinsic_helper(
     ty_context, ty_atomic_ref, ty_val, op_str  # pylint: disable=unused-argument
 ):
-    sig = types.void(ty_atomic_ref, ty_val)
+    sig = ty_atomic_ref.dtype(ty_atomic_ref, ty_val)
 
     def gen(context, builder, sig, args):
         atomic_ref_ty = sig.args[0]
@@ -126,7 +126,7 @@ def gen(context, builder, sig, args):
             args[1],
         ]
 
-        builder.call(func, fn_args)
+        return builder.call(func, fn_args)
 
     return sig, gen
 
@@ -145,7 +145,7 @@ def gen(context, builder, sig, args):
         args_lst[1] = builder.fneg(args[1])
         args = tuple(args_lst)
 
-        gen_fn(context, builder, sig, args)
+        return gen_fn(context, builder, sig, args)
 
     return gen
 

numba_dpex/kernel_api/atomic_ref.py

@@ -68,7 +68,7 @@ def fetch_add(self, val):
         Returns: The original value of the object referenced by the AtomicRef.
 
         """
-        old = self._ref[self._index]
+        old = self._ref[self._index].copy()
         self._ref[self._index] += val
         return old
 
@@ -84,7 +84,7 @@ def fetch_sub(self, val):
         Returns: The original value of the object referenced by the AtomicRef.
 
         """
-        old = self._ref[self._index]
+        old = self._ref[self._index].copy()
         self._ref[self._index] -= val
         return old
 
@@ -100,7 +100,7 @@ def fetch_min(self, val):
         Returns: The original value of the object referenced by the AtomicRef.
 
         """
-        old = self._ref[self._index]
+        old = self._ref[self._index].copy()
         self._ref[self._index] = min(old, val)
         return old
 
@@ -116,7 +116,7 @@ def fetch_max(self, val):
         Returns: The original value of the object referenced by the AtomicRef.
 
         """
-        old = self._ref[self._index]
+        old = self._ref[self._index].copy()
         self._ref[self._index] = max(old, val)
         return old
 
@@ -132,7 +132,7 @@ def fetch_and(self, val):
         Returns: The original value of the object referenced by the AtomicRef.
 
         """
-        old = self._ref[self._index]
+        old = self._ref[self._index].copy()
         self._ref[self._index] &= val
         return old
 
@@ -148,7 +148,7 @@ def fetch_or(self, val):
         Returns: The original value of the object referenced by the AtomicRef.
 
         """
-        old = self._ref[self._index]
+        old = self._ref[self._index].copy()
         self._ref[self._index] |= val
         return old
 
@@ -164,7 +164,7 @@ def fetch_xor(self, val):
         Returns: The original value of the object referenced by the AtomicRef.
 
         """
-        old = self._ref[self._index]
+        old = self._ref[self._index].copy()
         self._ref[self._index] ^= val
         return old
 
@@ -197,7 +197,7 @@ def exchange(self, val):
         Returns: The original value of the object referenced by the AtomicRef.
 
         """
-        old = self._ref[self._index]
+        old = self._ref[self._index].copy()
         self._ref[self._index] = val
         return old
 

numba_dpex/tests/experimental/kernel_api_overloads/spv_overloads/test_atomic_fetch_phi.py

@@ -72,6 +72,36 @@ def _kernel(a, b, ref_index):
         assert b[ref_index] == b[ref_index + 1]
 
 
+def test_fetch_phi_retval(fetch_phi_fn):
+    """A test for all fetch_phi atomic functions."""
+
+    @dpex_exp.kernel
+    def _kernel(a, b, c):
+        i = dpex.get_global_id(0)
+        v = AtomicRef(b, index=i)
+        c[i] = getattr(v, fetch_phi_fn)(a[i])
+
+    N = 10
+    a = dpnp.arange(N, dtype=dpnp.int32)
+    b = dpnp.ones(N, dtype=dpnp.int32)
+    c = dpnp.zeros(N, dtype=dpnp.int32)
+    a_copy = dpnp.copy(a)
+    b_copy = dpnp.copy(b)
+    c_copy = dpnp.copy(c)
+
+    dpex_exp.call_kernel(_kernel, dpex.Range(10), a, b, c)
+
+    # Verify if the value returned by fetch_phi kernel
+    # stored into `c` is same as the value returned
+    # by fetch_phi python stored into `c_copy`
+    for i in range(a.size):
+        v = AtomicRef(b_copy, index=i)
+        c_copy[i] = getattr(v, fetch_phi_fn)(a_copy[i])
+
+    for i in range(a.size):
+        assert c[i] == c_copy[i]
+
+
 def test_fetch_phi_diff_types(fetch_phi_fn):
     """A negative test that verifies that a TypingError is raised if
     AtomicRef type and value to be added are of different types.