numba-cuda 0.22.0__cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

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

@@ -0,0 +1,370 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import math
+import itertools
+
+import numpy as np
+
+from numba.cuda.testing import unittest, CUDATestCase
+from numba.cuda import types
+from numba import cuda
+from numba.cuda.tests.cudapy.complex_usecases import (
+    real_usecase,
+    imag_usecase,
+    conjugate_usecase,
+    phase_usecase,
+    polar_as_complex_usecase,
+    rect_usecase,
+    isnan_usecase,
+    isinf_usecase,
+    isfinite_usecase,
+    exp_usecase,
+    log_usecase,
+    log_base_usecase,
+    log10_usecase,
+    sqrt_usecase,
+    asin_usecase,
+    acos_usecase,
+    atan_usecase,
+    cos_usecase,
+    sin_usecase,
+    tan_usecase,
+    acosh_usecase,
+    asinh_usecase,
+    atanh_usecase,
+    cosh_usecase,
+    sinh_usecase,
+    tanh_usecase,
+)
+from numba.cuda.np import numpy_support
+
+
+def compile_scalar_func(pyfunc, argtypes, restype):
+    # First compile a scalar device function
+    assert not any(isinstance(tp, types.Array) for tp in argtypes)
+    assert not isinstance(restype, types.Array)
+    device_func = cuda.jit(restype(*argtypes), device=True)(pyfunc)
+
+    kernel_types = [
+        types.Array(tp, 1, "C") for tp in [restype] + list(argtypes)
+    ]
+
+    if len(argtypes) == 1:
+
+        def kernel_func(out, a):
+            i = cuda.grid(1)
+            if i < out.shape[0]:
+                out[i] = device_func(a[i])
+    elif len(argtypes) == 2:
+
+        def kernel_func(out, a, b):
+            i = cuda.grid(1)
+            if i < out.shape[0]:
+                out[i] = device_func(a[i], b[i])
+    else:
+        assert 0
+
+    kernel = cuda.jit(tuple(kernel_types))(kernel_func)
+
+    def kernel_wrapper(values):
+        n = len(values)
+        inputs = [
+            np.empty(n, dtype=numpy_support.as_dtype(tp)) for tp in argtypes
+        ]
+        output = np.empty(n, dtype=numpy_support.as_dtype(restype))
+        for i, vs in enumerate(values):
+            for v, inp in zip(vs, inputs):
+                inp[i] = v
+        args = [output] + inputs
+        kernel[int(math.ceil(n / 256)), 256](*args)
+        return list(output)
+
+    return kernel_wrapper
+
+
+class BaseComplexTest(CUDATestCase):
+    def basic_values(self):
+        reals = [
+            -0.0,
+            +0.0,
+            1,
+            -1,
+            +1.5,
+            -3.5,
+            float("-inf"),
+            float("+inf"),
+            float("nan"),
+        ]
+        return [complex(x, y) for x, y in itertools.product(reals, reals)]
+
+    def more_values(self):
+        reals = [
+            0.0,
+            +0.0,
+            1,
+            -1,
+            -math.pi,
+            +math.pi,
+            float("-inf"),
+            float("+inf"),
+            float("nan"),
+        ]
+        return [complex(x, y) for x, y in itertools.product(reals, reals)]
+
+    def non_nan_values(self):
+        reals = [
+            -0.0,
+            +0.0,
+            1,
+            -1,
+            -math.pi,
+            +math.pi,
+            float("inf"),
+            float("-inf"),
+        ]
+        return [complex(x, y) for x, y in itertools.product(reals, reals)]
+
+    def run_func(self, pyfunc, sigs, values, ulps=1, ignore_sign_on_zero=False):
+        for sig in sigs:
+            if isinstance(sig, types.Type):
+                sig = (sig,)
+            if isinstance(sig, tuple):
+                # Assume return type is the type of first argument
+                sig = sig[0](*sig)
+            prec = (
+                "single"
+                if sig.args[0] in (types.float32, types.complex64)
+                else "double"
+            )
+            cudafunc = compile_scalar_func(pyfunc, sig.args, sig.return_type)
+            ok_values = []
+            expected_list = []
+            for args in values:
+                if not isinstance(args, (list, tuple)):
+                    args = (args,)
+                try:
+                    expected_list.append(pyfunc(*args))
+                    ok_values.append(args)
+                except ValueError as e:
+                    self.assertIn("math domain error", str(e))
+                    continue
+            got_list = cudafunc(ok_values)
+            for got, expected, args in zip(got_list, expected_list, ok_values):
+                msg = "for input %r with prec %r" % (args, prec)
+                self.assertPreciseEqual(
+                    got,
+                    expected,
+                    prec=prec,
+                    ulps=ulps,
+                    ignore_sign_on_zero=ignore_sign_on_zero,
+                    msg=msg,
+                )
+
+    run_unary = run_func
+    run_binary = run_func
+
+
+class TestComplex(BaseComplexTest):
+    def check_real_image(self, pyfunc):
+        values = self.basic_values()
+        self.run_unary(
+            pyfunc,
+            [
+                tp.underlying_float(tp)
+                for tp in (types.complex64, types.complex128)
+            ],
+            values,
+        )
+
+    def test_real(self):
+        self.check_real_image(real_usecase)
+
+    def test_imag(self):
+        self.check_real_image(imag_usecase)
+
+    def test_conjugate(self):
+        pyfunc = conjugate_usecase
+        values = self.basic_values()
+        self.run_unary(pyfunc, [types.complex64, types.complex128], values)
+
+
+class TestCMath(BaseComplexTest):
+    """
+    Tests for cmath module support.
+    """
+
+    def check_predicate_func(self, pyfunc):
+        self.run_unary(
+            pyfunc,
+            [types.boolean(tp) for tp in (types.complex128, types.complex64)],
+            self.basic_values(),
+        )
+
+    def check_unary_func(
+        self,
+        pyfunc,
+        ulps=1,
+        values=None,
+        returns_float=False,
+        ignore_sign_on_zero=False,
+    ):
+        if returns_float:
+
+            def sig(tp):
+                return tp.underlying_float(tp)
+        else:
+
+            def sig(tp):
+                return tp(tp)
+
+        self.run_unary(
+            pyfunc,
+            [sig(types.complex128)],
+            values or self.more_values(),
+            ulps=ulps,
+            ignore_sign_on_zero=ignore_sign_on_zero,
+        )
+        # Avoid discontinuities around pi when in single precision.
+        self.run_unary(
+            pyfunc,
+            [sig(types.complex64)],
+            values or self.basic_values(),
+            ulps=ulps,
+            ignore_sign_on_zero=ignore_sign_on_zero,
+        )
+
+    # Conversions
+
+    def test_phase(self):
+        self.check_unary_func(phase_usecase, returns_float=True)
+
+    def test_polar(self):
+        self.check_unary_func(polar_as_complex_usecase)
+
+    def test_rect(self):
+        def do_test(tp, seed_values):
+            values = [
+                (z.real, z.imag)
+                for z in seed_values
+                if not math.isinf(z.imag) or z.real == 0
+            ]
+            float_type = tp.underlying_float
+            self.run_binary(rect_usecase, [tp(float_type, float_type)], values)
+
+        do_test(types.complex128, self.more_values())
+        # Avoid discontinuities around pi when in single precision.
+        do_test(types.complex64, self.basic_values())
+
+    # Classification
+
+    def test_isnan(self):
+        self.check_predicate_func(isnan_usecase)
+
+    def test_isinf(self):
+        self.check_predicate_func(isinf_usecase)
+
+    def test_isfinite(self):
+        self.check_predicate_func(isfinite_usecase)
+
+    # Power and logarithms
+
+    def test_exp(self):
+        self.check_unary_func(exp_usecase, ulps=2)
+
+    def test_log(self):
+        self.check_unary_func(log_usecase)
+
+    def test_log_base(self):
+        values = list(itertools.product(self.more_values(), self.more_values()))
+        value_types = [
+            (types.complex128, types.complex128),
+            (types.complex64, types.complex64),
+        ]
+        self.run_binary(log_base_usecase, value_types, values, ulps=3)
+
+    def test_log10(self):
+        self.check_unary_func(log10_usecase)
+
+    def test_sqrt(self):
+        self.check_unary_func(sqrt_usecase)
+
+    # Trigonometric functions
+
+    def test_acos(self):
+        self.check_unary_func(acos_usecase, ulps=2)
+
+    def test_asin(self):
+        self.check_unary_func(asin_usecase, ulps=2)
+
+    def test_atan(self):
+        self.check_unary_func(
+            atan_usecase, ulps=2, values=self.non_nan_values()
+        )
+
+    def test_cos(self):
+        self.check_unary_func(cos_usecase, ulps=2)
+
+    def test_sin(self):
+        # See test_sinh.
+        self.check_unary_func(sin_usecase, ulps=2)
+
+    def test_tan(self):
+        self.check_unary_func(tan_usecase, ulps=2, ignore_sign_on_zero=True)
+
+    # Hyperbolic functions
+
+    def test_acosh(self):
+        self.check_unary_func(acosh_usecase)
+
+    def test_asinh(self):
+        self.check_unary_func(asinh_usecase, ulps=2)
+
+    def test_atanh(self):
+        self.check_unary_func(atanh_usecase, ulps=2, ignore_sign_on_zero=True)
+
+    def test_cosh(self):
+        self.check_unary_func(cosh_usecase, ulps=2)
+
+    def test_sinh(self):
+        self.check_unary_func(sinh_usecase, ulps=2)
+
+    def test_tanh(self):
+        self.check_unary_func(tanh_usecase, ulps=2, 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()

numba_cuda/numba/cuda/tests/cudapy/test_complex_kernel.py

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numpy as np
+from numba import cuda
+from numba.cuda.testing import unittest, CUDATestCase
+
+
+class TestCudaComplex(CUDATestCase):
+    def test_cuda_complex_arg(self):
+        @cuda.jit("void(complex128[:], complex128)")
+        def foo(a, b):
+            i = cuda.grid(1)
+            a[i] += b
+
+        a = np.arange(5, dtype=np.complex128)
+        a0 = a.copy()
+        foo[1, a.shape](a, 2j)
+        self.assertTrue(np.allclose(a, a0 + 2j))
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/cudapy/test_const_string.py

@@ -0,0 +1,142 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import re
+import numpy as np
+from numba import cuda
+from numba.cuda.testing import unittest, skip_on_cudasim, CUDATestCase
+from llvmlite import ir
+
+
+@skip_on_cudasim("This is testing CUDA backend code generation")
+class TestConstStringCodegen(unittest.TestCase):
+    def test_const_string(self):
+        # These imports are incompatible with CUDASIM
+        from numba.cuda.descriptor import cuda_target
+        from numba.cuda.cudadrv.nvvm import compile_ir
+
+        targetctx = cuda_target.target_context
+        mod = targetctx.create_module("")
+        textstring = "A Little Brown Fox"
+        gv0 = targetctx.insert_const_string(mod, textstring)
+        # Insert the same const string a second time - the first should be
+        # reused.
+        targetctx.insert_const_string(mod, textstring)
+
+        res = re.findall(
+            r"@\"__conststring__.*internal.*constant.*\["
+            r"19\s+x\s+i8\]",
+            str(mod),
+        )
+        # Ensure that the const string was only inserted once
+        self.assertEqual(len(res), 1)
+
+        fnty = ir.FunctionType(ir.IntType(8).as_pointer(), [])
+
+        # Using insert_const_string
+        fn = ir.Function(mod, fnty, "test_insert_const_string")
+        builder = ir.IRBuilder(fn.append_basic_block())
+        res = builder.addrspacecast(
+            gv0, ir.PointerType(ir.IntType(8)), "generic"
+        )
+        builder.ret(res)
+
+        matches = re.findall(
+            r"@\"__conststring__.*internal.*constant.*\["
+            r"19\s+x\s+i8\]",
+            str(mod),
+        )
+        self.assertEqual(len(matches), 1)
+
+        # Using insert_string_const_addrspace
+        fn = ir.Function(mod, fnty, "test_insert_string_const_addrspace")
+        builder = ir.IRBuilder(fn.append_basic_block())
+        res = targetctx.insert_string_const_addrspace(builder, textstring)
+        builder.ret(res)
+
+        matches = re.findall(
+            r"@\"__conststring__.*internal.*constant.*\["
+            r"19\s+x\s+i8\]",
+            str(mod),
+        )
+        self.assertEqual(len(matches), 1)
+
+        ptx = compile_ir(str(mod)).decode("ascii")
+        matches = list(re.findall(r"\.const.*__conststring__", ptx))
+
+        self.assertEqual(len(matches), 1)
+
+
+# Inspired by the reproducer from Issue #7041.
+class TestConstString(CUDATestCase):
+    def test_assign_const_unicode_string(self):
+        @cuda.jit
+        def str_assign(arr):
+            i = cuda.grid(1)
+            if i < len(arr):
+                arr[i] = "XYZ"
+
+        n_strings = 8
+        arr = np.zeros(n_strings + 1, dtype="<U12")
+        str_assign[1, n_strings](arr)
+
+        # Expected result, e.g.:
+        #     ['XYZ' 'XYZ' 'XYZ' 'XYZ' 'XYZ' 'XYZ' 'XYZ' 'XYZ' '']
+        expected = np.zeros_like(arr)
+        expected[:-1] = "XYZ"
+        expected[-1] = ""
+        np.testing.assert_equal(arr, expected)
+
+    def test_assign_const_byte_string(self):
+        @cuda.jit
+        def bytes_assign(arr):
+            i = cuda.grid(1)
+            if i < len(arr):
+                arr[i] = b"XYZ"
+
+        n_strings = 8
+        arr = np.zeros(n_strings + 1, dtype="S12")
+        bytes_assign[1, n_strings](arr)
+
+        # Expected result, e.g.:
+        #     [b'XYZ' b'XYZ' b'XYZ' b'XYZ' b'XYZ' b'XYZ' b'XYZ' b'XYZ' b'']
+        expected = np.zeros_like(arr)
+        expected[:-1] = b"XYZ"
+        expected[-1] = b""
+        np.testing.assert_equal(arr, expected)
+
+    def test_assign_const_string_in_record(self):
+        @cuda.jit
+        def f(a):
+            a[0]["x"] = 1
+            a[0]["y"] = "ABC"
+            a[1]["x"] = 2
+            a[1]["y"] = "XYZ"
+
+        dt = np.dtype([("x", np.int32), ("y", np.dtype("<U12"))])
+        a = np.zeros(2, dt)
+
+        f[1, 1](a)
+
+        reference = np.asarray([(1, "ABC"), (2, "XYZ")], dtype=dt)
+        np.testing.assert_array_equal(reference, a)
+
+    def test_assign_const_bytes_in_record(self):
+        @cuda.jit
+        def f(a):
+            a[0]["x"] = 1
+            a[0]["y"] = b"ABC"
+            a[1]["x"] = 2
+            a[1]["y"] = b"XYZ"
+
+        dt = np.dtype([("x", np.float32), ("y", np.dtype("S12"))])
+        a = np.zeros(2, dt)
+
+        f[1, 1](a)
+
+        reference = np.asarray([(1, b"ABC"), (2, b"XYZ")], dtype=dt)
+        np.testing.assert_array_equal(reference, a)
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/cudapy/test_constmem.py

@@ -0,0 +1,178 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numpy as np
+
+from numba import cuda
+from numba.cuda import complex64, int32, float64
+from numba.cuda.testing import unittest, CUDATestCase
+from numba.cuda.core.config import ENABLE_CUDASIM
+
+CONST_EMPTY = np.array([])
+CONST1D = np.arange(10, dtype=np.float64) / 2.0
+CONST2D = np.asfortranarray(np.arange(100, dtype=np.int32).reshape(10, 10))
+CONST3D = (np.arange(5 * 5 * 5, dtype=np.complex64).reshape(5, 5, 5) + 1j) / 2j
+CONST3BYTES = np.arange(3, dtype=np.uint8)
+
+CONST_RECORD_EMPTY = np.array([], dtype=[("x", float), ("y", int)])
+CONST_RECORD = np.array([(1.0, 2), (3.0, 4)], dtype=[("x", float), ("y", int)])
+CONST_RECORD_ALIGN = np.array(
+    [(1, 2, 3, 0xDEADBEEF, 8), (4, 5, 6, 0xBEEFDEAD, 10)],
+    dtype=np.dtype(
+        dtype=[
+            ("a", np.uint8),
+            ("b", np.uint8),
+            ("x", np.uint8),
+            ("y", np.uint32),
+            ("z", np.uint8),
+        ],
+        align=True,
+    ),
+)
+
+
+def cuconstEmpty(A):
+    C = cuda.const.array_like(CONST_EMPTY)
+    i = cuda.grid(1)
+    A[i] = len(C)
+
+
+def cuconst(A):
+    C = cuda.const.array_like(CONST1D)
+    i = cuda.grid(1)
+
+    # +1 or it'll be loaded & stored as a u32
+    A[i] = C[i] + 1.0
+
+
+def cuconst2d(A):
+    C = cuda.const.array_like(CONST2D)
+    i, j = cuda.grid(2)
+    A[i, j] = C[i, j]
+
+
+def cuconst3d(A):
+    C = cuda.const.array_like(CONST3D)
+    i = cuda.threadIdx.x
+    j = cuda.threadIdx.y
+    k = cuda.threadIdx.z
+    A[i, j, k] = C[i, j, k]
+
+
+def cuconstRecEmpty(A):
+    C = cuda.const.array_like(CONST_RECORD_EMPTY)
+    i = cuda.grid(1)
+    A[i] = len(C)
+
+
+def cuconstRec(A, B):
+    C = cuda.const.array_like(CONST_RECORD)
+    i = cuda.grid(1)
+    A[i] = C[i]["x"]
+    B[i] = C[i]["y"]
+
+
+def cuconstRecAlign(A, B, C, D, E):
+    Z = cuda.const.array_like(CONST_RECORD_ALIGN)
+    i = cuda.grid(1)
+    A[i] = Z[i]["a"]
+    B[i] = Z[i]["b"]
+    C[i] = Z[i]["x"]
+    D[i] = Z[i]["y"]
+    E[i] = Z[i]["z"]
+
+
+def cuconstAlign(z):
+    a = cuda.const.array_like(CONST3BYTES)
+    b = cuda.const.array_like(CONST1D)
+    i = cuda.grid(1)
+    z[i] = a[i] + b[i]
+
+
+class TestCudaConstantMemory(CUDATestCase):
+    def test_const_array(self):
+        sig = (float64[:],)
+        jcuconst = cuda.jit(sig)(cuconst)
+        A = np.zeros_like(CONST1D)
+        jcuconst[2, 5](A)
+        self.assertTrue(np.all(A == CONST1D + 1))
+
+        if not ENABLE_CUDASIM:
+            self.assertRegex(
+                jcuconst.inspect_asm(sig),
+                r"ld\.const\.[bf]64",
+                "load as a 64-bit value",
+            )
+
+    def test_const_empty(self):
+        jcuconstEmpty = cuda.jit("void(int64[:])")(cuconstEmpty)
+        A = np.full(1, fill_value=-1, dtype=np.int64)
+        jcuconstEmpty[1, 1](A)
+        self.assertTrue(np.all(A == 0))
+
+    def test_const_align(self):
+        jcuconstAlign = cuda.jit("void(float64[:])")(cuconstAlign)
+        A = np.full(3, fill_value=np.nan, dtype=float)
+        jcuconstAlign[1, 3](A)
+        self.assertTrue(np.all(A == (CONST3BYTES + CONST1D[:3])))
+
+    def test_const_array_2d(self):
+        sig = (int32[:, :],)
+        jcuconst2d = cuda.jit(sig)(cuconst2d)
+        A = np.zeros_like(CONST2D, order="C")
+        jcuconst2d[(2, 2), (5, 5)](A)
+        self.assertTrue(np.all(A == CONST2D))
+
+        if not ENABLE_CUDASIM:
+            self.assertRegex(
+                jcuconst2d.inspect_asm(sig),
+                r"ld\.const\.[bu]32",
+                "load as a 32-bit value",
+            )
+
+    def test_const_array_3d(self):
+        sig = (complex64[:, :, :],)
+        jcuconst3d = cuda.jit(sig)(cuconst3d)
+        A = np.zeros_like(CONST3D, order="F")
+        jcuconst3d[1, (5, 5, 5)](A)
+        self.assertTrue(np.all(A == CONST3D))
+
+        if not ENABLE_CUDASIM:
+            asm = jcuconst3d.inspect_asm(sig)
+            complex_load = r"ld\.const\.v2\.[bf]32"
+            description = "Load the complex as a vector of 2x 32-bits"
+            self.assertRegex(asm, complex_load, description)
+
+    def test_const_record_empty(self):
+        jcuconstRecEmpty = cuda.jit("void(int64[:])")(cuconstRecEmpty)
+        A = np.full(1, fill_value=-1, dtype=np.int64)
+        jcuconstRecEmpty[1, 1](A)
+        self.assertTrue(np.all(A == 0))
+
+    def test_const_record(self):
+        A = np.zeros(2, dtype=float)
+        B = np.zeros(2, dtype=int)
+        jcuconst = cuda.jit(cuconstRec).specialize(A, B)
+
+        jcuconst[2, 1](A, B)
+        np.testing.assert_allclose(A, CONST_RECORD["x"])
+        np.testing.assert_allclose(B, CONST_RECORD["y"])
+
+    def test_const_record_align(self):
+        A = np.zeros(2, dtype=np.float64)
+        B = np.zeros(2, dtype=np.float64)
+        C = np.zeros(2, dtype=np.float64)
+        D = np.zeros(2, dtype=np.float64)
+        E = np.zeros(2, dtype=np.float64)
+        jcuconst = cuda.jit(cuconstRecAlign).specialize(A, B, C, D, E)
+
+        jcuconst[2, 1](A, B, C, D, E)
+        np.testing.assert_allclose(A, CONST_RECORD_ALIGN["a"])
+        np.testing.assert_allclose(B, CONST_RECORD_ALIGN["b"])
+        np.testing.assert_allclose(C, CONST_RECORD_ALIGN["x"])
+        np.testing.assert_allclose(D, CONST_RECORD_ALIGN["y"])
+        np.testing.assert_allclose(E, CONST_RECORD_ALIGN["z"])
+
+
+if __name__ == "__main__":
+    unittest.main()