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

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

@@ -0,0 +1,438 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numpy as np
+
+from numba.cuda import vectorize, guvectorize
+from numba import cuda
+from numba.cuda.testing import unittest, CUDATestCase, ForeignArray
+from numba.cuda.testing import skip_on_cudasim, skip_if_external_memmgr
+from numba.cuda.tests.support import linux_only, override_config
+from unittest.mock import call, patch
+
+
+@skip_on_cudasim("CUDA Array Interface is not supported in the simulator")
+class TestCudaArrayInterface(CUDATestCase):
+    def assertPointersEqual(self, a, b):
+        self.assertEqual(
+            a.device_ctypes_pointer.value, b.device_ctypes_pointer.value
+        )
+
+    def test_as_cuda_array(self):
+        h_arr = np.arange(10)
+        self.assertFalse(cuda.is_cuda_array(h_arr))
+        d_arr = cuda.to_device(h_arr)
+        self.assertTrue(cuda.is_cuda_array(d_arr))
+        my_arr = ForeignArray(d_arr)
+        self.assertTrue(cuda.is_cuda_array(my_arr))
+        wrapped = cuda.as_cuda_array(my_arr)
+        self.assertTrue(cuda.is_cuda_array(wrapped))
+        # Their values must equal the original array
+        np.testing.assert_array_equal(wrapped.copy_to_host(), h_arr)
+        np.testing.assert_array_equal(d_arr.copy_to_host(), h_arr)
+        # d_arr and wrapped must be the same buffer
+        self.assertPointersEqual(wrapped, d_arr)
+
+    def get_stream_value(self, stream):
+        return stream.handle.value
+
+    @skip_if_external_memmgr("Ownership not relevant with external memmgr")
+    def test_ownership(self):
+        # Get the deallocation queue
+        ctx = cuda.current_context()
+        deallocs = ctx.memory_manager.deallocations
+        # Flush all deallocations
+        deallocs.clear()
+        self.assertEqual(len(deallocs), 0)
+        # Make new device array
+        d_arr = cuda.to_device(np.arange(100))
+        # Convert it
+        cvted = cuda.as_cuda_array(d_arr)
+        # Drop reference to the original object such that
+        # only `cvted` has a reference to it.
+        del d_arr
+        # There shouldn't be any new deallocations
+        self.assertEqual(len(deallocs), 0)
+        # Try to access the memory and verify its content
+        np.testing.assert_equal(cvted.copy_to_host(), np.arange(100))
+        # Drop last reference to the memory
+        del cvted
+        self.assertEqual(len(deallocs), 1)
+        # Flush
+        deallocs.clear()
+
+    def test_kernel_arg(self):
+        h_arr = np.arange(10)
+        d_arr = cuda.to_device(h_arr)
+        my_arr = ForeignArray(d_arr)
+        wrapped = cuda.as_cuda_array(my_arr)
+
+        @cuda.jit
+        def mutate(arr, val):
+            i = cuda.grid(1)
+            if i >= len(arr):
+                return
+            arr[i] += val
+
+        val = 7
+        mutate.forall(wrapped.size)(wrapped, val)
+
+        np.testing.assert_array_equal(wrapped.copy_to_host(), h_arr + val)
+        np.testing.assert_array_equal(d_arr.copy_to_host(), h_arr + val)
+
+    def test_ufunc_arg(self):
+        @vectorize(["f8(f8, f8)"], target="cuda")
+        def vadd(a, b):
+            return a + b
+
+        # Case 1: use custom array as argument
+        h_arr = np.random.random(10)
+        arr = ForeignArray(cuda.to_device(h_arr))
+        val = 6
+        out = vadd(arr, val)
+        np.testing.assert_array_equal(out.copy_to_host(), h_arr + val)
+
+        # Case 2: use custom array as return
+        out = ForeignArray(cuda.device_array(h_arr.shape))
+        returned = vadd(h_arr, val, out=out)
+        np.testing.assert_array_equal(returned.copy_to_host(), h_arr + val)
+
+    def test_gufunc_arg(self):
+        @guvectorize(["(f8, f8, f8[:])"], "(),()->()", target="cuda")
+        def vadd(inp, val, out):
+            out[0] = inp + val
+
+        # Case 1: use custom array as argument
+        h_arr = np.random.random(10)
+        arr = ForeignArray(cuda.to_device(h_arr))
+        val = np.float64(7)
+        out = vadd(arr, val)
+        np.testing.assert_array_equal(out.copy_to_host(), h_arr + val)
+
+        # Case 2: use custom array as return
+        out = ForeignArray(cuda.device_array(h_arr.shape))
+        returned = vadd(h_arr, val, out=out)
+        np.testing.assert_array_equal(returned.copy_to_host(), h_arr + val)
+        self.assertPointersEqual(returned, out._arr)
+
+    def test_array_views(self):
+        """Views created via array interface support:
+        - Strided slices
+        - Strided slices
+        """
+        h_arr = np.random.random(10)
+        c_arr = cuda.to_device(h_arr)
+
+        arr = cuda.as_cuda_array(c_arr)
+
+        # __getitem__ interface accesses expected data
+
+        # Direct views
+        np.testing.assert_array_equal(arr.copy_to_host(), h_arr)
+        np.testing.assert_array_equal(arr[:].copy_to_host(), h_arr)
+
+        # Slicing
+        np.testing.assert_array_equal(arr[:5].copy_to_host(), h_arr[:5])
+
+        # Strided view
+        np.testing.assert_array_equal(arr[::2].copy_to_host(), h_arr[::2])
+
+        # View of strided array
+        arr_strided = cuda.as_cuda_array(c_arr[::2])
+        np.testing.assert_array_equal(arr_strided.copy_to_host(), h_arr[::2])
+
+        # A strided-view-of-array and view-of-strided-array have the same
+        # shape, strides, itemsize, and alloc_size
+        self.assertEqual(arr[::2].shape, arr_strided.shape)
+        self.assertEqual(arr[::2].strides, arr_strided.strides)
+        self.assertEqual(arr[::2].dtype.itemsize, arr_strided.dtype.itemsize)
+        self.assertEqual(arr[::2].alloc_size, arr_strided.alloc_size)
+        self.assertEqual(
+            arr[::2].nbytes, arr_strided.size * arr_strided.dtype.itemsize
+        )
+
+        # __setitem__ interface propagates into external array
+
+        # Writes to a slice
+        arr[:5] = np.pi
+        np.testing.assert_array_equal(
+            c_arr.copy_to_host(), np.concatenate((np.full(5, np.pi), h_arr[5:]))
+        )
+
+        # Writes to a slice from a view
+        arr[:5] = arr[5:]
+        np.testing.assert_array_equal(
+            c_arr.copy_to_host(), np.concatenate((h_arr[5:], h_arr[5:]))
+        )
+
+        # Writes through a view
+        arr[:] = cuda.to_device(h_arr)
+        np.testing.assert_array_equal(c_arr.copy_to_host(), h_arr)
+
+        # Writes to a strided slice
+        arr[::2] = np.pi
+        np.testing.assert_array_equal(
+            c_arr.copy_to_host()[::2],
+            np.full(5, np.pi),
+        )
+        np.testing.assert_array_equal(c_arr.copy_to_host()[1::2], h_arr[1::2])
+
+    def test_negative_strided_issue(self):
+        # issue #3705
+        h_arr = np.random.random(10)
+        c_arr = cuda.to_device(h_arr)
+
+        def base_offset(orig, sliced):
+            return sliced["data"][0] - orig["data"][0]
+
+        h_ai = h_arr.__array_interface__
+        c_ai = c_arr.__cuda_array_interface__
+
+        h_ai_sliced = h_arr[::-1].__array_interface__
+        c_ai_sliced = c_arr[::-1].__cuda_array_interface__
+
+        # Check data offset is correct
+        self.assertEqual(
+            base_offset(h_ai, h_ai_sliced),
+            base_offset(c_ai, c_ai_sliced),
+        )
+        # Check shape and strides are correct
+        self.assertEqual(h_ai_sliced["shape"], c_ai_sliced["shape"])
+        self.assertEqual(h_ai_sliced["strides"], c_ai_sliced["strides"])
+
+    def test_negative_strided_copy_to_host(self):
+        # issue #3705
+        h_arr = np.random.random(10)
+        c_arr = cuda.to_device(h_arr)
+        sliced = c_arr[::-1]
+        with self.assertRaises(NotImplementedError) as raises:
+            sliced.copy_to_host()
+        expected_msg = "D->H copy not implemented for negative strides"
+        self.assertIn(expected_msg, str(raises.exception))
+
+    def test_masked_array(self):
+        h_arr = np.random.random(10)
+        h_mask = np.random.randint(2, size=10, dtype="bool")
+        c_arr = cuda.to_device(h_arr)
+        c_mask = cuda.to_device(h_mask)
+
+        # Manually create a masked CUDA Array Interface dictionary
+        masked_cuda_array_interface = c_arr.__cuda_array_interface__.copy()
+        masked_cuda_array_interface["mask"] = c_mask
+
+        with self.assertRaises(NotImplementedError) as raises:
+            cuda.from_cuda_array_interface(masked_cuda_array_interface)
+        expected_msg = "Masked arrays are not supported"
+        self.assertIn(expected_msg, str(raises.exception))
+
+    def test_zero_size_array(self):
+        # for #4175
+        c_arr = cuda.device_array(0)
+        self.assertEqual(c_arr.__cuda_array_interface__["data"][0], 0)
+
+        @cuda.jit
+        def add_one(arr):
+            x = cuda.grid(1)
+            N = arr.shape[0]
+            if x < N:
+                arr[x] += 1
+
+        d_arr = ForeignArray(c_arr)
+        add_one[1, 10](d_arr)  # this should pass
+
+    def test_strides(self):
+        # for #4175
+        # First, test C-contiguous array
+        c_arr = cuda.device_array((2, 3, 4))
+        self.assertEqual(c_arr.__cuda_array_interface__["strides"], None)
+
+        # Second, test non C-contiguous array
+        c_arr = c_arr[:, 1, :]
+        self.assertNotEqual(c_arr.__cuda_array_interface__["strides"], None)
+
+    def test_consuming_strides(self):
+        hostarray = np.arange(10).reshape(2, 5)
+        devarray = cuda.to_device(hostarray)
+        face = devarray.__cuda_array_interface__
+        self.assertIsNone(face["strides"])
+        got = cuda.from_cuda_array_interface(face).copy_to_host()
+        np.testing.assert_array_equal(got, hostarray)
+        self.assertTrue(got.flags["C_CONTIGUOUS"])
+        # Try non-NULL strides
+        face["strides"] = hostarray.strides
+        self.assertIsNotNone(face["strides"])
+        got = cuda.from_cuda_array_interface(face).copy_to_host()
+        np.testing.assert_array_equal(got, hostarray)
+        self.assertTrue(got.flags["C_CONTIGUOUS"])
+
+    def test_produce_no_stream(self):
+        c_arr = cuda.device_array(10)
+        self.assertIsNone(c_arr.__cuda_array_interface__["stream"])
+
+        mapped_arr = cuda.mapped_array(10)
+        self.assertIsNone(mapped_arr.__cuda_array_interface__["stream"])
+
+    @linux_only
+    def test_produce_managed_no_stream(self):
+        managed_arr = cuda.managed_array(10)
+        self.assertIsNone(managed_arr.__cuda_array_interface__["stream"])
+
+    def test_produce_stream(self):
+        s = cuda.stream()
+        c_arr = cuda.device_array(10, stream=s)
+        cai_stream = c_arr.__cuda_array_interface__["stream"]
+        stream_value = self.get_stream_value(s)
+        self.assertEqual(stream_value, cai_stream)
+
+        s = cuda.stream()
+        mapped_arr = cuda.mapped_array(10, stream=s)
+        cai_stream = mapped_arr.__cuda_array_interface__["stream"]
+        stream_value = self.get_stream_value(s)
+        self.assertEqual(stream_value, cai_stream)
+
+    @linux_only
+    def test_produce_managed_stream(self):
+        s = cuda.stream()
+        managed_arr = cuda.managed_array(10, stream=s)
+        cai_stream = managed_arr.__cuda_array_interface__["stream"]
+        stream_value = self.get_stream_value(s)
+        self.assertEqual(stream_value, cai_stream)
+
+    def test_consume_no_stream(self):
+        # Create a foreign array with no stream
+        f_arr = ForeignArray(cuda.device_array(10))
+
+        # Ensure that the imported array has no default stream
+        c_arr = cuda.as_cuda_array(f_arr)
+        self.assertEqual(c_arr.stream, 0)
+
+    def test_consume_stream(self):
+        # Create a foreign array with a stream
+        s = cuda.stream()
+        f_arr = ForeignArray(cuda.device_array(10, stream=s))
+
+        # Ensure that an imported array has the stream as its default stream
+        c_arr = cuda.as_cuda_array(f_arr)
+        self.assertTrue(c_arr.stream.external)
+        stream_value = self.get_stream_value(s)
+        imported_stream_value = self.get_stream_value(c_arr.stream)
+        self.assertEqual(stream_value, imported_stream_value)
+
+    def test_consume_no_sync(self):
+        # Create a foreign array with no stream
+        f_arr = ForeignArray(cuda.device_array(10))
+
+        with patch.object(
+            cuda.cudadrv.driver.Stream, "synchronize", return_value=None
+        ) as mock_sync:
+            cuda.as_cuda_array(f_arr)
+
+        # Ensure the synchronize method of a stream was not called
+        mock_sync.assert_not_called()
+
+    def test_consume_sync(self):
+        # Create a foreign array with a stream
+        s = cuda.stream()
+        f_arr = ForeignArray(cuda.device_array(10, stream=s))
+
+        with patch.object(
+            cuda.cudadrv.driver.Stream, "synchronize", return_value=None
+        ) as mock_sync:
+            cuda.as_cuda_array(f_arr)
+
+        # Ensure the synchronize method of a stream was called
+        mock_sync.assert_called_once_with()
+
+    def test_consume_sync_disabled(self):
+        # Create a foreign array with a stream
+        s = cuda.stream()
+        f_arr = ForeignArray(cuda.device_array(10, stream=s))
+
+        # Set sync to false before testing. The test suite should generally be
+        # run with sync enabled, but stash the old value just in case it is
+        # not.
+        with override_config("CUDA_ARRAY_INTERFACE_SYNC", False):
+            with patch.object(
+                cuda.cudadrv.driver.Stream, "synchronize", return_value=None
+            ) as mock_sync:
+                cuda.as_cuda_array(f_arr)
+
+            # Ensure the synchronize method of a stream was not called
+            mock_sync.assert_not_called()
+
+    def test_launch_no_sync(self):
+        # Create a foreign array with no stream
+        f_arr = ForeignArray(cuda.device_array(10))
+
+        @cuda.jit
+        def f(x):
+            pass
+
+        with patch.object(
+            cuda.cudadrv.driver.Stream, "synchronize", return_value=None
+        ) as mock_sync:
+            f[1, 1](f_arr)
+
+        # Ensure the synchronize method of a stream was not called
+        mock_sync.assert_not_called()
+
+    def test_launch_sync(self):
+        # Create a foreign array with a stream
+        s = cuda.stream()
+        f_arr = ForeignArray(cuda.device_array(10, stream=s))
+
+        @cuda.jit
+        def f(x):
+            pass
+
+        with patch.object(
+            cuda.cudadrv.driver.Stream, "synchronize", return_value=None
+        ) as mock_sync:
+            f[1, 1](f_arr)
+
+        # Ensure the synchronize method of a stream was called
+        mock_sync.assert_called_once_with()
+
+    def test_launch_sync_two_streams(self):
+        # Create two foreign arrays with streams
+        s1 = cuda.stream()
+        s2 = cuda.stream()
+        f_arr1 = ForeignArray(cuda.device_array(10, stream=s1))
+        f_arr2 = ForeignArray(cuda.device_array(10, stream=s2))
+
+        @cuda.jit
+        def f(x, y):
+            pass
+
+        with patch.object(
+            cuda.cudadrv.driver.Stream, "synchronize", return_value=None
+        ) as mock_sync:
+            f[1, 1](f_arr1, f_arr2)
+
+        # Ensure that synchronize was called twice
+        mock_sync.assert_has_calls([call(), call()])
+
+    def test_launch_sync_disabled(self):
+        # Create two foreign arrays with streams
+        s1 = cuda.stream()
+        s2 = cuda.stream()
+        f_arr1 = ForeignArray(cuda.device_array(10, stream=s1))
+        f_arr2 = ForeignArray(cuda.device_array(10, stream=s2))
+
+        with override_config("CUDA_ARRAY_INTERFACE_SYNC", False):
+
+            @cuda.jit
+            def f(x, y):
+                pass
+
+            with patch.object(
+                cuda.cudadrv.driver.Stream, "synchronize", return_value=None
+            ) as mock_sync:
+                f[1, 1](f_arr1, f_arr2)
+
+            # Ensure that synchronize was not called
+            mock_sync.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,94 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba import cuda
+import numpy as np
+from numba.cuda.testing import CUDATestCase
+from numba.cuda.tests.support import override_config
+import unittest
+
+
+class TestCudaJitNoTypes(CUDATestCase):
+    """
+    Tests the jit decorator with no types provided.
+    """
+
+    def test_device_array(self):
+        @cuda.jit
+        def foo(x, y):
+            i = cuda.grid(1)
+            y[i] = x[i]
+
+        x = np.arange(10)
+        y = np.empty_like(x)
+
+        dx = cuda.to_device(x)
+        dy = cuda.to_device(y)
+
+        foo[10, 1](dx, dy)
+
+        dy.copy_to_host(y)
+
+        self.assertTrue(np.all(x == y))
+
+    def test_device_jit(self):
+        @cuda.jit(device=True)
+        def mapper(args):
+            a, b, c = args
+            return a + b + c
+
+        @cuda.jit(device=True)
+        def reducer(a, b):
+            return a + b
+
+        @cuda.jit
+        def driver(A, B):
+            i = cuda.grid(1)
+            if i < B.size:
+                args = A[i], A[i] + B[i], B[i]
+                B[i] = reducer(mapper(args), 1)
+
+        A = np.arange(100, dtype=np.float32)
+        B = np.arange(100, dtype=np.float32)
+
+        Acopy = A.copy()
+        Bcopy = B.copy()
+
+        driver[1, 100](A, B)
+
+        np.testing.assert_allclose(Acopy + Acopy + Bcopy + Bcopy + 1, B)
+
+    def test_device_jit_2(self):
+        @cuda.jit(device=True)
+        def inner(arg):
+            return arg + 1
+
+        @cuda.jit
+        def outer(argin, argout):
+            argout[0] = inner(argin[0]) + inner(2)
+
+        a = np.zeros(1)
+        b = np.zeros(1)
+
+        stream = cuda.stream()
+        d_a = cuda.to_device(a, stream)
+        d_b = cuda.to_device(b, stream)
+
+        outer[1, 1, stream](d_a, d_b)
+
+        d_b.copy_to_host(b, stream)
+
+        self.assertEqual(b[0], (a[0] + 1) + (2 + 1))
+
+    def test_jit_debug_simulator(self):
+        # Ensure that the jit decorator accepts the debug kwarg when the
+        # simulator is in use - see Issue #6615.
+        with override_config("ENABLE_CUDASIM", 1):
+
+            @cuda.jit(debug=True, opt=False)
+            def f(x):
+                pass
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,101 @@
+# 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, vectorize, guvectorize
+from numba.cuda.np.numpy_support import from_dtype
+from numba.cuda.testing import CUDATestCase, skip_on_cudasim
+import unittest
+
+
+class TestCudaDateTime(CUDATestCase):
+    def test_basic_datetime_kernel(self):
+        @cuda.jit
+        def foo(start, end, delta):
+            for i in range(cuda.grid(1), delta.size, cuda.gridsize(1)):
+                delta[i] = end[i] - start[i]
+
+        arr1 = np.arange("2005-02", "2006-02", dtype="datetime64[D]")
+        arr2 = arr1 + np.random.randint(0, 10000, arr1.size)
+        delta = np.zeros_like(arr1, dtype="timedelta64[D]")
+
+        foo[1, 32](arr1, arr2, delta)
+
+        self.assertPreciseEqual(delta, arr2 - arr1)
+
+    def test_scalar_datetime_kernel(self):
+        @cuda.jit
+        def foo(dates, target, delta, matches, outdelta):
+            for i in range(cuda.grid(1), matches.size, cuda.gridsize(1)):
+                matches[i] = dates[i] == target
+                outdelta[i] = dates[i] - delta
+
+        arr1 = np.arange("2005-02", "2006-02", dtype="datetime64[D]")
+        target = arr1[5]  # datetime
+        delta = arr1[6] - arr1[5]  # timedelta
+        matches = np.zeros_like(arr1, dtype=np.bool_)
+        outdelta = np.zeros_like(arr1, dtype="datetime64[D]")
+
+        foo[1, 32](arr1, target, delta, matches, outdelta)
+        where = matches.nonzero()
+
+        self.assertEqual(list(where), [5])
+        self.assertPreciseEqual(outdelta, arr1 - delta)
+
+    @skip_on_cudasim("ufunc API unsupported in the simulator")
+    def test_ufunc(self):
+        datetime_t = from_dtype(np.dtype("datetime64[D]"))
+
+        @vectorize([(datetime_t, datetime_t)], target="cuda")
+        def timediff(start, end):
+            return end - start
+
+        arr1 = np.arange("2005-02", "2006-02", dtype="datetime64[D]")
+        arr2 = arr1 + np.random.randint(0, 10000, arr1.size)
+
+        delta = timediff(arr1, arr2)
+
+        self.assertPreciseEqual(delta, arr2 - arr1)
+
+    @skip_on_cudasim("ufunc API unsupported in the simulator")
+    def test_gufunc(self):
+        datetime_t = from_dtype(np.dtype("datetime64[D]"))
+        timedelta_t = from_dtype(np.dtype("timedelta64[D]"))
+
+        @guvectorize(
+            [(datetime_t, datetime_t, timedelta_t[:])],
+            "(),()->()",
+            target="cuda",
+        )
+        def timediff(start, end, out):
+            out[0] = end - start
+
+        arr1 = np.arange("2005-02", "2006-02", dtype="datetime64[D]")
+        arr2 = arr1 + np.random.randint(0, 10000, arr1.size)
+
+        delta = timediff(arr1, arr2)
+
+        self.assertPreciseEqual(delta, arr2 - arr1)
+
+    @skip_on_cudasim("no .copy_to_host() in the simulator")
+    def test_datetime_view_as_int64(self):
+        arr = np.arange("2005-02", "2006-02", dtype="datetime64[D]")
+        darr = cuda.to_device(arr)
+        viewed = darr.view(np.int64)
+        self.assertPreciseEqual(arr.view(np.int64), viewed.copy_to_host())
+        self.assertEqual(viewed.gpu_data, darr.gpu_data)
+
+    @skip_on_cudasim("no .copy_to_host() in the simulator")
+    def test_timedelta_view_as_int64(self):
+        arr = np.arange("2005-02", "2006-02", dtype="datetime64[D]")
+        arr = arr - (arr - 1)
+        self.assertEqual(arr.dtype, np.dtype("timedelta64[D]"))
+        darr = cuda.to_device(arr)
+        viewed = darr.view(np.int64)
+        self.assertPreciseEqual(arr.view(np.int64), viewed.copy_to_host())
+        self.assertEqual(viewed.gpu_data, darr.gpu_data)
+
+
+if __name__ == "__main__":
+    unittest.main()