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

numba_cuda/numba/cuda/tests/cudasim/test_cudasim_issues.py

@@ -0,0 +1,111 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import threading
+
+import numpy as np
+
+from numba import cuda
+from numba.cuda.testing import CUDATestCase, skip_unless_cudasim
+import numba.cuda.simulator as simulator
+import unittest
+
+
+class TestCudaSimIssues(CUDATestCase):
+    def test_record_access(self):
+        backyard_type = [
+            ("statue", np.float64),
+            ("newspaper", np.float64, (6,)),
+        ]
+
+        goose_type = [
+            ("garden", np.float64, (12,)),
+            ("town", np.float64, (42,)),
+            ("backyard", backyard_type),
+        ]
+
+        goose_np_type = np.dtype(goose_type, align=True)
+
+        @cuda.jit
+        def simple_kernel(f):
+            f.garden[0] = 45.0
+            f.backyard.newspaper[3] = 2.0
+            f.backyard.newspaper[3] = f.backyard.newspaper[3] + 3.0
+
+        item = np.recarray(1, dtype=goose_np_type)
+        simple_kernel[1, 1](item[0])
+        np.testing.assert_equal(item[0]["garden"][0], 45)
+        np.testing.assert_equal(item[0]["backyard"]["newspaper"][3], 5)
+
+    def test_recarray_setting(self):
+        recordwith2darray = np.dtype(
+            [("i", np.int32), ("j", np.float32, (3, 2))]
+        )
+        rec = np.recarray(2, dtype=recordwith2darray)
+        rec[0]["i"] = 45
+
+        @cuda.jit
+        def simple_kernel(f):
+            f[1] = f[0]
+
+        simple_kernel[1, 1](rec)
+        np.testing.assert_equal(rec[0]["i"], rec[1]["i"])
+
+    def test_cuda_module_in_device_function(self):
+        """
+        Discovered in https://github.com/numba/numba/issues/1837.
+        When the `cuda` module is referenced in a device function,
+        it does not have the kernel API (e.g. cuda.threadIdx, cuda.shared)
+        """
+        from numba.cuda.tests.cudasim import support
+
+        inner = support.cuda_module_in_device_function
+
+        @cuda.jit
+        def outer(out):
+            tid = inner()
+            if tid < out.size:
+                out[tid] = tid
+
+        arr = np.zeros(10, dtype=np.int32)
+        outer[1, 11](arr)
+        expected = np.arange(arr.size, dtype=np.int32)
+        np.testing.assert_equal(expected, arr)
+
+    @skip_unless_cudasim("Only works on CUDASIM")
+    def test_deadlock_on_exception(self):
+        def assert_no_blockthreads():
+            blockthreads = []
+            for t in threading.enumerate():
+                if not isinstance(t, simulator.kernel.BlockThread):
+                    continue
+
+                # join blockthreads with a short timeout to allow aborted
+                # threads to exit
+                t.join(1)
+                if t.is_alive():
+                    self.fail("Blocked kernel thread: %s" % t)
+
+            self.assertListEqual(blockthreads, [])
+
+        @simulator.jit
+        def assign_with_sync(x, y):
+            i = cuda.grid(1)
+            y[i] = x[i]
+
+            cuda.syncthreads()
+            cuda.syncthreads()
+
+        x = np.arange(3)
+        y = np.empty(3)
+        assign_with_sync[1, 3](x, y)
+        np.testing.assert_array_equal(x, y)
+        assert_no_blockthreads()
+
+        with self.assertRaises(IndexError):
+            assign_with_sync[1, 6](x, y)
+        assert_no_blockthreads()
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause

numba_cuda/numba/cuda/tests/data/cta_barrier.cu

@@ -0,0 +1,28 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#include <cooperative_groups.h>
+#include <cuda/barrier>
+
+namespace cg = cooperative_groups;
+
+__device__ void _wait_on_tile(cuda::barrier<cuda::thread_scope_block> &tile)
+{
+    auto token = tile.arrive();
+    tile.wait(std::move(token));
+}
+
+extern "C"
+__device__ int cta_barrier(int *ret) {
+    auto cta = cg::this_thread_block();
+    cg::thread_block_tile<32> tile = cg::tiled_partition<32>(cta);
+    __shared__ cuda::barrier<cuda::thread_scope_block> barrier;
+    if (threadIdx.x == 0) {
+        init(&barrier, blockDim.x);
+    }
+
+    _wait_on_tile(barrier);
+    return 0;
+}

numba_cuda/numba/cuda/tests/data/cuda_include.cu

@@ -0,0 +1,10 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+// Not all CUDA includes are safe to include in device code compiled by NVRTC,
+// because it does not have paths to all system include directories. Headers
+// such as cuda_device_runtime_api.h are safe to use in NVRTC without adding
+// additional includes.
+#include <cuda_device_runtime_api.h>

numba_cuda/numba/cuda/tests/data/error.cu

@@ -0,0 +1,12 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+extern "C" __device__
+int bar(int* out, int a) {
+  // Explicitly placed to generate an error
+  SYNTAX ERROR
+  *out = a * 2;
+  return 0;
+}

numba_cuda/numba/cuda/tests/data/include/add.cuh

@@ -0,0 +1,8 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+// Templated addition function: myadd
+template <typename T>
+__device__ T myadd(T a, T b) { return a + b; }

numba_cuda/numba/cuda/tests/data/jitlink.cu

@@ -0,0 +1,28 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+// Compile with:
+//
+//   nvcc -gencode arch=compute_50,code=compute_50 -rdc true -ptx jitlink.cu
+//
+// using the oldest supported toolkit version (10.2 at the time of writing).
+
+extern "C" __device__
+int bar(int *out, int a)
+{
+  *out = a * 2;
+  return 0;
+}
+
+
+// The out argument is necessary due to Numba's CUDA calling convention, which
+// always reserves the first parameter for a pointer to a returned value, even
+// if there is no return value.
+extern "C" __device__
+int array_mutator(void *out, int *a)
+{
+  a[0] = a[1];
+  return 0;
+}

numba_cuda/numba/cuda/tests/data/jitlink.ptx

@@ -0,0 +1,49 @@
+//
+// Generated by NVIDIA NVVM Compiler
+//
+// Compiler Build ID: CL-27506705
+// Cuda compilation tools, release 10.2, V10.2.89
+// Based on LLVM 3.4svn
+//
+
+.version 6.5
+.target sm_50
+.address_size 64
+
+	// .globl	bar
+
+.visible .func  (.param .b32 func_retval0) bar(
+	.param .b64 bar_param_0,
+	.param .b32 bar_param_1
+)
+{
+	.reg .b32 	%r<4>;
+	.reg .b64 	%rd<2>;
+
+
+	ld.param.u64 	%rd1, [bar_param_0];
+	ld.param.u32 	%r1, [bar_param_1];
+	shl.b32 	%r2, %r1, 1;
+	st.u32 	[%rd1], %r2;
+	mov.u32 	%r3, 0;
+	st.param.b32	[func_retval0+0], %r3;
+	ret;
+}
+
+	// .globl	array_mutator
+.visible .func  (.param .b32 func_retval0) array_mutator(
+	.param .b64 array_mutator_param_0,
+	.param .b64 array_mutator_param_1
+)
+{
+	.reg .b32 	%r<3>;
+	.reg .b64 	%rd<2>;
+
+
+	ld.param.u64 	%rd1, [array_mutator_param_1];
+	ld.u32 	%r1, [%rd1+4];
+	st.u32 	[%rd1], %r1;
+	mov.u32 	%r2, 0;
+	st.param.b32	[func_retval0+0], %r2;
+	ret;
+}

numba_cuda/numba/cuda/tests/data/warn.cu

@@ -0,0 +1,12 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+extern "C" __device__
+int bar(int* out, int a) {
+  // Explicitly placed to generate a warning for testing the NVRTC program log
+  int unused;
+  *out = a * 2;
+  return 0;
+}

numba_cuda/numba/cuda/tests/doc_examples/__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__))

numba_cuda/numba/cuda/tests/doc_examples/ffi/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause

numba_cuda/numba/cuda/tests/doc_examples/ffi/functions.cu

@@ -0,0 +1,54 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+// magictoken.ex_mul_f32_f32.begin
+// Foreign function example: multiplication of a pair of floats
+
+extern "C" __device__ int
+mul_f32_f32(
+  float* return_value,
+  float x,
+  float y)
+{
+  // Compute result and store in caller-provided slot
+  *return_value = x * y;
+
+  // Signal that no Python exception occurred
+  return 0;
+}
+// magictoken.ex_mul_f32_f32.end
+
+
+// magictoken.ex_sum_reduce_proto.begin
+extern "C"
+__device__ int
+sum_reduce(
+  float* return_value,
+  float* array,
+  int n
+);
+// magictoken.ex_sum_reduce_proto.end
+
+
+// Performs a simple reduction on an array passed by pointer using the
+// ffi.from_buffer() method. Implements the prototype above.
+extern "C"
+__device__ int
+sum_reduce(
+  float* return_value,
+  float* array,
+  int n
+)
+{
+  double sum = 0.0;
+
+  for (size_t i = 0; i < n; ++i) {
+    sum += array[i];
+  }
+
+  *return_value = (float)sum;
+
+  return 0;
+}

numba_cuda/numba/cuda/tests/doc_examples/ffi/include/mul.cuh

@@ -0,0 +1,8 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+// Templated multiplication function: mymul
+template <typename T>
+__device__ T mymul(T a, T b) { return a * b; }

numba_cuda/numba/cuda/tests/doc_examples/ffi/saxpy.cu

@@ -0,0 +1,14 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#include <add.cuh> // In numba/cuda/tests/data/include
+#include <mul.cuh> // In numba/cuda/tests/doc_examples/ffi/include
+
+extern "C"
+__device__ int saxpy(float *ret, float a, float x, float y)
+{
+    *ret = myadd(mymul(a, x), y);
+    return 0;
+}

numba_cuda/numba/cuda/tests/doc_examples/test_cg.py

@@ -0,0 +1,86 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+# Contents in this file are referenced from the sphinx-generated docs.
+# "magictoken" is used for markers as beginning and ending of example text.
+
+import unittest
+from numba.cuda.testing import (
+    CUDATestCase,
+    skip_on_cudasim,
+    skip_if_cudadevrt_missing,
+    skip_unless_cc_60,
+)
+
+
+@skip_if_cudadevrt_missing
+@skip_unless_cc_60
+@skip_on_cudasim("cudasim doesn't support cuda import at non-top-level")
+class TestCooperativeGroups(CUDATestCase):
+    def test_ex_grid_sync(self):
+        # magictoken.ex_grid_sync_kernel.begin
+        from numba import cuda
+        from numba.cuda import int32
+        import numpy as np
+
+        sig = (int32[:, ::1],)
+
+        @cuda.jit(sig)
+        def sequential_rows(M):
+            col = cuda.grid(1)
+            g = cuda.cg.this_grid()
+
+            rows = M.shape[0]
+            cols = M.shape[1]
+
+            for row in range(1, rows):
+                opposite = cols - col - 1
+                # Each row's elements are one greater than the previous row
+                M[row, col] = M[row - 1, opposite] + 1
+                # Wait until all threads have written their column element,
+                # and that the write is visible to all other threads
+                g.sync()
+
+        # magictoken.ex_grid_sync_kernel.end
+
+        # magictoken.ex_grid_sync_data.begin
+        # Empty input data
+        A = np.zeros((1024, 1024), dtype=np.int32)
+        # A somewhat arbitrary choice (one warp), but generally smaller block sizes
+        # allow more blocks to be launched (noting that other limitations on
+        # occupancy apply such as shared memory size)
+        blockdim = 32
+        griddim = A.shape[1] // blockdim
+        # magictoken.ex_grid_sync_data.end
+
+        # Skip this test if the grid size used in the example is too large for
+        # a cooperative launch on the current GPU
+        mb = sequential_rows.overloads[sig].max_cooperative_grid_blocks(
+            blockdim
+        )
+        if mb < griddim:
+            self.skipTest("Device does not support a large enough coop grid")
+
+        # magictoken.ex_grid_sync_launch.begin
+        # Kernel launch - this is implicitly a cooperative launch
+        sequential_rows[griddim, blockdim](A)
+
+        # What do the results look like?
+        # print(A)
+        #
+        # [[   0    0    0 ...    0    0    0]
+        #  [   1    1    1 ...    1    1    1]
+        #  [   2    2    2 ...    2    2    2]
+        #  ...
+        #  [1021 1021 1021 ... 1021 1021 1021]
+        #  [1022 1022 1022 ... 1022 1022 1022]
+        #  [1023 1023 1023 ... 1023 1023 1023]]
+        # magictoken.ex_grid_sync_launch.end
+
+        # Sanity check - are the results what we expect?
+        reference = np.tile(np.arange(1024), (1024, 1)).T
+        np.testing.assert_equal(A, reference)
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/doc_examples/test_cpointer.py

@@ -0,0 +1,68 @@
+# 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, skip_on_cudasim
+from numba.cuda.tests.support import captured_stdout
+
+
+@skip_on_cudasim("cudasim doesn't support cuda import at non-top-level")
+class TestCPointer(CUDATestCase):
+    """
+    Test simple vector addition
+    """
+
+    def setUp(self):
+        # Prevent output from this test showing
+        # up when running the test suite
+        self._captured_stdout = captured_stdout()
+        self._captured_stdout.__enter__()
+        super().setUp()
+
+    def tearDown(self):
+        # No exception type, value, or traceback
+        self._captured_stdout.__exit__(None, None, None)
+        super().tearDown()
+
+    def test_ex_cpointer(self):
+        # ex_cpointer.sig.begin
+        import numpy as np
+        from numba import cuda
+        from numba.cuda import types
+
+        # The first kernel argument is a pointer to a uint8 array.
+        # The second argument holds the length as a uint32.
+        # The return type of a kernel is always void.
+        sig = types.void(types.CPointer(types.uint8), types.uint32)
+        # ex_cpointer.sig.end
+
+        # ex_cpointer.kernel.begin
+        @cuda.jit(sig)
+        def add_one(x, n):
+            i = cuda.grid(1)
+            if i < n:
+                x[i] += 1
+
+        # ex_cpointer.kernel.end
+
+        # ex_cpointer.launch.begin
+        x = cuda.to_device(np.arange(10, dtype=np.uint8))
+
+        # Print initial values of x
+        print(x.copy_to_host())  # [0 1 2 3 4 5 6 7 8 9]
+
+        # Obtain a pointer to the data from from the CUDA Array Interface
+        x_ptr = x.__cuda_array_interface__["data"][0]
+        x_len = len(x)
+
+        # Launch the kernel with the pointer and length
+        add_one[1, 32](x_ptr, x_len)
+
+        # Demonstrate that the data was updated by the kernel
+        print(x.copy_to_host())  # [ 1  2  3  4  5  6  7  8  9 10]
+        # ex_cpointer.launch.end
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/doc_examples/test_cpu_gpu_compat.py

@@ -0,0 +1,81 @@
+# 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,
+    skip_on_cudasim,
+    skip_on_standalone_numba_cuda,
+)
+from numba.cuda.tests.support import captured_stdout
+import numpy as np
+
+
+@skip_on_cudasim("cudasim doesn't support cuda import at non-top-level")
+class TestCpuGpuCompat(CUDATestCase):
+    """
+    Test compatibility of CPU and GPU functions
+    """
+
+    def setUp(self):
+        # Prevent output from this test showing up when running the test suite
+        self._captured_stdout = captured_stdout()
+        self._captured_stdout.__enter__()
+        super().setUp()
+
+    def tearDown(self):
+        # No exception type, value, or traceback
+        self._captured_stdout.__exit__(None, None, None)
+        super().tearDown()
+
+    @skip_on_standalone_numba_cuda
+    def test_ex_cpu_gpu_compat(self):
+        # ex_cpu_gpu_compat.import.begin
+        from math import pi
+
+        import numba
+        from numba import cuda
+        # ex_cpu_gpu_compat.import.end
+
+        # ex_cpu_gpu_compat.allocate.begin
+        X = cuda.to_device([1, 10, 234])
+        Y = cuda.to_device([2, 2, 4014])
+        Z = cuda.to_device([3, 14, 2211])
+        results = cuda.to_device([0.0, 0.0, 0.0])
+        # ex_cpu_gpu_compat.allocate.end
+
+        # ex_cpu_gpu_compat.define.begin
+        @numba.jit
+        def business_logic(x, y, z):
+            return 4 * z * (2 * x - (4 * y) / 2 * pi)
+
+        # ex_cpu_gpu_compat.define.end
+
+        # ex_cpu_gpu_compat.cpurun.begin
+        print(business_logic(1, 2, 3))  # -126.79644737231007
+        # ex_cpu_gpu_compat.cpurun.end
+
+        # ex_cpu_gpu_compat.usegpu.begin
+        @cuda.jit
+        def f(res, xarr, yarr, zarr):
+            tid = cuda.grid(1)
+            if tid < len(xarr):
+                # The function decorated with numba.jit may be directly reused
+                res[tid] = business_logic(xarr[tid], yarr[tid], zarr[tid])
+
+        # ex_cpu_gpu_compat.usegpu.end
+
+        # ex_cpu_gpu_compat.launch.begin
+        f.forall(len(X))(results, X, Y, Z)
+        print(results)
+        # [-126.79644737231007, 416.28324559588634, -218912930.2987788]
+        # ex_cpu_gpu_compat.launch.end
+
+        expect = [business_logic(x, y, z) for x, y, z in zip(X, Y, Z)]
+
+        np.testing.assert_equal(expect, results.copy_to_host())
+
+
+if __name__ == "__main__":
+    unittest.main()