numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

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

@@ -0,0 +1,276 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.testing import CUDATestCase, skip_on_cudasim
+from numba.cuda.core.compiler import CompilerBase
+from numba.cuda.flags import Flags
+from numba.cuda.core.compiler_machinery import PassManager
+from numba.cuda.core import ir_utils
+from numba.cuda import types
+from numba.cuda.core import ir, bytecode
+from numba.cuda import compiler
+from numba.cuda.core.untyped_passes import (
+    ExtractByteCode,
+    TranslateByteCode,
+    FixupArgs,
+    IRProcessing,
+)
+from numba.cuda.core.typed_passes import (
+    NopythonTypeInference,
+    DeadCodeElimination,
+)
+
+
+# global constant for testing find_const
+GLOBAL_B = 11
+
+
+@skip_on_cudasim("Requires CUDA target")
+class TestIrUtils(CUDATestCase):
+    """
+    Tests ir handling utility functions like find_callname.
+    """
+
+    def test_dead_code_elimination(self):
+        class Tester(CompilerBase):
+            @classmethod
+            def mk_pipeline(
+                cls,
+                args,
+                return_type=None,
+                flags=None,
+                locals=None,
+                library=None,
+                typing_context=None,
+                target_context=None,
+            ):
+                if locals is None:
+                    locals = {}
+                if not flags:
+                    flags = Flags()
+
+                if typing_context is None:
+                    from numba.cuda.descriptor import cuda_target
+
+                    typing_context = cuda_target.typing_context
+                if target_context is None:
+                    from numba.cuda.descriptor import cuda_target
+
+                    target_context = cuda_target.target_context
+                return cls(
+                    typing_context,
+                    target_context,
+                    library,
+                    args,
+                    return_type,
+                    flags,
+                    locals,
+                )
+
+            def compile_to_ir(self, func, DCE=False):
+                """
+                Compile and return IR
+                """
+                func_id = bytecode.FunctionIdentity.from_function(func)
+                self.state.func_id = func_id
+                ExtractByteCode().run_pass(self.state)
+                state = self.state
+
+                name = "DCE_testing"
+                pm = PassManager(name)
+                pm.add_pass(TranslateByteCode, "analyzing bytecode")
+                pm.add_pass(FixupArgs, "fix up args")
+                pm.add_pass(IRProcessing, "processing IR")
+                pm.add_pass(NopythonTypeInference, "nopython frontend")
+                if DCE is True:
+                    pm.add_pass(DeadCodeElimination, "DCE after typing")
+                pm.finalize()
+                pm.run(state)
+                return state.func_ir
+
+        def check_initial_ir(the_ir):
+            # dead stuff:
+            # a const int value 0xdead
+            # an assign of above into to variable `dead`
+            # a const int above 0xdeaddead
+            # an assign of said int to variable `deaddead`
+            # this is 2 statements to remove
+
+            self.assertEqual(len(the_ir.blocks), 1)
+            block = the_ir.blocks[0]
+            deads = []
+            for x in block.find_insts(ir.Assign):
+                if isinstance(getattr(x, "target", None), ir.Var):
+                    if "dead" in getattr(x.target, "name", ""):
+                        deads.append(x)
+
+            self.assertEqual(len(deads), 2)
+            for d in deads:
+                # check the ir.Const is the definition and the value is expected
+                const_val = the_ir.get_definition(d.value)
+                self.assertTrue(
+                    int("0x%s" % d.target.name, 16), const_val.value
+                )
+
+            return deads
+
+        def check_dce_ir(the_ir):
+            self.assertEqual(len(the_ir.blocks), 1)
+            block = the_ir.blocks[0]
+            deads = []
+            consts = []
+            for x in block.find_insts(ir.Assign):
+                if isinstance(getattr(x, "target", None), ir.Var):
+                    if "dead" in getattr(x.target, "name", ""):
+                        deads.append(x)
+                if isinstance(getattr(x, "value", None), ir.Const):
+                    consts.append(x)
+            self.assertEqual(len(deads), 0)
+
+            # check the consts to make sure there's no reference to 0xdead or
+            # 0xdeaddead
+            for x in consts:
+                self.assertTrue(x.value.value not in [0xDEAD, 0xDEADDEAD])
+
+        def foo(x):
+            y = x + 1
+            dead = 0xDEAD  # noqa
+            z = y + 2
+            deaddead = 0xDEADDEAD  # noqa
+            ret = z * z
+            return ret
+
+        test_pipeline = Tester.mk_pipeline((types.intp,))
+        no_dce = test_pipeline.compile_to_ir(foo)
+        removed = check_initial_ir(no_dce)
+
+        test_pipeline = Tester.mk_pipeline((types.intp,))
+        w_dce = test_pipeline.compile_to_ir(foo, DCE=True)
+        check_dce_ir(w_dce)
+
+        # check that the count of initial - removed = dce
+        self.assertEqual(
+            len(no_dce.blocks[0].body) - len(removed), len(w_dce.blocks[0].body)
+        )
+
+    @skip_on_cudasim("Skipping ir utils tests on CUDA simulator")
+    def test_find_const_global(self):
+        """
+        Test find_const() for values in globals (ir.Global) and freevars
+        (ir.FreeVar) that are considered constants for compilation.
+        """
+        FREEVAR_C = 12
+
+        def foo(a):
+            b = GLOBAL_B
+            c = FREEVAR_C
+            return a + b + c
+
+        f_ir = compiler.run_frontend(foo)
+        block = f_ir.blocks[0]
+        const_b = None
+        const_c = None
+
+        for inst in block.body:
+            if isinstance(inst, ir.Assign) and inst.target.name == "b":
+                const_b = ir_utils.guard(ir_utils.find_const, f_ir, inst.target)
+            if isinstance(inst, ir.Assign) and inst.target.name == "c":
+                const_c = ir_utils.guard(ir_utils.find_const, f_ir, inst.target)
+
+        self.assertEqual(const_b, GLOBAL_B)
+        self.assertEqual(const_c, FREEVAR_C)
+
+    @skip_on_cudasim("Skipping ir utils tests on CUDA simulator")
+    def test_flatten_labels(self):
+        """tests flatten_labels"""
+
+        def foo(a):
+            acc = 0
+            if a > 3:
+                acc += 1
+                if a > 19:
+                    return 53
+            elif a < 1000:
+                if a >= 12:
+                    acc += 1
+                for x in range(10):
+                    acc -= 1
+                    if acc < 2:
+                        break
+                else:
+                    acc += 7
+            else:
+                raise ValueError("some string")
+            # prevents inline of return on py310
+            py310_defeat1 = 1  # noqa
+            py310_defeat2 = 2  # noqa
+            py310_defeat3 = 3  # noqa
+            py310_defeat4 = 4  # noqa
+            return acc
+
+        def bar(a):
+            acc = 0
+            z = 12
+            if a > 3:
+                acc += 1
+                z += 12
+                if a > 19:
+                    z += 12
+                    return 53
+            elif a < 1000:
+                if a >= 12:
+                    z += 12
+                    acc += 1
+                for x in range(10):
+                    z += 12
+                    acc -= 1
+                    if acc < 2:
+                        break
+                else:
+                    z += 12
+                    acc += 7
+            else:
+                raise ValueError("some string")
+            py310_defeat1 = 1  # noqa
+            py310_defeat2 = 2  # noqa
+            py310_defeat3 = 3  # noqa
+            py310_defeat4 = 4  # noqa
+            return acc
+
+        def baz(a):
+            acc = 0
+            if a > 3:
+                acc += 1
+                if a > 19:
+                    return 53
+                else:  # extra control flow in comparison to foo
+                    return 55
+            elif a < 1000:
+                if a >= 12:
+                    acc += 1
+                for x in range(10):
+                    acc -= 1
+                    if acc < 2:
+                        break
+                else:
+                    acc += 7
+            else:
+                raise ValueError("some string")
+            py310_defeat1 = 1  # noqa
+            py310_defeat2 = 2  # noqa
+            py310_defeat3 = 3  # noqa
+            py310_defeat4 = 4  # noqa
+            return acc
+
+        def get_flat_cfg(func):
+            func_ir = ir_utils.compile_to_numba_ir(func, dict())
+            flat_blocks = ir_utils.flatten_labels(func_ir.blocks)
+            self.assertEqual(max(flat_blocks.keys()) + 1, len(func_ir.blocks))
+            return ir_utils.compute_cfg_from_blocks(flat_blocks)
+
+        foo_cfg = get_flat_cfg(foo)
+        bar_cfg = get_flat_cfg(bar)
+        baz_cfg = get_flat_cfg(baz)
+
+        self.assertEqual(foo_cfg, bar_cfg)
+        self.assertNotEqual(foo_cfg, baz_cfg)

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

@@ -0,0 +1,101 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba import cuda
+from numba.cuda.testing import unittest, CUDATestCase
+
+import numpy as np
+
+
+class TestIterators(CUDATestCase):
+    def test_enumerate(self):
+        @cuda.jit
+        def enumerator(x, error):
+            count = 0
+
+            for i, v in enumerate(x):
+                if count != i:
+                    error[0] = 1
+                if v != x[i]:
+                    error[0] = 2
+
+                count += 1
+
+            if count != len(x):
+                error[0] = 3
+
+        x = np.asarray((10, 9, 8, 7, 6))
+        error = np.zeros(1, dtype=np.int32)
+
+        enumerator[1, 1](x, error)
+        self.assertEqual(error[0], 0)
+
+    def _test_twoarg_function(self, f):
+        x = np.asarray((10, 9, 8, 7, 6))
+        y = np.asarray((1, 2, 3, 4, 5))
+        error = np.zeros(1, dtype=np.int32)
+
+        f[1, 1](x, y, error)
+        self.assertEqual(error[0], 0)
+
+    def test_zip(self):
+        @cuda.jit
+        def zipper(x, y, error):
+            i = 0
+
+            for xv, yv in zip(x, y):
+                if xv != x[i]:
+                    error[0] = 1
+                if yv != y[i]:
+                    error[0] = 2
+
+                i += 1
+
+            if i != len(x):
+                error[0] = 3
+
+        self._test_twoarg_function(zipper)
+
+    def test_enumerate_zip(self):
+        @cuda.jit
+        def enumerator_zipper(x, y, error):
+            count = 0
+
+            for i, (xv, yv) in enumerate(zip(x, y)):
+                if i != count:
+                    error[0] = 1
+                if xv != x[i]:
+                    error[0] = 2
+                if yv != y[i]:
+                    error[0] = 3
+
+                count += 1
+
+            if count != len(x):
+                error[0] = 4
+
+        self._test_twoarg_function(enumerator_zipper)
+
+    def test_zip_enumerate(self):
+        @cuda.jit
+        def zipper_enumerator(x, y, error):
+            count = 0
+
+            for (i, xv), yv in zip(enumerate(x), y):
+                if i != count:
+                    error[0] = 1
+                if xv != x[i]:
+                    error[0] = 2
+                if yv != y[i]:
+                    error[0] = 3
+
+                count += 1
+
+            if count != len(x):
+                error[0] = 4
+
+        self._test_twoarg_function(zipper_enumerator)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,68 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Test basic language features
+
+"""
+
+import numpy as np
+from numba import cuda
+from numba.cuda import float64
+from numba.cuda.testing import unittest, CUDATestCase
+
+
+class TestLang(CUDATestCase):
+    def test_enumerate(self):
+        tup = (1.0, 2.5, 3.0)
+
+        @cuda.jit("void(float64[:])")
+        def foo(a):
+            for i, v in enumerate(tup):
+                a[i] = v
+
+        a = np.zeros(len(tup))
+        foo[1, 1](a)
+        self.assertTrue(np.all(a == tup))
+
+    def test_zip(self):
+        t1 = (1, 2, 3)
+        t2 = (4.5, 5.6, 6.7)
+
+        @cuda.jit("void(float64[:])")
+        def foo(a):
+            c = 0
+            for i, j in zip(t1, t2):
+                c += i + j
+            a[0] = c
+
+        a = np.zeros(1)
+        foo[1, 1](a)
+        b = np.array(t1)
+        c = np.array(t2)
+        self.assertTrue(np.all(a == (b + c).sum()))
+
+    def test_issue_872(self):
+        """
+        Ensure that typing and lowering of CUDA kernel API primitives works in
+        more than one block. Was originally to ensure that macro expansion works
+        for more than one block (issue #872), but macro expansion has been
+        replaced by a "proper" implementation of all kernel API functions.
+        """
+
+        @cuda.jit("void(float64[:,:])")
+        def cuda_kernel_api_in_multiple_blocks(ary):
+            for i in range(2):
+                tx = cuda.threadIdx.x
+            for j in range(3):
+                ty = cuda.threadIdx.y
+            sm = cuda.shared.array((2, 3), float64)
+            sm[tx, ty] = 1.0
+            ary[tx, ty] = sm[tx, ty]
+
+        a = np.zeros((2, 3))
+        cuda_kernel_api_in_multiple_blocks[1, (2, 3)](a)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import numpy as np
+from numba import cuda
+from numba.cuda import float64, void
+from numba.cuda.testing import unittest, CUDATestCase
+from numba.cuda.core import config
+
+# NOTE: CUDA kernel does not return any value
+
+if config.ENABLE_CUDASIM:
+    tpb = 4
+else:
+    tpb = 16
+SM_SIZE = tpb, tpb
+
+
+class TestCudaLaplace(CUDATestCase):
+    def test_laplace_small(self):
+        @cuda.jit(float64(float64, float64), device=True, inline="always")
+        def get_max(a, b):
+            if a > b:
+                return a
+            else:
+                return b
+
+        @cuda.jit(void(float64[:, :], float64[:, :], float64[:, :]))
+        def jocabi_relax_core(A, Anew, error):
+            err_sm = cuda.shared.array(SM_SIZE, dtype=float64)
+
+            ty = cuda.threadIdx.x
+            tx = cuda.threadIdx.y
+            bx = cuda.blockIdx.x
+            by = cuda.blockIdx.y
+
+            n = A.shape[0]
+            m = A.shape[1]
+
+            i, j = cuda.grid(2)
+
+            err_sm[ty, tx] = 0
+            if j >= 1 and j < n - 1 and i >= 1 and i < m - 1:
+                Anew[j, i] = 0.25 * (
+                    A[j, i + 1] + A[j, i - 1] + A[j - 1, i] + A[j + 1, i]
+                )
+                err_sm[ty, tx] = Anew[j, i] - A[j, i]
+
+            cuda.syncthreads()
+
+            # max-reduce err_sm vertically
+            t = tpb // 2
+            while t > 0:
+                if ty < t:
+                    err_sm[ty, tx] = get_max(err_sm[ty, tx], err_sm[ty + t, tx])
+                t //= 2
+                cuda.syncthreads()
+
+            # max-reduce err_sm horizontally
+            t = tpb // 2
+            while t > 0:
+                if tx < t and ty == 0:
+                    err_sm[ty, tx] = get_max(err_sm[ty, tx], err_sm[ty, tx + t])
+                t //= 2
+                cuda.syncthreads()
+
+            if tx == 0 and ty == 0:
+                error[by, bx] = err_sm[0, 0]
+
+        if config.ENABLE_CUDASIM:
+            NN, NM = 4, 4
+            iter_max = 20
+        else:
+            NN, NM = 256, 256
+            iter_max = 1000
+
+        A = np.zeros((NN, NM), dtype=np.float64)
+        Anew = np.zeros((NN, NM), dtype=np.float64)
+
+        n = NN
+
+        tol = 1.0e-6
+        error = 1.0
+
+        for j in range(n):
+            A[j, 0] = 1.0
+            Anew[j, 0] = 1.0
+
+        iter = 0
+
+        blockdim = (tpb, tpb)
+        griddim = (NN // blockdim[0], NM // blockdim[1])
+
+        error_grid = np.zeros(griddim)
+
+        stream = cuda.stream()
+
+        dA = cuda.to_device(A, stream)  # to device and don't come back
+        dAnew = cuda.to_device(Anew, stream)  # to device and don't come back
+        derror_grid = cuda.to_device(error_grid, stream)
+
+        while error > tol and iter < iter_max:
+            self.assertTrue(error_grid.dtype == np.float64)
+
+            jocabi_relax_core[griddim, blockdim, stream](dA, dAnew, derror_grid)
+
+            derror_grid.copy_to_host(error_grid, stream=stream)
+
+            # error_grid is available on host
+            stream.synchronize()
+
+            error = np.abs(error_grid).max()
+
+            # swap dA and dAnew
+            tmp = dA
+            dA = dAnew
+            dAnew = tmp
+
+            iter += 1
+
+
+if __name__ == "__main__":
+    unittest.main()