numba-cuda 0.22.0__cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

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

@@ -0,0 +1,185 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import pytest
+
+from llvmlite import ir
+from numba.cuda.cudadrv import nvrtc, nvvm, runtime
+from numba.cuda.testing import unittest
+from numba.cuda.cudadrv.nvvm import LibDevice, NvvmError, NVVM
+from numba.cuda.testing import skip_on_cudasim
+
+
+@skip_on_cudasim("NVVM Driver unsupported in the simulator")
+class TestNvvmDriver(unittest.TestCase):
+    def get_nvvmir(self):
+        versions = NVVM().get_ir_version()
+        data_layout = NVVM().data_layout
+        return nvvmir_generic.format(data_layout=data_layout, v=versions)
+
+    def test_nvvm_compile_simple(self):
+        nvvmir = self.get_nvvmir()
+        ptx = nvvm.compile_ir(nvvmir).decode("utf8")
+        self.assertTrue("simple" in ptx)
+        self.assertTrue("ave" in ptx)
+
+    def test_nvvm_compile_nullary_option(self):
+        # Tests compilation with an option that doesn't take an argument
+        # ("-gen-lto") - all other NVVM options are of the form
+        # "-<name>=<value>"
+
+        nvvmir = self.get_nvvmir()
+        arch = "compute_%d%d" % nvrtc.get_lowest_supported_cc()
+        ltoir = nvvm.compile_ir(nvvmir, opt=3, gen_lto=None, arch=arch)
+
+        # Verify we correctly passed the option by checking if we got LTOIR
+        # from NVVM (by looking for the expected magic number for LTOIR)
+        self.assertEqual(ltoir[:4], b"\xed\x43\x4e\x7f")
+
+    def test_nvvm_bad_option(self):
+        # Ensure that unsupported / non-existent options are reported as such
+        # to the user / caller
+        msg = "-made-up-option=2 is an unsupported option"
+        with self.assertRaisesRegex(NvvmError, msg):
+            nvvm.compile_ir("", made_up_option=2)
+
+    def test_nvvm_from_llvm(self):
+        m = ir.Module("test_nvvm_from_llvm")
+        m.triple = "nvptx64-nvidia-cuda"
+        nvvm.add_ir_version(m)
+        fty = ir.FunctionType(ir.VoidType(), [ir.IntType(32)])
+        kernel = ir.Function(m, fty, name="mycudakernel")
+        bldr = ir.IRBuilder(kernel.append_basic_block("entry"))
+        bldr.ret_void()
+        nvvm.set_cuda_kernel(kernel)
+
+        m.data_layout = NVVM().data_layout
+        ptx = nvvm.compile_ir(str(m)).decode("utf8")
+        self.assertTrue("mycudakernel" in ptx)
+        self.assertTrue(".address_size 64" in ptx)
+
+    def test_used_list(self):
+        # Construct a module
+        m = ir.Module("test_used_list")
+        m.triple = "nvptx64-nvidia-cuda"
+        m.data_layout = NVVM().data_layout
+        nvvm.add_ir_version(m)
+
+        # Add a function and mark it as a kernel
+        fty = ir.FunctionType(ir.VoidType(), [ir.IntType(32)])
+        kernel = ir.Function(m, fty, name="mycudakernel")
+        bldr = ir.IRBuilder(kernel.append_basic_block("entry"))
+        bldr.ret_void()
+        nvvm.set_cuda_kernel(kernel)
+
+        # Verify that the used list was correctly constructed
+        used_lines = [
+            line for line in str(m).splitlines() if "llvm.used" in line
+        ]
+        msg = 'Expected exactly one @"llvm.used" array'
+        self.assertEqual(len(used_lines), 1, msg)
+
+        used_line = used_lines[0]
+        # Kernel should be referenced in the used list
+        self.assertIn("mycudakernel", used_line)
+        # Check linkage of the used list
+        self.assertIn("appending global", used_line)
+        # Ensure used list is in the metadata section
+        self.assertIn('section "llvm.metadata"', used_line)
+
+    def test_nvvm_ir_verify_fail(self):
+        if runtime.get_version() >= (12, 5):
+            self.skipTest("Bad triple doesn't fail verify on CUDA >= 12.5")
+        m = ir.Module("test_bad_ir")
+        m.triple = "unknown-unknown-unknown"
+        m.data_layout = NVVM().data_layout
+        nvvm.add_ir_version(m)
+        with self.assertRaisesRegex(NvvmError, "Invalid target triple"):
+            nvvm.compile_ir(str(m))
+
+    def _test_nvvm_support(self, arch):
+        compute_xx = "compute_{0}{1}".format(*arch)
+        nvvmir = self.get_nvvmir()
+        ptx = nvvm.compile_ir(
+            nvvmir, arch=compute_xx, ftz=1, prec_sqrt=0, prec_div=0
+        ).decode("utf8")
+        self.assertIn(".target sm_{0}{1}".format(*arch), ptx)
+        self.assertIn("simple", ptx)
+        self.assertIn("ave", ptx)
+
+    def test_nvvm_support(self):
+        """Test supported CC by NVVM"""
+        for arch in nvrtc.get_supported_ccs():
+            self._test_nvvm_support(arch=arch)
+
+    def test_nvvm_warning(self):
+        m = ir.Module("test_nvvm_warning")
+        m.triple = "nvptx64-nvidia-cuda"
+        m.data_layout = NVVM().data_layout
+        nvvm.add_ir_version(m)
+
+        fty = ir.FunctionType(ir.VoidType(), [])
+        kernel = ir.Function(m, fty, name="inlinekernel")
+        builder = ir.IRBuilder(kernel.append_basic_block("entry"))
+        builder.ret_void()
+        nvvm.set_cuda_kernel(kernel)
+
+        # Add the noinline attribute to trigger NVVM to generate a warning
+        kernel.attributes.add("noinline")
+
+        code = str(m)
+        with pytest.warns(Warning, match="overriding noinline attribute"):
+            nvvm.compile_ir(code)
+
+
+@skip_on_cudasim("NVVM Driver unsupported in the simulator")
+class TestLibDevice(unittest.TestCase):
+    def test_libdevice_load(self):
+        # Test that constructing LibDevice gives a bitcode file
+        libdevice = LibDevice()
+        self.assertEqual(libdevice.bc[:4], b"BC\xc0\xde")
+
+
+nvvmir_generic = """\
+target triple="nvptx64-nvidia-cuda"
+target datalayout = "{data_layout}"
+
+define i32 @ave(i32 %a, i32 %b) {{
+entry:
+%add = add nsw i32 %a, %b
+%div = sdiv i32 %add, 2
+ret i32 %div
+}}
+
+define void @simple(i32* %data) {{
+entry:
+%0 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
+%1 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x()
+%mul = mul i32 %0, %1
+%2 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
+%add = add i32 %mul, %2
+%call = call i32 @ave(i32 %add, i32 %add)
+%idxprom = sext i32 %add to i64
+%arrayidx = getelementptr inbounds i32, i32* %data, i64 %idxprom
+store i32 %call, i32* %arrayidx, align 4
+ret void
+}}
+
+declare i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() nounwind readnone
+
+declare i32 @llvm.nvvm.read.ptx.sreg.ntid.x() nounwind readnone
+
+declare i32 @llvm.nvvm.read.ptx.sreg.tid.x() nounwind readnone
+
+!nvvmir.version = !{{!1}}
+!1 = !{{i32 {v[0]}, i32 {v[1]}, i32 {v[2]}, i32 {v[3]}}}
+
+!nvvm.annotations = !{{!2}}
+!2 = !{{void (i32*)* @simple, !"kernel", i32 1}}
+
+@"llvm.used" = appending global [1 x i8*] [i8* bitcast (void (i32*)* @simple to i8*)], section "llvm.metadata"
+"""  # noqa: E501
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numpy as np
+import platform
+
+from numba import cuda
+from numba.cuda.testing import unittest, CUDATestCase
+
+
+class TestPinned(CUDATestCase):
+    def _run_copies(self, A):
+        A0 = np.copy(A)
+
+        stream = cuda.stream()
+        ptr = cuda.to_device(A, copy=False, stream=stream)
+        ptr.copy_to_device(A, stream=stream)
+        ptr.copy_to_host(A, stream=stream)
+        stream.synchronize()
+
+        self.assertTrue(np.allclose(A, A0))
+
+    def test_pinned(self):
+        machine = platform.machine()
+        if machine.startswith("arm") or machine.startswith("aarch64"):
+            count = 262144  # 2MB
+        else:
+            count = 2097152  # 16MB
+        A = np.arange(count)
+        with cuda.pinned(A):
+            self._run_copies(A)
+
+    def test_unpinned(self):
+        A = np.arange(2 * 1024 * 1024)  # 16 MB
+        self._run_copies(A)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import unittest
+from numba.cuda.testing import CUDATestCase
+from numba import cuda
+from numba.cuda.testing import skip_on_cudasim
+
+
+@skip_on_cudasim("CUDA Profiler unsupported in the simulator")
+class TestProfiler(CUDATestCase):
+    def test_profiling(self):
+        with cuda.profiling():
+            a = cuda.device_array(10)
+            del a
+
+        with cuda.profiling():
+            a = cuda.device_array(100)
+            del a
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,38 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import threading
+from numba import cuda
+from numba.cuda.cudadrv.driver import driver
+from numba.cuda.testing import unittest, CUDATestCase
+from queue import Queue
+
+
+class TestResetDevice(CUDATestCase):
+    def test_reset_device(self):
+        def newthread(exception_queue):
+            try:
+                devices = range(driver.get_device_count())
+                for _ in range(2):
+                    for d in devices:
+                        cuda.select_device(d)
+                        cuda.close()
+            except Exception as e:
+                exception_queue.put(e)
+
+        # Do test on a separate thread so that we don't affect
+        # the current context in the main thread.
+
+        exception_queue = Queue()
+        t = threading.Thread(target=newthread, args=(exception_queue,))
+        t.start()
+        t.join()
+
+        exceptions = []
+        while not exception_queue.empty():
+            exceptions.append(exception_queue.get())
+        self.assertEqual(exceptions, [])
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,48 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import concurrent.futures
+import multiprocessing
+import os
+from numba.cuda.testing import unittest
+
+
+def set_visible_devices_and_check():
+    from numba import cuda
+    import os
+
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    return len(cuda.gpus.lst)
+
+
+class TestVisibleDevices(unittest.TestCase):
+    def test_visible_devices_set_after_import(self):
+        # See Issue #6149. This test checks that we can set
+        # CUDA_VISIBLE_DEVICES after importing Numba and have the value
+        # reflected in the available list of GPUs. Prior to the fix for this
+        # issue, Numba made a call to runtime.get_version() on import that
+        # initialized the driver and froze the list of available devices before
+        # CUDA_VISIBLE_DEVICES could be set by the user.
+
+        # Avoid importing cuda at the top level so that
+        # set_visible_devices_and_check gets to import it first in its process
+        from numba import cuda
+
+        if len(cuda.gpus.lst) in (0, 1):
+            self.skipTest("This test requires multiple GPUs")
+
+        if os.environ.get("CUDA_VISIBLE_DEVICES"):
+            msg = "Cannot test when CUDA_VISIBLE_DEVICES already set"
+            self.skipTest(msg)
+
+        with concurrent.futures.ProcessPoolExecutor(
+            mp_context=multiprocessing.get_context("spawn")
+        ) as exe:
+            future = exe.submit(set_visible_devices_and_check)
+
+        visible_gpu_count = future.result()
+        assert visible_gpu_count == 1
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+#
+# Test does not work on some cards.
+#
+import threading
+from queue import Queue
+
+import numpy as np
+from numba import cuda
+from numba.cuda.testing import unittest, CUDATestCase
+
+
+def newthread(exception_queue):
+    try:
+        cuda.select_device(0)
+        stream = cuda.stream()
+        A = np.arange(100)
+        dA = cuda.to_device(A, stream=stream)
+        stream.synchronize()
+        del dA
+        del stream
+        cuda.synchronize()
+    except Exception as e:
+        exception_queue.put(e)
+
+
+class TestSelectDevice(CUDATestCase):
+    def test_select_device(self):
+        exception_queue = Queue()
+        for i in range(10):
+            t = threading.Thread(target=newthread, args=(exception_queue,))
+            t.start()
+            t.join()
+
+        exceptions = []
+        while not exception_queue.empty():
+            exceptions.append(exception_queue.get())
+        self.assertEqual(exceptions, [])
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,127 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import asyncio
+import functools
+import threading
+import numpy as np
+from numba import cuda
+from numba.cuda.testing import unittest, CUDATestCase, skip_on_cudasim
+
+
+def with_asyncio_loop(f):
+    @functools.wraps(f)
+    def runner(*args, **kwds):
+        loop = asyncio.new_event_loop()
+        loop.set_debug(True)
+        try:
+            return loop.run_until_complete(f(*args, **kwds))
+        finally:
+            loop.close()
+
+    return runner
+
+
+@unittest.skip("Disabled temporarily due to Issue #317")
+@skip_on_cudasim("CUDA Driver API unsupported in the simulator")
+class TestCudaStream(CUDATestCase):
+    def test_add_callback(self):
+        def callback(stream, status, event):
+            event.set()
+
+        stream = cuda.stream()
+        callback_event = threading.Event()
+        stream.add_callback(callback, callback_event)
+        self.assertTrue(callback_event.wait(1.0))
+
+    def test_add_callback_with_default_arg(self):
+        callback_event = threading.Event()
+
+        def callback(stream, status, arg):
+            self.assertIsNone(arg)
+            callback_event.set()
+
+        stream = cuda.stream()
+        stream.add_callback(callback)
+        self.assertTrue(callback_event.wait(1.0))
+
+    @with_asyncio_loop
+    async def test_async_done(self):
+        stream = cuda.stream()
+        await stream.async_done()
+
+    @with_asyncio_loop
+    async def test_parallel_tasks(self):
+        async def async_cuda_fn(value_in: float) -> float:
+            stream = cuda.stream()
+            h_src, h_dst = cuda.pinned_array(8), cuda.pinned_array(8)
+            h_src[:] = value_in
+            d_ary = cuda.to_device(h_src, stream=stream)
+            d_ary.copy_to_host(h_dst, stream=stream)
+            done_result = await stream.async_done()
+            self.assertEqual(done_result, stream)
+            return h_dst.mean()
+
+        values_in = [1, 2, 3, 4]
+        tasks = [asyncio.create_task(async_cuda_fn(v)) for v in values_in]
+        values_out = await asyncio.gather(*tasks)
+        self.assertTrue(np.allclose(values_in, values_out))
+
+    @with_asyncio_loop
+    async def test_multiple_async_done(self):
+        stream = cuda.stream()
+        done_aws = [stream.async_done() for _ in range(4)]
+        done = await asyncio.gather(*done_aws)
+        for d in done:
+            self.assertEqual(d, stream)
+
+    @with_asyncio_loop
+    async def test_multiple_async_done_multiple_streams(self):
+        streams = [cuda.stream() for _ in range(4)]
+        done_aws = [stream.async_done() for stream in streams]
+        done = await asyncio.gather(*done_aws)
+
+        # Ensure we got the four original streams in done
+        self.assertSetEqual(set(done), set(streams))
+
+    @with_asyncio_loop
+    async def test_cancelled_future(self):
+        stream = cuda.stream()
+        done1, done2 = stream.async_done(), stream.async_done()
+        done1.cancel()
+        await done2
+        self.assertTrue(done1.cancelled())
+        self.assertTrue(done2.done())
+
+
+@skip_on_cudasim("CUDA Driver API unsupported in the simulator")
+class TestFailingStream(CUDATestCase):
+    # This test can only be run in isolation because it corrupts the CUDA
+    # context, which cannot be recovered from within the same process. It is
+    # left here so that it can be run manually for debugging / testing purposes
+    # - or may be re-enabled if in future there is infrastructure added for
+    # running tests in a separate process (a subprocess cannot be used because
+    # CUDA will have been initialized before the fork, so it cannot be used in
+    # the child process).
+    @unittest.skip
+    @with_asyncio_loop
+    async def test_failed_stream(self):
+        ctx = cuda.current_context()
+        module = ctx.create_module_ptx("""
+            .version 6.5
+            .target sm_30
+            .address_size 64
+            .visible .entry failing_kernel() { trap; }
+        """)
+        failing_kernel = module.get_function("failing_kernel")
+
+        stream = cuda.stream()
+        failing_kernel.configure((1,), (1,), stream=stream).__call__()
+        done = stream.async_done()
+        with self.assertRaises(Exception):
+            await done
+        self.assertIsNotNone(done.exception())
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.tests import load_testsuite
+import os
+
+
+def load_tests(loader, tests, pattern):
+    return load_testsuite(loader, os.path.dirname(__file__))