Repository: numba/numba

Issue:
Numba 0.56 does not support `atomic.add` on arrays of complexs anymore
## Reporting a bug

- [x] I have tried using the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/main/CHANGE_LOG).
- [x] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

The following code snippet uses an `atomic.add` on an array of complex numbers. 
It runs without any issues with `numba 0.55.2`, but fails with `numba 0.56`

```python
from numba import cuda
import numpy as np

@cuda.jit
def atomic_add_one(values):
    """
    Adds 1 to each element values, with an atomic add

    Parameters:
    -----------
    values: 1darray of `complex128`
    """
    tid = cuda.threadIdx.x
    bid = cuda.blockIdx.x
    bdim = cuda.blockDim.x
    i = (bid * bdim) + tid
    cuda.atomic.add(values.real, i, 1)

arr = np.random.rand(16384) + 1.j*np.random.rand(16384)
print(arr.dtype)

atomic_add_one[256,64](arr)
```



Discussion:
Thanks for the reproducer - I can reproduce this issue, and I'm looking into it.

Resolving patch:
diff --git a/numba/core/typing/arraydecl.py b/numba/core/typing/arraydecl.py
--- a/numba/core/typing/arraydecl.py
+++ b/numba/core/typing/arraydecl.py
@@ -280,7 +280,23 @@ def resolve_T(self, ary):
             retty = ary.copy(layout=layout)
         return retty
 
-    # NOTE: .imag and .real are implemented with @overload_attribute
+    def resolve_real(self, ary):
+        return self._resolve_real_imag(ary, attr='real')
+
+    def resolve_imag(self, ary):
+        return self._resolve_real_imag(ary, attr='imag')
+
+    def _resolve_real_imag(self, ary, attr):
+        if ary.dtype in types.complex_domain:
+            return ary.copy(dtype=ary.dtype.underlying_float, layout='A')
+        elif ary.dtype in types.number_domain:
+            res = ary.copy(dtype=ary.dtype)
+            if attr == 'imag':
+                res = res.copy(readonly=True)
+            return res
+        else:
+            msg = "cannot access .{} of array of {}"
+            raise TypingError(msg.format(attr, ary.dtype))
 
     @bound_function("array.transpose")
     def resolve_transpose(self, ary, args, kws):
diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py
--- a/numba/np/arrayobj.py
+++ b/numba/np/arrayobj.py
@@ -26,7 +26,7 @@
                                  RefType)
 from numba.core.typing import signature
 from numba.core.extending import (register_jitable, overload, overload_method,
-                                  intrinsic, overload_attribute)
+                                  intrinsic)
 from numba.misc import quicksort, mergesort
 from numba.cpython import slicing
 from numba.cpython.unsafe.tuple import tuple_setitem, build_full_slice_tuple
@@ -2618,127 +2618,87 @@ def array_flags_f_contiguous(context, builder, typ, value):
 # ------------------------------------------------------------------------------
 # .real / .imag
 
-@overload_attribute(types.Array, 'real')
-def ol_array_real(arr):
-    typ = arr
+@lower_getattr(types.Array, "real")
+def array_real_part(context, builder, typ, value):
     if typ.dtype in types.complex_domain:
-        def impl(arr):
-            return _array_real_attr(arr)
-        return impl
+        return array_complex_attr(context, builder, typ, value, attr='real')
     elif typ.dtype in types.number_domain:
-        def impl(arr):
-            return arr
-        return impl
+        # as an identity function
+        return impl_ret_borrowed(context, builder, typ, value)
     else:
-        msg = f"cannot access .real of array of {arr.dtype}"
-        raise errors.TypingError(msg)
+        raise NotImplementedError('unsupported .real for {}'.format(type.dtype))
 
 
-@intrinsic
-def _force_readonly(tyctx, arr):
-    """Takes an array, and returns it with the readonly bit set through
-    typing"""
-    sig = arr.copy(readonly=True)(arr)
-
-    def impl(cgctx, builder, sig, llargs):
-        arrayty = make_array(arr)
-        array = arrayty(cgctx, builder, llargs[0])
-        return impl_ret_borrowed(cgctx, builder, sig.return_type,
-                                 array._getvalue())
-    return sig, impl
-
-
-@overload_attribute(types.Array, 'imag')
-def ol_array_imag(arr):
-    typ = arr
+@lower_getattr(types.Array, "imag")
+def array_imag_part(context, builder, typ, value):
     if typ.dtype in types.complex_domain:
-        def impl(arr):
-            return _array_imag_attr(arr)
-        return impl
+        return array_complex_attr(context, builder, typ, value, attr='imag')
     elif typ.dtype in types.number_domain:
-        def impl(arr):
-            # .imag on a real domain dtype is a readonly zeros array
-            tmp = np.zeros_like(arr)
-            rotmp = _force_readonly(tmp)
-            return rotmp
-        return impl
+        # return a readonly zero array
+        sig = signature(typ.copy(readonly=True), typ)
+        arrtype, shapes = _parse_empty_like_args(context, builder, sig, [value])
+        ary = _empty_nd_impl(context, builder, arrtype, shapes)
+        cgutils.memset(builder, ary.data, builder.mul(ary.itemsize,
+                                                      ary.nitems), 0)
+        return impl_ret_new_ref(context, builder, sig.return_type,
+                                ary._getvalue())
     else:
-        msg = f"cannot access .imag of array of {arr.dtype}"
-        raise errors.TypingError(msg)
+        raise NotImplementedError('unsupported .imag for {}'.format(type.dtype))
 
 
-def _generate_real_imag_attr(attr):
-    """
-    Generates intrinsics for obtaining the real/imaginary part of complex array
+def array_complex_attr(context, builder, typ, value, attr):
     """
+    Given a complex array, it's memory layout is:
 
-    @intrinsic
-    def intrin_typing(tyctx, arr):
-        sig = arr.copy(dtype=arr.dtype.underlying_float, layout='A')(arr)
+        R C R C R C
+        ^   ^   ^
 
-        def codegen(cgctx, builder, sig, llargs):
-            value = llargs[0]
-            typ = sig.args[0]
-            """
-            Given a complex array, it's memory layout is:
+    (`R` indicates a float for the real part;
+     `C` indicates a float for the imaginary part;
+     the `^` indicates the start of each element)
 
-                R C R C R C
-                ^   ^   ^
+    To get the real part, we can simply change the dtype and itemsize to that
+    of the underlying float type.  The new layout is:
 
-            (`R` indicates a float for the real part;
-            `C` indicates a float for the imaginary part;
-            the `^` indicates the start of each element)
+        R x R x R x
+        ^   ^   ^
 
-            To get the real part, we can simply change the dtype and itemsize to
-            that of the underlying float type.  The new layout is:
+    (`x` indicates unused)
 
-                R x R x R x
-                ^   ^   ^
+    A load operation will use the dtype to determine the number of bytes to
+    load.
 
-            (`x` indicates unused)
+    To get the imaginary part, we shift the pointer by 1 float offset and
+    change the dtype and itemsize.  The new layout is:
 
-            A load operation will use the dtype to determine the number of bytes
-            to load.
+        x C x C x C
+          ^   ^   ^
+    """
+    if attr not in ['real', 'imag'] or typ.dtype not in types.complex_domain:
+        raise NotImplementedError("cannot get attribute `{}`".format(attr))
 
-            To get the imaginary part, we shift the pointer by 1 float offset
-            and change the dtype and itemsize.  The new layout is:
+    arrayty = make_array(typ)
+    array = arrayty(context, builder, value)
 
-                x C x C x C
-                ^   ^   ^
-            """
-            if (attr not in ['real', 'imag'] or
-                    typ.dtype not in types.complex_domain):
-                raise NotImplementedError(f"cannot get attribute `{attr}`")
-
-            arrayty = make_array(typ)
-            array = arrayty(cgctx, builder, value)
-
-            # sizeof underlying float type
-            flty = typ.dtype.underlying_float
-            sizeof_flty = cgctx.get_abi_sizeof(cgctx.get_data_type(flty))
-            itemsize = array.itemsize.type(sizeof_flty)
-
-            # cast data pointer to float type
-            llfltptrty = cgctx.get_value_type(flty).as_pointer()
-            dataptr = builder.bitcast(array.data, llfltptrty)
-
-            # add offset
-            if attr == 'imag':
-                dataptr = builder.gep(dataptr, [ir.IntType(32)(1)])
-
-            # make result
-            resultty = typ.copy(dtype=flty, layout='A')
-            result = make_array(resultty)(cgctx, builder)
-            repl = dict(data=dataptr, itemsize=itemsize)
-            cgutils.copy_struct(result, array, repl)
-            return impl_ret_borrowed(cgctx, builder, resultty,
-                                     result._getvalue())
-        return sig, codegen
-    return intrin_typing
-
-
-_array_real_attr = _generate_real_imag_attr("real")
-_array_imag_attr = _generate_real_imag_attr("imag")
+    # sizeof underlying float type
+    flty = typ.dtype.underlying_float
+    sizeof_flty = context.get_abi_sizeof(context.get_data_type(flty))
+    itemsize = array.itemsize.type(sizeof_flty)
+
+    # cast data pointer to float type
+    llfltptrty = context.get_value_type(flty).as_pointer()
+    dataptr = builder.bitcast(array.data, llfltptrty)
+
+    # add offset
+    if attr == 'imag':
+        dataptr = builder.gep(dataptr, [ir.IntType(32)(1)])
+
+    # make result
+    resultty = typ.copy(dtype=flty, layout='A')
+    result = make_array(resultty)(context, builder)
+    repl = dict(data=dataptr, itemsize=itemsize)
+    cgutils.copy_struct(result, array, repl)
+    return impl_ret_borrowed(context, builder, resultty, result._getvalue())
 
 
 @overload_method(types.Array, 'conj')


Test patch:
diff --git a/numba/cuda/tests/cudapy/test_complex.py b/numba/cuda/tests/cudapy/test_complex.py
--- a/numba/cuda/tests/cudapy/test_complex.py
+++ b/numba/cuda/tests/cudapy/test_complex.py
@@ -259,5 +259,38 @@ def test_tanh(self):
                               ignore_sign_on_zero=True)
 
 
+class TestAtomicOnComplexComponents(CUDATestCase):
+    # Based on the reproducer from Issue #8309. array.real and array.imag could
+    # not be used because they required returning an array from a generated
+    # function, and even if this was permitted, they could not be resolved from
+    # the atomic lowering when they were overloads.
+    #
+    # See https://github.com/numba/numba/issues/8309
+
+    def test_atomic_on_real(self):
+        @cuda.jit
+        def atomic_add_one(values):
+            i = cuda.grid(1)
+            cuda.atomic.add(values.real, i, 1)
+
+        N = 32
+        arr1 = np.arange(N) + np.arange(N) * 1j
+        arr2 = arr1.copy()
+        atomic_add_one[1, N](arr2)
+        np.testing.assert_equal(arr1 + 1, arr2)
+
+    def test_atomic_on_imag(self):
+        @cuda.jit
+        def atomic_add_one_j(values):
+            i = cuda.grid(1)
+            cuda.atomic.add(values.imag, i, 1)
+
+        N = 32
+        arr1 = np.arange(N) + np.arange(N) * 1j
+        arr2 = arr1.copy()
+        atomic_add_one_j[1, N](arr2)
+        np.testing.assert_equal(arr1 + 1j, arr2)
+
+
 if __name__ == '__main__':
     unittest.main()