numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

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

@@ -0,0 +1,247 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from contextlib import contextmanager
+
+import numpy as np
+
+from numba import cuda
+from numba.cuda.testing import (
+    unittest,
+    skip_on_cudasim,
+    skip_if_external_memmgr,
+    CUDATestCase,
+)
+from numba.cuda.tests.support import captured_stderr
+from numba.cuda.core import config
+
+
+@skip_on_cudasim("not supported on CUDASIM")
+class TestDeallocation(CUDATestCase):
+    @skip_if_external_memmgr("Deallocation specific to Numba memory management")
+    def test_max_pending_count(self):
+        # get deallocation manager and flush it
+        deallocs = cuda.current_context().memory_manager.deallocations
+        deallocs.clear()
+        self.assertEqual(len(deallocs), 0)
+        # deallocate to maximum count
+        for i in range(config.CUDA_DEALLOCS_COUNT):
+            cuda.to_device(np.arange(1))
+            self.assertEqual(len(deallocs), i + 1)
+        # one more to trigger .clear()
+        cuda.to_device(np.arange(1))
+        self.assertEqual(len(deallocs), 0)
+
+    @skip_if_external_memmgr("Deallocation specific to Numba memory management")
+    def test_max_pending_bytes(self):
+        # get deallocation manager and flush it
+        ctx = cuda.current_context()
+        deallocs = ctx.memory_manager.deallocations
+        deallocs.clear()
+        self.assertEqual(len(deallocs), 0)
+
+        mi = ctx.get_memory_info()
+
+        max_pending = 10**6  # 1MB
+        old_ratio = config.CUDA_DEALLOCS_RATIO
+        try:
+            # change to a smaller ratio
+            config.CUDA_DEALLOCS_RATIO = max_pending / mi.total
+            # due to round off error (floor is used in calculating
+            # _max_pending_bytes) it can be off by 1.
+            self.assertAlmostEqual(
+                deallocs._max_pending_bytes, max_pending, delta=1
+            )
+
+            # allocate half the max size
+            # this will not trigger deallocation
+            cuda.to_device(np.ones(max_pending // 2, dtype=np.int8))
+            self.assertEqual(len(deallocs), 1)
+
+            # allocate another remaining
+            # this will not trigger deallocation
+            cuda.to_device(
+                np.ones(
+                    deallocs._max_pending_bytes - deallocs._size, dtype=np.int8
+                )
+            )
+            self.assertEqual(len(deallocs), 2)
+
+            # another byte to trigger .clear()
+            cuda.to_device(np.ones(1, dtype=np.int8))
+            self.assertEqual(len(deallocs), 0)
+        finally:
+            # restore old ratio
+            config.CUDA_DEALLOCS_RATIO = old_ratio
+
+    @skip_if_external_memmgr("Deallocation specific to Numba memory management")
+    def test_defer_cleanup(self):
+        harr = np.arange(5)
+        darr1 = cuda.to_device(harr)
+        deallocs = cuda.current_context().memory_manager.deallocations
+        deallocs.clear()
+        self.assertEqual(len(deallocs), 0)
+        with cuda.defer_cleanup():
+            darr2 = cuda.to_device(harr)
+            del darr1
+            self.assertEqual(len(deallocs), 1)
+            del darr2
+            self.assertEqual(len(deallocs), 2)
+            deallocs.clear()
+            self.assertEqual(len(deallocs), 2)
+
+        deallocs.clear()
+        self.assertEqual(len(deallocs), 0)
+
+    @skip_if_external_memmgr("Deallocation specific to Numba memory management")
+    def test_nested_defer_cleanup(self):
+        harr = np.arange(5)
+        darr1 = cuda.to_device(harr)
+        deallocs = cuda.current_context().memory_manager.deallocations
+        deallocs.clear()
+        self.assertEqual(len(deallocs), 0)
+        with cuda.defer_cleanup():
+            with cuda.defer_cleanup():
+                darr2 = cuda.to_device(harr)
+                del darr1
+                self.assertEqual(len(deallocs), 1)
+                del darr2
+                self.assertEqual(len(deallocs), 2)
+                deallocs.clear()
+                self.assertEqual(len(deallocs), 2)
+            deallocs.clear()
+            self.assertEqual(len(deallocs), 2)
+
+        deallocs.clear()
+        self.assertEqual(len(deallocs), 0)
+
+    @skip_if_external_memmgr("Deallocation specific to Numba memory management")
+    def test_exception(self):
+        harr = np.arange(5)
+        darr1 = cuda.to_device(harr)
+        deallocs = cuda.current_context().memory_manager.deallocations
+        deallocs.clear()
+        self.assertEqual(len(deallocs), 0)
+
+        class CustomError(Exception):
+            pass
+
+        with self.assertRaises(CustomError):
+            with cuda.defer_cleanup():
+                darr2 = cuda.to_device(harr)
+                del darr2
+                self.assertEqual(len(deallocs), 1)
+                deallocs.clear()
+                self.assertEqual(len(deallocs), 1)
+                raise CustomError
+        deallocs.clear()
+        self.assertEqual(len(deallocs), 0)
+        del darr1
+        self.assertEqual(len(deallocs), 1)
+        deallocs.clear()
+        self.assertEqual(len(deallocs), 0)
+
+    @contextmanager
+    def check_ignored_exception(self, ctx):
+        with captured_stderr() as cap:
+            yield
+            ctx.deallocations.clear()
+        self.assertFalse(cap.getvalue())
+
+    def test_del_stream(self):
+        ctx = cuda.current_context()
+        stream = ctx.create_stream()
+        with self.check_ignored_exception(ctx):
+            del stream
+
+    def test_del_event(self):
+        ctx = cuda.current_context()
+        event = ctx.create_event()
+        with self.check_ignored_exception(ctx):
+            del event
+
+    def test_del_pinned_memory(self):
+        ctx = cuda.current_context()
+        mem = ctx.memhostalloc(32)
+        with self.check_ignored_exception(ctx):
+            del mem
+
+    def test_del_mapped_memory(self):
+        ctx = cuda.current_context()
+        mem = ctx.memhostalloc(32, mapped=True)
+        with self.check_ignored_exception(ctx):
+            del mem
+
+    def test_del_device_memory(self):
+        ctx = cuda.current_context()
+        mem = ctx.memalloc(32)
+        with self.check_ignored_exception(ctx):
+            del mem
+
+    def test_del_managed_memory(self):
+        ctx = cuda.current_context()
+        mem = ctx.memallocmanaged(32)
+        with self.check_ignored_exception(ctx):
+            del mem
+
+    def test_del_pinned_contextmanager(self):
+        # Check that temporarily pinned memory is unregistered immediately,
+        # such that it can be re-pinned at any time
+        class PinnedException(Exception):
+            pass
+
+        arr = np.zeros(1)
+        ctx = cuda.current_context()
+        ctx.deallocations.clear()
+        with self.check_ignored_exception(ctx):
+            with cuda.pinned(arr):
+                pass
+            with cuda.pinned(arr):
+                pass
+            # Should also work inside a `defer_cleanup` block
+            with cuda.defer_cleanup():
+                with cuda.pinned(arr):
+                    pass
+                with cuda.pinned(arr):
+                    pass
+            # Should also work when breaking out of the block due to an
+            # exception
+            try:
+                with cuda.pinned(arr):
+                    raise PinnedException
+            except PinnedException:
+                with cuda.pinned(arr):
+                    pass
+
+    def test_del_mapped_contextmanager(self):
+        # Check that temporarily mapped memory is unregistered immediately,
+        # such that it can be re-mapped at any time
+        class MappedException(Exception):
+            pass
+
+        arr = np.zeros(1)
+        ctx = cuda.current_context()
+        ctx.deallocations.clear()
+        with self.check_ignored_exception(ctx):
+            with cuda.mapped(arr):
+                pass
+            with cuda.mapped(arr):
+                pass
+            # Should also work inside a `defer_cleanup` block
+            with cuda.defer_cleanup():
+                with cuda.mapped(arr):
+                    pass
+                with cuda.mapped(arr):
+                    pass
+            # Should also work when breaking out of the block due to an
+            # exception
+            try:
+                with cuda.mapped(arr):
+                    raise MappedException
+            except MappedException:
+                with cuda.mapped(arr):
+                    pass
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,100 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import os
+import sys
+import subprocess
+import threading
+from numba import cuda
+from numba.cuda.testing import (
+    unittest,
+    CUDATestCase,
+    skip_on_cudasim,
+    skip_under_cuda_memcheck,
+)
+from numba.cuda.tests.support import captured_stdout
+
+
+class TestCudaDetect(CUDATestCase):
+    def test_cuda_detect(self):
+        # exercise the code path
+        with captured_stdout() as out:
+            cuda.detect()
+        output = out.getvalue()
+        self.assertIn("Found", output)
+        self.assertIn("CUDA devices", output)
+
+
+class TestSupportedVersion(CUDATestCase):
+    def test_is_supported_version(self):
+        # Exercise the `cuda.is_supported_version()` API.
+        #
+        # Assume for the purpose of the test that we're running on a supported
+        # toolkit version; if not, there's not much point in running the test
+        # suite.
+        self.assertTrue(cuda.is_supported_version())
+
+
+@skip_under_cuda_memcheck("Hangs cuda-memcheck")
+class TestCUDAFindLibs(CUDATestCase):
+    def run_cmd(self, cmdline, env):
+        popen = subprocess.Popen(
+            cmdline, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env
+        )
+
+        # finish in 5 minutes or kill it
+        timeout = threading.Timer(5 * 60.0, popen.kill)
+        try:
+            timeout.start()
+            out, err = popen.communicate()
+            # the process should exit with an error
+            return out.decode(), err.decode()
+        finally:
+            timeout.cancel()
+        return None, None
+
+    def run_test_in_separate_process(self, envvar, envvar_value):
+        env_copy = os.environ.copy()
+        env_copy[envvar] = str(envvar_value)
+        code = """if 1:
+            from numba import cuda
+            @cuda.jit('(int64,)')
+            def kernel(x):
+                pass
+            kernel(1,)
+            """
+        cmdline = [sys.executable, "-c", code]
+        return self.run_cmd(cmdline, env_copy)
+
+    @skip_on_cudasim("Simulator does not hit device library search code path")
+    @unittest.skipIf(not sys.platform.startswith("linux"), "linux only")
+    def test_cuda_find_lib_errors(self):
+        """
+        This tests that driver discovery attempts to load from typical system
+        locations and fails gracefully if pointed at an invalid directory.
+        """
+        # one of these is likely to exist on linux, it's also unlikely that
+        # someone has extracted the contents of libdevice into here!
+        locs = ["lib", "lib64"]
+
+        looking_for = None
+        for l in locs:
+            looking_for = os.path.join(os.path.sep, l)
+            if os.path.exists(looking_for):
+                break
+
+        # This is the testing part, the test will only run if there's a valid
+        # path in which to look
+        if looking_for is not None:
+            # We no longer support NUMBA_CUDA_DRIVER. Still run the subprocess
+            # to ensure importing and a trivial kernel launch path work, but
+            # do not set any Numba-specific driver env vars.
+            out, err = self.run_test_in_separate_process(
+                "DUMMY_UNUSED", looking_for
+            )
+            self.assertTrue(out is not None)
+            self.assertTrue(err is not None)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,198 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import ctypes
+import numpy as np
+
+from numba import cuda
+from numba.cuda.core import config
+from numba.cuda.testing import unittest, CUDATestCase, skip_on_cudasim
+from numba.cuda.tests.support import linux_only
+
+if not config.ENABLE_CUDASIM:
+
+    class DeviceOnlyEMMPlugin(cuda.HostOnlyCUDAMemoryManager):
+        """
+        Dummy EMM Plugin implementation for testing. It memorises which plugin
+        API methods have been called so that the tests can check that Numba
+        called into the plugin as expected.
+        """
+
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+
+            # For tracking our dummy allocations
+            self.allocations = {}
+            self.count = 0
+
+            # For tracking which methods have been called
+            self.initialized = False
+            self.memalloc_called = False
+            self.reset_called = False
+            self.get_memory_info_called = False
+            self.get_ipc_handle_called = False
+
+        def memalloc(self, size):
+            # We maintain a list of allocations and keep track of them, so that
+            # we can test that the finalizers of objects returned by memalloc
+            # get called.
+
+            # Numba should have initialized the memory manager when preparing
+            # the context for use, prior to any memalloc call.
+            if not self.initialized:
+                raise RuntimeError("memalloc called before initialize")
+            self.memalloc_called = True
+
+            # Create an allocation and record it
+            self.count += 1
+            alloc_count = self.count
+            self.allocations[alloc_count] = size
+
+            # The finalizer deletes the record from our internal dict of
+            # allocations.
+            finalizer_allocs = self.allocations
+
+            def finalizer():
+                del finalizer_allocs[alloc_count]
+
+            # We use an AutoFreePointer so that the finalizer will be run when
+            # the reference count drops to zero.
+            ptr = ctypes.c_void_p(alloc_count)
+            return cuda.cudadrv.driver.AutoFreePointer(
+                ptr, size, finalizer=finalizer
+            )
+
+        def initialize(self):
+            # No special initialization needed.
+            self.initialized = True
+
+        def reset(self):
+            # We remove all allocations on reset, just as a real EMM Plugin
+            # would do. Note that our finalizers in memalloc don't check
+            # whether the allocations are still alive, so running them after
+            # reset will detect any allocations that are floating around at
+            # exit time; however, the atexit finalizer for weakref will only
+            # print a traceback, not terminate the interpreter abnormally.
+            self.reset_called = True
+
+        def get_memory_info(self):
+            # Return some dummy memory information
+            self.get_memory_info_called = True
+            return cuda.MemoryInfo(free=32, total=64)
+
+        def get_ipc_handle(self, memory):
+            # The dummy IPC handle is only a string, so it is important that
+            # the tests don't try to do too much with it (e.g. open / close
+            # it).
+            self.get_ipc_handle_called = True
+            return "Dummy IPC handle for alloc %s" % memory.device_pointer_value
+
+        @property
+        def interface_version(self):
+            # The expected version for an EMM Plugin.
+            return 1
+
+    class BadVersionEMMPlugin(DeviceOnlyEMMPlugin):
+        """A plugin that claims to implement a different interface version"""
+
+        @property
+        def interface_version(self):
+            return 2
+
+
+@skip_on_cudasim("EMM Plugins not supported on CUDA simulator")
+class TestDeviceOnlyEMMPlugin(CUDATestCase):
+    """
+    Tests that the API of an EMM Plugin that implements device allocations
+    only is used correctly by Numba.
+    """
+
+    def setUp(self):
+        super().setUp()
+        # Always start afresh with a new context and memory manager
+        ctx = cuda.current_context()
+        ctx.reset()
+        self._initial_memory_manager = ctx.memory_manager
+        ctx.memory_manager = DeviceOnlyEMMPlugin(context=ctx)
+
+    def tearDown(self):
+        super().tearDown()
+        ctx = cuda.current_context()
+        ctx.reset()
+        ctx.memory_manager = self._initial_memory_manager
+
+    def test_memalloc(self):
+        mgr = cuda.current_context().memory_manager
+
+        # Allocate an array and check that memalloc was called with the correct
+        # size.
+        arr_1 = np.arange(10)
+        d_arr_1 = cuda.device_array_like(arr_1)
+        self.assertTrue(mgr.memalloc_called)
+
+        self.assertEqual(mgr.count, 1)
+        self.assertEqual(mgr.allocations[1], arr_1.nbytes)
+
+        # Allocate again, with a different size, and check that it is also
+        # correct.
+        arr_2 = np.arange(5)
+        d_arr_2 = cuda.device_array_like(arr_2)
+        self.assertEqual(mgr.count, 2)
+        self.assertEqual(mgr.allocations[2], arr_2.nbytes)
+
+        # Remove the first array, and check that our finalizer was called for
+        # the first array only.
+        del d_arr_1
+        self.assertNotIn(1, mgr.allocations)
+        self.assertIn(2, mgr.allocations)
+
+        # Remove the second array and check that its finalizer was also
+        # called.
+        del d_arr_2
+        self.assertNotIn(2, mgr.allocations)
+
+    def test_initialized_in_context(self):
+        # If we have a CUDA context, it should already have initialized its
+        # memory manager.
+        self.assertTrue(cuda.current_context().memory_manager.initialized)
+
+    def test_reset(self):
+        ctx = cuda.current_context()
+        ctx.reset()
+        self.assertTrue(ctx.memory_manager.reset_called)
+
+    def test_get_memory_info(self):
+        ctx = cuda.current_context()
+        meminfo = ctx.get_memory_info()
+        self.assertTrue(ctx.memory_manager.get_memory_info_called)
+        self.assertEqual(meminfo.free, 32)
+        self.assertEqual(meminfo.total, 64)
+
+    @linux_only
+    def test_get_ipc_handle(self):
+        # We don't attempt to close the IPC handle in this test because Numba
+        # will be expecting a real IpcHandle object to have been returned from
+        # get_ipc_handle, and it would cause problems to do so.
+        arr = np.arange(2)
+        d_arr = cuda.device_array_like(arr)
+        ipch = d_arr.get_ipc_handle()
+        ctx = cuda.current_context()
+        self.assertTrue(ctx.memory_manager.get_ipc_handle_called)
+        self.assertIn("Dummy IPC handle for alloc 1", ipch._ipc_handle)
+
+
+@skip_on_cudasim("EMM Plugins not supported on CUDA simulator")
+class TestBadEMMPluginVersion(CUDATestCase):
+    """
+    Ensure that Numba rejects EMM Plugins with incompatible version
+    numbers.
+    """
+
+    def test_bad_plugin_version(self):
+        with self.assertRaises(RuntimeError) as raises:
+            cuda.set_memory_manager(BadVersionEMMPlugin)
+        self.assertIn("version 1 required", str(raises.exception))
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,53 @@
+# 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
+from cuda.core.experimental import Device
+from numba.cuda.testing import skip_on_cudasim
+
+
+class TestCudaEvent(CUDATestCase):
+    def test_event_elapsed(self):
+        N = 32
+        dary = cuda.device_array(N, dtype=np.double)
+        evtstart = cuda.event()
+        evtend = cuda.event()
+
+        evtstart.record()
+        cuda.to_device(np.arange(N, dtype=np.double), to=dary)
+        evtend.record()
+        evtend.wait()
+        evtend.synchronize()
+        # Exercise the code path
+        evtstart.elapsed_time(evtend)
+
+    def test_event_elapsed_stream(self):
+        stream = cuda.stream()
+        self.event_elapsed_inner(stream)
+
+    @skip_on_cudasim("Testing cuda.core events requires driver")
+    def test_event_elapsed_cuda_core_stream(self):
+        dev = Device()
+        dev.set_current()
+        stream = dev.create_stream()
+        self.event_elapsed_inner(stream)
+
+    def event_elapsed_inner(self, stream):
+        N = 32
+        dary = cuda.device_array(N, dtype=np.double)
+        evtstart = cuda.event()
+        evtend = cuda.event()
+
+        evtstart.record(stream=stream)
+        cuda.to_device(np.arange(N, dtype=np.double), to=dary, stream=stream)
+        evtend.record(stream=stream)
+        evtend.wait(stream=stream)
+        evtend.synchronize()
+        # Exercise the code path
+        evtstart.elapsed_time(evtend)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,72 @@
+# 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 import driver
+from numba import cuda
+from numba.cuda.testing import unittest, CUDATestCase
+
+
+class TestHostAlloc(CUDATestCase):
+    def tearDown(self):
+        cuda.current_context().reset()
+
+    def test_host_alloc_driver(self):
+        n = 32
+        mem = cuda.current_context().memhostalloc(n, mapped=True)
+
+        dtype = np.dtype(np.uint8)
+        ary = np.ndarray(shape=n // dtype.itemsize, dtype=dtype, buffer=mem)
+
+        magic = 0xAB
+        driver.device_memset(mem, magic, n)
+
+        self.assertTrue(np.all(ary == magic))
+
+        ary.fill(n)
+
+        recv = np.empty_like(ary)
+
+        driver.device_to_host(recv, mem, ary.size)
+
+        self.assertTrue(np.all(ary == recv))
+        self.assertTrue(np.all(recv == n))
+
+    def test_host_alloc_pinned(self):
+        ary = cuda.pinned_array(10, dtype=np.uint32)
+        ary.fill(123)
+        self.assertTrue(all(ary == 123))
+        devary = cuda.to_device(ary)
+        driver.device_memset(devary, 0, driver.device_memory_size(devary))
+        self.assertTrue(all(ary == 123))
+        devary.copy_to_host(ary)
+        self.assertTrue(all(ary == 0))
+
+    def test_host_alloc_mapped(self):
+        ary = cuda.mapped_array(10, dtype=np.uint32)
+        ary.fill(123)
+        self.assertTrue(all(ary == 123))
+        driver.device_memset(ary, 0, driver.device_memory_size(ary))
+        self.assertTrue(all(ary == 0))
+        self.assertTrue(sum(ary != 0) == 0)
+
+    def test_host_operators(self):
+        for ary in [
+            cuda.mapped_array(10, dtype=np.uint32),
+            cuda.pinned_array(10, dtype=np.uint32),
+        ]:
+            ary[:] = range(10)
+            self.assertTrue(sum(ary + 1) == 55)
+            self.assertTrue(sum((ary + 1) * 2 - 1) == 100)
+            self.assertTrue(sum(ary < 5) == 5)
+            self.assertTrue(sum(ary <= 5) == 6)
+            self.assertTrue(sum(ary > 6) == 3)
+            self.assertTrue(sum(ary >= 6) == 4)
+            self.assertTrue(sum(ary**2) == 285)
+            self.assertTrue(sum(ary // 2) == 20)
+            self.assertTrue(sum(ary / 2.0) == 22.5)
+            self.assertTrue(sum(ary % 2) == 5)
+
+
+if __name__ == "__main__":
+    unittest.main()