numba-cuda 0.22.1__cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

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

@@ -0,0 +1,180 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numpy as np
+import ctypes
+from numba.cuda.cudadrv.devicearray import (
+    DeviceRecord,
+    from_record_like,
+    auto_device,
+)
+from numba.cuda.testing import unittest, CUDATestCase
+from numba.cuda.testing import skip_on_cudasim
+from numba.cuda.np import numpy_support
+from numba import cuda
+
+N_CHARS = 5
+
+recordtype = np.dtype(
+    [
+        ("a", np.float64),
+        ("b", np.int32),
+        ("c", np.complex64),
+        ("d", (np.str_, N_CHARS)),
+    ],
+    align=True,
+)
+
+recordwitharray = np.dtype([("g", np.int32), ("h", np.float32, 2)], align=True)
+
+recwithmat = np.dtype([("i", np.int32), ("j", np.float32, (3, 3))])
+
+recwithrecwithmat = np.dtype([("x", np.int32), ("y", recwithmat)])
+
+
+@skip_on_cudasim("Device Record API unsupported in the simulator")
+class TestCudaDeviceRecord(CUDATestCase):
+    """
+    Tests the DeviceRecord class with np.void host types.
+    """
+
+    def setUp(self):
+        super().setUp()
+        self._create_data(np.zeros)
+
+    def _create_data(self, array_ctor):
+        self.dtype = np.dtype([("a", np.int32), ("b", np.float32)], align=True)
+        self.hostz = array_ctor(1, self.dtype)[0]
+        self.hostnz = array_ctor(1, self.dtype)[0]
+        self.hostnz["a"] = 10
+        self.hostnz["b"] = 11.0
+
+    def _check_device_record(self, reference, rec):
+        self.assertEqual(rec.shape, tuple())
+        self.assertEqual(rec.strides, tuple())
+        self.assertEqual(rec.dtype, reference.dtype)
+        self.assertEqual(rec.alloc_size, reference.dtype.itemsize)
+        self.assertIsNotNone(rec.gpu_data)
+        self.assertNotEqual(rec.device_ctypes_pointer, ctypes.c_void_p(0))
+
+        numba_type = numpy_support.from_dtype(reference.dtype)
+        self.assertEqual(rec._numba_type_, numba_type)
+
+    def test_device_record_interface(self):
+        hostrec = self.hostz.copy()
+        devrec = DeviceRecord(self.dtype)
+        self._check_device_record(hostrec, devrec)
+
+    def test_device_record_copy(self):
+        hostrec = self.hostz.copy()
+        devrec = DeviceRecord(self.dtype)
+        devrec.copy_to_device(hostrec)
+
+        # Copy back and check values are all zeros
+        hostrec2 = self.hostnz.copy()
+        devrec.copy_to_host(hostrec2)
+        np.testing.assert_equal(self.hostz, hostrec2)
+
+        # Copy non-zero values to GPU and back and check values
+        hostrec3 = self.hostnz.copy()
+        devrec.copy_to_device(hostrec3)
+
+        hostrec4 = self.hostz.copy()
+        devrec.copy_to_host(hostrec4)
+        np.testing.assert_equal(hostrec4, self.hostnz)
+
+    def test_from_record_like(self):
+        # Create record from host record
+        hostrec = self.hostz.copy()
+        devrec = from_record_like(hostrec)
+        self._check_device_record(hostrec, devrec)
+
+        # Create record from device record and check for distinct data
+        devrec2 = from_record_like(devrec)
+        self._check_device_record(devrec, devrec2)
+        self.assertNotEqual(devrec.gpu_data, devrec2.gpu_data)
+
+    def test_auto_device(self):
+        # Create record from host record
+        hostrec = self.hostnz.copy()
+        devrec, new_gpu_obj = auto_device(hostrec)
+        self._check_device_record(hostrec, devrec)
+        self.assertTrue(new_gpu_obj)
+
+        # Copy data back and check it is equal to auto_device arg
+        hostrec2 = self.hostz.copy()
+        devrec.copy_to_host(hostrec2)
+        np.testing.assert_equal(hostrec2, hostrec)
+
+
+class TestCudaDeviceRecordWithRecord(TestCudaDeviceRecord):
+    """
+    Tests the DeviceRecord class with np.record host types
+    """
+
+    def setUp(self):
+        CUDATestCase.setUp(self)
+        self._create_data(np.recarray)
+
+
+@skip_on_cudasim("Structured array attr access not supported in simulator")
+class TestRecordDtypeWithStructArrays(CUDATestCase):
+    """
+    Test operation of device arrays on structured arrays.
+    """
+
+    def _createSampleArrays(self):
+        self.sample1d = cuda.device_array(3, dtype=recordtype)
+        self.samplerec1darr = cuda.device_array(1, dtype=recordwitharray)[0]
+        self.samplerecmat = cuda.device_array(1, dtype=recwithmat)[0]
+
+    def setUp(self):
+        super().setUp()
+        self._createSampleArrays()
+
+        ary = self.sample1d
+        for i in range(ary.size):
+            x = i + 1
+            ary[i]["a"] = x / 2
+            ary[i]["b"] = x
+            ary[i]["c"] = x * 1j
+            ary[i]["d"] = str(x) * N_CHARS
+
+    def test_structured_array1(self):
+        ary = self.sample1d
+        for i in range(self.sample1d.size):
+            x = i + 1
+            self.assertEqual(ary[i]["a"], x / 2)
+            self.assertEqual(ary[i]["b"], x)
+            self.assertEqual(ary[i]["c"], x * 1j)
+            self.assertEqual(ary[i]["d"], str(x) * N_CHARS)
+
+    def test_structured_array2(self):
+        ary = self.samplerec1darr
+        ary["g"] = 2
+        ary["h"][0] = 3.0
+        ary["h"][1] = 4.0
+        self.assertEqual(ary["g"], 2)
+        self.assertEqual(ary["h"][0], 3.0)
+        self.assertEqual(ary["h"][1], 4.0)
+
+    def test_structured_array3(self):
+        ary = self.samplerecmat
+        mat = np.array(
+            [[5.0, 10.0, 15.0], [20.0, 25.0, 30.0], [35.0, 40.0, 45.0]],
+            dtype=np.float32,
+        ).reshape(3, 3)
+        ary["j"][:] = mat
+        np.testing.assert_equal(ary["j"], mat)
+
+    def test_structured_array4(self):
+        arr = np.zeros(1, dtype=recwithrecwithmat)
+        d_arr = cuda.to_device(arr)
+        d_arr[0]["y"]["i"] = 1
+        self.assertEqual(d_arr[0]["y"]["i"], 1)
+        d_arr[0]["y"]["j"][0, 0] = 2.0
+        self.assertEqual(d_arr[0]["y"]["j"][0, 0], 2.0)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,360 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from ctypes import c_int, sizeof
+
+from numba.cuda.cudadrv.driver import (
+    host_to_device,
+    device_to_host,
+    driver,
+    launch_kernel,
+)
+
+from numba import cuda
+from numba.cuda.cudadrv import devices, driver as _driver
+from numba.cuda.testing import unittest, CUDATestCase
+from numba.cuda.testing import skip_on_cudasim
+import contextlib
+
+from cuda.core.experimental import Device
+
+ptx1 = """
+    .version 1.4
+    .target sm_10, map_f64_to_f32
+
+    .entry _Z10helloworldPi (
+    .param .u64 __cudaparm__Z10helloworldPi_A)
+    {
+    .reg .u32 %r<3>;
+    .reg .u64 %rd<6>;
+    .loc	14	4	0
+$LDWbegin__Z10helloworldPi:
+    .loc	14	6	0
+    cvt.s32.u16 	%r1, %tid.x;
+    ld.param.u64 	%rd1, [__cudaparm__Z10helloworldPi_A];
+    cvt.u64.u16 	%rd2, %tid.x;
+    mul.lo.u64 	%rd3, %rd2, 4;
+    add.u64 	%rd4, %rd1, %rd3;
+    st.global.s32 	[%rd4+0], %r1;
+    .loc	14	7	0
+    exit;
+$LDWend__Z10helloworldPi:
+    } // _Z10helloworldPi
+"""
+
+ptx2 = """
+.version 3.0
+.target sm_20
+.address_size 64
+
+    .file	1 "/tmp/tmpxft_000012c7_00000000-9_testcuda.cpp3.i"
+    .file	2 "testcuda.cu"
+
+.entry _Z10helloworldPi(
+    .param .u64 _Z10helloworldPi_param_0
+)
+{
+    .reg .s32 	%r<3>;
+    .reg .s64 	%rl<5>;
+
+
+    ld.param.u64 	%rl1, [_Z10helloworldPi_param_0];
+    cvta.to.global.u64 	%rl2, %rl1;
+    .loc 2 6 1
+    mov.u32 	%r1, %tid.x;
+    mul.wide.u32 	%rl3, %r1, 4;
+    add.s64 	%rl4, %rl2, %rl3;
+    st.global.u32 	[%rl4], %r1;
+    .loc 2 7 2
+    ret;
+}
+"""
+
+
+@skip_on_cudasim("CUDA Driver API unsupported in the simulator")
+class TestCudaDriver(CUDATestCase):
+    def setUp(self):
+        super().setUp()
+        self.assertTrue(len(devices.gpus) > 0)
+        self.context = devices.get_context()
+        device = self.context.device
+        ccmajor, _ = device.compute_capability
+        if ccmajor >= 2:
+            self.ptx = ptx2
+        else:
+            self.ptx = ptx1
+
+    def tearDown(self):
+        super().tearDown()
+        del self.context
+
+    def test_cuda_driver_basic(self):
+        module = self.context.create_module_ptx(self.ptx)
+        function = module.get_function("_Z10helloworldPi")
+
+        array = (c_int * 100)()
+
+        memory = self.context.memalloc(sizeof(array))
+        host_to_device(memory, array, sizeof(array))
+
+        ptr = memory.device_ctypes_pointer
+        stream = 0
+
+        stream = _driver.binding.CUstream(stream)
+
+        launch_kernel(
+            function.handle,  # Kernel
+            1,
+            1,
+            1,  # gx, gy, gz
+            100,
+            1,
+            1,  # bx, by, bz
+            0,  # dynamic shared mem
+            stream,  # stream
+            [ptr],
+        )  # arguments
+
+        device_to_host(array, memory, sizeof(array))
+        for i, v in enumerate(array):
+            self.assertEqual(i, v)
+
+        module.unload()
+
+    def test_cuda_driver_stream_operations(self):
+        module = self.context.create_module_ptx(self.ptx)
+        function = module.get_function("_Z10helloworldPi")
+
+        array = (c_int * 100)()
+
+        stream = self.context.create_stream()
+
+        with stream.auto_synchronize():
+            memory = self.context.memalloc(sizeof(array))
+            host_to_device(memory, array, sizeof(array), stream=stream)
+
+            ptr = memory.device_ctypes_pointer
+
+            stream_handle = stream.handle
+            stream_handle = stream_handle.value
+
+            launch_kernel(
+                function.handle,  # Kernel
+                1,
+                1,
+                1,  # gx, gy, gz
+                100,
+                1,
+                1,  # bx, by, bz
+                0,  # dynamic shared mem
+                stream_handle,  # stream
+                [ptr],
+            )  # arguments
+
+        device_to_host(array, memory, sizeof(array), stream=stream)
+
+        for i, v in enumerate(array):
+            self.assertEqual(i, v)
+
+    def test_cuda_core_stream_operations(self):
+        module = self.context.create_module_ptx(self.ptx)
+        function = module.get_function("_Z10helloworldPi")
+        array = (c_int * 100)()
+        dev = Device()
+        dev.set_current()
+        stream = dev.create_stream()
+
+        @contextlib.contextmanager
+        def auto_synchronize(stream):
+            try:
+                yield stream
+            finally:
+                stream.sync()
+
+        with auto_synchronize(stream):
+            memory = self.context.memalloc(sizeof(array))
+            host_to_device(memory, array, sizeof(array), stream=stream)
+
+            ptr = memory.device_ctypes_pointer
+
+            launch_kernel(
+                function.handle,  # Kernel
+                1,
+                1,
+                1,  # gx, gy, gz
+                100,
+                1,
+                1,  # bx, by, bz
+                0,  # dynamic shared mem
+                stream.handle,  # stream
+                [ptr],
+            )
+
+            device_to_host(array, memory, sizeof(array), stream=stream)
+        for i, v in enumerate(array):
+            self.assertEqual(i, v)
+
+    def test_cuda_core_stream_launch_user_facing(self):
+        @cuda.jit
+        def kernel(a):
+            idx = cuda.grid(1)
+            if idx < len(a):
+                a[idx] = idx
+
+        dev = Device()
+        dev.set_current()
+        stream = dev.create_stream()
+
+        ary = cuda.to_device([0] * 100, stream=stream)
+        stream.sync()
+
+        kernel[1, 100, stream](ary)
+        stream.sync()
+
+        result = ary.copy_to_host(stream=stream)
+        for i, v in enumerate(result):
+            self.assertEqual(i, v)
+
+    def test_cuda_driver_default_stream(self):
+        # Test properties of the default stream
+        ds = self.context.get_default_stream()
+        self.assertIn("Default CUDA stream", repr(ds))
+        self.assertEqual(0, int(ds))
+        # bool(stream) is the check that is done in memcpy to decide if async
+        # version should be used. So the default (0) stream should be true-ish
+        # even though 0 is usually false-ish in Python.
+        self.assertTrue(ds)
+        self.assertFalse(ds.external)
+
+    def test_cuda_driver_legacy_default_stream(self):
+        # Test properties of the legacy default stream
+        ds = self.context.get_legacy_default_stream()
+        self.assertIn("Legacy default CUDA stream", repr(ds))
+        self.assertEqual(1, int(ds))
+        self.assertTrue(ds)
+        self.assertFalse(ds.external)
+
+    def test_cuda_driver_per_thread_default_stream(self):
+        # Test properties of the per-thread default stream
+        ds = self.context.get_per_thread_default_stream()
+        self.assertIn("Per-thread default CUDA stream", repr(ds))
+        self.assertEqual(2, int(ds))
+        self.assertTrue(ds)
+        self.assertFalse(ds.external)
+
+    def test_cuda_driver_stream(self):
+        # Test properties of non-default streams
+        s = self.context.create_stream()
+        self.assertIn("CUDA stream", repr(s))
+        self.assertNotIn("Default", repr(s))
+        self.assertNotIn("External", repr(s))
+        self.assertNotEqual(0, int(s))
+        self.assertTrue(s)
+        self.assertFalse(s.external)
+
+    def test_cuda_driver_external_stream(self):
+        # Test properties of a stream created from an external stream object.
+        # We use the driver API directly to create a stream, to emulate an
+        # external library creating a stream
+        handle = driver.cuStreamCreate(0)
+        ptr = int(handle)
+        s = self.context.create_external_stream(ptr)
+
+        self.assertIn("External CUDA stream", repr(s))
+        # Ensure neither "Default" nor "default"
+        self.assertNotIn("efault", repr(s))
+        self.assertEqual(ptr, int(s))
+        self.assertTrue(s)
+        self.assertTrue(s.external)
+
+    def test_cuda_driver_occupancy(self):
+        module = self.context.create_module_ptx(self.ptx)
+        function = module.get_function("_Z10helloworldPi")
+
+        value = self.context.get_active_blocks_per_multiprocessor(
+            function, 128, 128
+        )
+        self.assertTrue(value > 0)
+
+        def b2d(bs):
+            return bs
+
+        grid, block = self.context.get_max_potential_block_size(
+            function, b2d, 128, 128
+        )
+        self.assertTrue(grid > 0)
+        self.assertTrue(block > 0)
+
+    def test_cuda_cache_config(self):
+        from numba import types
+        import numpy as np
+
+        sig = (types.float32[::1], types.float32[::1])
+
+        @cuda.jit(sig)
+        def add_one(r, x):
+            i = cuda.grid(1)
+            if i < len(r):
+                r[i] = x[i] + 1
+
+        kernel = add_one.overloads[sig]
+        cufunc = kernel._codelibrary.get_cufunc()
+
+        configs_to_test = [
+            ("prefer_shared", dict(prefer_shared=True)),
+            ("prefer_cache", dict(prefer_cache=True)),
+            ("prefer_equal", dict(prefer_equal=True)),
+            ("default", dict()),
+        ]
+
+        for name, kwargs in configs_to_test:
+            with self.subTest(config=name):
+                try:
+                    cufunc.cache_config(**kwargs)
+                except Exception as e:
+                    self.fail(f"cache_config({name}) failed: {e}")
+
+        x = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32)
+        r = np.zeros_like(x)
+
+        d_x = cuda.to_device(x)
+        d_r = cuda.to_device(r)
+
+        cufunc.cache_config(prefer_shared=True)
+        add_one[1, 5](d_r, d_x)
+
+        result = d_r.copy_to_host()
+        expected = x + 1
+
+        np.testing.assert_array_almost_equal(
+            result,
+            expected,
+            err_msg="Kernel produced incorrect results after cache_config",
+        )
+
+
+class TestDevice(CUDATestCase):
+    def test_device_get_uuid(self):
+        # A device UUID looks like:
+        #
+        #     GPU-e6489c45-5b68-3b03-bab7-0e7c8e809643
+        #
+        # To test, we construct an RE that matches this form and verify that
+        # the returned UUID matches.
+        #
+        # Device UUIDs may not conform to parts of the UUID specification (RFC
+        # 4122) pertaining to versions and variants, so we do not extract and
+        # validate the values of these bits.
+
+        h = "[0-9a-f]{%d}"
+        h4 = h % 4
+        h8 = h % 8
+        h12 = h % 12
+        uuid_format = f"^GPU-{h8}-{h4}-{h4}-{h4}-{h12}$"
+
+        dev = devices.get_context().device
+        self.assertRegex(dev.uuid, uuid_format)
+
+
+if __name__ == "__main__":
+    unittest.main()