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

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

@@ -0,0 +1,80 @@
+# 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 TestReduction(CUDATestCase):
+    """
+    Test shared memory reduction
+    """
+
+    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_reduction(self):
+        # ex_reduction.import.begin
+        import numpy as np
+        from numba import cuda
+        from numba.cuda.types import int32
+        # ex_reduction.import.end
+
+        # ex_reduction.allocate.begin
+        # generate data
+        a = cuda.to_device(np.arange(1024))
+        nelem = len(a)
+        # ex_reduction.allocate.end
+
+        # ex_reduction.kernel.begin
+        @cuda.jit
+        def array_sum(data):
+            tid = cuda.threadIdx.x
+            size = len(data)
+            if tid < size:
+                i = cuda.grid(1)
+
+                # Declare an array in shared memory
+                shr = cuda.shared.array(nelem, int32)
+                shr[tid] = data[i]
+
+                # Ensure writes to shared memory are visible
+                # to all threads before reducing
+                cuda.syncthreads()
+
+                s = 1
+                while s < cuda.blockDim.x:
+                    if tid % (2 * s) == 0:
+                        # Stride by `s` and add
+                        shr[tid] += shr[tid + s]
+                    s *= 2
+                    cuda.syncthreads()
+
+                # After the loop, the zeroth  element contains the sum
+                if tid == 0:
+                    data[tid] = shr[tid]
+
+        # ex_reduction.kernel.end
+
+        # ex_reduction.launch.begin
+        array_sum[1, nelem](a)
+        print(a[0])  # 523776
+        print(sum(np.arange(1024)))  # 523776
+        # ex_reduction.launch.end
+
+        np.testing.assert_equal(a[0], sum(np.arange(1024)))
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,206 @@
+# 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 TestSessionization(CUDATestCase):
+    """
+    Test click stream sessionization
+    """
+
+    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_sessionize(self):
+        # ex_sessionize.import.begin
+        import numpy as np
+        from numba import cuda
+
+        # Set the timeout to one hour
+        session_timeout = np.int64(np.timedelta64("3600", "s"))
+        # ex_sessionize.import.end
+
+        # ex_sessionize.allocate.begin
+        # Generate data
+        ids = cuda.to_device(
+            np.array(
+                [
+                    1,
+                    1,
+                    1,
+                    1,
+                    1,
+                    1,
+                    2,
+                    2,
+                    2,
+                    3,
+                    3,
+                    3,
+                    3,
+                    3,
+                    3,
+                    3,
+                    3,
+                    3,
+                    3,
+                    4,
+                    4,
+                    4,
+                    4,
+                    4,
+                    4,
+                    4,
+                    4,
+                    4,
+                ]
+            )
+        )
+        sec = cuda.to_device(
+            np.array(
+                [
+                    1,
+                    2,
+                    3,
+                    5000,
+                    5001,
+                    5002,
+                    1,
+                    2,
+                    3,
+                    1,
+                    2,
+                    5000,
+                    5001,
+                    10000,
+                    10001,
+                    10002,
+                    10003,
+                    15000,
+                    150001,
+                    1,
+                    5000,
+                    50001,
+                    15000,
+                    20000,
+                    25000,
+                    25001,
+                    25002,
+                    25003,
+                ],
+                dtype="datetime64[ns]",
+            ).astype("int64")  # Cast to int64 for compatibility
+        )
+        # Create a vector to hold the results
+        results = cuda.to_device(np.zeros(len(ids)))
+        # ex_sessionize.allocate.end
+
+        # ex_sessionize.kernel.begin
+        @cuda.jit
+        def sessionize(user_id, timestamp, results):
+            gid = cuda.grid(1)
+            size = len(user_id)
+
+            if gid >= size:
+                return
+
+            # Determine session boundaries
+            is_first_datapoint = gid == 0
+            if not is_first_datapoint:
+                new_user = user_id[gid] != user_id[gid - 1]
+                timed_out = (
+                    timestamp[gid] - timestamp[gid - 1] > session_timeout
+                )
+                is_sess_boundary = new_user or timed_out
+            else:
+                is_sess_boundary = True
+
+            # Determine session labels
+            if is_sess_boundary:
+                # This thread marks the start of a session
+                results[gid] = gid
+
+                # Make sure all session boundaries are written
+                # before populating the session id
+                grid = cuda.cg.this_grid()
+                grid.sync()
+
+                look_ahead = 1
+                # Check elements 'forward' of this one
+                # until a new session boundary is found
+                while results[gid + look_ahead] == 0:
+                    results[gid + look_ahead] = gid
+                    look_ahead += 1
+                    # Avoid out-of-bounds accesses by the last thread
+                    if gid + look_ahead == size - 1:
+                        results[gid + look_ahead] = gid
+                        break
+
+        # ex_sessionize.kernel.end
+
+        # ex_sessionize.launch.begin
+        sessionize.forall(len(ids))(ids, sec, results)
+
+        print(results.copy_to_host())
+        # array([ 0.,  0.,  0.,  3.,  3.,  3.,
+        #         6.,  6.,  6.,  9.,  9., 11.,
+        #         11., 13., 13., 13., 13., 17.,
+        #         18., 19., 20., 21., 21., 23.,
+        #         24., 24., 24., 24.])
+        # ex_sessionize.launch.end
+
+        expect = [
+            0,
+            0,
+            0,
+            3,
+            3,
+            3,
+            6,
+            6,
+            6,
+            9,
+            9,
+            11,
+            11,
+            13,
+            13,
+            13,
+            13,
+            17,
+            18,
+            19,
+            20,
+            21,
+            21,
+            23,
+            24,
+            24,
+            24,
+            24,
+        ]
+        np.testing.assert_equal(expect, results.copy_to_host())
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,53 @@
+# 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 TestUFunc(CUDATestCase):
+    """
+    Test calling a UFunc
+    """
+
+    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_cuda_ufunc_call(self):
+        # ex_cuda_ufunc.begin
+        import numpy as np
+        from numba import cuda
+
+        # A kernel calling a ufunc (sin, in this case)
+        @cuda.jit
+        def f(r, x):
+            # Compute sin(x) with result written to r
+            np.sin(x, r)
+
+        # Declare input and output arrays
+        x = np.arange(10, dtype=np.float32) - 5
+        r = np.zeros_like(x)
+
+        # Launch kernel that calls the ufunc
+        f[1, 1](r, x)
+
+        # A quick sanity check demonstrating equality of the sine computed by
+        # the sin ufunc inside the kernel, and NumPy's sin ufunc
+        np.testing.assert_allclose(r, np.sin(x))
+        # ex_cuda_ufunc.end
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,76 @@
+# 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 TestVecAdd(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_vecadd(self):
+        # ex_vecadd.import.begin
+        import numpy as np
+        from numba import cuda
+        # ex_vecadd.import.end
+
+        # ex_vecadd.kernel.begin
+        @cuda.jit
+        def f(a, b, c):
+            # like threadIdx.x + (blockIdx.x * blockDim.x)
+            tid = cuda.grid(1)
+            size = len(c)
+
+            if tid < size:
+                c[tid] = a[tid] + b[tid]
+
+        # ex_vecadd.kernel.end
+
+        # Seed RNG for test repeatability
+        np.random.seed(1)
+
+        # ex_vecadd.allocate.begin
+        N = 100000
+        a = cuda.to_device(np.random.random(N))
+        b = cuda.to_device(np.random.random(N))
+        c = cuda.device_array_like(a)
+        # ex_vecadd.allocate.end
+
+        # ex_vecadd.forall.begin
+        f.forall(len(a))(a, b, c)
+        print(c.copy_to_host())
+        # ex_vecadd.forall.end
+
+        # ex_vecadd.launch.begin
+        # Enough threads per block for several warps per block
+        nthreads = 256
+        # Enough blocks to cover the entire vector depending on its length
+        nblocks = (len(a) // nthreads) + 1
+        f[nblocks, nthreads](a, b, c)
+        print(c.copy_to_host())
+        # ex_vecadd.launch.end
+
+        np.testing.assert_equal(
+            c.copy_to_host(), a.copy_to_host() + b.copy_to_host()
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

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