import unittest from unittest.case import TestCase import warnings import numpy as np from numba import objmode from numba.core import ir, compiler from numba.core import errors from numba.core.compiler import ( CompilerBase, ReconstructSSA, ) from numba.core.compiler_machinery import ( FunctionPass, PassManager, register_pass, ) from numba.core.untyped_passes import ( TranslateByteCode, IRProcessing, ) from numba import njit class TestIR(unittest.TestCase): 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(unittest.TestCase): 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]) 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) 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: 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(TestCase): 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()