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

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

@@ -0,0 +1,402 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba import cuda
+from numba.cuda import types
+from numba.cuda import HAS_NUMBA
+
+if HAS_NUMBA:
+    from numba.core.errors import TypingError
+    from numba import njit
+    import numba
+else:
+    from numba.cuda.core.errors import TypingError
+from numba.cuda.extending import overload, overload_attribute
+from numba.cuda.typing.typeof import typeof
+from numba.core.typing.typeof import typeof as cpu_typeof
+from numba.cuda.testing import (
+    CUDATestCase,
+    skip_on_cudasim,
+    unittest,
+    skip_on_standalone_numba_cuda,
+)
+import numpy as np
+
+
+# Dummy function definitions to overload
+
+
+def generic_func_1():
+    pass
+
+
+def cuda_func_1():
+    pass
+
+
+def generic_func_2():
+    pass
+
+
+def cuda_func_2():
+    pass
+
+
+def generic_calls_generic():
+    pass
+
+
+def generic_calls_cuda():
+    pass
+
+
+def cuda_calls_generic():
+    pass
+
+
+def cuda_calls_cuda():
+    pass
+
+
+def target_overloaded():
+    pass
+
+
+def generic_calls_target_overloaded():
+    pass
+
+
+def cuda_calls_target_overloaded():
+    pass
+
+
+def target_overloaded_calls_target_overloaded():
+    pass
+
+
+def default_values_and_kwargs():
+    pass
+
+
+# To recognise which functions are resolved for a call, we identify each with a
+# prime number. Each function called multiplies a value by its prime (starting
+# with the value 1), and we can check that the result is as expected based on
+# the final value after all multiplications.
+
+GENERIC_FUNCTION_1 = 2
+CUDA_FUNCTION_1 = 3
+GENERIC_FUNCTION_2 = 5
+CUDA_FUNCTION_2 = 7
+GENERIC_CALLS_GENERIC = 11
+GENERIC_CALLS_CUDA = 13
+CUDA_CALLS_GENERIC = 17
+CUDA_CALLS_CUDA = 19
+GENERIC_TARGET_OL = 23
+CUDA_TARGET_OL = 29
+GENERIC_CALLS_TARGET_OL = 31
+CUDA_CALLS_TARGET_OL = 37
+GENERIC_TARGET_OL_CALLS_TARGET_OL = 41
+CUDA_TARGET_OL_CALLS_TARGET_OL = 43
+
+
+# Overload implementations
+
+
+@overload(generic_func_1, target="generic")
+def ol_generic_func_1(x):
+    def impl(x):
+        x[0] *= GENERIC_FUNCTION_1
+
+    return impl
+
+
+@overload(cuda_func_1, target="cuda")
+def ol_cuda_func_1(x):
+    def impl(x):
+        x[0] *= CUDA_FUNCTION_1
+
+    return impl
+
+
+@overload(generic_func_2, target="generic")
+def ol_generic_func_2(x):
+    def impl(x):
+        x[0] *= GENERIC_FUNCTION_2
+
+    return impl
+
+
+@overload(cuda_func_2, target="cuda")
+def ol_cuda_func(x):
+    def impl(x):
+        x[0] *= CUDA_FUNCTION_2
+
+    return impl
+
+
+@overload(generic_calls_generic, target="generic")
+def ol_generic_calls_generic(x):
+    def impl(x):
+        x[0] *= GENERIC_CALLS_GENERIC
+        generic_func_1(x)
+
+    return impl
+
+
+@overload(generic_calls_cuda, target="generic")
+def ol_generic_calls_cuda(x):
+    def impl(x):
+        x[0] *= GENERIC_CALLS_CUDA
+        cuda_func_1(x)
+
+    return impl
+
+
+@overload(cuda_calls_generic, target="cuda")
+def ol_cuda_calls_generic(x):
+    def impl(x):
+        x[0] *= CUDA_CALLS_GENERIC
+        generic_func_1(x)
+
+    return impl
+
+
+@overload(cuda_calls_cuda, target="cuda")
+def ol_cuda_calls_cuda(x):
+    def impl(x):
+        x[0] *= CUDA_CALLS_CUDA
+        cuda_func_1(x)
+
+    return impl
+
+
+@overload(target_overloaded, target="generic")
+def ol_target_overloaded_generic(x):
+    def impl(x):
+        x[0] *= GENERIC_TARGET_OL
+
+    return impl
+
+
+@overload(target_overloaded, target="cuda")
+def ol_target_overloaded_cuda(x):
+    def impl(x):
+        x[0] *= CUDA_TARGET_OL
+
+    return impl
+
+
+@overload(generic_calls_target_overloaded, target="generic")
+def ol_generic_calls_target_overloaded(x):
+    def impl(x):
+        x[0] *= GENERIC_CALLS_TARGET_OL
+        target_overloaded(x)
+
+    return impl
+
+
+@overload(cuda_calls_target_overloaded, target="cuda")
+def ol_cuda_calls_target_overloaded(x):
+    def impl(x):
+        x[0] *= CUDA_CALLS_TARGET_OL
+        target_overloaded(x)
+
+    return impl
+
+
+@overload(target_overloaded_calls_target_overloaded, target="generic")
+def ol_generic_calls_target_overloaded_generic(x):
+    def impl(x):
+        x[0] *= GENERIC_TARGET_OL_CALLS_TARGET_OL
+        target_overloaded(x)
+
+    return impl
+
+
+@overload(target_overloaded_calls_target_overloaded, target="cuda")
+def ol_generic_calls_target_overloaded_cuda(x):
+    def impl(x):
+        x[0] *= CUDA_TARGET_OL_CALLS_TARGET_OL
+        target_overloaded(x)
+
+    return impl
+
+
+@overload(default_values_and_kwargs)
+def ol_default_values_and_kwargs(out, x, y=5, z=6):
+    def impl(out, x, y=5, z=6):
+        out[0], out[1] = x + y, z
+
+    return impl
+
+
+@skip_on_cudasim("Overloading not supported in cudasim")
+class TestOverload(CUDATestCase):
+    def check_overload(self, kernel, expected):
+        x = np.ones(1, dtype=np.int32)
+        cuda.jit(kernel)[1, 1](x)
+        self.assertEqual(x[0], expected)
+
+    @skip_on_standalone_numba_cuda
+    def check_overload_cpu(self, kernel, expected):
+        x = np.ones(1, dtype=np.int32)
+        njit(kernel)(x)
+        self.assertEqual(x[0], expected)
+
+    def test_generic(self):
+        def kernel(x):
+            generic_func_1(x)
+
+        expected = GENERIC_FUNCTION_1
+        self.check_overload(kernel, expected)
+
+    def test_cuda(self):
+        def kernel(x):
+            cuda_func_1(x)
+
+        expected = CUDA_FUNCTION_1
+        self.check_overload(kernel, expected)
+
+    def test_generic_and_cuda(self):
+        def kernel(x):
+            generic_func_1(x)
+            cuda_func_1(x)
+
+        expected = GENERIC_FUNCTION_1 * CUDA_FUNCTION_1
+        self.check_overload(kernel, expected)
+
+    def test_call_two_generic_calls(self):
+        def kernel(x):
+            generic_func_1(x)
+            generic_func_2(x)
+
+        expected = GENERIC_FUNCTION_1 * GENERIC_FUNCTION_2
+        self.check_overload(kernel, expected)
+
+    def test_call_two_cuda_calls(self):
+        def kernel(x):
+            cuda_func_1(x)
+            cuda_func_2(x)
+
+        expected = CUDA_FUNCTION_1 * CUDA_FUNCTION_2
+        self.check_overload(kernel, expected)
+
+    def test_generic_calls_generic(self):
+        def kernel(x):
+            generic_calls_generic(x)
+
+        expected = GENERIC_CALLS_GENERIC * GENERIC_FUNCTION_1
+        self.check_overload(kernel, expected)
+
+    def test_generic_calls_cuda(self):
+        def kernel(x):
+            generic_calls_cuda(x)
+
+        expected = GENERIC_CALLS_CUDA * CUDA_FUNCTION_1
+        self.check_overload(kernel, expected)
+
+    def test_cuda_calls_generic(self):
+        def kernel(x):
+            cuda_calls_generic(x)
+
+        expected = CUDA_CALLS_GENERIC * GENERIC_FUNCTION_1
+        self.check_overload(kernel, expected)
+
+    def test_cuda_calls_cuda(self):
+        def kernel(x):
+            cuda_calls_cuda(x)
+
+        expected = CUDA_CALLS_CUDA * CUDA_FUNCTION_1
+        self.check_overload(kernel, expected)
+
+    def test_call_target_overloaded(self):
+        def kernel(x):
+            target_overloaded(x)
+
+        expected = CUDA_TARGET_OL
+        self.check_overload(kernel, expected)
+
+    def test_generic_calls_target_overloaded(self):
+        def kernel(x):
+            generic_calls_target_overloaded(x)
+
+        expected = GENERIC_CALLS_TARGET_OL * CUDA_TARGET_OL
+        self.check_overload(kernel, expected)
+
+    def test_cuda_calls_target_overloaded(self):
+        def kernel(x):
+            cuda_calls_target_overloaded(x)
+
+        expected = CUDA_CALLS_TARGET_OL * CUDA_TARGET_OL
+        self.check_overload(kernel, expected)
+
+    def test_target_overloaded_calls_target_overloaded(self):
+        def kernel(x):
+            target_overloaded_calls_target_overloaded(x)
+
+        # Check the CUDA overloads are used on CUDA
+        expected = CUDA_TARGET_OL_CALLS_TARGET_OL * CUDA_TARGET_OL
+        self.check_overload(kernel, expected)
+
+    @skip_on_standalone_numba_cuda
+    def test_target_overloaded_calls_target_overloaded_cpu(self):
+        def kernel(x):
+            target_overloaded_calls_target_overloaded(x)
+
+        # Check that the CPU overloads are used on the CPU
+        expected = GENERIC_TARGET_OL_CALLS_TARGET_OL * GENERIC_TARGET_OL
+        self.check_overload_cpu(kernel, expected)
+
+    @skip_on_standalone_numba_cuda
+    def test_overload_attribute_target(self):
+        MyDummy, MyDummyType = self.make_dummy_type()
+        mydummy_type_cpu = cpu_typeof(MyDummy())  # For @njit (cpu)
+        mydummy_type = typeof(MyDummy())  # For @cuda.jit (CUDA)
+
+        @overload_attribute(MyDummyType, "cuda_only", target="cuda")
+        def ov_dummy_cuda_attr(obj):
+            def imp(obj):
+                return 42
+
+            return imp
+
+        # Ensure that we cannot use the CUDA target-specific attribute on the
+        # CPU, and that an appropriate typing error is raised
+
+        # A different error is produced prior to version 0.60
+        # (the fixes in #9454 improved the message)
+        # https://github.com/numba/numba/pull/9454
+        if HAS_NUMBA and numba.version_info[:2] < (0, 60):
+            msg = 'resolving type of attribute "cuda_only" of "x"'
+        else:
+            msg = "Unknown attribute 'cuda_only'"
+
+        with self.assertRaisesRegex(TypingError, msg):
+
+            @njit(types.int64(mydummy_type_cpu))
+            def illegal_target_attr_use(x):
+                return x.cuda_only
+
+        # Ensure that the CUDA target-specific attribute is usable and works
+        # correctly when the target is CUDA - note eager compilation via
+        # signature
+        @cuda.jit(types.void(types.int64[::1], mydummy_type))
+        def cuda_target_attr_use(res, dummy):
+            res[0] = dummy.cuda_only
+
+    def test_default_values_and_kwargs(self):
+        """
+        Test default values and kwargs.
+        """
+
+        @cuda.jit()
+        def kernel(a, b, out):
+            default_values_and_kwargs(out, a, z=b)
+
+        out = np.empty(2, dtype=np.int64)
+        kernel[1, 1](1, 2, out)
+        self.assertEqual(tuple(out), (6, 2))
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import math
+import numpy as np
+from numba import cuda
+from numba.cuda import float64, int8, int32, void
+from numba.cuda.testing import unittest, CUDATestCase
+
+
+def cu_mat_power(A, power, power_A):
+    y, x = cuda.grid(2)
+
+    m, n = power_A.shape
+    if x >= n or y >= m:
+        return
+
+    power_A[y, x] = math.pow(A[y, x], int32(power))
+
+
+def cu_mat_power_binop(A, power, power_A):
+    y, x = cuda.grid(2)
+
+    m, n = power_A.shape
+    if x >= n or y >= m:
+        return
+
+    power_A[y, x] = A[y, x] ** power
+
+
+def vec_pow(r, x, y):
+    i = cuda.grid(1)
+
+    if i < len(r):
+        r[i] = pow(x[i], y[i])
+
+
+def vec_pow_binop(r, x, y):
+    i = cuda.grid(1)
+
+    if i < len(r):
+        r[i] = x[i] ** y[i]
+
+
+def vec_pow_inplace_binop(r, x):
+    i = cuda.grid(1)
+
+    if i < len(r):
+        r[i] **= x[i]
+
+
+def random_complex(N):
+    np.random.seed(123)
+    return np.random.random(1) + np.random.random(1) * 1j
+
+
+class TestCudaPowi(CUDATestCase):
+    def test_powi(self):
+        dec = cuda.jit(void(float64[:, :], int8, float64[:, :]))
+        kernel = dec(cu_mat_power)
+
+        power = 2
+        A = np.arange(10, dtype=np.float64).reshape(2, 5)
+        Aout = np.empty_like(A)
+        kernel[1, A.shape](A, power, Aout)
+        self.assertTrue(np.allclose(Aout, A**power))
+
+    def test_powi_binop(self):
+        dec = cuda.jit(void(float64[:, :], int8, float64[:, :]))
+        kernel = dec(cu_mat_power_binop)
+
+        power = 2
+        A = np.arange(10, dtype=np.float64).reshape(2, 5)
+        Aout = np.empty_like(A)
+        kernel[1, A.shape](A, power, Aout)
+        self.assertTrue(np.allclose(Aout, A**power))
+
+    # Relative tolerance kwarg is provided because 1.0e-7 (the default for
+    # assert_allclose) is a bit tight for single precision.
+    def _test_cpow(self, dtype, func, rtol=1.0e-7):
+        N = 32
+        x = random_complex(N).astype(dtype)
+        y = random_complex(N).astype(dtype)
+        r = np.zeros_like(x)
+
+        cfunc = cuda.jit(func)
+        cfunc[1, N](r, x, y)
+        np.testing.assert_allclose(r, x**y, rtol=rtol)
+
+        # Checks special cases
+        x = np.asarray([0.0j, 1.0j], dtype=dtype)
+        y = np.asarray([0.0j, 1.0], dtype=dtype)
+        r = np.zeros_like(x)
+
+        cfunc[1, 2](r, x, y)
+        np.testing.assert_allclose(r, x**y, rtol=rtol)
+
+    def test_cpow_complex64_pow(self):
+        self._test_cpow(np.complex64, vec_pow, rtol=3.0e-7)
+
+    def test_cpow_complex64_binop(self):
+        self._test_cpow(np.complex64, vec_pow_binop, rtol=3.0e-7)
+
+    def test_cpow_complex128_pow(self):
+        self._test_cpow(np.complex128, vec_pow)
+
+    def test_cpow_complex128_binop(self):
+        self._test_cpow(np.complex128, vec_pow_binop)
+
+    def _test_cpow_inplace_binop(self, dtype, rtol=1.0e-7):
+        N = 32
+        x = random_complex(N).astype(dtype)
+        y = random_complex(N).astype(dtype)
+        r = x**y
+
+        cfunc = cuda.jit(vec_pow_inplace_binop)
+        cfunc[1, N](x, y)
+        np.testing.assert_allclose(x, r, rtol=rtol)
+
+    def test_cpow_complex64_inplace_binop(self):
+        self._test_cpow_inplace_binop(np.complex64, rtol=3.0e-7)
+
+    def test_cpow_complex128_inplace_binop(self):
+        self._test_cpow_inplace_binop(np.complex128, rtol=3.0e-7)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,193 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.testing import CUDATestCase, skip_on_cudasim
+import numpy as np
+import subprocess
+import sys
+import unittest
+
+
+cuhello_usecase = """\
+from numba import cuda
+
+@cuda.jit
+def cuhello():
+    i = cuda.grid(1)
+    print(i, 999)
+    print(-42)
+
+cuhello[2, 3]()
+cuda.synchronize()
+"""
+
+
+printfloat_usecase = """\
+from numba import cuda
+
+@cuda.jit
+def printfloat():
+    i = cuda.grid(1)
+    print(i, 23, 34.75, 321)
+
+printfloat[1, 1]()
+cuda.synchronize()
+"""
+
+
+printbool_usecase = """\
+from numba import cuda
+
+@cuda.jit
+def printbool(x):
+    print(True)
+    print(False)
+    print(x == 0)
+
+printbool[1, 1](0)
+printbool[1, 1](1)
+cuda.synchronize()
+"""
+
+
+printstring_usecase = """\
+from numba import cuda
+
+@cuda.jit
+def printstring():
+    i = cuda.grid(1)
+    print(i, "hop!", 999)
+
+printstring[1, 3]()
+cuda.synchronize()
+"""
+
+
+printdim3_usecase = """\
+from numba import cuda
+
+@cuda.jit
+def printdim3():
+    print(cuda.threadIdx)
+
+printdim3[1, (2, 2, 2)]()
+cuda.synchronize()
+"""
+
+
+printempty_usecase = """\
+from numba import cuda
+
+@cuda.jit
+def printempty():
+    print()
+
+printempty[1, 1]()
+cuda.synchronize()
+"""
+
+
+print_too_many_usecase = """\
+from numba import cuda
+import numpy as np
+
+@cuda.jit
+def print_too_many(r):
+    print(r[0], r[1], r[2], r[3], r[4], r[5], r[6], r[7], r[8], r[9], r[10],
+          r[11], r[12], r[13], r[14], r[15], r[16], r[17], r[18], r[19], r[20],
+          r[21], r[22], r[23], r[24], r[25], r[26], r[27], r[28], r[29], r[30],
+          r[31], r[32])
+
+print_too_many[1, 1](np.arange(33))
+cuda.synchronize()
+"""
+
+print_bfloat16_usecase = """\
+from numba import cuda
+from numba.cuda import config
+
+@cuda.jit
+def print_bfloat16():
+    # 0.9375 is a dyadic rational, it's integer significand can expand within 7 digits.
+    # printing this should not give any rounding error.
+    a = cuda.bfloat16(0.9375)
+    print(a, a, a)
+
+print_bfloat16[1, 1]()
+cuda.synchronize()
+"""
+
+
+class TestPrint(CUDATestCase):
+    # Note that in these tests we generally strip the output to avoid dealing
+    # with platform-specific line ending issues, e.g. '\r\n' vs '\n' etc.
+
+    def run_code(self, code):
+        """Runs code in a subprocess and returns the captured output"""
+        cmd = [sys.executable, "-c", code]
+        cp = subprocess.run(cmd, timeout=60, capture_output=True, check=True)
+        return cp.stdout.decode(), cp.stderr.decode()
+
+    def test_cuhello(self):
+        output, _ = self.run_code(cuhello_usecase)
+        actual = [line.strip() for line in output.splitlines()]
+        expected = ["-42"] * 6 + ["%d 999" % i for i in range(6)]
+        # The output of GPU threads is intermingled, but each print()
+        # call is still atomic
+        self.assertEqual(sorted(actual), expected)
+
+    def test_printfloat(self):
+        output, _ = self.run_code(printfloat_usecase)
+        # CUDA and the simulator use different formats for float formatting
+        expected_cases = ["0 23 34.750000 321", "0 23 34.75 321"]
+        self.assertIn(output.strip(), expected_cases)
+
+    def test_bool(self):
+        output, _ = self.run_code(printbool_usecase)
+        expected = "True\r?\nFalse\r?\nTrue\r?\nTrue\r?\nFalse\r?\nFalse"
+        self.assertRegex(output.strip(), expected)
+
+    def test_printempty(self):
+        output, _ = self.run_code(printempty_usecase)
+        self.assertEqual(output.strip(), "")
+
+    def test_string(self):
+        output, _ = self.run_code(printstring_usecase)
+        lines = [line.strip() for line in output.splitlines(True)]
+        expected = ["%d hop! 999" % i for i in range(3)]
+        self.assertEqual(sorted(lines), expected)
+
+    def test_dim3(self):
+        output, _ = self.run_code(printdim3_usecase)
+        lines = [line.strip() for line in output.splitlines(True)]
+        expected = [str(i) for i in np.ndindex(2, 2, 2)]
+        self.assertEqual(sorted(lines), expected)
+
+    @skip_on_cudasim("bfloat16 on host is not yet supported.")
+    def test_bfloat16(self):
+        output, _ = self.run_code(print_bfloat16_usecase)
+        self.assertEqual(output.strip(), "0.937500 0.937500 0.937500")
+
+    @skip_on_cudasim("cudasim can print unlimited output")
+    def test_too_many_args(self):
+        # Tests that we emit the format string and warn when there are more
+        # than 32 arguments, in common with CUDA C/C++ printf - this is due to
+        # a limitation in CUDA vprintf, see:
+        # https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#limitations
+
+        output, errors = self.run_code(print_too_many_usecase)
+
+        # Check that the format string was printed instead of formatted garbage
+        expected_fmt_string = " ".join(["%lld" for _ in range(33)])
+        self.assertIn(expected_fmt_string, output)
+
+        # Check for the expected warning about formatting more than 32 items
+        warn_msg = (
+            "CUDA print() cannot print more than 32 items. The raw "
+            "format string will be emitted by the kernel instead."
+        )
+        self.assertIn(warn_msg, errors)
+
+
+if __name__ == "__main__":
+    unittest.main()