numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

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

@@ -0,0 +1,598 @@
+# 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
+import warnings
+import numpy as np
+
+from numba.cuda import HAS_NUMBA
+
+if HAS_NUMBA:
+    from numba import objmode
+    from numba import njit
+from numba.cuda.core import ir
+from numba.cuda import compiler
+from numba.cuda.core import errors
+from numba.cuda.core.compiler import (
+    CompilerBase,
+)
+from numba.cuda.core.compiler_machinery import (
+    FunctionPass,
+    PassManager,
+    register_pass,
+)
+from numba.cuda.core.untyped_passes import (
+    TranslateByteCode,
+    IRProcessing,
+    ReconstructSSA,
+)
+from numba.cuda.testing import skip_on_standalone_numba_cuda
+
+
+def requires_fully_vendored_ir(fn):
+    return unittest.skip(
+        "Cannot run until the dependency on the un-vendored IR module is removed"
+    )(fn)
+
+
+class TestIR(CUDATestCase):
+    def test_IRScope(self):
+        filename = "<?>"
+        top = ir.Scope(parent=None, loc=ir.Loc(filename=filename, line=1))
+        local = ir.Scope(parent=top, loc=ir.Loc(filename=filename, line=2))
+
+        apple = local.define("apple", loc=ir.Loc(filename=filename, line=3))
+        self.assertIs(local.get("apple"), apple)
+        self.assertEqual(len(local.localvars), 1)
+
+        orange = top.define("orange", loc=ir.Loc(filename=filename, line=4))
+        self.assertEqual(len(local.localvars), 1)
+        self.assertEqual(len(top.localvars), 1)
+        self.assertIs(top.get("orange"), orange)
+        self.assertIs(local.get("orange"), orange)
+
+        more_orange = local.define(
+            "orange", loc=ir.Loc(filename=filename, line=5)
+        )
+        self.assertIs(top.get("orange"), orange)
+        self.assertIsNot(local.get("orange"), not orange)
+        self.assertIs(local.get("orange"), more_orange)
+
+        try:
+            local.define("orange", loc=ir.Loc(filename=filename, line=5))
+        except ir.RedefinedError:
+            pass
+        else:
+            self.fail("Expecting an %s" % ir.RedefinedError)
+
+
+class CheckEquality(CUDATestCase):
+    var_a = ir.Var(None, "a", ir.unknown_loc)
+    var_b = ir.Var(None, "b", ir.unknown_loc)
+    var_c = ir.Var(None, "c", ir.unknown_loc)
+    var_d = ir.Var(None, "d", ir.unknown_loc)
+    var_e = ir.Var(None, "e", ir.unknown_loc)
+    loc1 = ir.Loc("mock", 1, 0)
+    loc2 = ir.Loc("mock", 2, 0)
+    loc3 = ir.Loc("mock", 3, 0)
+
+    def check(self, base, same=[], different=[]):
+        for s in same:
+            self.assertTrue(base == s)
+        for d in different:
+            self.assertTrue(base != d)
+
+
+class TestIRMeta(CheckEquality):
+    """
+    Tests IR node meta, like Loc and Scope
+    """
+
+    def test_loc(self):
+        a = ir.Loc("file", 1, 0)
+        b = ir.Loc("file", 1, 0)
+        c = ir.Loc("pile", 1, 0)
+        d = ir.Loc("file", 2, 0)
+        e = ir.Loc("file", 1, 1)
+        self.check(
+            a,
+            same=[
+                b,
+            ],
+            different=[c, d, e],
+        )
+
+        f = ir.Loc("file", 1, 0, maybe_decorator=False)
+        g = ir.Loc("file", 1, 0, maybe_decorator=True)
+        self.check(a, same=[f, g])
+
+    def test_scope(self):
+        parent1 = ir.Scope(None, self.loc1)
+        parent2 = ir.Scope(None, self.loc1)
+        parent3 = ir.Scope(None, self.loc2)
+        self.check(
+            parent1,
+            same=[
+                parent2,
+                parent3,
+            ],
+        )
+
+        a = ir.Scope(parent1, self.loc1)
+        b = ir.Scope(parent1, self.loc1)
+        c = ir.Scope(parent1, self.loc2)
+        d = ir.Scope(parent3, self.loc1)
+        self.check(a, same=[b, c, d])
+
+        # parent1 and parent2 are equal, so children referring to either parent
+        # should be equal
+        e = ir.Scope(parent2, self.loc1)
+        self.check(
+            a,
+            same=[
+                e,
+            ],
+        )
+
+
+class TestIRNodes(CheckEquality):
+    """
+    Tests IR nodes
+    """
+
+    def test_terminator(self):
+        # terminator base class inst should always be equal
+        t1 = ir.Terminator()
+        t2 = ir.Terminator()
+        self.check(t1, same=[t2])
+
+    def test_jump(self):
+        a = ir.Jump(1, self.loc1)
+        b = ir.Jump(1, self.loc1)
+        c = ir.Jump(1, self.loc2)
+        d = ir.Jump(2, self.loc1)
+        self.check(a, same=[b, c], different=[d])
+
+    def test_return(self):
+        a = ir.Return(self.var_a, self.loc1)
+        b = ir.Return(self.var_a, self.loc1)
+        c = ir.Return(self.var_a, self.loc2)
+        d = ir.Return(self.var_b, self.loc1)
+        self.check(a, same=[b, c], different=[d])
+
+    def test_raise(self):
+        a = ir.Raise(self.var_a, self.loc1)
+        b = ir.Raise(self.var_a, self.loc1)
+        c = ir.Raise(self.var_a, self.loc2)
+        d = ir.Raise(self.var_b, self.loc1)
+        self.check(a, same=[b, c], different=[d])
+
+    def test_staticraise(self):
+        a = ir.StaticRaise(AssertionError, None, self.loc1)
+        b = ir.StaticRaise(AssertionError, None, self.loc1)
+        c = ir.StaticRaise(AssertionError, None, self.loc2)
+        e = ir.StaticRaise(AssertionError, ("str",), self.loc1)
+        d = ir.StaticRaise(RuntimeError, None, self.loc1)
+        self.check(a, same=[b, c], different=[d, e])
+
+    def test_branch(self):
+        a = ir.Branch(self.var_a, 1, 2, self.loc1)
+        b = ir.Branch(self.var_a, 1, 2, self.loc1)
+        c = ir.Branch(self.var_a, 1, 2, self.loc2)
+        d = ir.Branch(self.var_b, 1, 2, self.loc1)
+        e = ir.Branch(self.var_a, 2, 2, self.loc1)
+        f = ir.Branch(self.var_a, 1, 3, self.loc1)
+        self.check(a, same=[b, c], different=[d, e, f])
+
+    def test_expr(self):
+        a = ir.Expr("some_op", self.loc1)
+        b = ir.Expr("some_op", self.loc1)
+        c = ir.Expr("some_op", self.loc2)
+        d = ir.Expr("some_other_op", self.loc1)
+        self.check(a, same=[b, c], different=[d])
+
+    def test_setitem(self):
+        a = ir.SetItem(self.var_a, self.var_b, self.var_c, self.loc1)
+        b = ir.SetItem(self.var_a, self.var_b, self.var_c, self.loc1)
+        c = ir.SetItem(self.var_a, self.var_b, self.var_c, self.loc2)
+        d = ir.SetItem(self.var_d, self.var_b, self.var_c, self.loc1)
+        e = ir.SetItem(self.var_a, self.var_d, self.var_c, self.loc1)
+        f = ir.SetItem(self.var_a, self.var_b, self.var_d, self.loc1)
+        self.check(a, same=[b, c], different=[d, e, f])
+
+    def test_staticsetitem(self):
+        a = ir.StaticSetItem(self.var_a, 1, self.var_b, self.var_c, self.loc1)
+        b = ir.StaticSetItem(self.var_a, 1, self.var_b, self.var_c, self.loc1)
+        c = ir.StaticSetItem(self.var_a, 1, self.var_b, self.var_c, self.loc2)
+        d = ir.StaticSetItem(self.var_d, 1, self.var_b, self.var_c, self.loc1)
+        e = ir.StaticSetItem(self.var_a, 2, self.var_b, self.var_c, self.loc1)
+        f = ir.StaticSetItem(self.var_a, 1, self.var_d, self.var_c, self.loc1)
+        g = ir.StaticSetItem(self.var_a, 1, self.var_b, self.var_d, self.loc1)
+        self.check(a, same=[b, c], different=[d, e, f, g])
+
+    def test_delitem(self):
+        a = ir.DelItem(self.var_a, self.var_b, self.loc1)
+        b = ir.DelItem(self.var_a, self.var_b, self.loc1)
+        c = ir.DelItem(self.var_a, self.var_b, self.loc2)
+        d = ir.DelItem(self.var_c, self.var_b, self.loc1)
+        e = ir.DelItem(self.var_a, self.var_c, self.loc1)
+        self.check(a, same=[b, c], different=[d, e])
+
+    def test_del(self):
+        a = ir.Del(self.var_a.name, self.loc1)
+        b = ir.Del(self.var_a.name, self.loc1)
+        c = ir.Del(self.var_a.name, self.loc2)
+        d = ir.Del(self.var_b.name, self.loc1)
+        self.check(a, same=[b, c], different=[d])
+
+    def test_setattr(self):
+        a = ir.SetAttr(self.var_a, "foo", self.var_b, self.loc1)
+        b = ir.SetAttr(self.var_a, "foo", self.var_b, self.loc1)
+        c = ir.SetAttr(self.var_a, "foo", self.var_b, self.loc2)
+        d = ir.SetAttr(self.var_c, "foo", self.var_b, self.loc1)
+        e = ir.SetAttr(self.var_a, "bar", self.var_b, self.loc1)
+        f = ir.SetAttr(self.var_a, "foo", self.var_c, self.loc1)
+        self.check(a, same=[b, c], different=[d, e, f])
+
+    def test_delattr(self):
+        a = ir.DelAttr(self.var_a, "foo", self.loc1)
+        b = ir.DelAttr(self.var_a, "foo", self.loc1)
+        c = ir.DelAttr(self.var_a, "foo", self.loc2)
+        d = ir.DelAttr(self.var_c, "foo", self.loc1)
+        e = ir.DelAttr(self.var_a, "bar", self.loc1)
+        self.check(a, same=[b, c], different=[d, e])
+
+    def test_assign(self):
+        a = ir.Assign(self.var_a, self.var_b, self.loc1)
+        b = ir.Assign(self.var_a, self.var_b, self.loc1)
+        c = ir.Assign(self.var_a, self.var_b, self.loc2)
+        d = ir.Assign(self.var_c, self.var_b, self.loc1)
+        e = ir.Assign(self.var_a, self.var_c, self.loc1)
+        self.check(a, same=[b, c], different=[d, e])
+
+    def test_print(self):
+        a = ir.Print((self.var_a,), self.var_b, self.loc1)
+        b = ir.Print((self.var_a,), self.var_b, self.loc1)
+        c = ir.Print((self.var_a,), self.var_b, self.loc2)
+        d = ir.Print((self.var_c,), self.var_b, self.loc1)
+        e = ir.Print((self.var_a,), self.var_c, self.loc1)
+        self.check(a, same=[b, c], different=[d, e])
+
+    def test_storemap(self):
+        a = ir.StoreMap(self.var_a, self.var_b, self.var_c, self.loc1)
+        b = ir.StoreMap(self.var_a, self.var_b, self.var_c, self.loc1)
+        c = ir.StoreMap(self.var_a, self.var_b, self.var_c, self.loc2)
+        d = ir.StoreMap(self.var_d, self.var_b, self.var_c, self.loc1)
+        e = ir.StoreMap(self.var_a, self.var_d, self.var_c, self.loc1)
+        f = ir.StoreMap(self.var_a, self.var_b, self.var_d, self.loc1)
+        self.check(a, same=[b, c], different=[d, e, f])
+
+    def test_yield(self):
+        a = ir.Yield(self.var_a, self.loc1, 0)
+        b = ir.Yield(self.var_a, self.loc1, 0)
+        c = ir.Yield(self.var_a, self.loc2, 0)
+        d = ir.Yield(self.var_b, self.loc1, 0)
+        e = ir.Yield(self.var_a, self.loc1, 1)
+        self.check(a, same=[b, c], different=[d, e])
+
+    def test_enterwith(self):
+        a = ir.EnterWith(self.var_a, 0, 1, self.loc1)
+        b = ir.EnterWith(self.var_a, 0, 1, self.loc1)
+        c = ir.EnterWith(self.var_a, 0, 1, self.loc2)
+        d = ir.EnterWith(self.var_b, 0, 1, self.loc1)
+        e = ir.EnterWith(self.var_a, 1, 1, self.loc1)
+        f = ir.EnterWith(self.var_a, 0, 2, self.loc1)
+        self.check(a, same=[b, c], different=[d, e, f])
+
+    def test_arg(self):
+        a = ir.Arg("foo", 0, self.loc1)
+        b = ir.Arg("foo", 0, self.loc1)
+        c = ir.Arg("foo", 0, self.loc2)
+        d = ir.Arg("bar", 0, self.loc1)
+        e = ir.Arg("foo", 1, self.loc1)
+        self.check(a, same=[b, c], different=[d, e])
+
+    def test_const(self):
+        a = ir.Const(1, self.loc1)
+        b = ir.Const(1, self.loc1)
+        c = ir.Const(1, self.loc2)
+        d = ir.Const(2, self.loc1)
+        self.check(a, same=[b, c], different=[d])
+
+    def test_global(self):
+        a = ir.Global("foo", 0, self.loc1)
+        b = ir.Global("foo", 0, self.loc1)
+        c = ir.Global("foo", 0, self.loc2)
+        d = ir.Global("bar", 0, self.loc1)
+        e = ir.Global("foo", 1, self.loc1)
+        self.check(a, same=[b, c], different=[d, e])
+
+    def test_var(self):
+        a = ir.Var(None, "foo", self.loc1)
+        b = ir.Var(None, "foo", self.loc1)
+        c = ir.Var(None, "foo", self.loc2)
+        d = ir.Var(ir.Scope(None, ir.unknown_loc), "foo", self.loc1)
+        e = ir.Var(None, "bar", self.loc1)
+        self.check(a, same=[b, c, d], different=[e])
+
+    def test_undefinedtype(self):
+        a = ir.UndefinedType()
+        b = ir.UndefinedType()
+        self.check(a, same=[b])
+
+    def test_loop(self):
+        a = ir.Loop(1, 3)
+        b = ir.Loop(1, 3)
+        c = ir.Loop(2, 3)
+        d = ir.Loop(1, 4)
+        self.check(a, same=[b], different=[c, d])
+
+    def test_with(self):
+        a = ir.With(1, 3)
+        b = ir.With(1, 3)
+        c = ir.With(2, 3)
+        d = ir.With(1, 4)
+        self.check(a, same=[b], different=[c, d])
+
+
+# used later
+_GLOBAL = 1234
+
+
+class TestIRCompounds(CheckEquality):
+    """
+    Tests IR concepts that have state
+    """
+
+    def test_varmap(self):
+        a = ir.VarMap()
+        a.define(self.var_a, "foo")
+        a.define(self.var_b, "bar")
+
+        b = ir.VarMap()
+        b.define(self.var_a, "foo")
+        b.define(self.var_b, "bar")
+
+        c = ir.VarMap()
+        c.define(self.var_a, "foo")
+        c.define(self.var_c, "bar")
+
+        self.check(a, same=[b], different=[c])
+
+    def test_block(self):
+        def gen_block():
+            parent = ir.Scope(None, self.loc1)
+            tmp = ir.Block(parent, self.loc2)
+            assign1 = ir.Assign(self.var_a, self.var_b, self.loc3)
+            assign2 = ir.Assign(self.var_a, self.var_c, self.loc3)
+            assign3 = ir.Assign(self.var_c, self.var_b, self.loc3)
+            tmp.append(assign1)
+            tmp.append(assign2)
+            tmp.append(assign3)
+            return tmp
+
+        a = gen_block()
+        b = gen_block()
+        c = gen_block().append(ir.Assign(self.var_a, self.var_b, self.loc3))
+
+        self.check(a, same=[b], different=[c])
+
+    @skip_on_cudasim("doesn't work in the simulator")
+    def test_functionir(self):
+        def run_frontend(x):
+            return compiler.run_frontend(x, emit_dels=True)
+
+        # this creates a function full of all sorts of things to ensure the IR
+        # is pretty involved, it then compares two instances of the compiled
+        # function IR to check the IR is the same invariant of objects, and then
+        # a tiny mutation is made to the IR in the second function and detection
+        # of this change is checked.
+        def gen():
+            _FREEVAR = 0xCAFE
+
+            def foo(a, b, c=12, d=1j, e=None):
+                f = a + b
+                a += _FREEVAR
+                g = np.zeros(c, dtype=np.complex64)
+                h = f + g
+                i = 1j / d
+                if np.abs(i) > 0:
+                    k = h / i
+                    l = np.arange(1, c + 1)
+                    if HAS_NUMBA:
+                        with objmode():
+                            print(e, k)
+                    m = np.sqrt(l - g)
+                    if np.abs(m[0]) < 1:
+                        n = 0
+                        for o in range(a):
+                            n += 0
+                            if np.abs(n) < 3:
+                                break
+                        n += m[2]
+                    p = g / l
+                    q = []
+                    for r in range(len(p)):
+                        q.append(p[r])
+                        if r > 4 + 1:
+                            if HAS_NUMBA:
+                                with objmode(s="intp", t="complex128"):
+                                    s = 123
+                                    t = 5
+                            if s > 122:
+                                t += s
+                        t += q[0] + _GLOBAL
+
+                return f + o + r + t + r + a + n
+
+            return foo
+
+        x = gen()
+        y = gen()
+        x_ir = run_frontend(x)
+        y_ir = run_frontend(y)
+
+        self.assertTrue(x_ir.equal_ir(y_ir))
+
+        def check_diffstr(string, pointing_at=[]):
+            lines = string.splitlines()
+            for item in pointing_at:
+                for l in lines:
+                    if l.startswith("->"):
+                        if item in l:
+                            break
+                else:
+                    raise AssertionError("Could not find %s " % item)
+
+        self.assertIn("IR is considered equivalent", x_ir.diff_str(y_ir))
+
+        # minor mutation, simply switch branch targets on last branch
+        for label in reversed(list(y_ir.blocks.keys())):
+            blk = y_ir.blocks[label]
+            if isinstance(blk.body[-1], ir.Branch):
+                ref = blk.body[-1]
+                ref.truebr, ref.falsebr = ref.falsebr, ref.truebr
+                break
+
+        check_diffstr(x_ir.diff_str(y_ir), ["branch"])
+
+        z = gen()
+        self.assertFalse(x_ir.equal_ir(y_ir))
+
+        z_ir = run_frontend(z)
+
+        change_set = set()
+        for label in reversed(list(z_ir.blocks.keys())):
+            blk = z_ir.blocks[label]
+            ref = blk.body[:-1]
+            idx = None
+            for i in range(len(ref) - 1):
+                # look for two adjacent Del
+                if isinstance(ref[i], ir.Del) and isinstance(
+                    ref[i + 1], ir.Del
+                ):
+                    idx = i
+                    break
+            if idx is not None:
+                b = blk.body
+                change_set.add(str(b[idx + 1]))
+                change_set.add(str(b[idx]))
+                b[idx], b[idx + 1] = b[idx + 1], b[idx]
+                break
+
+        # ensure that a mutation occurred.
+        self.assertTrue(change_set)
+
+        self.assertFalse(x_ir.equal_ir(z_ir))
+        self.assertEqual(len(change_set), 2)
+        for item in change_set:
+            self.assertTrue(item.startswith("del "))
+        check_diffstr(x_ir.diff_str(z_ir), change_set)
+
+        def foo(a, b):
+            c = a * 2
+            d = c + b
+            e = np.sqrt(d)
+            return e
+
+        def bar(a, b):  # same as foo
+            c = a * 2
+            d = c + b
+            e = np.sqrt(d)
+            return e
+
+        def baz(a, b):
+            c = a * 2
+            d = b + c
+            e = np.sqrt(d + 1)
+            return e
+
+        foo_ir = run_frontend(foo)
+        bar_ir = run_frontend(bar)
+        self.assertTrue(foo_ir.equal_ir(bar_ir))
+        self.assertIn("IR is considered equivalent", foo_ir.diff_str(bar_ir))
+
+        baz_ir = run_frontend(baz)
+        self.assertFalse(foo_ir.equal_ir(baz_ir))
+        tmp = foo_ir.diff_str(baz_ir)
+        self.assertIn("Other block contains more statements", tmp)
+        check_diffstr(tmp, ["c + b", "b + c"])
+
+
+class TestIRPedanticChecks(CUDATestCase):
+    @skip_on_standalone_numba_cuda
+    def test_var_in_scope_assumption(self):
+        # Create a pass that clears ir.Scope in ir.Block
+        @register_pass(mutates_CFG=False, analysis_only=False)
+        class RemoveVarInScope(FunctionPass):
+            _name = "_remove_var_in_scope"
+
+            def __init__(self):
+                FunctionPass.__init__(self)
+
+            # implement method to do the work, "state" is the internal compiler
+            # state from the CompilerBase instance.
+            def run_pass(self, state):
+                func_ir = state.func_ir
+                # walk the blocks
+                for blk in func_ir.blocks.values():
+                    oldscope = blk.scope
+                    # put in an empty Scope
+                    blk.scope = ir.Scope(
+                        parent=oldscope.parent, loc=oldscope.loc
+                    )
+                return True
+
+        # Create a pass that always fails, to stop the compiler
+        @register_pass(mutates_CFG=False, analysis_only=False)
+        class FailPass(FunctionPass):
+            _name = "_fail"
+
+            def __init__(self, *args, **kwargs):
+                FunctionPass.__init__(self)
+
+            def run_pass(self, state):
+                # This is unreachable. SSA pass should have raised before this
+                # pass when run with `error.NumbaPedanticWarning`s raised as
+                # errors.
+                raise AssertionError("unreachable")
+
+        class MyCompiler(CompilerBase):
+            def define_pipelines(self):
+                pm = PassManager("testing pm")
+                pm.add_pass(TranslateByteCode, "analyzing bytecode")
+                pm.add_pass(IRProcessing, "processing IR")
+                pm.add_pass(RemoveVarInScope, "_remove_var_in_scope")
+                pm.add_pass(ReconstructSSA, "ssa")
+                pm.add_pass(FailPass, "_fail")
+                pm.finalize()
+                return [pm]
+
+        @njit(pipeline_class=MyCompiler)
+        def dummy(x):
+            # To trigger SSA and the pedantic check, this function must have
+            # multiple assignments to the same variable in different blocks.
+            a = 1
+            b = 2
+            if a < b:
+                a = 2
+            else:
+                b = 3
+            return a, b
+
+        with warnings.catch_warnings():
+            # Make NumbaPedanticWarning an error
+            warnings.simplefilter("error", errors.NumbaPedanticWarning)
+            # Catch NumbaIRAssumptionWarning
+            with self.assertRaises(errors.NumbaIRAssumptionWarning) as raises:
+                dummy(1)
+            # Verify the error message
+            self.assertRegex(
+                str(raises.exception),
+                r"variable '[a-z]' is not in scope",
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()