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

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

@@ -0,0 +1,193 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from __future__ import print_function
+
+import os
+
+import cffi
+
+import numpy as np
+
+from numba import cuda
+from numba.cuda import int32
+from numba.cuda import config
+from numba.cuda.types import CPointer
+from numba.cuda.testing import (
+    unittest,
+    CUDATestCase,
+    skip_on_cudasim,
+    skip_unless_cc_60,
+    skip_if_cudadevrt_missing,
+)
+from numba.cuda.typing import signature
+
+ffi = cffi.FFI()
+
+
+@cuda.jit
+def this_grid(A):
+    cuda.cg.this_grid()
+    A[0] = 1.0
+
+
+@cuda.jit
+def sync_group(A):
+    g = cuda.cg.this_grid()
+    g.sync()
+    A[0] = 1.0
+
+
+@cuda.jit
+def no_sync(A):
+    A[0] = cuda.grid(1)
+
+
+def sequential_rows(M):
+    # The grid writes rows one at a time. Each thread reads an element from
+    # the previous row written by its "opposite" thread.
+    #
+    # A failure to sync the grid at each row would result in an incorrect
+    # result as some threads could run ahead of threads in other blocks, or
+    # fail to see the update to the previous row from their opposite thread.
+
+    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
+        M[row, col] = M[row - 1, opposite] + 1
+        g.sync()
+
+
+@skip_if_cudadevrt_missing
+class TestCudaCooperativeGroups(CUDATestCase):
+    @skip_unless_cc_60
+    def test_this_grid(self):
+        A = np.full(1, fill_value=np.nan)
+        this_grid[1, 1](A)
+
+        # Ensure the kernel executed beyond the call to cuda.this_grid()
+        self.assertFalse(np.isnan(A[0]), "Value was not set")
+
+    @skip_unless_cc_60
+    @skip_on_cudasim(
+        "Simulator doesn't differentiate between normal and cooperative kernels"
+    )
+    def test_this_grid_is_cooperative(self):
+        A = np.full(1, fill_value=np.nan)
+        this_grid[1, 1](A)
+
+        # this_grid should have been determined to be cooperative
+        for key, overload in this_grid.overloads.items():
+            self.assertTrue(overload.cooperative)
+
+    @skip_unless_cc_60
+    def test_sync_group(self):
+        A = np.full(1, fill_value=np.nan)
+        sync_group[1, 1](A)
+
+        # Ensure the kernel executed beyond the call to cuda.sync_group()
+        self.assertFalse(np.isnan(A[0]), "Value was not set")
+
+    @skip_unless_cc_60
+    @skip_on_cudasim(
+        "Simulator doesn't differentiate between normal and cooperative kernels"
+    )
+    def test_sync_group_is_cooperative(self):
+        A = np.full(1, fill_value=np.nan)
+        sync_group[1, 1](A)
+        # sync_group should have been determined to be cooperative
+        for key, overload in sync_group.overloads.items():
+            self.assertTrue(overload.cooperative)
+
+    @skip_on_cudasim("Simulator does not implement linking")
+    def test_false_cooperative_doesnt_link_cudadevrt(self):
+        """
+        We should only mark a kernel as cooperative and link cudadevrt if the
+        kernel uses grid sync. Here we ensure that one that doesn't use grid
+        synsync isn't marked as such.
+        """
+        A = np.full(1, fill_value=np.nan)
+        no_sync[1, 1](A)
+
+        for key, overload in no_sync.overloads.items():
+            self.assertFalse(overload.cooperative)
+            for link in overload._codelibrary._linking_files:
+                self.assertNotIn("cudadevrt", link)
+
+    @skip_unless_cc_60
+    def test_sync_at_matrix_row(self):
+        if config.ENABLE_CUDASIM:
+            # Use a small matrix to compute using a single block in a
+            # reasonable amount of time
+            shape = (32, 32)
+        else:
+            shape = (1024, 1024)
+        A = np.zeros(shape, dtype=np.int32)
+        blockdim = 32
+        griddim = A.shape[1] // blockdim
+
+        sig = (int32[:, ::1],)
+        c_sequential_rows = cuda.jit(sig)(sequential_rows)
+
+        overload = c_sequential_rows.overloads[sig]
+        mb = overload.max_cooperative_grid_blocks(blockdim)
+        if griddim > mb:
+            unittest.skip("GPU cannot support enough cooperative grid blocks")
+
+        c_sequential_rows[griddim, blockdim](A)
+
+        reference = np.tile(np.arange(shape[0]), (shape[1], 1)).T
+        np.testing.assert_equal(A, reference)
+
+    @skip_unless_cc_60
+    def test_max_cooperative_grid_blocks(self):
+        # The maximum number of blocks will vary based on the device so we
+        # can't test for an expected value, but we can check that the function
+        # doesn't error, and that varying the number of dimensions of the block
+        # whilst keeping the total number of threads constant doesn't change
+        # the maximum to validate some of the logic.
+        sig = (int32[:, ::1],)
+        c_sequential_rows = cuda.jit(sig)(sequential_rows)
+        overload = c_sequential_rows.overloads[sig]
+        blocks1d = overload.max_cooperative_grid_blocks(256)
+        blocks2d = overload.max_cooperative_grid_blocks((16, 16))
+        blocks3d = overload.max_cooperative_grid_blocks((16, 4, 4))
+        self.assertEqual(blocks1d, blocks2d)
+        self.assertEqual(blocks1d, blocks3d)
+
+    @skip_on_cudasim("External code unsupported on cudasim")
+    @skip_unless_cc_60
+    def test_external_cooperative_func(self):
+        cudapy_test_path = os.path.dirname(__file__)
+        tests_path = os.path.dirname(cudapy_test_path)
+        data_path = os.path.join(tests_path, "data")
+        src = os.path.join(data_path, "cta_barrier.cu")
+
+        sig = signature(
+            CPointer(int32),
+        )
+        cta_barrier = cuda.declare_device(
+            "cta_barrier", sig=sig, link=[src], use_cooperative=True
+        )
+
+        @cuda.jit("void()")
+        def kernel():
+            cta_barrier()
+
+        overload = kernel.overloads[()]
+        block_size = 32
+        grid_size = overload.max_cooperative_grid_blocks(block_size)
+
+        kernel[grid_size, block_size]()
+
+        overload = kernel.overloads[()]
+        self.assertTrue(overload.cooperative)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,131 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Copyright (c) 2017 Intel Corporation
+# SPDX-License-Identifier: BSD-2-Clause
+#
+
+from numba.cuda import types, config
+from numba.cuda.core import ir
+from numba.cuda import compiler
+from numba.cuda.core.annotations import type_annotations
+from numba.cuda.core.ir_utils import (
+    copy_propagate,
+    apply_copy_propagate,
+    get_name_var_table,
+)
+from numba.cuda.core.typed_passes import type_inference_stage
+from numba.cuda.testing import CUDATestCase, skip_on_cudasim
+import unittest
+
+
+def _test_will_propagate(b, z, w):
+    x = 3
+    x1 = x
+    if b > 0:
+        y = z + w  # noqa: F821
+    else:
+        y = 0  # noqa: F841
+    a = 2 * x1
+    return a < b
+
+
+def _test_wont_propagate(b, z, w):
+    x = 3
+    if b > 0:
+        y = z + w  # noqa: F841
+        x = 1
+    else:
+        y = 0  # noqa: F841
+    a = 2 * x
+    return a < b
+
+
+def _in_list_var(list_var, var):
+    for i in list_var:
+        if i.name == var:
+            return True
+    return False
+
+
+def _find_assign(func_ir, var):
+    for label, block in func_ir.blocks.items():
+        for i, inst in enumerate(block.body):
+            if isinstance(inst, ir.Assign) and inst.target.name != var:
+                all_var = inst.list_vars()
+                if _in_list_var(all_var, var):
+                    return True
+
+    return False
+
+
+@skip_on_cudasim("cudasim doesn't support run_frontend")
+class TestCopyPropagate(CUDATestCase):
+    def test1(self):
+        from numba.cuda.descriptor import cuda_target
+
+        typingctx = cuda_target.typing_context
+        targetctx = cuda_target.target_context
+        test_ir = compiler.run_frontend(_test_will_propagate)
+        typingctx.refresh()
+        targetctx.refresh()
+        args = (types.int64, types.int64, types.int64)
+        typemap, return_type, calltypes, _ = type_inference_stage(
+            typingctx, targetctx, test_ir, args, None
+        )
+        _ = type_annotations.TypeAnnotation(
+            func_ir=test_ir,
+            typemap=typemap,
+            calltypes=calltypes,
+            lifted=(),
+            lifted_from=None,
+            args=args,
+            return_type=return_type,
+            html_output=config.HTML,
+        )
+        in_cps, out_cps = copy_propagate(test_ir.blocks, typemap)
+        _ = apply_copy_propagate(
+            test_ir.blocks,
+            in_cps,
+            get_name_var_table(test_ir.blocks),
+            typemap,
+            calltypes,
+        )
+
+        self.assertFalse(_find_assign(test_ir, "x1"))
+
+    def test2(self):
+        from numba.cuda.descriptor import cuda_target
+
+        typingctx = cuda_target.typing_context
+        targetctx = cuda_target.target_context
+        test_ir = compiler.run_frontend(_test_wont_propagate)
+        typingctx.refresh()
+        targetctx.refresh()
+        args = (types.int64, types.int64, types.int64)
+        typemap, return_type, calltypes, _ = type_inference_stage(
+            typingctx, targetctx, test_ir, args, None
+        )
+        _ = type_annotations.TypeAnnotation(
+            func_ir=test_ir,
+            typemap=typemap,
+            calltypes=calltypes,
+            lifted=(),
+            lifted_from=None,
+            args=args,
+            return_type=return_type,
+            html_output=config.HTML,
+        )
+        in_cps, out_cps = copy_propagate(test_ir.blocks, typemap)
+        _ = apply_copy_propagate(
+            test_ir.blocks,
+            in_cps,
+            get_name_var_table(test_ir.blocks),
+            typemap,
+            calltypes,
+        )
+
+        self.assertTrue(_find_assign(test_ir, "x"))
+
+
+if __name__ == "__main__":
+    unittest.main()