numba-cuda 0.22.0__cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

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

@@ -0,0 +1,76 @@
+# 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 import float32, void
+from numba.cuda.testing import unittest, CUDATestCase
+from numba.cuda.core import config
+
+# Ensure the test takes a reasonable amount of time in the simulator
+if config.ENABLE_CUDASIM:
+    bpg, tpb = 2, 8
+else:
+    bpg, tpb = 50, 32
+
+n = bpg * tpb
+SM_SIZE = (tpb, tpb)
+
+
+class TestCudaMatMul(CUDATestCase):
+    def test_func(self):
+        @cuda.jit(void(float32[:, ::1], float32[:, ::1], float32[:, ::1]))
+        def cu_square_matrix_mul(A, B, C):
+            sA = cuda.shared.array(shape=SM_SIZE, dtype=float32)
+            sB = cuda.shared.array(shape=(tpb, tpb), dtype=float32)
+
+            tx = cuda.threadIdx.x
+            ty = cuda.threadIdx.y
+            bx = cuda.blockIdx.x
+            by = cuda.blockIdx.y
+            bw = cuda.blockDim.x
+            bh = cuda.blockDim.y
+
+            x = tx + bx * bw
+            y = ty + by * bh
+
+            acc = float32(0)  # forces all the math to be f32
+            for i in range(bpg):
+                if x < n and y < n:
+                    sA[ty, tx] = A[y, tx + i * tpb]
+                    sB[ty, tx] = B[ty + i * tpb, x]
+
+                cuda.syncthreads()
+
+                if x < n and y < n:
+                    for j in range(tpb):
+                        acc += sA[ty, j] * sB[j, tx]
+
+                cuda.syncthreads()
+
+            if x < n and y < n:
+                C[y, x] = acc
+
+        np.random.seed(42)
+        A = np.array(np.random.random((n, n)), dtype=np.float32)
+        B = np.array(np.random.random((n, n)), dtype=np.float32)
+        C = np.empty_like(A)
+
+        stream = cuda.stream()
+        with stream.auto_synchronize():
+            dA = cuda.to_device(A, stream)
+            dB = cuda.to_device(B, stream)
+            dC = cuda.to_device(C, stream)
+            cu_square_matrix_mul[(bpg, bpg), (tpb, tpb), stream](dA, dB, dC)
+            dC.copy_to_host(C, stream)
+
+        # Host compute
+        Cans = np.dot(A, B)
+
+        # Check result
+        np.testing.assert_allclose(C, Cans, rtol=1e-5)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,78 @@
+# 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 import float64
+from numba.cuda.testing import unittest, CUDATestCase, skip_on_cudasim
+
+
+def builtin_max(A, B, C):
+    i = cuda.grid(1)
+
+    if i >= len(C):
+        return
+
+    C[i] = float64(max(A[i], B[i]))
+
+
+def builtin_min(A, B, C):
+    i = cuda.grid(1)
+
+    if i >= len(C):
+        return
+
+    C[i] = float64(min(A[i], B[i]))
+
+
+@skip_on_cudasim("Tests PTX emission")
+class TestCudaMinMax(CUDATestCase):
+    def _run(
+        self,
+        kernel,
+        numpy_equivalent,
+        ptx_instruction,
+        dtype_left,
+        dtype_right,
+        n=5,
+    ):
+        kernel = cuda.jit(kernel)
+
+        c = np.zeros(n, dtype=np.float64)
+        a = np.arange(n, dtype=dtype_left) + 0.5
+        b = np.full(n, fill_value=2, dtype=dtype_right)
+
+        kernel[1, c.shape](a, b, c)
+        np.testing.assert_allclose(c, numpy_equivalent(a, b))
+
+        ptx = next(p for p in kernel.inspect_asm().values())
+        self.assertIn(ptx_instruction, ptx)
+
+    def test_max_f8f8(self):
+        self._run(builtin_max, np.maximum, "max.f64", np.float64, np.float64)
+
+    def test_max_f4f8(self):
+        self._run(builtin_max, np.maximum, "max.f64", np.float32, np.float64)
+
+    def test_max_f8f4(self):
+        self._run(builtin_max, np.maximum, "max.f64", np.float64, np.float32)
+
+    def test_max_f4f4(self):
+        self._run(builtin_max, np.maximum, "max.f32", np.float32, np.float32)
+
+    def test_min_f8f8(self):
+        self._run(builtin_min, np.minimum, "min.f64", np.float64, np.float64)
+
+    def test_min_f4f8(self):
+        self._run(builtin_min, np.minimum, "min.f64", np.float32, np.float64)
+
+    def test_min_f8f4(self):
+        self._run(builtin_min, np.minimum, "min.f64", np.float64, np.float32)
+
+    def test_min_f4f4(self):
+        self._run(builtin_min, np.minimum, "min.f32", np.float32, np.float32)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import math
+from numba import cuda
+from numba.cuda.testing import unittest, CUDATestCase
+
+
+class TestCudaMonteCarlo(CUDATestCase):
+    def test_montecarlo(self):
+        """Just make sure we can compile this"""
+
+        @cuda.jit(
+            "void(double[:], double[:], double, double, double, double[:])"
+        )
+        def step(last, paths, dt, c0, c1, normdist):
+            i = cuda.grid(1)
+            if i >= paths.shape[0]:
+                return
+            noise = normdist[i]
+            paths[i] = last[i] * math.exp(c0 * dt + c1 * noise)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,145 @@
+# 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 skip_on_cudasim, CUDATestCase
+import threading
+import unittest
+
+
+class TestMultiGPUContext(CUDATestCase):
+    @unittest.skipIf(len(cuda.gpus) < 2, "need more than 1 gpus")
+    def test_multigpu_context(self):
+        @cuda.jit("void(float64[:], float64[:])")
+        def copy_plus_1(inp, out):
+            i = cuda.grid(1)
+            if i < out.size:
+                out[i] = inp[i] + 1
+
+        def check(inp, out):
+            np.testing.assert_equal(inp + 1, out)
+
+        N = 32
+        A = np.arange(N, dtype=np.float64)
+        B = np.arange(N, dtype=np.float64)
+
+        with cuda.gpus[0]:
+            copy_plus_1[1, N](A, B)
+
+        check(A, B)
+
+        copy_plus_1[1, N](A, B)
+        check(A, B)
+
+        with cuda.gpus[0]:
+            A0 = np.arange(N, dtype=np.float64)
+            B0 = np.arange(N, dtype=np.float64)
+            copy_plus_1[1, N](A0, B0)
+
+            with cuda.gpus[1]:
+                A1 = np.arange(N, dtype=np.float64)
+                B1 = np.arange(N, dtype=np.float64)
+                copy_plus_1[1, N](A1, B1)
+
+        check(A0, B0)
+        check(A1, B1)
+
+        A = np.arange(N, dtype=np.float64)
+        B = np.arange(N, dtype=np.float64)
+        copy_plus_1[1, N](A, B)
+        check(A, B)
+
+    @skip_on_cudasim("Simulator does not support multiple threads")
+    def test_multithreaded(self):
+        def work(gpu, dA, results, ridx):
+            try:
+                with gpu:
+                    arr = dA.copy_to_host()
+
+            except Exception as e:
+                results[ridx] = e
+
+            else:
+                results[ridx] = np.all(arr == np.arange(10))
+
+        dA = cuda.to_device(np.arange(10))
+
+        nthreads = 10
+        results = [None] * nthreads
+        threads = [
+            threading.Thread(
+                target=work, args=(cuda.gpus.current, dA, results, i)
+            )
+            for i in range(nthreads)
+        ]
+        for th in threads:
+            th.start()
+
+        for th in threads:
+            th.join()
+
+        for r in results:
+            if isinstance(r, BaseException):
+                raise r
+            else:
+                self.assertTrue(r)
+
+    @unittest.skipIf(len(cuda.gpus) < 2, "need more than 1 gpus")
+    def test_with_context(self):
+        @cuda.jit
+        def vector_add_scalar(arr, val):
+            i = cuda.grid(1)
+            if i < arr.size:
+                arr[i] += val
+
+        hostarr = np.arange(10, dtype=np.float32)
+        with cuda.gpus[0]:
+            arr1 = cuda.to_device(hostarr)
+
+        with cuda.gpus[1]:
+            arr2 = cuda.to_device(hostarr)
+
+        with cuda.gpus[0]:
+            vector_add_scalar[1, 10](arr1, 1)
+
+        with cuda.gpus[1]:
+            vector_add_scalar[1, 10](arr2, 2)
+
+        with cuda.gpus[0]:
+            np.testing.assert_equal(arr1.copy_to_host(), (hostarr + 1))
+
+        with cuda.gpus[1]:
+            np.testing.assert_equal(arr2.copy_to_host(), (hostarr + 2))
+
+    @unittest.skipIf(len(cuda.gpus) < 2, "need more than 1 gpus")
+    def test_with_context_peer_copy(self):
+        # Peer access is not always possible - for example, with one GPU in TCC
+        # mode and one in WDDM - if that is the case, this test would fail so
+        # we need to skip it.
+        with cuda.gpus[0]:
+            ctx = cuda.current_context()
+            if not ctx.can_access_peer(1):
+                self.skipTest("Peer access between GPUs disabled")
+
+        # 1. Create a range in an array
+        hostarr = np.arange(10, dtype=np.float32)
+
+        # 2. Copy range array from host -> GPU 0
+        with cuda.gpus[0]:
+            arr1 = cuda.to_device(hostarr)
+
+        # 3. Initialize a zero-filled array on GPU 1
+        with cuda.gpus[1]:
+            arr2 = cuda.to_device(np.zeros_like(hostarr))
+
+        with cuda.gpus[0]:
+            # 4. Copy range from GPU 0 -> GPU 1
+            arr2.copy_to_device(arr1)
+
+            # 5. Copy range from GPU 1 -> host and check contents
+            np.testing.assert_equal(arr2.copy_to_host(), hostarr)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import os
+import multiprocessing as mp
+import pytest
+import concurrent.futures
+
+import numpy as np
+
+from numba import cuda
+from numba.cuda.testing import skip_on_cudasim, CUDATestCase
+import unittest
+
+
+@skip_on_cudasim("disabled for cudasim")
+class TestMultiprocessing(CUDATestCase):
+    @unittest.skipUnless(hasattr(mp, "get_context"), "requires mp.get_context")
+    @unittest.skipUnless(os.name == "posix", "requires Unix")
+    def test_fork(self):
+        """
+        Test fork detection.
+        """
+        from numba.cuda.cudadrv.error import CudaDriverError
+
+        cuda.current_context()  # force cuda initialize
+        with concurrent.futures.ProcessPoolExecutor(
+            mp_context=mp.get_context("fork")
+        ) as exe:
+            future = exe.submit(cuda.to_device, np.arange(1))
+
+        with pytest.raises(
+            CudaDriverError, match="CUDA initialized before forking"
+        ):
+            future.result()
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,82 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import threading
+import multiprocessing
+import numpy as np
+from numba import cuda
+from numba.cuda.testing import (
+    skip_on_cudasim,
+    skip_under_cuda_memcheck,
+    CUDATestCase,
+)
+import unittest
+
+from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
+
+
+has_mp_get_context = hasattr(multiprocessing, "get_context")
+
+
+def check_concurrent_compiling():
+    @cuda.jit
+    def foo(x):
+        x[0] += 1
+
+    def use_foo(x):
+        foo[1, 1](x)
+        return x
+
+    arrays = [cuda.to_device(np.arange(10)) for i in range(10)]
+    expected = np.arange(10)
+    expected[0] += 1
+    with ThreadPoolExecutor(max_workers=4) as e:
+        for ary in e.map(use_foo, arrays):
+            np.testing.assert_equal(ary, expected)
+
+
+@skip_under_cuda_memcheck("Hangs cuda-memcheck")
+@skip_on_cudasim("disabled for cudasim")
+class TestMultiThreadCompiling(CUDATestCase):
+    def test_concurrent_compiling(self):
+        check_concurrent_compiling()
+
+    @unittest.skipIf(not has_mp_get_context, "no multiprocessing.get_context")
+    def test_spawn_concurrent_compilation(self):
+        # force CUDA context init
+        cuda.get_current_device()
+
+        with ProcessPoolExecutor(
+            # use "spawn" to avoid inheriting the CUDA context
+            mp_context=multiprocessing.get_context("spawn")
+        ) as exe:
+            future = exe.submit(check_concurrent_compiling)
+        future.result()
+
+    def test_invalid_context_error_with_d2h(self):
+        def d2h(arr, out):
+            out[:] = arr.copy_to_host()
+
+        arr = np.arange(1, 4)
+        out = np.zeros_like(arr)
+
+        with ThreadPoolExecutor() as exe:
+            exe.submit(d2h, cuda.to_device(arr), out)
+
+        np.testing.assert_equal(arr, out)
+
+    def test_invalid_context_error_with_d2d(self):
+        def d2d(dst, src):
+            dst.copy_to_device(src)
+
+        arr = np.arange(100)
+        common = cuda.to_device(arr)
+        darr = cuda.to_device(np.zeros(common.shape, dtype=common.dtype))
+        th = threading.Thread(target=d2d, args=[darr, common])
+        th.start()
+        th.join()
+        np.testing.assert_equal(darr.copy_to_host(), arr)
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/cudapy/test_nondet.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 import float32, void
+from numba.cuda.testing import unittest, CUDATestCase
+
+
+def generate_input(n):
+    A = np.array(np.arange(n * n).reshape(n, n), dtype=np.float32)
+    B = np.array(np.arange(n) + 0, dtype=A.dtype)
+    return A, B
+
+
+class TestCudaNonDet(CUDATestCase):
+    def test_for_pre(self):
+        """Test issue with loop not running due to bad sign-extension at the for
+        loop precondition.
+        """
+
+        @cuda.jit(void(float32[:, :], float32[:, :], float32[:]))
+        def diagproduct(c, a, b):
+            startX, startY = cuda.grid(2)
+            gridX = cuda.gridDim.x * cuda.blockDim.x
+            gridY = cuda.gridDim.y * cuda.blockDim.y
+            height = c.shape[0]
+            width = c.shape[1]
+
+            for x in range(startX, width, (gridX)):
+                for y in range(startY, height, (gridY)):
+                    c[y, x] = a[y, x] * b[x]
+
+        N = 8
+
+        A, B = generate_input(N)
+
+        F = np.empty(A.shape, dtype=A.dtype)
+
+        blockdim = (32, 8)
+        griddim = (1, 1)
+
+        dA = cuda.to_device(A)
+        dB = cuda.to_device(B)
+        dF = cuda.to_device(F, copy=False)
+        diagproduct[griddim, blockdim](dF, dA, dB)
+
+        E = np.dot(A, np.diag(B))
+        np.testing.assert_array_almost_equal(dF.copy_to_host(), E)
+
+
+if __name__ == "__main__":
+    unittest.main()