numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

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

@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numpy as np
+from numba.cuda.cudadrv.driver import device_memset, driver
+from numba import cuda
+from numba.cuda.testing import unittest, CUDATestCase
+from numba.cuda.testing import skip_on_cudasim, skip_on_arm
+from numba.cuda.tests.support import linux_only
+
+
+@skip_on_cudasim("CUDA Driver API unsupported in the simulator")
+@linux_only
+@skip_on_arm("Managed Alloc support is experimental/untested on ARM")
+class TestManagedAlloc(CUDATestCase):
+    def tearDown(self):
+        super().tearDown()
+        cuda.current_context().reset()
+
+    def get_total_gpu_memory(self):
+        # We use a driver function to directly get the total GPU memory because
+        # an EMM plugin may report something different (or not implement
+        # get_memory_info at all).
+        free, total = driver.cuMemGetInfo()
+        return total
+
+    def skip_if_cc_major_lt(self, min_required, reason):
+        """
+        Skip the current test if the compute capability of the device is
+        less than `min_required`.
+        """
+        ctx = cuda.current_context()
+        cc_major = ctx.device.compute_capability[0]
+        if cc_major < min_required:
+            self.skipTest(reason)
+
+    # CUDA Unified Memory comes in two flavors. For GPUs in the Kepler and
+    # Maxwell generations, managed memory allocations work as opaque,
+    # contiguous segments that can either be on the device or the host. For
+    # GPUs in the Pascal or later generations, managed memory operates on a
+    # per-page basis, so we can have arrays larger than GPU memory, where only
+    # part of them is resident on the device at one time. To ensure that this
+    # test works correctly on all supported GPUs, we'll select the size of our
+    # memory such that we only oversubscribe the GPU memory if we're on a
+    # Pascal or newer GPU (compute capability at least 6.0).
+
+    def test_managed_alloc_driver_undersubscribe(self):
+        msg = "Managed memory unsupported prior to CC 3.0"
+        self.skip_if_cc_major_lt(3, msg)
+        # We keep the allocation small so that it doesn't hang on GPUs
+        # with large memory (H100)
+        self._test_managed_alloc_driver(0.1)
+
+    # This test is skipped by default because it is easy to hang the machine
+    # for a very long time or get OOM killed if the GPU memory size is >50% of
+    # the system memory size. Even if the system does have more than 2x the RAM
+    # of the GPU, this test runs for a very long time (in comparison to the
+    # rest of the tests in the suite).
+    #
+    # However, it is left in here for manual testing as required.
+
+    @unittest.skip
+    def test_managed_alloc_driver_oversubscribe(self):
+        msg = "Oversubscription of managed memory unsupported prior to CC 6.0"
+        self.skip_if_cc_major_lt(6, msg)
+        self._test_managed_alloc_driver(2.0)
+
+    def test_managed_alloc_driver_host_attach(self):
+        msg = "Host attached managed memory is not accessible prior to CC 6.0"
+        self.skip_if_cc_major_lt(6, msg)
+        # Only test with a small array (0.01 * memory size) to keep the test
+        # quick.
+        self._test_managed_alloc_driver(0.01, attach_global=False)
+
+    def _test_managed_alloc_driver(self, memory_factor, attach_global=True):
+        # Verify that we can allocate and operate on managed
+        # memory through the CUDA driver interface.
+
+        total_mem_size = self.get_total_gpu_memory()
+        n_bytes = int(memory_factor * total_mem_size)
+
+        ctx = cuda.current_context()
+        mem = ctx.memallocmanaged(n_bytes, attach_global=attach_global)
+
+        dtype = np.dtype(np.uint8)
+        n_elems = n_bytes // dtype.itemsize
+        ary = np.ndarray(shape=n_elems, dtype=dtype, buffer=mem)
+
+        magic = 0xAB
+        device_memset(mem, magic, n_bytes)
+        ctx.synchronize()
+
+        # Note that this assertion operates on the CPU, so this
+        # test effectively drives both the CPU and the GPU on
+        # managed memory.
+
+        self.assertTrue(np.all(ary == magic))
+
+    def _test_managed_array(self, attach_global=True):
+        # Check the managed_array interface on both host and device.
+
+        ary = cuda.managed_array(100, dtype=np.double)
+        ary.fill(123.456)
+        self.assertTrue(all(ary == 123.456))
+
+        @cuda.jit("void(double[:])")
+        def kernel(x):
+            i = cuda.grid(1)
+            if i < x.shape[0]:
+                x[i] = 1.0
+
+        kernel[10, 10](ary)
+        cuda.current_context().synchronize()
+
+        self.assertTrue(all(ary == 1.0))
+
+    def test_managed_array_attach_global(self):
+        self._test_managed_array()
+
+    def test_managed_array_attach_host(self):
+        self._test_managed_array()
+        msg = "Host attached managed memory is not accessible prior to CC 6.0"
+        self.skip_if_cc_major_lt(6, msg)
+        self._test_managed_array(attach_global=False)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,301 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import unittest
+import threading
+
+import numpy as np
+
+from numba import cuda
+from numba.cuda import config
+from numba.cuda.cudadrv.linkable_code import CUSource
+from numba.cuda.testing import (
+    CUDATestCase,
+    skip_on_cudasim,
+)
+
+if not config.ENABLE_CUDASIM:
+    from cuda.bindings.driver import cuModuleGetGlobal, cuMemcpyHtoD
+
+    from cuda.bindings.driver import CUmodule as cu_module_type
+
+
+def wipe_all_modules_in_context():
+    """Cleans all modules reference held by current context.
+    This simulates the behavior on interpreter shutdown.
+
+    TODO: This is a temp solution until
+    https://github.com/NVIDIA/numba-cuda/issues/171 is implemented.
+    """
+    ctx = cuda.current_context()
+    ctx.reset()
+
+
+def get_hashable_handle_value(handle):
+    return handle
+
+
+@skip_on_cudasim("Module loading not implemented in the simulator")
+class TestModuleCallbacksBasic(CUDATestCase):
+    def test_basic(self):
+        counter = 0
+
+        def setup(handle):
+            self.assertTrue(isinstance(handle, cu_module_type))
+            nonlocal counter
+            counter += 1
+
+        def teardown(handle):
+            self.assertTrue(isinstance(handle, cu_module_type))
+            nonlocal counter
+            counter -= 1
+
+        lib = CUSource("", setup_callback=setup, teardown_callback=teardown)
+
+        @cuda.jit(link=[lib])
+        def kernel():
+            pass
+
+        self.assertEqual(counter, 0)
+        kernel[1, 1]()
+        self.assertEqual(counter, 1)
+        kernel[1, 1]()  # cached
+        self.assertEqual(counter, 1)
+
+        wipe_all_modules_in_context()
+        del kernel
+        self.assertEqual(counter, 0)
+
+    def test_different_argtypes(self):
+        counter = 0
+        setup_seen = set()
+        teardown_seen = set()
+
+        def setup(handle):
+            nonlocal counter, setup_seen
+            counter += 1
+            setup_seen.add(get_hashable_handle_value(handle))
+
+        def teardown(handle):
+            nonlocal counter
+            counter -= 1
+            teardown_seen.add(get_hashable_handle_value(handle))
+
+        lib = CUSource("", setup_callback=setup, teardown_callback=teardown)
+
+        @cuda.jit(link=[lib])
+        def kernel(arg):
+            pass
+
+        self.assertEqual(counter, 0)
+        kernel[1, 1](42)  # (int64)->() : module 1
+        self.assertEqual(counter, 1)
+        kernel[1, 1](100)  # (int64)->() : module 1, cached
+        self.assertEqual(counter, 1)
+        kernel[1, 1](3.14)  # (float64)->() : module 2
+        self.assertEqual(counter, 2)
+
+        wipe_all_modules_in_context()
+        del kernel
+        self.assertEqual(counter, 0)
+
+        self.assertEqual(len(setup_seen), 2)
+        self.assertEqual(len(teardown_seen), 2)
+
+    def test_two_kernels(self):
+        counter = 0
+        setup_seen = set()
+        teardown_seen = set()
+
+        def setup(handle):
+            nonlocal counter, setup_seen
+            counter += 1
+            setup_seen.add(get_hashable_handle_value(handle))
+
+        def teardown(handle):
+            nonlocal counter, teardown_seen
+            counter -= 1
+            teardown_seen.add(get_hashable_handle_value(handle))
+
+        lib = CUSource("", setup_callback=setup, teardown_callback=teardown)
+
+        @cuda.jit(link=[lib])
+        def kernel():
+            pass
+
+        @cuda.jit(link=[lib])
+        def kernel2():
+            pass
+
+        kernel[1, 1]()
+        self.assertEqual(counter, 1)
+        kernel2[1, 1]()
+        self.assertEqual(counter, 2)
+
+        wipe_all_modules_in_context()
+        del kernel
+        self.assertEqual(counter, 0)
+
+        self.assertEqual(len(setup_seen), 2)
+        self.assertEqual(len(teardown_seen), 2)
+
+
+@skip_on_cudasim("Module loading not implemented in the simulator")
+class TestModuleCallbacksAPICompleteness(CUDATestCase):
+    def test_api(self):
+        def setup(handle):
+            pass
+
+        def teardown(handle):
+            pass
+
+        api_combo = [
+            (setup, teardown),
+            (setup, None),
+            (None, teardown),
+            (None, None),
+        ]
+
+        for setup, teardown in api_combo:
+            with self.subTest(setup=setup, teardown=teardown):
+                lib = CUSource(
+                    "", setup_callback=setup, teardown_callback=teardown
+                )
+
+                @cuda.jit(link=[lib])
+                def kernel():
+                    pass
+
+                kernel[1, 1]()
+
+
+@skip_on_cudasim("Module loading not implemented in the simulator")
+class TestModuleCallbacks(CUDATestCase):
+    def setUp(self):
+        super().setUp()
+
+        module = """
+__device__ int num = 0;
+extern "C"
+__device__ int get_num(int &retval) {
+    retval = num;
+    return 0;
+}
+"""
+
+        def set_forty_two(handle):
+            # Initialize 42 to global variable `num`
+            res, dptr, size = cuModuleGetGlobal(
+                get_hashable_handle_value(handle), "num".encode()
+            )
+
+            arr = np.array([42], np.int32)
+            cuMemcpyHtoD(dptr, arr.ctypes.data, size)
+
+        self.lib = CUSource(
+            module, setup_callback=set_forty_two, teardown_callback=None
+        )
+
+    def test_decldevice_arg(self):
+        get_num = cuda.declare_device("get_num", "int32()", link=[self.lib])
+
+        @cuda.jit
+        def kernel(arr):
+            arr[0] = get_num()
+
+        arr = np.zeros(1, np.int32)
+        kernel[1, 1](arr)
+        self.assertEqual(arr[0], 42)
+
+    def test_jitarg(self):
+        get_num = cuda.declare_device("get_num", "int32()")
+
+        @cuda.jit(link=[self.lib])
+        def kernel(arr):
+            arr[0] = get_num()
+
+        arr = np.zeros(1, np.int32)
+        kernel[1, 1](arr)
+        self.assertEqual(arr[0], 42)
+
+
+@skip_on_cudasim("Module loading not implemented in the simulator")
+class TestMultithreadedCallbacks(CUDATestCase):
+    def test_concurrent_initialization(self):
+        seen_mods = set()
+        max_seen_mods = 0
+
+        def setup(mod):
+            nonlocal seen_mods, max_seen_mods
+            seen_mods.add(get_hashable_handle_value(mod))
+            max_seen_mods = max(max_seen_mods, len(seen_mods))
+
+        def teardown(mod):
+            nonlocal seen_mods
+            # Raises an error if the module is not found in the seen_mods
+            seen_mods.remove(get_hashable_handle_value(mod))
+
+        lib = CUSource("", setup_callback=setup, teardown_callback=teardown)
+
+        @cuda.jit(link=[lib])
+        def kernel():
+            pass
+
+        def concurrent_compilation_launch(kernel):
+            kernel[1, 1]()
+
+        threads = [
+            threading.Thread(
+                target=concurrent_compilation_launch, args=(kernel,)
+            )
+            for _ in range(4)
+        ]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        wipe_all_modules_in_context()
+        self.assertEqual(len(seen_mods), 0)
+        self.assertEqual(max_seen_mods, 1)  # one moduled shared across threads
+
+    def test_concurrent_initialization_different_args(self):
+        seen_mods = set()
+        max_seen_mods = 0
+
+        def setup(mod):
+            nonlocal seen_mods, max_seen_mods
+            seen_mods.add(get_hashable_handle_value(mod))
+            max_seen_mods = max(max_seen_mods, len(seen_mods))
+
+        def teardown(mod):
+            nonlocal seen_mods
+            seen_mods.remove(get_hashable_handle_value(mod))
+
+        lib = CUSource("", setup_callback=setup, teardown_callback=teardown)
+
+        @cuda.jit(link=[lib])
+        def kernel(a):
+            pass
+
+        def concurrent_compilation_launch():
+            kernel[1, 1](42)  # (int64)->() : module 1
+            kernel[1, 1](9)  # (int64)->() : module 1 from cache
+            kernel[1, 1](3.14)  # (float64)->() : module 2
+
+        threads = [
+            threading.Thread(target=concurrent_compilation_launch)
+            for _ in range(4)
+        ]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        wipe_all_modules_in_context()
+        assert len(seen_mods) == 0
+        self.assertEqual(max_seen_mods, 2)  # two modules shared across threads
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,174 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import pytest
+from numba.cuda.testing import unittest
+from numba.cuda.testing import skip_on_cudasim
+from numba.cuda.testing import CUDATestCase
+from numba.cuda import get_current_device
+from numba.cuda.cudadrv.driver import _Linker, _have_nvjitlink
+
+from numba import cuda
+from numba.cuda import config
+
+import os
+import io
+import contextlib
+
+
+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"
+    )
+
+
+@unittest.skipIf(
+    not TEST_BIN_DIR or not _have_nvjitlink(),
+    "nvJitLink not installed or new enough (>12.3)",
+)
+@skip_on_cudasim("Linking unsupported in the simulator")
+class TestLinker(CUDATestCase):
+    def test_nvjitlink_add_file_guess_ext_linkable_code(self):
+        files = (
+            test_device_functions_a,
+            test_device_functions_cubin,
+            test_device_functions_cu,
+            test_device_functions_fatbin,
+            test_device_functions_o,
+            test_device_functions_ptx,
+        )
+        for file in files:
+            with self.subTest(file=file):
+                linker = _Linker(cc=get_current_device().compute_capability)
+                linker.add_file_guess_ext(file)
+
+    def test_nvjitlink_test_add_file_guess_ext_invalid_input(self):
+        with open(test_device_functions_cubin, "rb") as f:
+            content = f.read()
+
+        linker = _Linker(cc=get_current_device().compute_capability)
+        with self.assertRaisesRegex(
+            TypeError, "Expected path to file or a LinkableCode"
+        ):
+            # Feeding raw data as bytes to add_file_guess_ext should raise,
+            # because there's no way to know what kind of file to treat it as
+            linker.add_file_guess_ext(content)
+
+    def test_nvjitlink_jit_with_linkable_code(self):
+        files = (
+            test_device_functions_a,
+            test_device_functions_cubin,
+            test_device_functions_cu,
+            test_device_functions_fatbin,
+            test_device_functions_o,
+            test_device_functions_ptx,
+        )
+        for lto in [True, False]:
+            for file in files:
+                with self.subTest(file=file):
+                    sig = "uint32(uint32, uint32)"
+                    add_from_numba = cuda.declare_device("add_from_numba", sig)
+
+                    @cuda.jit(link=[file], lto=lto)
+                    def kernel(result):
+                        result[0] = add_from_numba(1, 2)
+
+                    result = cuda.device_array(1)
+                    kernel[1, 1](result)
+                    assert result[0] == 3
+
+    def test_nvjitlink_jit_with_linkable_code_lto_dump_assembly(self):
+        files = [
+            test_device_functions_cu,
+            test_device_functions_ltoir,
+            test_device_functions_fatbin_multi,
+        ]
+
+        config.DUMP_ASSEMBLY = True
+
+        for file in files:
+            with self.subTest(file=file):
+                f = io.StringIO()
+                with contextlib.redirect_stdout(f):
+                    sig = "uint32(uint32, uint32)"
+                    add_from_numba = cuda.declare_device("add_from_numba", sig)
+
+                    @cuda.jit(link=[file], lto=True)
+                    def kernel(result):
+                        result[0] = add_from_numba(1, 2)
+
+                    result = cuda.device_array(1)
+                    kernel[1, 1](result)
+                    assert result[0] == 3
+
+                self.assertTrue("ASSEMBLY (AFTER LTO)" in f.getvalue())
+
+        config.DUMP_ASSEMBLY = False
+
+    def test_nvjitlink_jit_with_linkable_code_lto_dump_assembly_warn(self):
+        files = [
+            test_device_functions_a,
+            test_device_functions_cubin,
+            test_device_functions_fatbin,
+            test_device_functions_o,
+            test_device_functions_ptx,
+        ]
+
+        config.DUMP_ASSEMBLY = True
+
+        for file in files:
+            with self.subTest(file=file):
+                sig = "uint32(uint32, uint32)"
+                add_from_numba = cuda.declare_device("add_from_numba", sig)
+
+                @cuda.jit(link=[file], lto=True)
+                def kernel(result):
+                    result[0] = add_from_numba(1, 2)
+
+                result = cuda.device_array(1)
+                func = kernel[1, 1]
+                with pytest.warns(
+                    UserWarning,
+                    match="it is not optimizable at link time, and `ignore_nonlto == True`",
+                ):
+                    func(result)
+                assert result[0] == 3
+
+        config.DUMP_ASSEMBLY = False
+
+    def test_nvjitlink_jit_with_invalid_linkable_code(self):
+        with open(test_device_functions_cubin, "rb") as f:
+            content = f.read()
+        with self.assertRaisesRegex(
+            TypeError, "Expected path to file or a LinkableCode"
+        ):
+
+            @cuda.jit("void()", link=[content])
+            def kernel():
+                pass
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,28 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.cudadrv import nvrtc
+from numba.cuda.testing import skip_on_cudasim
+
+import unittest
+
+
+@skip_on_cudasim("NVVM Driver unsupported in the simulator")
+class TestArchOption(unittest.TestCase):
+    def test_get_arch_option(self):
+        # Test returning the nearest lowest arch.
+        self.assertEqual(nvrtc.get_arch_option(7, 5), "compute_75")
+        self.assertEqual(nvrtc.get_arch_option(7, 7), "compute_75")
+        self.assertEqual(nvrtc.get_arch_option(8, 5), "compute_80")
+        self.assertEqual(nvrtc.get_arch_option(9, 1), "compute_90")
+        # Test known arch.
+        supported_ccs = nvrtc.get_supported_ccs()
+        for cc in supported_ccs:
+            self.assertEqual(nvrtc.get_arch_option(*cc), "compute_%d%d" % cc)
+        self.assertEqual(
+            nvrtc.get_arch_option(1000, 0), "compute_%d%d" % supported_ccs[-1]
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()