numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

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

@@ -0,0 +1,206 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import pytest
+import numpy as np
+from numba import cuda
+from numba.cuda.cudadrv import driver
+from numba.cuda.testing import (
+    unittest,
+    CUDATestCase,
+    skip_on_cudasim,
+)
+from numba.cuda.tests.support import (
+    linux_only,
+    override_config,
+    run_in_subprocess,
+)
+from numba.cuda.core.errors import (
+    NumbaPerformanceWarning,
+    NumbaInvalidConfigWarning,
+)
+from numba.cuda.core import config
+import warnings
+
+
+@skip_on_cudasim("cudasim does not raise performance warnings")
+class TestWarnings(CUDATestCase):
+    def test_float16_warn_if_lto_missing(self):
+        fp16_kernel_invocation = """
+import math
+from numba import cuda
+
+@cuda.jit
+def kernel():
+    x = cuda.types.float16(1.0)
+    y = math.sin(x)
+
+kernel[1,1]()
+kernel[1,1]()
+"""
+        performance_warning = "float16 relies on LTO for performance"
+        expected_warning_count = 0 if driver._have_nvjitlink() else 1
+        _, err = run_in_subprocess(fp16_kernel_invocation)
+        self.assertEqual(
+            err.decode().count(performance_warning), expected_warning_count
+        )
+
+    def test_inefficient_launch_configuration(self):
+        @cuda.jit
+        def kernel():
+            pass
+
+        with override_config("CUDA_LOW_OCCUPANCY_WARNINGS", 1):
+            with pytest.warns(
+                NumbaPerformanceWarning, match="Grid size .+ low occupancy"
+            ):
+                func = kernel[1, 1]
+        func()
+
+    def test_efficient_launch_configuration(self):
+        @cuda.jit
+        def kernel():
+            pass
+
+        with override_config("CUDA_LOW_OCCUPANCY_WARNINGS", 1):
+            with warnings.catch_warnings(record=True) as w:
+                kernel[256, 256]()
+
+        self.assertEqual(len(w), 0)
+
+    def test_warn_on_host_array(self):
+        @cuda.jit
+        def foo(r, x):
+            r[0] = x + 1
+
+        N = 10
+        arr_f32 = np.zeros(N, dtype=np.float32)
+        func = foo[1, N]
+        with override_config("CUDA_WARN_ON_IMPLICIT_COPY", 1):
+            with pytest.warns(
+                NumbaPerformanceWarning,
+                match="Host array used in CUDA kernel will incur.+copy overhead",
+            ):
+                func(arr_f32, N)
+
+    def test_pinned_warn_on_host_array(self):
+        @cuda.jit
+        def foo(r, x):
+            r[0] = x + 1
+
+        N = 10
+        ary = cuda.pinned_array(N, dtype=np.float32)
+
+        func = foo[1, N]
+        with override_config("CUDA_WARN_ON_IMPLICIT_COPY", 1):
+            with pytest.warns(
+                NumbaPerformanceWarning,
+                match="Host array used in CUDA kernel will incur.+copy overhead",
+            ):
+                func(ary, N)
+
+    def test_nowarn_on_mapped_array(self):
+        @cuda.jit
+        def foo(r, x):
+            r[0] = x + 1
+
+        N = 10
+        ary = cuda.mapped_array(N, dtype=np.float32)
+
+        with override_config("CUDA_WARN_ON_IMPLICIT_COPY", 1):
+            with warnings.catch_warnings(record=True) as w:
+                foo[1, N](ary, N)
+
+        self.assertEqual(len(w), 0)
+
+    @linux_only
+    def test_nowarn_on_managed_array(self):
+        @cuda.jit
+        def foo(r, x):
+            r[0] = x + 1
+
+        N = 10
+        ary = cuda.managed_array(N, dtype=np.float32)
+
+        with override_config("CUDA_WARN_ON_IMPLICIT_COPY", 1):
+            with warnings.catch_warnings(record=True) as w:
+                foo[1, N](ary, N)
+
+        self.assertEqual(len(w), 0)
+
+    def test_nowarn_on_device_array(self):
+        @cuda.jit
+        def foo(r, x):
+            r[0] = x + 1
+
+        N = 10
+        ary = cuda.device_array(N, dtype=np.float32)
+
+        with override_config("CUDA_WARN_ON_IMPLICIT_COPY", 1):
+            with warnings.catch_warnings(record=True) as w:
+                foo[1, N](ary, N)
+
+        self.assertEqual(len(w), 0)
+
+    def test_warn_on_debug_and_opt(self):
+        with pytest.warns(
+            NumbaInvalidConfigWarning, match="not supported by CUDA"
+        ):
+            cuda.jit(debug=True, opt=True)
+
+    def test_warn_on_debug_and_opt_default(self):
+        with pytest.warns(
+            NumbaInvalidConfigWarning, match="not supported by CUDA"
+        ):
+            cuda.jit(debug=True)
+
+    def test_no_warn_on_debug_and_no_opt(self):
+        with warnings.catch_warnings(record=True) as w:
+            cuda.jit(debug=True, opt=False)
+
+        self.assertEqual(len(w), 0)
+
+    def test_no_warn_with_no_debug_and_opt_kwargs(self):
+        with warnings.catch_warnings(record=True) as w:
+            cuda.jit()
+
+        self.assertEqual(len(w), 0)
+
+    def test_no_warn_on_debug_and_opt_with_config(self):
+        with override_config("CUDA_DEBUGINFO_DEFAULT", 1):
+            with override_config("OPT", config._OptLevel(0)):
+                with warnings.catch_warnings(record=True) as w:
+                    cuda.jit()
+
+            self.assertEqual(len(w), 0)
+
+            with warnings.catch_warnings(record=True) as w:
+                cuda.jit(opt=False)
+
+            self.assertEqual(len(w), 0)
+
+        with override_config("OPT", config._OptLevel(0)):
+            with warnings.catch_warnings(record=True) as w:
+                cuda.jit(debug=True)
+
+            self.assertEqual(len(w), 0)
+
+    def test_warn_on_debug_and_opt_with_config(self):
+        with override_config("CUDA_DEBUGINFO_DEFAULT", 1):
+            for opt in (1, 2, 3, "max"):
+                with override_config("OPT", config._OptLevel(opt)):
+                    with pytest.warns(
+                        NumbaInvalidConfigWarning, match="not supported by CUDA"
+                    ):
+                        cuda.jit()
+
+        for opt in (1, 2, 3, "max"):
+            with override_config("OPT", config._OptLevel(opt)):
+                with pytest.warns(
+                    NumbaInvalidConfigWarning, match="not supported by CUDA"
+                ):
+                    cuda.jit(debug=True)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,331 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import re
+
+import numpy as np
+from numba import cuda
+from numba.cuda import int32, int64, float32, float64
+from numba.cuda.testing import unittest, CUDATestCase, skip_on_cudasim
+from numba.cuda.compiler import compile_ptx
+from numba.cuda.core import config
+
+
+def useful_syncwarp(ary):
+    i = cuda.grid(1)
+    if i == 0:
+        ary[0] = 42
+    cuda.syncwarp(0xFFFFFFFF)
+    ary[i] = ary[0]
+
+
+def use_shfl_sync_idx(ary, idx):
+    i = cuda.grid(1)
+    val = cuda.shfl_sync(0xFFFFFFFF, i, idx)
+    ary[i] = val
+
+
+def use_shfl_sync_up(ary, delta):
+    i = cuda.grid(1)
+    val = cuda.shfl_up_sync(0xFFFFFFFF, i, delta)
+    ary[i] = val
+
+
+def use_shfl_sync_down(ary, delta):
+    i = cuda.grid(1)
+    val = cuda.shfl_down_sync(0xFFFFFFFF, i, delta)
+    ary[i] = val
+
+
+def use_shfl_sync_xor(ary, xor):
+    i = cuda.grid(1)
+    val = cuda.shfl_xor_sync(0xFFFFFFFF, i, xor)
+    ary[i] = val
+
+
+def use_shfl_sync_with_val(ary, into):
+    i = cuda.grid(1)
+    val = cuda.shfl_sync(0xFFFFFFFF, into, 0)
+    ary[i] = val
+
+
+def use_vote_sync_all(ary_in, ary_out):
+    i = cuda.grid(1)
+    pred = cuda.all_sync(0xFFFFFFFF, ary_in[i])
+    ary_out[i] = pred
+
+
+def use_vote_sync_any(ary_in, ary_out):
+    i = cuda.grid(1)
+    pred = cuda.any_sync(0xFFFFFFFF, ary_in[i])
+    ary_out[i] = pred
+
+
+def use_vote_sync_eq(ary_in, ary_out):
+    i = cuda.grid(1)
+    pred = cuda.eq_sync(0xFFFFFFFF, ary_in[i])
+    ary_out[i] = pred
+
+
+def use_vote_sync_ballot(ary):
+    i = cuda.threadIdx.x
+    ballot = cuda.ballot_sync(0xFFFFFFFF, True)
+    ary[i] = ballot
+
+
+def use_match_any_sync(ary_in, ary_out):
+    i = cuda.grid(1)
+    ballot = cuda.match_any_sync(0xFFFFFFFF, ary_in[i])
+    ary_out[i] = ballot
+
+
+def use_match_all_sync(ary_in, ary_out):
+    i = cuda.grid(1)
+    ballot, pred = cuda.match_all_sync(0xFFFFFFFF, ary_in[i])
+    ary_out[i] = ballot if pred else 0
+
+
+def use_independent_scheduling(arr):
+    i = cuda.threadIdx.x
+    if i % 4 == 0:
+        ballot = cuda.ballot_sync(0x11111111, True)
+    elif i % 4 == 1:
+        ballot = cuda.ballot_sync(0x22222222, True)
+    elif i % 4 == 2:
+        ballot = cuda.ballot_sync(0x44444444, True)
+    elif i % 4 == 3:
+        ballot = cuda.ballot_sync(0x88888888, True)
+    arr[i] = ballot
+
+
+def _safe_cc_check(cc):
+    if config.ENABLE_CUDASIM:
+        return True
+    else:
+        return cuda.get_current_device().compute_capability >= cc
+
+
+@skip_on_cudasim("Warp Operations are not yet implemented on cudasim")
+class TestCudaWarpOperations(CUDATestCase):
+    def test_useful_syncwarp(self):
+        compiled = cuda.jit("void(int32[:])")(useful_syncwarp)
+        nelem = 32
+        ary = np.empty(nelem, dtype=np.int32)
+        compiled[1, nelem](ary)
+        self.assertTrue(np.all(ary == 42))
+
+    def test_shfl_sync_idx(self):
+        compiled = cuda.jit("void(int32[:], int32)")(use_shfl_sync_idx)
+        nelem = 32
+        idx = 4
+        ary = np.empty(nelem, dtype=np.int32)
+        compiled[1, nelem](ary, idx)
+        self.assertTrue(np.all(ary == idx))
+
+    def test_shfl_sync_up(self):
+        compiled = cuda.jit("void(int32[:], int32)")(use_shfl_sync_up)
+        nelem = 32
+        delta = 4
+        ary = np.empty(nelem, dtype=np.int32)
+        exp = np.arange(nelem, dtype=np.int32)
+        exp[delta:] -= delta
+        compiled[1, nelem](ary, delta)
+        self.assertTrue(np.all(ary == exp))
+
+    def test_shfl_sync_down(self):
+        compiled = cuda.jit("void(int32[:], int32)")(use_shfl_sync_down)
+        nelem = 32
+        delta = 4
+        ary = np.empty(nelem, dtype=np.int32)
+        exp = np.arange(nelem, dtype=np.int32)
+        exp[:-delta] += delta
+        compiled[1, nelem](ary, delta)
+        self.assertTrue(np.all(ary == exp))
+
+    def test_shfl_sync_xor(self):
+        compiled = cuda.jit("void(int32[:], int32)")(use_shfl_sync_xor)
+        nelem = 32
+        xor = 16
+        ary = np.empty(nelem, dtype=np.int32)
+        exp = np.arange(nelem, dtype=np.int32) ^ xor
+        compiled[1, nelem](ary, xor)
+        self.assertTrue(np.all(ary == exp))
+
+    def test_shfl_sync_const_mode_val(self):
+        # Test `mode` argument is constant in shfl_sync calls.
+        # Related to https://github.com/NVIDIA/numba-cuda/pull/231
+        subtest = [
+            (use_shfl_sync_idx, 4),
+            (use_shfl_sync_up, 4),
+            (use_shfl_sync_down, 4),
+            (use_shfl_sync_xor, 16),
+        ]
+
+        args_re = r"\((.*)\)"
+        m = re.compile(args_re)
+
+        for func, value in subtest:
+            with self.subTest(func=func.__name__):
+                compiled = cuda.jit("void(int32[:], int32)")(func)
+                nelem = 32
+                ary = np.empty(nelem, dtype=np.int32)
+                compiled[1, nelem](ary, value)
+                irs = next(iter(compiled.inspect_llvm().values()))
+
+                for ir in irs.split("\n"):
+                    if "call" in ir and "llvm.nvvm.shfl.sync.i32" in ir:
+                        args = m.search(ir).group(0)
+                        arglist = args.split(",")
+                        mode_arg = arglist[1]
+                        self.assertNotIn("%", mode_arg)
+
+    def test_shfl_sync_const_mode_val_sm100(self):
+        # Test shfl_sync compiles with cc=(10, 0)
+        subtest = [
+            use_shfl_sync_idx,
+            use_shfl_sync_up,
+            use_shfl_sync_down,
+            use_shfl_sync_xor,
+        ]
+
+        for func in subtest:
+            with self.subTest(func=func.__name__):
+                compile_ptx(func, (int32[:], int32), cc=(10, 0))
+
+    def test_shfl_sync_types(self):
+        types = int32, int64, float32, float64
+        values = (
+            np.int32(-1),
+            np.int64(1 << 42),
+            np.float32(np.pi),
+            np.float64(np.pi),
+        )
+        for typ, val in zip(types, values):
+            with self.subTest(typ=typ):
+                compiled = cuda.jit((typ[:], typ))(use_shfl_sync_with_val)
+                nelem = 32
+                ary = np.empty(nelem, dtype=val.dtype)
+                compiled[1, nelem](ary, val)
+                self.assertTrue(np.all(ary == val))
+
+    def test_vote_sync_all(self):
+        compiled = cuda.jit("void(int32[:], int32[:])")(use_vote_sync_all)
+        nelem = 32
+        ary_in = np.ones(nelem, dtype=np.int32)
+        ary_out = np.empty(nelem, dtype=np.int32)
+        compiled[1, nelem](ary_in, ary_out)
+        self.assertTrue(np.all(ary_out == 1))
+        ary_in[-1] = 0
+        compiled[1, nelem](ary_in, ary_out)
+        self.assertTrue(np.all(ary_out == 0))
+
+    def test_vote_sync_any(self):
+        compiled = cuda.jit("void(int32[:], int32[:])")(use_vote_sync_any)
+        nelem = 32
+        ary_in = np.zeros(nelem, dtype=np.int32)
+        ary_out = np.empty(nelem, dtype=np.int32)
+        compiled[1, nelem](ary_in, ary_out)
+        self.assertTrue(np.all(ary_out == 0))
+        ary_in[2] = 1
+        ary_in[5] = 1
+        compiled[1, nelem](ary_in, ary_out)
+        self.assertTrue(np.all(ary_out == 1))
+
+    def test_vote_sync_eq(self):
+        compiled = cuda.jit("void(int32[:], int32[:])")(use_vote_sync_eq)
+        nelem = 32
+        ary_in = np.zeros(nelem, dtype=np.int32)
+        ary_out = np.empty(nelem, dtype=np.int32)
+        compiled[1, nelem](ary_in, ary_out)
+        self.assertTrue(np.all(ary_out == 1))
+        ary_in[1] = 1
+        compiled[1, nelem](ary_in, ary_out)
+        self.assertTrue(np.all(ary_out == 0))
+        ary_in[:] = 1
+        compiled[1, nelem](ary_in, ary_out)
+        self.assertTrue(np.all(ary_out == 1))
+
+    def test_vote_sync_ballot(self):
+        compiled = cuda.jit("void(uint32[:])")(use_vote_sync_ballot)
+        nelem = 32
+        ary = np.empty(nelem, dtype=np.uint32)
+        compiled[1, nelem](ary)
+        self.assertTrue(np.all(ary == np.uint32(0xFFFFFFFF)))
+
+    @unittest.skipUnless(
+        _safe_cc_check((7, 0)), "Matching requires at least Volta Architecture"
+    )
+    def test_match_any_sync(self):
+        compiled = cuda.jit("void(int32[:], int32[:])")(use_match_any_sync)
+        nelem = 10
+        ary_in = np.arange(nelem, dtype=np.int32) % 2
+        ary_out = np.empty(nelem, dtype=np.int32)
+        exp = np.tile((0b0101010101, 0b1010101010), 5)
+        compiled[1, nelem](ary_in, ary_out)
+        self.assertTrue(np.all(ary_out == exp))
+
+    @unittest.skipUnless(
+        _safe_cc_check((7, 0)), "Matching requires at least Volta Architecture"
+    )
+    def test_match_all_sync(self):
+        compiled = cuda.jit("void(int32[:], int32[:])")(use_match_all_sync)
+        nelem = 10
+        ary_in = np.zeros(nelem, dtype=np.int32)
+        ary_out = np.empty(nelem, dtype=np.int32)
+        compiled[1, nelem](ary_in, ary_out)
+        self.assertTrue(np.all(ary_out == 0b1111111111))
+        ary_in[1] = 4
+        compiled[1, nelem](ary_in, ary_out)
+        self.assertTrue(np.all(ary_out == 0))
+
+    @unittest.skipUnless(
+        _safe_cc_check((7, 0)),
+        "Independent scheduling requires at least Volta Architecture",
+    )
+    def test_independent_scheduling(self):
+        compiled = cuda.jit("void(uint32[:])")(use_independent_scheduling)
+        arr = np.empty(32, dtype=np.uint32)
+        exp = np.tile((0x11111111, 0x22222222, 0x44444444, 0x88888888), 8)
+        compiled[1, 32](arr)
+        self.assertTrue(np.all(arr == exp))
+
+    def test_activemask(self):
+        @cuda.jit
+        def use_activemask(x):
+            i = cuda.grid(1)
+            if (i % 2) == 0:
+                # Even numbered threads fill in even numbered array entries
+                # with binary "...01010101"
+                x[i] = cuda.activemask()
+            else:
+                # Odd numbered threads fill in odd numbered array entries
+                # with binary "...10101010"
+                x[i] = cuda.activemask()
+
+        out = np.zeros(32, dtype=np.uint32)
+        use_activemask[1, 32](out)
+
+        # 0x5 = 0101: The pattern from even-numbered threads
+        # 0xA = 1010: The pattern from odd-numbered threads
+        expected = np.tile((0x55555555, 0xAAAAAAAA), 16)
+        np.testing.assert_equal(expected, out)
+
+    def test_lanemask_lt(self):
+        @cuda.jit
+        def use_lanemask_lt(x):
+            i = cuda.grid(1)
+            x[i] = cuda.lanemask_lt()
+
+        out = np.zeros(32, dtype=np.uint32)
+        use_lanemask_lt[1, 32](out)
+
+        # A string of 1s that grows from the LSB for each entry:
+        # 0, 1, 3, 7, F, 1F, 3F, 7F, FF, 1FF, etc.
+        # or in binary:
+        # ...0001, ....0011, ...0111, etc.
+        expected = np.asarray([(2**i) - 1 for i in range(32)], dtype=np.uint32)
+        np.testing.assert_equal(expected, out)
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/cudasim/__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/cudasim/support.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba import cuda
+
+
+@cuda.jit(device=True)
+def cuda_module_in_device_function():
+    return cuda.threadIdx.x

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;
+}