numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

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

@@ -0,0 +1,161 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numbers
+import weakref
+
+from numba import cuda
+from numba.cuda.testing import unittest, CUDATestCase, skip_on_cudasim
+from numba.cuda.cudadrv import driver
+
+
+class TestContextStack(CUDATestCase):
+    def setUp(self):
+        super().setUp()
+        # Reset before testing
+        cuda.current_context().reset()
+
+    def test_gpus_len(self):
+        self.assertGreater(len(cuda.gpus), 0)
+
+    def test_gpus_iter(self):
+        gpulist = list(cuda.gpus)
+        self.assertGreater(len(gpulist), 0)
+
+    def test_gpus_cudevice_indexing(self):
+        """Test that CUdevice objects can be used to index into cuda.gpus"""
+        # When using the CUDA Python bindings, the device ids are CUdevice
+        # objects, otherwise they are integers. We test that the device id is
+        # usable as an index into cuda.gpus.
+        device_ids = [device.id for device in cuda.list_devices()]
+        for device_id in device_ids:
+            with cuda.gpus[device_id]:
+                self.assertEqual(cuda.gpus.current.id, device_id)
+
+
+class TestContextAPI(CUDATestCase):
+    def tearDown(self):
+        super().tearDown()
+        cuda.current_context().reset()
+
+    def test_context_memory(self):
+        try:
+            mem = cuda.current_context().get_memory_info()
+        except NotImplementedError:
+            self.skipTest("EMM Plugin does not implement get_memory_info()")
+
+        self.assertIsInstance(mem.free, numbers.Number)
+        self.assertEqual(mem.free, mem[0])
+
+        self.assertIsInstance(mem.total, numbers.Number)
+        self.assertEqual(mem.total, mem[1])
+
+        self.assertLessEqual(mem.free, mem.total)
+
+    @unittest.skipIf(len(cuda.gpus) < 2, "need more than 1 gpus")
+    @skip_on_cudasim("CUDA HW required")
+    def test_forbidden_context_switch(self):
+        # Cannot switch context inside a `cuda.require_context`
+        @cuda.require_context
+        def switch_gpu():
+            with cuda.gpus[1]:
+                pass
+
+        with cuda.gpus[0]:
+            with self.assertRaises(RuntimeError) as raises:
+                switch_gpu()
+
+            self.assertIn("Cannot switch CUDA-context.", str(raises.exception))
+
+    @unittest.skipIf(len(cuda.gpus) < 2, "need more than 1 gpus")
+    def test_accepted_context_switch(self):
+        def switch_gpu():
+            with cuda.gpus[1]:
+                return cuda.current_context().device.id
+
+        with cuda.gpus[0]:
+            devid = switch_gpu()
+        self.assertEqual(int(devid), 1)
+
+
+@skip_on_cudasim("CUDA HW required")
+class Test3rdPartyContext(CUDATestCase):
+    def tearDown(self):
+        super().tearDown()
+        cuda.current_context().reset()
+
+    def test_attached_primary(self, extra_work=lambda: None):
+        # Emulate primary context creation by 3rd party
+        the_driver = driver.driver
+        dev = driver.binding.CUdevice(0)
+        binding_hctx = the_driver.cuDevicePrimaryCtxRetain(dev)
+        hctx = driver.drvapi.cu_context(int(binding_hctx))
+        try:
+            ctx = driver.Context(weakref.proxy(self), hctx)
+            ctx.push()
+            # Check that the context from numba matches the created primary
+            # context.
+            my_ctx = cuda.current_context()
+            self.assertEqual(my_ctx.handle.value, ctx.handle.value)
+
+            extra_work()
+        finally:
+            ctx.pop()
+            the_driver.cuDevicePrimaryCtxRelease(dev)
+
+    def test_attached_non_primary(self):
+        # Emulate non-primary context creation by 3rd party
+        the_driver = driver.driver
+        flags = 0
+        dev = driver.binding.CUdevice(0)
+
+        result, version = driver.binding.cuDriverGetVersion()
+        self.assertEqual(
+            result,
+            driver.binding.CUresult.CUDA_SUCCESS,
+            "Error getting CUDA driver version",
+        )
+
+        # CUDA 13's cuCtxCreate has an optional parameter prepended. The
+        # version of cuCtxCreate in use depends on the cuda.bindings major
+        # version rather than the installed driver version on the machine
+        # we're running on.
+        from cuda import bindings
+
+        bindings_version = int(bindings.__version__.split(".")[0])
+        if bindings_version in (11, 12):
+            args = (flags, dev)
+        else:
+            args = (None, flags, dev)
+
+        hctx = the_driver.cuCtxCreate(*args)
+        try:
+            cuda.current_context()
+        except RuntimeError as e:
+            # Expecting an error about non-primary CUDA context
+            self.assertIn(
+                "Numba cannot operate on non-primary CUDA context ", str(e)
+            )
+        else:
+            self.fail("No RuntimeError raised")
+        finally:
+            the_driver.cuCtxDestroy(hctx)
+
+    def test_cudajit_in_attached_primary_context(self):
+        def do():
+            from numba import cuda
+
+            @cuda.jit
+            def foo(a):
+                for i in range(a.size):
+                    a[i] = i
+
+            a = cuda.device_array(10)
+            foo[1, 1](a)
+            self.assertEqual(list(a.copy_to_host()), list(range(10)))
+
+        self.test_attached_primary(do)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,397 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from itertools import product
+
+import numpy as np
+
+from numba import cuda
+from numba.cuda.testing import unittest, CUDATestCase, skip_on_cudasim
+from unittest.mock import patch
+
+
+class CudaArrayIndexing(CUDATestCase):
+    def test_index_1d(self):
+        arr = np.arange(10)
+        darr = cuda.to_device(arr)
+        (x,) = arr.shape
+        for i in range(-x, x):
+            self.assertEqual(arr[i], darr[i])
+        with self.assertRaises(IndexError):
+            darr[-x - 1]
+        with self.assertRaises(IndexError):
+            darr[x]
+
+    def test_index_2d(self):
+        arr = np.arange(3 * 4).reshape(3, 4)
+        darr = cuda.to_device(arr)
+        x, y = arr.shape
+        for i in range(-x, x):
+            for j in range(-y, y):
+                self.assertEqual(arr[i, j], darr[i, j])
+        with self.assertRaises(IndexError):
+            darr[-x - 1, 0]
+        with self.assertRaises(IndexError):
+            darr[x, 0]
+        with self.assertRaises(IndexError):
+            darr[0, -y - 1]
+        with self.assertRaises(IndexError):
+            darr[0, y]
+
+    def test_index_3d(self):
+        arr = np.arange(3 * 4 * 5).reshape(3, 4, 5)
+        darr = cuda.to_device(arr)
+        x, y, z = arr.shape
+        for i in range(-x, x):
+            for j in range(-y, y):
+                for k in range(-z, z):
+                    self.assertEqual(arr[i, j, k], darr[i, j, k])
+        with self.assertRaises(IndexError):
+            darr[-x - 1, 0, 0]
+        with self.assertRaises(IndexError):
+            darr[x, 0, 0]
+        with self.assertRaises(IndexError):
+            darr[0, -y - 1, 0]
+        with self.assertRaises(IndexError):
+            darr[0, y, 0]
+        with self.assertRaises(IndexError):
+            darr[0, 0, -z - 1]
+        with self.assertRaises(IndexError):
+            darr[0, 0, z]
+
+
+class CudaArrayStridedSlice(CUDATestCase):
+    def test_strided_index_1d(self):
+        arr = np.arange(10)
+        darr = cuda.to_device(arr)
+        for i in range(arr.size):
+            np.testing.assert_equal(arr[i::2], darr[i::2].copy_to_host())
+
+    def test_strided_index_2d(self):
+        arr = np.arange(6 * 7).reshape(6, 7)
+        darr = cuda.to_device(arr)
+
+        for i in range(arr.shape[0]):
+            for j in range(arr.shape[1]):
+                np.testing.assert_equal(
+                    arr[i::2, j::2], darr[i::2, j::2].copy_to_host()
+                )
+
+    def test_strided_index_3d(self):
+        arr = np.arange(6 * 7 * 8).reshape(6, 7, 8)
+        darr = cuda.to_device(arr)
+
+        for i in range(arr.shape[0]):
+            for j in range(arr.shape[1]):
+                for k in range(arr.shape[2]):
+                    np.testing.assert_equal(
+                        arr[i::2, j::2, k::2],
+                        darr[i::2, j::2, k::2].copy_to_host(),
+                    )
+
+
+class CudaArraySlicing(CUDATestCase):
+    def test_prefix_1d(self):
+        arr = np.arange(5)
+        darr = cuda.to_device(arr)
+        for i in range(arr.size):
+            expect = arr[i:]
+            got = darr[i:].copy_to_host()
+            self.assertTrue(np.all(expect == got))
+
+    def test_prefix_2d(self):
+        arr = np.arange(3**2).reshape(3, 3)
+        darr = cuda.to_device(arr)
+        for i in range(arr.shape[0]):
+            for j in range(arr.shape[1]):
+                expect = arr[i:, j:]
+                sliced = darr[i:, j:]
+                self.assertEqual(expect.shape, sliced.shape)
+                self.assertEqual(expect.strides, sliced.strides)
+                got = sliced.copy_to_host()
+                self.assertTrue(np.all(expect == got))
+
+    def test_select_3d_first_two_dim(self):
+        arr = np.arange(3 * 4 * 5).reshape(3, 4, 5)
+        darr = cuda.to_device(arr)
+        # Select first dimension
+        for i in range(arr.shape[0]):
+            expect = arr[i]
+            sliced = darr[i]
+            self.assertEqual(expect.shape, sliced.shape)
+            self.assertEqual(expect.strides, sliced.strides)
+            got = sliced.copy_to_host()
+            self.assertTrue(np.all(expect == got))
+        # Select second dimension
+        for i in range(arr.shape[0]):
+            for j in range(arr.shape[1]):
+                expect = arr[i, j]
+                sliced = darr[i, j]
+                self.assertEqual(expect.shape, sliced.shape)
+                self.assertEqual(expect.strides, sliced.strides)
+                got = sliced.copy_to_host()
+                self.assertTrue(np.all(expect == got))
+
+    def test_select_f(self):
+        a = np.arange(5 * 6 * 7).reshape(5, 6, 7, order="F")
+        da = cuda.to_device(a)
+
+        for i in range(a.shape[0]):
+            for j in range(a.shape[1]):
+                self.assertTrue(
+                    np.array_equal(da[i, j, :].copy_to_host(), a[i, j, :])
+                )
+            for j in range(a.shape[2]):
+                self.assertTrue(
+                    np.array_equal(da[i, :, j].copy_to_host(), a[i, :, j])
+                )
+        for i in range(a.shape[1]):
+            for j in range(a.shape[2]):
+                self.assertTrue(
+                    np.array_equal(da[:, i, j].copy_to_host(), a[:, i, j])
+                )
+
+    def test_select_c(self):
+        a = np.arange(5 * 6 * 7).reshape(5, 6, 7, order="C")
+        da = cuda.to_device(a)
+
+        for i in range(a.shape[0]):
+            for j in range(a.shape[1]):
+                self.assertTrue(
+                    np.array_equal(da[i, j, :].copy_to_host(), a[i, j, :])
+                )
+            for j in range(a.shape[2]):
+                self.assertTrue(
+                    np.array_equal(da[i, :, j].copy_to_host(), a[i, :, j])
+                )
+        for i in range(a.shape[1]):
+            for j in range(a.shape[2]):
+                self.assertTrue(
+                    np.array_equal(da[:, i, j].copy_to_host(), a[:, i, j])
+                )
+
+    def test_prefix_select(self):
+        arr = np.arange(5 * 7).reshape(5, 7, order="F")
+
+        darr = cuda.to_device(arr)
+        self.assertTrue(np.all(darr[:1, 1].copy_to_host() == arr[:1, 1]))
+
+    def test_negative_slicing_1d(self):
+        arr = np.arange(10)
+        darr = cuda.to_device(arr)
+        for i, j in product(range(-10, 10), repeat=2):
+            np.testing.assert_array_equal(arr[i:j], darr[i:j].copy_to_host())
+
+    def test_negative_slicing_2d(self):
+        arr = np.arange(12).reshape(3, 4)
+        darr = cuda.to_device(arr)
+        for x, y, w, s in product(range(-4, 4), repeat=4):
+            np.testing.assert_array_equal(
+                arr[x:y, w:s], darr[x:y, w:s].copy_to_host()
+            )
+
+    def test_empty_slice_1d(self):
+        arr = np.arange(5)
+        darr = cuda.to_device(arr)
+        for i in range(darr.shape[0]):
+            np.testing.assert_array_equal(darr[i:i].copy_to_host(), arr[i:i])
+        # empty slice of empty slice
+        self.assertFalse(darr[:0][:0].copy_to_host().size > 0)
+        # out-of-bound slice just produces empty slices
+        np.testing.assert_array_equal(darr[:0][:1].copy_to_host(), arr[:0][:1])
+        np.testing.assert_array_equal(
+            darr[:0][-1:].copy_to_host(), arr[:0][-1:]
+        )
+
+    def test_empty_slice_2d(self):
+        arr = np.arange(5 * 7).reshape(5, 7)
+        darr = cuda.to_device(arr)
+        np.testing.assert_array_equal(darr[:0].copy_to_host(), arr[:0])
+        np.testing.assert_array_equal(darr[3, :0].copy_to_host(), arr[3, :0])
+        # empty slice of empty slice
+        self.assertFalse(darr[:0][:0].copy_to_host().size > 0)
+        # out-of-bound slice just produces empty slices
+        np.testing.assert_array_equal(darr[:0][:1].copy_to_host(), arr[:0][:1])
+        np.testing.assert_array_equal(
+            darr[:0][-1:].copy_to_host(), arr[:0][-1:]
+        )
+
+
+class CudaArraySetting(CUDATestCase):
+    """
+    Most of the slicing logic is tested in the cases above, so these
+    tests focus on the setting logic.
+    """
+
+    def test_scalar(self):
+        arr = np.arange(5 * 7).reshape(5, 7)
+        darr = cuda.to_device(arr)
+        arr[2, 2] = 500
+        darr[2, 2] = 500
+        np.testing.assert_array_equal(darr.copy_to_host(), arr)
+
+    def test_rank(self):
+        arr = np.arange(5 * 7).reshape(5, 7)
+        darr = cuda.to_device(arr)
+        arr[2] = 500
+        darr[2] = 500
+        np.testing.assert_array_equal(darr.copy_to_host(), arr)
+
+    def test_broadcast(self):
+        arr = np.arange(5 * 7).reshape(5, 7)
+        darr = cuda.to_device(arr)
+        arr[:, 2] = 500
+        darr[:, 2] = 500
+        np.testing.assert_array_equal(darr.copy_to_host(), arr)
+
+    def test_array_assign_column(self):
+        arr = np.arange(5 * 7).reshape(5, 7)
+        darr = cuda.to_device(arr)
+        _400 = np.full(shape=7, fill_value=400)
+        arr[2] = _400
+        darr[2] = _400
+        np.testing.assert_array_equal(darr.copy_to_host(), arr)
+
+    def test_array_assign_row(self):
+        arr = np.arange(5 * 7).reshape(5, 7)
+        darr = cuda.to_device(arr)
+        _400 = np.full(shape=5, fill_value=400)
+        arr[:, 2] = _400
+        darr[:, 2] = _400
+        np.testing.assert_array_equal(darr.copy_to_host(), arr)
+
+    def test_array_assign_subarray(self):
+        arr = np.arange(5 * 6 * 7).reshape(5, 6, 7)
+        darr = cuda.to_device(arr)
+        _400 = np.full(shape=(6, 7), fill_value=400)
+        arr[2] = _400
+        darr[2] = _400
+        np.testing.assert_array_equal(darr.copy_to_host(), arr)
+
+    def test_array_assign_deep_subarray(self):
+        arr = np.arange(5 * 6 * 7 * 8).reshape(5, 6, 7, 8)
+        darr = cuda.to_device(arr)
+        _400 = np.full(shape=(5, 6, 8), fill_value=400)
+        arr[:, :, 2] = _400
+        darr[:, :, 2] = _400
+        np.testing.assert_array_equal(darr.copy_to_host(), arr)
+
+    def test_array_assign_all(self):
+        arr = np.arange(5 * 7).reshape(5, 7)
+        darr = cuda.to_device(arr)
+        _400 = np.full(shape=(5, 7), fill_value=400)
+        arr[:] = _400
+        darr[:] = _400
+        np.testing.assert_array_equal(darr.copy_to_host(), arr)
+
+    def test_strides(self):
+        arr = np.ones(20)
+        darr = cuda.to_device(arr)
+        arr[::2] = 500
+        darr[::2] = 500
+        np.testing.assert_array_equal(darr.copy_to_host(), arr)
+
+    def test_incompatible_highdim(self):
+        darr = cuda.to_device(np.arange(5 * 7))
+
+        with self.assertRaises(ValueError) as e:
+            darr[:] = np.ones(shape=(1, 2, 3))
+
+        self.assertIn(
+            member=str(e.exception),
+            container=[
+                "Can't assign 3-D array to 1-D self",  # device
+                "could not broadcast input array from shape (2,3) "
+                "into shape (35,)",  # simulator, NP >= 1.20
+            ],
+        )
+
+    def test_incompatible_shape(self):
+        darr = cuda.to_device(np.arange(5))
+
+        with self.assertRaises(ValueError) as e:
+            darr[:] = [1, 3]
+
+        self.assertIn(
+            member=str(e.exception),
+            container=[
+                "Can't copy sequence with size 2 to array axis 0 with "
+                "dimension 5",  # device
+                "could not broadcast input array from shape (2,) into "
+                "shape (5,)",  # simulator, NP >= 1.20
+            ],
+        )
+
+    @skip_on_cudasim("cudasim does not use streams and operates synchronously")
+    def test_sync(self):
+        # There should be a synchronization when no stream is supplied
+        darr = cuda.to_device(np.arange(5))
+
+        with patch.object(
+            cuda.cudadrv.driver.Stream, "synchronize", return_value=None
+        ) as mock_sync:
+            darr[0] = 10
+
+        mock_sync.assert_called_once()
+
+    @skip_on_cudasim("cudasim does not use streams and operates synchronously")
+    def test_no_sync_default_stream(self):
+        # There should not be a synchronization when the array has a default
+        # stream, whether it is the default stream, the legacy default stream,
+        # the per-thread default stream, or another stream.
+        streams = (
+            cuda.stream(),
+            cuda.default_stream(),
+            cuda.legacy_default_stream(),
+            cuda.per_thread_default_stream(),
+        )
+
+        for stream in streams:
+            darr = cuda.to_device(np.arange(5), stream=stream)
+
+            with patch.object(
+                cuda.cudadrv.driver.Stream, "synchronize", return_value=None
+            ) as mock_sync:
+                darr[0] = 10
+
+            mock_sync.assert_not_called()
+
+    @skip_on_cudasim("cudasim does not use streams and operates synchronously")
+    def test_no_sync_supplied_stream(self):
+        # There should not be a synchronization when a stream is supplied for
+        # the setitem call, whether it is the default stream, the legacy default
+        # stream, the per-thread default stream, or another stream.
+        streams = (
+            cuda.stream(),
+            cuda.default_stream(),
+            cuda.legacy_default_stream(),
+            cuda.per_thread_default_stream(),
+        )
+
+        for stream in streams:
+            darr = cuda.to_device(np.arange(5))
+
+            with patch.object(
+                cuda.cudadrv.driver.Stream, "synchronize", return_value=None
+            ) as mock_sync:
+                darr.setitem(0, 10, stream=stream)
+
+            mock_sync.assert_not_called()
+
+    @unittest.skip("Requires PR #6367")
+    def test_issue_6505(self):
+        # On Windows, the writes to ary_v would not be visible prior to the
+        # assertion, due to the assignment being done with a kernel launch that
+        # returns asynchronously - there should now be a sync after the kernel
+        # launch to ensure that the writes are always visible.
+        ary = cuda.mapped_array(2, dtype=np.int32)
+        ary[:] = 0
+
+        ary_v = ary.view("u1")
+        ary_v[1] = 1
+        ary_v[5] = 1
+        self.assertEqual(sum(ary), 512)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,24 @@
+# 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
+
+
+class TestCudaAutoContext(CUDATestCase):
+    def test_auto_context(self):
+        """A problem was revealed by a customer that the use cuda.to_device
+        does not create a CUDA context.
+        This tests the problem
+        """
+        A = np.arange(10, dtype=np.float32)
+        newA = np.empty_like(A)
+        dA = cuda.to_device(A)
+
+        dA.copy_to_host(newA)
+        self.assertTrue(np.allclose(A, newA))
+
+
+if __name__ == "__main__":
+    unittest.main()