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

numba_cuda/numba/cuda/tests/cudadrv/test_init.py

@@ -0,0 +1,138 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import concurrent.futures
+import multiprocessing as mp
+import os
+
+from numba import cuda
+from numba.cuda.cudadrv.driver import CudaAPIError, driver
+from numba.cuda.cudadrv.error import CudaSupportError
+from numba.cuda.testing import skip_on_cudasim, unittest, CUDATestCase
+
+
+# A mock of cuInit that always raises a CudaAPIError
+def cuInit_raising(arg):
+    raise CudaAPIError(999, "CUDA_ERROR_UNKNOWN")
+
+
+# Test code to run in a child that patches driver.cuInit to a variant that
+# always raises. We can't use mock.patch.object here because driver.cuInit is
+# not assigned until we attempt to initialize - mock.patch.object cannot locate
+# the non-existent original method, and so fails. Instead we patch
+# driver.cuInit with our raising version prior to any attempt to initialize.
+def cuInit_raising_test():
+    driver.cuInit = cuInit_raising
+
+    try:
+        # A CUDA operation that forces initialization of the device
+        cuda.device_array(1)
+    except CudaSupportError as e:
+        success = True
+        msg = e.msg
+    else:
+        success = False
+        msg = None
+
+    return success, msg
+
+
+# Similar to cuInit_raising_test above, but for testing that the string
+# returned by cuda_error() is as expected.
+def initialization_error_test():
+    driver.cuInit = cuInit_raising
+
+    try:
+        # A CUDA operation that forces initialization of the device
+        cuda.device_array(1)
+    except CudaSupportError:
+        success = True
+    else:
+        success = False
+
+    return success, cuda.cuda_error()
+
+
+# For testing the path where Driver.__init__() catches a CudaSupportError
+def cuda_disabled_test():
+    try:
+        # A CUDA operation that forces initialization of the device
+        cuda.device_array(1)
+    except CudaSupportError as e:
+        success = True
+        msg = e.msg
+    else:
+        success = False
+        msg = None
+
+    return success, msg
+
+
+# Similar to cuda_disabled_test, but checks cuda.cuda_error() instead of the
+# exception raised on initialization
+def cuda_disabled_error_test():
+    try:
+        # A CUDA operation that forces initialization of the device
+        cuda.device_array(1)
+    except CudaSupportError:
+        success = True
+    else:
+        success = False
+
+    return success, cuda.cuda_error()
+
+
+@skip_on_cudasim("CUDA Simulator does not initialize driver")
+class TestInit(CUDATestCase):
+    def _test_init_failure(self, target, expected):
+        # Run the initialization failure test in a separate subprocess
+        with concurrent.futures.ProcessPoolExecutor(
+            mp_context=mp.get_context("spawn")
+        ) as exe:
+            # should complete within 30s
+            success, msg = exe.submit(target).result(timeout=30)
+
+        # Ensure the child process raised an exception during initialization
+        # before checking the message
+        if not success:
+            assert "CudaSupportError not raised" in msg
+
+        assert expected in msg
+
+    def test_init_failure_raising(self):
+        expected = "Error at driver init: CUDA_ERROR_UNKNOWN (999)"
+        self._test_init_failure(cuInit_raising_test, expected)
+
+    def test_init_failure_error(self):
+        expected = "CUDA_ERROR_UNKNOWN (999)"
+        self._test_init_failure(initialization_error_test, expected)
+
+    def _test_cuda_disabled(self, target):
+        # Uses _test_init_failure to launch the test in a separate subprocess
+        # with CUDA disabled.
+        cuda_disabled = os.environ.get("NUMBA_DISABLE_CUDA")
+        os.environ["NUMBA_DISABLE_CUDA"] = "1"
+        try:
+            expected = "CUDA is disabled due to setting NUMBA_DISABLE_CUDA=1"
+            self._test_init_failure(cuda_disabled_test, expected)
+        finally:
+            if cuda_disabled is not None:
+                os.environ["NUMBA_DISABLE_CUDA"] = cuda_disabled
+            else:
+                os.environ.pop("NUMBA_DISABLE_CUDA")
+
+    def test_cuda_disabled_raising(self):
+        self._test_cuda_disabled(cuda_disabled_test)
+
+    def test_cuda_disabled_error(self):
+        self._test_cuda_disabled(cuda_disabled_error_test)
+
+    def test_init_success(self):
+        # Here we assume that initialization is successful (because many bad
+        # things will happen with the test suite if it is not) and check that
+        # there is no error recorded.
+        self.assertIsNone(cuda.cuda_error())
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/cudadrv/test_inline_ptx.py

@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from llvmlite import ir
+
+from numba.cuda.cudadrv import nvvm
+from numba.cuda.testing import unittest, CUDATestCase
+from numba.cuda.testing import skip_on_cudasim
+
+
+@skip_on_cudasim("Inline PTX cannot be used in the simulator")
+class TestCudaInlineAsm(CUDATestCase):
+    def test_inline_rsqrt(self):
+        mod = ir.Module(__name__)
+        mod.triple = "nvptx64-nvidia-cuda"
+        nvvm.add_ir_version(mod)
+        fnty = ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.FloatType())])
+        fn = ir.Function(mod, fnty, "cu_rsqrt")
+        bldr = ir.IRBuilder(fn.append_basic_block("entry"))
+
+        rsqrt_approx_fnty = ir.FunctionType(ir.FloatType(), [ir.FloatType()])
+        inlineasm = ir.InlineAsm(
+            rsqrt_approx_fnty,
+            "rsqrt.approx.f32 $0, $1;",
+            "=f,f",
+            side_effect=True,
+        )
+        val = bldr.load(fn.args[0])
+        res = bldr.call(inlineasm, [val])
+
+        bldr.store(res, fn.args[0])
+        bldr.ret_void()
+
+        # generate ptx
+        mod.data_layout = nvvm.NVVM().data_layout
+        nvvm.set_cuda_kernel(fn)
+        nvvmir = str(mod)
+        ptx = nvvm.compile_ir(nvvmir)
+        self.assertTrue("rsqrt.approx.f32" in str(ptx))
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/cudadrv/test_is_fp16.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba import cuda
+from numba.cuda.testing import CUDATestCase, skip_on_cudasim, skip_unless_cc_53
+
+
+class TestIsFP16Supported(CUDATestCase):
+    def test_is_fp16_supported(self):
+        self.assertTrue(cuda.is_float16_supported())
+
+    @skip_on_cudasim("fp16 not available in sim")
+    @skip_unless_cc_53
+    def test_device_supports_float16(self):
+        self.assertTrue(cuda.get_current_device().supports_float16)

numba_cuda/numba/cuda/tests/cudadrv/test_linkable_code.py

@@ -0,0 +1,58 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import os
+from numba import cuda
+from numba.cuda.cudadrv.linkable_code import LinkableCode
+from numba.cuda.testing import CUDATestCase, skip_on_cudasim
+import unittest
+
+TEST_BIN_DIR = os.getenv("NUMBA_CUDA_TEST_BIN_DIR")
+if TEST_BIN_DIR:
+    test_device_functions_a = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.a"
+    )
+    test_device_functions_cubin = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.cubin"
+    )
+    test_device_functions_cu = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.cu"
+    )
+    test_device_functions_fatbin = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.fatbin"
+    )
+    test_device_functions_fatbin_multi = os.path.join(
+        TEST_BIN_DIR, "test_device_functions_multi.fatbin"
+    )
+    test_device_functions_o = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.o"
+    )
+    test_device_functions_ptx = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.ptx"
+    )
+    test_device_functions_ltoir = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.ltoir"
+    )
+
+
+class TestLinkableCode(CUDATestCase):
+    @skip_on_cudasim(reason="Simulator does not support linkable code")
+    @unittest.skipIf(not TEST_BIN_DIR, "necessary binaries not generated.")
+    def test_linkable_code_from_path_or_obj(self):
+        files_kind = [
+            (test_device_functions_a, cuda.Archive),
+            (test_device_functions_cubin, cuda.Cubin),
+            (test_device_functions_cu, cuda.CUSource),
+            (test_device_functions_fatbin, cuda.Fatbin),
+            (test_device_functions_o, cuda.Object),
+            (test_device_functions_ptx, cuda.PTXSource),
+            (test_device_functions_ltoir, cuda.LTOIR),
+        ]
+
+        for path, kind in files_kind:
+            obj = LinkableCode.from_path_or_obj(path)
+            assert isinstance(obj, kind)
+
+            # test identity of from_path_or_obj
+            obj2 = LinkableCode.from_path_or_obj(obj)
+            assert obj2 is obj

numba_cuda/numba/cuda/tests/cudadrv/test_linker.py

@@ -0,0 +1,348 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numpy as np
+import pytest
+from numba.cuda.testing import unittest
+from numba.cuda.testing import (
+    skip_on_cudasim,
+    skip_if_cuda_includes_missing,
+    skip_if_nvjitlink_missing,
+)
+from numba.cuda.testing import CUDATestCase, test_data_dir
+from numba.cuda.cudadrv.driver import CudaAPIError, _Linker, LinkerError
+from numba.cuda import require_context
+from numba import cuda
+from numba.cuda import void, float64, int64, int32, float32
+from numba.cuda.typing.typeof import typeof
+
+CONST1D = np.arange(10, dtype=np.float64)
+
+
+def simple_const_mem(A):
+    C = cuda.const.array_like(CONST1D)
+    i = cuda.grid(1)
+
+    A[i] = C[i] + 1.0
+
+
+def func_with_lots_of_registers(x, a, b, c, d, e, f):
+    a1 = 1.0
+    a2 = 1.0
+    a3 = 1.0
+    a4 = 1.0
+    a5 = 1.0
+    b1 = 1.0
+    b2 = 1.0
+    b3 = 1.0
+    b4 = 1.0
+    b5 = 1.0
+    c1 = 1.0
+    c2 = 1.0
+    c3 = 1.0
+    c4 = 1.0
+    c5 = 1.0
+    d1 = 10
+    d2 = 10
+    d3 = 10
+    d4 = 10
+    d5 = 10
+    for i in range(a):
+        a1 += b
+        a2 += c
+        a3 += d
+        a4 += e
+        a5 += f
+        b1 *= b
+        b2 *= c
+        b3 *= d
+        b4 *= e
+        b5 *= f
+        c1 /= b
+        c2 /= c
+        c3 /= d
+        c4 /= e
+        c5 /= f
+        d1 <<= b
+        d2 <<= c
+        d3 <<= d
+        d4 <<= e
+        d5 <<= f
+    x[cuda.grid(1)] = a1 + a2 + a3 + a4 + a5
+    x[cuda.grid(1)] += b1 + b2 + b3 + b4 + b5
+    x[cuda.grid(1)] += c1 + c2 + c3 + c4 + c5
+    x[cuda.grid(1)] += d1 + d2 + d3 + d4 + d5
+
+
+def simple_smem(ary, dty):
+    sm = cuda.shared.array(100, dty)
+    i = cuda.grid(1)
+    if i == 0:
+        for j in range(100):
+            sm[j] = j
+    cuda.syncthreads()
+    ary[i] = sm[i]
+
+
+def coop_smem2d(ary):
+    i, j = cuda.grid(2)
+    sm = cuda.shared.array((10, 20), float32)
+    sm[i, j] = (i + 1) / (j + 1)
+    cuda.syncthreads()
+    ary[i, j] = sm[i, j]
+
+
+def simple_maxthreads(ary):
+    i = cuda.grid(1)
+    ary[i] = i
+
+
+LMEM_SIZE = 1000
+
+
+def simple_lmem(A, B, dty):
+    C = cuda.local.array(LMEM_SIZE, dty)
+    for i in range(C.shape[0]):
+        C[i] = A[i]
+    for i in range(C.shape[0]):
+        B[i] = C[i]
+
+
+@skip_on_cudasim("Linking unsupported in the simulator")
+class TestLinker(CUDATestCase):
+    @require_context
+    def test_linker_basic(self):
+        """Simply go through the constructor and destructor"""
+        linker = _Linker.new(cc=(7, 5))
+        del linker
+
+    def _test_linking(self, eager):
+        global bar  # must be a global; other it is recognized as a freevar
+        bar = cuda.declare_device("bar", "int32(int32)")
+
+        link = str(test_data_dir / "jitlink.ptx")
+
+        if eager:
+            args = ["void(int32[:], int32[:])"]
+        else:
+            args = []
+
+        @cuda.jit(*args, link=[link])
+        def foo(x, y):
+            i = cuda.grid(1)
+            x[i] += bar(y[i])
+
+        A = np.array([123], dtype=np.int32)
+        B = np.array([321], dtype=np.int32)
+
+        foo[1, 1](A, B)
+
+        self.assertTrue(A[0] == 123 + 2 * 321)
+
+    def test_linking_lazy_compile(self):
+        self._test_linking(eager=False)
+
+    def test_linking_eager_compile(self):
+        self._test_linking(eager=True)
+
+    def test_linking_cu(self):
+        bar = cuda.declare_device("bar", "int32(int32)")
+
+        link = str(test_data_dir / "jitlink.cu")
+
+        @cuda.jit(link=[link])
+        def kernel(r, x):
+            i = cuda.grid(1)
+
+            if i < len(r):
+                r[i] = bar(x[i])
+
+        x = np.arange(10, dtype=np.int32)
+        r = np.zeros_like(x)
+
+        kernel[1, 32](r, x)
+
+        # Matches the operation of bar() in jitlink.cu
+        expected = x * 2
+        np.testing.assert_array_equal(r, expected)
+
+    def test_linking_cu_log_warning(self):
+        bar = cuda.declare_device("bar", "int32(int32)")
+
+        link = str(test_data_dir / "warn.cu")
+
+        with pytest.warns(UserWarning) as w:
+
+            @cuda.jit("void(int32)", link=[link])
+            def kernel(x):
+                bar(x)
+
+        nvrtc_log_warnings = [
+            wi for wi in w if "NVRTC log messages" in str(wi.message)
+        ]
+        self.assertEqual(
+            len(nvrtc_log_warnings), 1, "Expected warnings from NVRTC"
+        )
+        # Check the warning refers to the log messages
+        self.assertIn("NVRTC log messages", str(nvrtc_log_warnings[0].message))
+        # Check the message pertaining to the unused variable is provided
+        self.assertIn(
+            "declared but never referenced", str(nvrtc_log_warnings[0].message)
+        )
+
+    def test_linking_cu_error(self):
+        bar = cuda.declare_device("bar", "int32(int32)")
+
+        link = str(test_data_dir / "error.cu")
+
+        from cuda.core.experimental._utils.cuda_utils import NVRTCError
+
+        errty = NVRTCError
+        with self.assertRaises(errty) as e:
+
+            @cuda.jit("void(int32)", link=[link])
+            def kernel(x):
+                bar(x)
+
+        msg = e.exception.args[0]
+        # Check the error message refers to the NVRTC compile
+        nvrtc_err_str = "NVRTC_ERROR_COMPILATION"
+        self.assertIn(nvrtc_err_str, msg)
+        # Check the expected error in the CUDA source is reported
+        self.assertIn('identifier "SYNTAX" is undefined', msg)
+        # Check the filename is reported correctly
+        self.assertIn('in the compilation of "error.cu"', msg)
+
+    def test_linking_unknown_filetype_error(self):
+        expected_err = "Don't know how to link file with extension .cuh"
+        with self.assertRaisesRegex(RuntimeError, expected_err):
+
+            @cuda.jit("void()", link=["header.cuh"])
+            def kernel():
+                pass
+
+    def test_linking_file_with_no_extension_error(self):
+        expected_err = "Don't know how to link file with no extension"
+        with self.assertRaisesRegex(RuntimeError, expected_err):
+
+            @cuda.jit("void()", link=["data"])
+            def kernel():
+                pass
+
+    @skip_if_cuda_includes_missing
+    def test_linking_cu_cuda_include(self):
+        link = str(test_data_dir / "cuda_include.cu")
+
+        # An exception will be raised when linking this kernel due to the
+        # compile failure if CUDA includes cannot be found by Nvrtc.
+        @cuda.jit("void()", link=[link])
+        def kernel():
+            pass
+
+    def test_try_to_link_nonexistent(self):
+        with self.assertRaises(LinkerError) as e:
+
+            @cuda.jit("void(int32[::1])", link=["nonexistent.a"])
+            def f(x):
+                x[0] = 0
+
+        self.assertIn("nonexistent.a not found", e.exception.args)
+
+    def test_set_registers_no_max(self):
+        """Ensure that the jitted kernel used in the test_set_registers_* tests
+        uses more than 57 registers - this ensures that test_set_registers_*
+        are really checking that they reduced the number of registers used from
+        something greater than the maximum."""
+        compiled = cuda.jit(func_with_lots_of_registers)
+        compiled = compiled.specialize(np.empty(32), *range(6))
+        self.assertGreater(compiled.get_regs_per_thread(), 57)
+
+    def test_set_registers_57(self):
+        compiled = cuda.jit(max_registers=57)(func_with_lots_of_registers)
+        compiled = compiled.specialize(np.empty(32), *range(6))
+        self.assertLessEqual(compiled.get_regs_per_thread(), 57)
+
+    def test_set_registers_38(self):
+        compiled = cuda.jit(max_registers=38)(func_with_lots_of_registers)
+        compiled = compiled.specialize(np.empty(32), *range(6))
+        self.assertLessEqual(compiled.get_regs_per_thread(), 38)
+
+    def test_set_registers_eager(self):
+        sig = void(float64[::1], int64, int64, int64, int64, int64, int64)
+        compiled = cuda.jit(sig, max_registers=38)(func_with_lots_of_registers)
+        self.assertLessEqual(compiled.get_regs_per_thread(), 38)
+
+    def test_get_const_mem_size(self):
+        sig = void(float64[::1])
+        compiled = cuda.jit(sig)(simple_const_mem)
+        const_mem_size = compiled.get_const_mem_size()
+        self.assertGreaterEqual(const_mem_size, CONST1D.nbytes)
+
+    def test_get_no_shared_memory(self):
+        compiled = cuda.jit(func_with_lots_of_registers)
+        compiled = compiled.specialize(np.empty(32), *range(6))
+        shared_mem_size = compiled.get_shared_mem_per_block()
+        self.assertEqual(shared_mem_size, 0)
+
+    def test_get_shared_mem_per_block(self):
+        sig = void(int32[::1], typeof(np.int32))
+        compiled = cuda.jit(sig)(simple_smem)
+        shared_mem_size = compiled.get_shared_mem_per_block()
+        self.assertEqual(shared_mem_size, 400)
+
+    def test_get_shared_mem_per_specialized(self):
+        compiled = cuda.jit(simple_smem)
+        compiled_specialized = compiled.specialize(
+            np.zeros(100, dtype=np.int32), np.float64
+        )
+        shared_mem_size = compiled_specialized.get_shared_mem_per_block()
+        self.assertEqual(shared_mem_size, 800)
+
+    def test_get_max_threads_per_block(self):
+        compiled = cuda.jit("void(float32[:,::1])")(coop_smem2d)
+        max_threads = compiled.get_max_threads_per_block()
+        self.assertGreater(max_threads, 0)
+
+    def test_max_threads_exceeded(self):
+        compiled = cuda.jit("void(int32[::1])")(simple_maxthreads)
+        max_threads = compiled.get_max_threads_per_block()
+        nelem = max_threads + 1
+        ary = np.empty(nelem, dtype=np.int32)
+        try:
+            compiled[1, nelem](ary)
+        except CudaAPIError as e:
+            self.assertIn("cuLaunchKernel", e.msg)
+
+    def test_get_local_mem_per_thread(self):
+        sig = void(int32[::1], int32[::1], typeof(np.int32))
+        compiled = cuda.jit(sig)(simple_lmem)
+        local_mem_size = compiled.get_local_mem_per_thread()
+        calc_size = np.dtype(np.int32).itemsize * LMEM_SIZE
+        self.assertGreaterEqual(local_mem_size, calc_size)
+
+    def test_get_local_mem_per_specialized(self):
+        compiled = cuda.jit(simple_lmem)
+        compiled_specialized = compiled.specialize(
+            np.zeros(LMEM_SIZE, dtype=np.int32),
+            np.zeros(LMEM_SIZE, dtype=np.int32),
+            np.float64,
+        )
+        local_mem_size = compiled_specialized.get_local_mem_per_thread()
+        calc_size = np.dtype(np.float64).itemsize * LMEM_SIZE
+        self.assertGreaterEqual(local_mem_size, calc_size)
+
+    @skip_if_nvjitlink_missing("nvJitLink not installed or new enough (>12.3)")
+    def test_link_for_different_cc(self):
+        linker = _Linker.new(cc=(7, 5), lto=True)
+        code = """
+__device__ int foo(int x) {
+    return x + 1;
+}
+"""
+        linker.add_cu(code, "foo")
+        ptx = linker.get_linked_ptx().decode()
+        assert "target sm_75" in ptx
+
+
+if __name__ == "__main__":
+    unittest.main()