Skip to content

Commit e6521ef

Browse files
authored
Merge pull request #8310 from gmarkall/issue-8309
CUDA: Fix Issue #8309 - atomics don't work on complex components
2 parents dbc0c3d + eafa2ac commit e6521ef

3 files changed

Lines changed: 111 additions & 102 deletions

File tree

‎numba/core/typing/arraydecl.py‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -280,7 +280,23 @@ def resolve_T(self, ary):
280280
retty = ary.copy(layout=layout)
281281
return retty
282282

283-
# NOTE: .imag and .real are implemented with @overload_attribute
283+
def resolve_real(self, ary):
284+
return self._resolve_real_imag(ary, attr='real')
285+
286+
def resolve_imag(self, ary):
287+
return self._resolve_real_imag(ary, attr='imag')
288+
289+
def _resolve_real_imag(self, ary, attr):
290+
if ary.dtype in types.complex_domain:
291+
return ary.copy(dtype=ary.dtype.underlying_float, layout='A')
292+
elif ary.dtype in types.number_domain:
293+
res = ary.copy(dtype=ary.dtype)
294+
if attr == 'imag':
295+
res = res.copy(readonly=True)
296+
return res
297+
else:
298+
msg = "cannot access .{} of array of {}"
299+
raise TypingError(msg.format(attr, ary.dtype))
284300

285301
@bound_function("array.transpose")
286302
def resolve_transpose(self, ary, args, kws):

‎numba/cuda/tests/cudapy/test_complex.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -259,5 +259,38 @@ def test_tanh(self):
259259
ignore_sign_on_zero=True)
260260

261261

262+
class TestAtomicOnComplexComponents(CUDATestCase):
263+
# Based on the reproducer from Issue #8309. array.real and array.imag could
264+
# not be used because they required returning an array from a generated
265+
# function, and even if this was permitted, they could not be resolved from
266+
# the atomic lowering when they were overloads.
267+
#
268+
# See https://github.com/numba/numba/issues/8309
269+
270+
def test_atomic_on_real(self):
271+
@cuda.jit
272+
def atomic_add_one(values):
273+
i = cuda.grid(1)
274+
cuda.atomic.add(values.real, i, 1)
275+
276+
N = 32
277+
arr1 = np.arange(N) + np.arange(N) * 1j
278+
arr2 = arr1.copy()
279+
atomic_add_one[1, N](arr2)
280+
np.testing.assert_equal(arr1 + 1, arr2)
281+
282+
def test_atomic_on_imag(self):
283+
@cuda.jit
284+
def atomic_add_one_j(values):
285+
i = cuda.grid(1)
286+
cuda.atomic.add(values.imag, i, 1)
287+
288+
N = 32
289+
arr1 = np.arange(N) + np.arange(N) * 1j
290+
arr2 = arr1.copy()
291+
atomic_add_one_j[1, N](arr2)
292+
np.testing.assert_equal(arr1 + 1j, arr2)
293+
294+
262295
if __name__ == '__main__':
263296
unittest.main()

‎numba/np/arrayobj.py‎

Lines changed: 61 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
RefType)
2727
from numba.core.typing import signature
2828
from numba.core.extending import (register_jitable, overload, overload_method,
29-
intrinsic, overload_attribute)
29+
intrinsic)
3030
from numba.misc import quicksort, mergesort
3131
from numba.cpython import slicing
3232
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):
26182618
# ------------------------------------------------------------------------------
26192619
# .real / .imag
26202620

2621-
@overload_attribute(types.Array, 'real')
2622-
def ol_array_real(arr):
2623-
typ = arr
2621+
@lower_getattr(types.Array, "real")
2622+
def array_real_part(context, builder, typ, value):
26242623
if typ.dtype in types.complex_domain:
2625-
def impl(arr):
2626-
return _array_real_attr(arr)
2627-
return impl
2624+
return array_complex_attr(context, builder, typ, value, attr='real')
26282625
elif typ.dtype in types.number_domain:
2629-
def impl(arr):
2630-
return arr
2631-
return impl
2626+
# as an identity function
2627+
return impl_ret_borrowed(context, builder, typ, value)
26322628
else:
2633-
msg = f"cannot access .real of array of {arr.dtype}"
2634-
raise errors.TypingError(msg)
2629+
raise NotImplementedError('unsupported .real for {}'.format(type.dtype))
26352630

26362631

2637-
@intrinsic
2638-
def _force_readonly(tyctx, arr):
2639-
"""Takes an array, and returns it with the readonly bit set through
2640-
typing"""
2641-
sig = arr.copy(readonly=True)(arr)
2642-
2643-
def impl(cgctx, builder, sig, llargs):
2644-
arrayty = make_array(arr)
2645-
array = arrayty(cgctx, builder, llargs[0])
2646-
return impl_ret_borrowed(cgctx, builder, sig.return_type,
2647-
array._getvalue())
2648-
return sig, impl
2649-
2650-
2651-
@overload_attribute(types.Array, 'imag')
2652-
def ol_array_imag(arr):
2653-
typ = arr
2632+
@lower_getattr(types.Array, "imag")
2633+
def array_imag_part(context, builder, typ, value):
26542634
if typ.dtype in types.complex_domain:
2655-
def impl(arr):
2656-
return _array_imag_attr(arr)
2657-
return impl
2635+
return array_complex_attr(context, builder, typ, value, attr='imag')
26582636
elif typ.dtype in types.number_domain:
2659-
def impl(arr):
2660-
# .imag on a real domain dtype is a readonly zeros array
2661-
tmp = np.zeros_like(arr)
2662-
rotmp = _force_readonly(tmp)
2663-
return rotmp
2664-
return impl
2637+
# return a readonly zero array
2638+
sig = signature(typ.copy(readonly=True), typ)
2639+
arrtype, shapes = _parse_empty_like_args(context, builder, sig, [value])
2640+
ary = _empty_nd_impl(context, builder, arrtype, shapes)
2641+
cgutils.memset(builder, ary.data, builder.mul(ary.itemsize,
2642+
ary.nitems), 0)
2643+
return impl_ret_new_ref(context, builder, sig.return_type,
2644+
ary._getvalue())
26652645
else:
2666-
msg = f"cannot access .imag of array of {arr.dtype}"
2667-
raise errors.TypingError(msg)
2646+
raise NotImplementedError('unsupported .imag for {}'.format(type.dtype))
26682647

26692648

2670-
def _generate_real_imag_attr(attr):
2671-
"""
2672-
Generates intrinsics for obtaining the real/imaginary part of complex array
2649+
def array_complex_attr(context, builder, typ, value, attr):
26732650
"""
2651+
Given a complex array, it's memory layout is:
26742652
2675-
@intrinsic
2676-
def intrin_typing(tyctx, arr):
2677-
sig = arr.copy(dtype=arr.dtype.underlying_float, layout='A')(arr)
2653+
R C R C R C
2654+
^ ^ ^
26782655
2679-
def codegen(cgctx, builder, sig, llargs):
2680-
value = llargs[0]
2681-
typ = sig.args[0]
2682-
"""
2683-
Given a complex array, it's memory layout is:
2656+
(`R` indicates a float for the real part;
2657+
`C` indicates a float for the imaginary part;
2658+
the `^` indicates the start of each element)
26842659
2685-
R C R C R C
2686-
^ ^ ^
2660+
To get the real part, we can simply change the dtype and itemsize to that
2661+
of the underlying float type. The new layout is:
26872662
2688-
(`R` indicates a float for the real part;
2689-
`C` indicates a float for the imaginary part;
2690-
the `^` indicates the start of each element)
2663+
R x R x R x
2664+
^ ^ ^
26912665
2692-
To get the real part, we can simply change the dtype and itemsize to
2693-
that of the underlying float type. The new layout is:
2666+
(`x` indicates unused)
26942667
2695-
R x R x R x
2696-
^ ^ ^
2668+
A load operation will use the dtype to determine the number of bytes to
2669+
load.
26972670
2698-
(`x` indicates unused)
2671+
To get the imaginary part, we shift the pointer by 1 float offset and
2672+
change the dtype and itemsize. The new layout is:
26992673
2700-
A load operation will use the dtype to determine the number of bytes
2701-
to load.
2674+
x C x C x C
2675+
^ ^ ^
2676+
"""
2677+
if attr not in ['real', 'imag'] or typ.dtype not in types.complex_domain:
2678+
raise NotImplementedError("cannot get attribute `{}`".format(attr))
27022679

2703-
To get the imaginary part, we shift the pointer by 1 float offset
2704-
and change the dtype and itemsize. The new layout is:
2680+
arrayty = make_array(typ)
2681+
array = arrayty(context, builder, value)
27052682

2706-
x C x C x C
2707-
^ ^ ^
2708-
"""
2709-
if (attr not in ['real', 'imag'] or
2710-
typ.dtype not in types.complex_domain):
2711-
raise NotImplementedError(f"cannot get attribute `{attr}`")
2712-
2713-
arrayty = make_array(typ)
2714-
array = arrayty(cgctx, builder, value)
2715-
2716-
# sizeof underlying float type
2717-
flty = typ.dtype.underlying_float
2718-
sizeof_flty = cgctx.get_abi_sizeof(cgctx.get_data_type(flty))
2719-
itemsize = array.itemsize.type(sizeof_flty)
2720-
2721-
# cast data pointer to float type
2722-
llfltptrty = cgctx.get_value_type(flty).as_pointer()
2723-
dataptr = builder.bitcast(array.data, llfltptrty)
2724-
2725-
# add offset
2726-
if attr == 'imag':
2727-
dataptr = builder.gep(dataptr, [ir.IntType(32)(1)])
2728-
2729-
# make result
2730-
resultty = typ.copy(dtype=flty, layout='A')
2731-
result = make_array(resultty)(cgctx, builder)
2732-
repl = dict(data=dataptr, itemsize=itemsize)
2733-
cgutils.copy_struct(result, array, repl)
2734-
return impl_ret_borrowed(cgctx, builder, resultty,
2735-
result._getvalue())
2736-
return sig, codegen
2737-
return intrin_typing
2738-
2739-
2740-
_array_real_attr = _generate_real_imag_attr("real")
2741-
_array_imag_attr = _generate_real_imag_attr("imag")
2683+
# sizeof underlying float type
2684+
flty = typ.dtype.underlying_float
2685+
sizeof_flty = context.get_abi_sizeof(context.get_data_type(flty))
2686+
itemsize = array.itemsize.type(sizeof_flty)
2687+
2688+
# cast data pointer to float type
2689+
llfltptrty = context.get_value_type(flty).as_pointer()
2690+
dataptr = builder.bitcast(array.data, llfltptrty)
2691+
2692+
# add offset
2693+
if attr == 'imag':
2694+
dataptr = builder.gep(dataptr, [ir.IntType(32)(1)])
2695+
2696+
# make result
2697+
resultty = typ.copy(dtype=flty, layout='A')
2698+
result = make_array(resultty)(context, builder)
2699+
repl = dict(data=dataptr, itemsize=itemsize)
2700+
cgutils.copy_struct(result, array, repl)
2701+
return impl_ret_borrowed(context, builder, resultty, result._getvalue())
27422702

27432703

27442704
@overload_method(types.Array, 'conj')

0 commit comments

Comments
 (0)