numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

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()

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

@@ -0,0 +1,141 @@
+# 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
+from numba.cuda.tests.support import skip_unless_cffi, override_config
+
+
+@skip_unless_cffi
+@skip_on_cudasim("cudasim doesn't support cuda import at non-top-level")
+class TestFFI(CUDATestCase):
+    def test_ex_linking_cu(self):
+        # magictoken.ex_linking_cu.begin
+        from numba import cuda
+        import numpy as np
+        import os
+
+        # Path to the source containing the foreign function
+        # (here assumed to be in a subdirectory called "ffi")
+        basedir = os.path.dirname(os.path.abspath(__file__))
+        functions_cu = os.path.join(basedir, "ffi", "functions.cu")
+
+        # Declaration of the foreign function
+        mul = cuda.declare_device(
+            "mul_f32_f32", "float32(float32, float32)", link=functions_cu
+        )
+
+        # A kernel that calls mul; functions.cu is linked automatically due to
+        # the call to mul.
+        @cuda.jit
+        def multiply_vectors(r, x, y):
+            i = cuda.grid(1)
+
+            if i < len(r):
+                r[i] = mul(x[i], y[i])
+
+        # Generate random data
+        N = 32
+        np.random.seed(1)
+        x = np.random.rand(N).astype(np.float32)
+        y = np.random.rand(N).astype(np.float32)
+        r = np.zeros_like(x)
+
+        # Run the kernel
+        multiply_vectors[1, 32](r, x, y)
+
+        # Sanity check - ensure the results match those expected
+        np.testing.assert_array_equal(r, x * y)
+        # magictoken.ex_linking_cu.end
+
+    def test_ex_from_buffer(self):
+        from numba import cuda
+        import os
+
+        basedir = os.path.dirname(os.path.abspath(__file__))
+        functions_cu = os.path.join(basedir, "ffi", "functions.cu")
+
+        # magictoken.ex_from_buffer_decl.begin
+        signature = "float32(CPointer(float32), int32)"
+        sum_reduce = cuda.declare_device(
+            "sum_reduce", signature, link=functions_cu
+        )
+        # magictoken.ex_from_buffer_decl.end
+
+        # magictoken.ex_from_buffer_kernel.begin
+        import cffi
+
+        ffi = cffi.FFI()
+
+        @cuda.jit
+        def reduction_caller(result, array):
+            array_ptr = ffi.from_buffer(array)
+            result[()] = sum_reduce(array_ptr, len(array))
+
+        # magictoken.ex_from_buffer_kernel.end
+
+        import numpy as np
+
+        x = np.arange(10).astype(np.float32)
+        r = np.ndarray((), dtype=np.float32)
+
+        reduction_caller[1, 1](r, x)
+
+        expected = np.sum(x)
+        actual = r[()]
+        np.testing.assert_allclose(expected, actual)
+
+    def test_ex_extra_includes(self):
+        import numpy as np
+        from numba import cuda
+        from numba.cuda import config
+        import os
+
+        basedir = os.path.dirname(os.path.abspath(__file__))
+        mul_dir = os.path.join(basedir, "ffi", "include")
+        saxpy_cu = os.path.join(basedir, "ffi", "saxpy.cu")
+
+        testdir = os.path.dirname(basedir)
+        add_dir = os.path.join(testdir, "data", "include")
+
+        includedir = ":".join([mul_dir, add_dir])
+        with override_config("CUDA_NVRTC_EXTRA_SEARCH_PATHS", includedir):
+            # magictoken.ex_extra_search_paths.begin
+            from numba.cuda import config
+
+            includedir = ":".join([mul_dir, add_dir])
+            config.CUDA_NVRTC_EXTRA_SEARCH_PATHS = includedir
+            # magictoken.ex_extra_search_paths.end
+
+            # magictoken.ex_extra_search_paths_kernel.begin
+            sig = "float32(float32, float32, float32)"
+            saxpy = cuda.declare_device("saxpy", sig=sig, link=saxpy_cu)
+
+            @cuda.jit
+            def vector_saxpy(a, x, y, res):
+                i = cuda.grid(1)
+                if i < len(res):
+                    res[i] = saxpy(a, x[i], y[i])
+
+            # magictoken.ex_extra_search_paths_kernel.end
+
+            size = 10_000
+            a = 3.0
+            X = np.ones((size,), dtype="float32")
+            Y = np.ones((size,), dtype="float32")
+            R = np.zeros((size,), dtype="float32")
+
+            block_size = 32
+            num_blocks = (size // block_size) + 1
+
+            vector_saxpy[num_blocks, block_size](a, X, Y, R)
+
+            expected = a * X + Y
+            np.testing.assert_equal(R, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()