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

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

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

@@ -0,0 +1,160 @@
+# 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_if_cudadevrt_missing,
+    skip_on_cudasim,
+    skip_unless_cc_60,
+)
+from numba.cuda.tests.support import captured_stdout
+
+
+@skip_if_cudadevrt_missing
+@skip_unless_cc_60
+@skip_on_cudasim("cudasim doesn't support cuda import at non-top-level")
+class TestLaplace(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_laplace(self):
+        # set True to regenerate the figures that
+        # accompany this example
+        plot = False
+
+        # ex_laplace.import.begin
+        import numpy as np
+        from numba import cuda
+        # ex_laplace.import.end
+
+        # ex_laplace.allocate.begin
+        # Use an odd problem size.
+        # This is so there can be an element truly in the "middle" for symmetry.
+        size = 1001
+        data = np.zeros(size)
+
+        # Middle element is made very hot
+        data[500] = 10000
+        buf_0 = cuda.to_device(data)
+
+        # This extra array is used for synchronization purposes
+        buf_1 = cuda.device_array_like(buf_0)
+
+        niter = 10000
+        # ex_laplace.allocate.end
+
+        if plot:
+            import matplotlib.pyplot as plt
+
+            fig, ax = plt.subplots(figsize=(16 * 0.66, 9 * 0.66))
+            plt.plot(
+                np.arange(len(buf_0)),
+                buf_0.copy_to_host(),
+                lw=3,
+                marker="*",
+                color="black",
+            )
+
+            plt.title("Initial State", fontsize=24)
+            plt.xlabel("Position", fontsize=24)
+            plt.ylabel("Temperature", fontsize=24)
+
+            ax.set_xticks(ax.get_xticks(), fontsize=16)
+            ax.set_yticks(ax.get_yticks(), fontsize=16)
+            plt.xlim(0, len(data))
+            plt.ylim(0, 10001)
+            plt.savefig("laplace_initial.svg")
+
+        # ex_laplace.kernel.begin
+        @cuda.jit
+        def solve_heat_equation(buf_0, buf_1, timesteps, k):
+            i = cuda.grid(1)
+
+            # Don't continue if our index is outside the domain
+            if i >= len(buf_0):
+                return
+
+            # Prepare to do a grid-wide synchronization later
+            grid = cuda.cg.this_grid()
+
+            for step in range(timesteps):
+                # Select the buffer from the previous timestep
+                if (step % 2) == 0:
+                    data = buf_0
+                    next_data = buf_1
+                else:
+                    data = buf_1
+                    next_data = buf_0
+
+                # Get the current temperature associated with this point
+                curr_temp = data[i]
+
+                # Apply formula from finite difference equation
+                if i == 0:
+                    # Left wall is held at T = 0
+                    next_temp = curr_temp + k * (data[i + 1] - (2 * curr_temp))
+                elif i == len(data) - 1:
+                    # Right wall is held at T = 0
+                    next_temp = curr_temp + k * (data[i - 1] - (2 * curr_temp))
+                else:
+                    # Interior points are a weighted average of their neighbors
+                    next_temp = curr_temp + k * (
+                        data[i - 1] - (2 * curr_temp) + data[i + 1]
+                    )
+
+                # Write new value to the next buffer
+                next_data[i] = next_temp
+
+                # Wait for every thread to write before moving on
+                grid.sync()
+
+        # ex_laplace.kernel.end
+
+        # ex_laplace.launch.begin
+        solve_heat_equation.forall(len(data))(buf_0, buf_1, niter, 0.25)
+        # ex_laplace.launch.end
+
+        results = buf_1.copy_to_host()
+        if plot:
+            fig, ax = plt.subplots(figsize=(16 * 0.66, 9 * 0.66))
+            plt.plot(
+                np.arange(len(results)),
+                results,
+                lw=3,
+                marker="*",
+                color="black",
+            )
+            plt.title(f"T = {niter}", fontsize=24)
+            plt.xlabel("Position", fontsize=24)
+            plt.ylabel("Temperature", fontsize=24)
+
+            ax.set_xticks(ax.get_xticks(), fontsize=16)
+            ax.set_yticks(ax.get_yticks(), fontsize=16)
+
+            plt.ylim(0, max(results))
+            plt.xlim(0, len(results))
+            plt.savefig("laplace_final.svg")
+
+        # Integral over the domain should be equal to its initial value.
+        # Note that this should match the initial value of data[500] above, but
+        # we don't assign it to a variable because that would make the example
+        # code look a bit oddly verbose.
+        np.testing.assert_allclose(results.sum(), 10000)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,180 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Matrix multiplication example via `cuda.jit`.
+
+Reference: https://stackoverflow.com/a/64198479/13697228 by @RobertCrovella
+
+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 captured_stdout
+
+
+@skip_on_cudasim("cudasim doesn't support cuda import at non-top-level")
+class TestMatMul(CUDATestCase):
+    """
+    Text matrix multiplication using simple, shared memory/square, and shared
+    memory/nonsquare cases.
+    """
+
+    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_matmul(self):
+        """Test of matrix multiplication on various cases."""
+        # magictoken.ex_import.begin
+        from numba import cuda
+        from numba.cuda import float32
+        import numpy as np
+        import math
+        # magictoken.ex_import.end
+
+        # magictoken.ex_matmul.begin
+        @cuda.jit
+        def matmul(A, B, C):
+            """Perform square matrix multiplication of C = A * B."""
+            i, j = cuda.grid(2)
+            if i < C.shape[0] and j < C.shape[1]:
+                tmp = 0.0
+                for k in range(A.shape[1]):
+                    tmp += A[i, k] * B[k, j]
+                C[i, j] = tmp
+
+        # magictoken.ex_matmul.end
+
+        # magictoken.ex_run_matmul.begin
+        x_h = np.arange(16).reshape([4, 4])
+        y_h = np.ones([4, 4])
+        z_h = np.zeros([4, 4])
+
+        x_d = cuda.to_device(x_h)
+        y_d = cuda.to_device(y_h)
+        z_d = cuda.to_device(z_h)
+
+        threadsperblock = (16, 16)
+        blockspergrid_x = math.ceil(z_h.shape[0] / threadsperblock[0])
+        blockspergrid_y = math.ceil(z_h.shape[1] / threadsperblock[1])
+        blockspergrid = (blockspergrid_x, blockspergrid_y)
+
+        matmul[blockspergrid, threadsperblock](x_d, y_d, z_d)
+        z_h = z_d.copy_to_host()
+        print(z_h)
+        print(x_h @ y_h)
+        # magictoken.ex_run_matmul.end
+
+        # magictoken.ex_fast_matmul.begin
+        # Controls threads per block and shared memory usage.
+        # The computation will be done on blocks of TPBxTPB elements.
+        # TPB should not be larger than 32 in this example
+        TPB = 16
+
+        @cuda.jit
+        def fast_matmul(A, B, C):
+            """
+            Perform matrix multiplication of C = A * B using CUDA shared memory.
+
+            Reference: https://stackoverflow.com/a/64198479/13697228 by @RobertCrovella
+            """
+            # Define an array in the shared memory
+            # The size and type of the arrays must be known at compile time
+            sA = cuda.shared.array(shape=(TPB, TPB), dtype=float32)
+            sB = cuda.shared.array(shape=(TPB, TPB), dtype=float32)
+
+            x, y = cuda.grid(2)
+
+            tx = cuda.threadIdx.x
+            ty = cuda.threadIdx.y
+            bpg = cuda.gridDim.x  # blocks per grid
+
+            # Each thread computes one element in the result matrix.
+            # The dot product is chunked into dot products of TPB-long vectors.
+            tmp = float32(0.0)
+            for i in range(bpg):
+                # Preload data into shared memory
+                sA[ty, tx] = 0
+                sB[ty, tx] = 0
+                if y < A.shape[0] and (tx + i * TPB) < A.shape[1]:
+                    sA[ty, tx] = A[y, tx + i * TPB]
+                if x < B.shape[1] and (ty + i * TPB) < B.shape[0]:
+                    sB[ty, tx] = B[ty + i * TPB, x]
+
+                # Wait until all threads finish preloading
+                cuda.syncthreads()
+
+                # Computes partial product on the shared memory
+                for j in range(TPB):
+                    tmp += sA[ty, j] * sB[j, tx]
+
+                # Wait until all threads finish computing
+                cuda.syncthreads()
+            if y < C.shape[0] and x < C.shape[1]:
+                C[y, x] = tmp
+
+        # magictoken.ex_fast_matmul.end
+
+        # magictoken.ex_run_fast_matmul.begin
+        x_h = np.arange(16).reshape([4, 4])
+        y_h = np.ones([4, 4])
+        z_h = np.zeros([4, 4])
+
+        x_d = cuda.to_device(x_h)
+        y_d = cuda.to_device(y_h)
+        z_d = cuda.to_device(z_h)
+
+        threadsperblock = (TPB, TPB)
+        blockspergrid_x = math.ceil(z_h.shape[0] / threadsperblock[0])
+        blockspergrid_y = math.ceil(z_h.shape[1] / threadsperblock[1])
+        blockspergrid = (blockspergrid_x, blockspergrid_y)
+
+        fast_matmul[blockspergrid, threadsperblock](x_d, y_d, z_d)
+        z_h = z_d.copy_to_host()
+        print(z_h)
+        print(x_h @ y_h)
+        # magictoken.ex_run_fast_matmul.end
+
+        # fast_matmul test(s)
+        msg = "fast_matmul incorrect for shared memory, square case."
+        self.assertTrue(np.all(z_h == x_h @ y_h), msg=msg)
+
+        # magictoken.ex_run_nonsquare.begin
+        x_h = np.arange(115).reshape([5, 23])
+        y_h = np.ones([23, 7])
+        z_h = np.zeros([5, 7])
+
+        x_d = cuda.to_device(x_h)
+        y_d = cuda.to_device(y_h)
+        z_d = cuda.to_device(z_h)
+
+        threadsperblock = (TPB, TPB)
+        grid_y_max = max(x_h.shape[0], y_h.shape[0])
+        grid_x_max = max(x_h.shape[1], y_h.shape[1])
+        blockspergrid_x = math.ceil(grid_x_max / threadsperblock[0])
+        blockspergrid_y = math.ceil(grid_y_max / threadsperblock[1])
+        blockspergrid = (blockspergrid_x, blockspergrid_y)
+
+        fast_matmul[blockspergrid, threadsperblock](x_d, y_d, z_d)
+        z_h = z_d.copy_to_host()
+        print(z_h)
+        print(x_h @ y_h)
+        # magictoken.ex_run_nonsquare.end
+
+        # nonsquare fast_matmul test(s)
+        msg = "fast_matmul incorrect for shared memory, non-square case."
+        self.assertTrue(np.all(z_h == x_h @ y_h), msg=msg)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,119 @@
+# 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
+
+
+@skip_on_cudasim("cudasim doesn't support cuda import at non-top-level")
+class TestMonteCarlo(CUDATestCase):
+    """
+    Test monte-carlo integration
+    """
+
+    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_montecarlo(self):
+        # ex_montecarlo.import.begin
+        import numba
+        import numpy as np
+        from numba import cuda
+        from numba.cuda.random import (
+            create_xoroshiro128p_states,
+            xoroshiro128p_uniform_float32,
+        )
+        # ex_montecarlo.import.end
+
+        # ex_montecarlo.define.begin
+        # number of samples, higher will lead to a more accurate answer
+        nsamps = 1000000
+        # ex_montecarlo.define.end
+
+        # ex_montecarlo.kernel.begin
+        @cuda.jit
+        def mc_integrator_kernel(out, rng_states, lower_lim, upper_lim):
+            """
+            kernel to draw random samples and evaluate the function to
+            be integrated at those sample values
+            """
+            size = len(out)
+
+            gid = cuda.grid(1)
+            if gid < size:
+                # draw a sample between 0 and 1 on this thread
+                samp = xoroshiro128p_uniform_float32(rng_states, gid)
+
+                # normalize this sample to the limit range
+                samp = samp * (upper_lim - lower_lim) + lower_lim
+
+                # evaluate the function to be
+                # integrated at the normalized
+                # value of the sample
+                y = func(samp)
+                out[gid] = y
+
+        # ex_montecarlo.kernel.end
+
+        # ex_montecarlo.callfunc.begin
+        @cuda.reduce
+        def sum_reduce(a, b):
+            return a + b
+
+        def mc_integrate(lower_lim, upper_lim, nsamps):
+            """
+            approximate the definite integral of `func` from
+            `lower_lim` to `upper_lim`
+            """
+            out = cuda.to_device(np.zeros(nsamps, dtype="float32"))
+            rng_states = create_xoroshiro128p_states(nsamps, seed=42)
+
+            # jit the function for use in CUDA kernels
+
+            mc_integrator_kernel.forall(nsamps)(
+                out, rng_states, lower_lim, upper_lim
+            )
+            # normalization factor to convert
+            # to the average: (b - a)/(N - 1)
+            factor = (upper_lim - lower_lim) / (nsamps - 1)
+
+            return sum_reduce(out) * factor
+
+        # ex_montecarlo.callfunc.end
+
+        # ex_montecarlo.launch.begin
+        # define a function to integrate
+        @numba.jit
+        def func(x):
+            return 1.0 / x
+
+        mc_integrate(1, 2, nsamps)  # array(0.6929643, dtype=float32)
+        mc_integrate(2, 3, nsamps)  # array(0.4054021, dtype=float32)
+        # ex_montecarlo.launch.end
+
+        # values computed independently using maple
+        np.testing.assert_allclose(
+            mc_integrate(1, 2, nsamps), 0.69315, atol=0.001
+        )
+        np.testing.assert_allclose(
+            mc_integrate(2, 3, nsamps), 0.4055, atol=0.001
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,66 @@
+# 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_on_cudasim("cudasim doesn't support cuda import at non-top-level")
+class TestRandom(CUDATestCase):
+    def test_ex_3d_grid(self):
+        # magictoken.ex_3d_grid.begin
+        from numba import cuda
+        from numba.cuda.random import (
+            create_xoroshiro128p_states,
+            xoroshiro128p_uniform_float32,
+        )
+        import numpy as np
+
+        @cuda.jit
+        def random_3d(arr, rng_states):
+            # Per-dimension thread indices and strides
+            startx, starty, startz = cuda.grid(3)
+            stridex, stridey, stridez = cuda.gridsize(3)
+
+            # Linearized thread index
+            tid = (startz * stridey * stridex) + (starty * stridex) + startx
+
+            # Use strided loops over the array to assign a random value to each entry
+            for i in range(startz, arr.shape[0], stridez):
+                for j in range(starty, arr.shape[1], stridey):
+                    for k in range(startx, arr.shape[2], stridex):
+                        arr[i, j, k] = xoroshiro128p_uniform_float32(
+                            rng_states, tid
+                        )
+
+        # Array dimensions
+        X, Y, Z = 701, 900, 719
+
+        # Block and grid dimensions
+        bx, by, bz = 8, 8, 8
+        gx, gy, gz = 16, 16, 16
+
+        # Total number of threads
+        nthreads = bx * by * bz * gx * gy * gz
+
+        # Initialize a state for each thread
+        rng_states = create_xoroshiro128p_states(nthreads, seed=1)
+
+        # Generate random numbers
+        arr = cuda.device_array((X, Y, Z), dtype=np.float32)
+        random_3d[(gx, gy, gz), (bx, by, bz)](arr, rng_states)
+        # magictoken.ex_3d_grid.end
+
+        # Some basic tests of the randomly-generated numbers
+        host_arr = arr.copy_to_host()
+        self.assertGreater(np.mean(host_arr), 0.49)
+        self.assertLess(np.mean(host_arr), 0.51)
+        self.assertTrue(np.all(host_arr <= 1.0))
+        self.assertTrue(np.all(host_arr >= 0.0))
+
+
+if __name__ == "__main__":
+    unittest.main()