632 lines
18 KiB
Python
632 lines
18 KiB
Python
|
"""
|
||
|
|
Tests for SSA reconstruction
|
||
|
|
"""
|
||
|
|
import sys
|
||
|
|
import copy
|
||
|
|
import logging
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from numba import njit, jit, types
|
||
|
|
from numba.core import errors, ir
|
||
|
|
from numba.core.compiler_machinery import FunctionPass, register_pass
|
||
|
|
from numba.core.compiler import DefaultPassBuilder, CompilerBase
|
||
|
|
from numba.core.untyped_passes import ReconstructSSA, PreserveIR
|
||
|
|
from numba.core.typed_passes import NativeLowering
|
||
|
|
from numba.extending import overload
|
||
|
|
from numba.tests.support import MemoryLeakMixin, TestCase, override_config
|
||
|
|
|
||
|
|
|
||
|
|
_DEBUG = False
|
||
|
|
|
||
|
|
if _DEBUG:
|
||
|
|
# Enable debug logger on SSA reconstruction
|
||
|
|
ssa_logger = logging.getLogger("numba.core.ssa")
|
||
|
|
ssa_logger.setLevel(level=logging.DEBUG)
|
||
|
|
ssa_logger.addHandler(logging.StreamHandler(sys.stderr))
|
||
|
|
|
||
|
|
|
||
|
|
class SSABaseTest(TestCase):
|
||
|
|
|
||
|
|
def check_func(self, func, *args):
|
||
|
|
got = func(*copy.deepcopy(args))
|
||
|
|
exp = func.py_func(*copy.deepcopy(args))
|
||
|
|
self.assertEqual(got, exp)
|
||
|
|
|
||
|
|
|
||
|
|
class TestSSA(SSABaseTest):
|
||
|
|
"""
|
||
|
|
Contains tests to help isolate problems in SSA
|
||
|
|
"""
|
||
|
|
|
||
|
|
def test_argument_name_reused(self):
|
||
|
|
@njit
|
||
|
|
def foo(x):
|
||
|
|
x += 1
|
||
|
|
return x
|
||
|
|
|
||
|
|
self.check_func(foo, 123)
|
||
|
|
|
||
|
|
def test_if_else_redefine(self):
|
||
|
|
@njit
|
||
|
|
def foo(x, y):
|
||
|
|
z = x * y
|
||
|
|
if x < y:
|
||
|
|
z = x
|
||
|
|
else:
|
||
|
|
z = y
|
||
|
|
return z
|
||
|
|
|
||
|
|
self.check_func(foo, 3, 2)
|
||
|
|
self.check_func(foo, 2, 3)
|
||
|
|
|
||
|
|
def test_sum_loop(self):
|
||
|
|
@njit
|
||
|
|
def foo(n):
|
||
|
|
c = 0
|
||
|
|
for i in range(n):
|
||
|
|
c += i
|
||
|
|
return c
|
||
|
|
|
||
|
|
self.check_func(foo, 0)
|
||
|
|
self.check_func(foo, 10)
|
||
|
|
|
||
|
|
def test_sum_loop_2vars(self):
|
||
|
|
@njit
|
||
|
|
def foo(n):
|
||
|
|
c = 0
|
||
|
|
d = n
|
||
|
|
for i in range(n):
|
||
|
|
c += i
|
||
|
|
d += n
|
||
|
|
return c, d
|
||
|
|
|
||
|
|
self.check_func(foo, 0)
|
||
|
|
self.check_func(foo, 10)
|
||
|
|
|
||
|
|
def test_sum_2d_loop(self):
|
||
|
|
@njit
|
||
|
|
def foo(n):
|
||
|
|
c = 0
|
||
|
|
for i in range(n):
|
||
|
|
for j in range(n):
|
||
|
|
c += j
|
||
|
|
c += i
|
||
|
|
return c
|
||
|
|
|
||
|
|
self.check_func(foo, 0)
|
||
|
|
self.check_func(foo, 10)
|
||
|
|
|
||
|
|
def check_undefined_var(self, should_warn):
|
||
|
|
@njit
|
||
|
|
def foo(n):
|
||
|
|
if n:
|
||
|
|
if n > 0:
|
||
|
|
c = 0
|
||
|
|
return c
|
||
|
|
else:
|
||
|
|
# variable c is not defined in this branch
|
||
|
|
c += 1
|
||
|
|
return c
|
||
|
|
|
||
|
|
if should_warn:
|
||
|
|
with self.assertWarns(errors.NumbaWarning) as warns:
|
||
|
|
# n=1 so we won't actually run the branch with the uninitialized
|
||
|
|
self.check_func(foo, 1)
|
||
|
|
self.assertIn("Detected uninitialized variable c",
|
||
|
|
str(warns.warning))
|
||
|
|
else:
|
||
|
|
self.check_func(foo, 1)
|
||
|
|
|
||
|
|
with self.assertRaises(UnboundLocalError):
|
||
|
|
foo.py_func(0)
|
||
|
|
|
||
|
|
def test_undefined_var(self):
|
||
|
|
with override_config('ALWAYS_WARN_UNINIT_VAR', 0):
|
||
|
|
self.check_undefined_var(should_warn=False)
|
||
|
|
with override_config('ALWAYS_WARN_UNINIT_VAR', 1):
|
||
|
|
self.check_undefined_var(should_warn=True)
|
||
|
|
|
||
|
|
def test_phi_propagation(self):
|
||
|
|
@njit
|
||
|
|
def foo(actions):
|
||
|
|
n = 1
|
||
|
|
|
||
|
|
i = 0
|
||
|
|
ct = 0
|
||
|
|
while n > 0 and i < len(actions):
|
||
|
|
n -= 1
|
||
|
|
|
||
|
|
while actions[i]:
|
||
|
|
if actions[i]:
|
||
|
|
if actions[i]:
|
||
|
|
n += 10
|
||
|
|
actions[i] -= 1
|
||
|
|
else:
|
||
|
|
if actions[i]:
|
||
|
|
n += 20
|
||
|
|
actions[i] += 1
|
||
|
|
|
||
|
|
ct += n
|
||
|
|
ct += n
|
||
|
|
return ct, n
|
||
|
|
|
||
|
|
self.check_func(foo, np.array([1, 2]))
|
||
|
|
|
||
|
|
def test_unhandled_undefined(self):
|
||
|
|
def function1(arg1, arg2, arg3, arg4, arg5):
|
||
|
|
# This function is auto-generated.
|
||
|
|
if arg1:
|
||
|
|
var1 = arg2
|
||
|
|
var2 = arg3
|
||
|
|
var3 = var2
|
||
|
|
var4 = arg1
|
||
|
|
return
|
||
|
|
else:
|
||
|
|
if arg2:
|
||
|
|
if arg4:
|
||
|
|
var5 = arg4 # noqa: F841
|
||
|
|
return
|
||
|
|
else:
|
||
|
|
var6 = var4
|
||
|
|
return
|
||
|
|
return var6
|
||
|
|
else:
|
||
|
|
if arg5:
|
||
|
|
if var1:
|
||
|
|
if arg5:
|
||
|
|
var1 = var6
|
||
|
|
return
|
||
|
|
else:
|
||
|
|
var7 = arg2 # noqa: F841
|
||
|
|
return arg2
|
||
|
|
return
|
||
|
|
else:
|
||
|
|
if var2:
|
||
|
|
arg5 = arg2
|
||
|
|
return arg1
|
||
|
|
else:
|
||
|
|
var6 = var3
|
||
|
|
return var4
|
||
|
|
return
|
||
|
|
return
|
||
|
|
else:
|
||
|
|
var8 = var1
|
||
|
|
return
|
||
|
|
return var8
|
||
|
|
var9 = var3 # noqa: F841
|
||
|
|
var10 = arg5 # noqa: F841
|
||
|
|
return var1
|
||
|
|
|
||
|
|
# The argument values is not critical for re-creating the bug
|
||
|
|
# because the bug is in compile-time.
|
||
|
|
expect = function1(2, 3, 6, 0, 7)
|
||
|
|
got = njit(function1)(2, 3, 6, 0, 7)
|
||
|
|
self.assertEqual(expect, got)
|
||
|
|
|
||
|
|
|
||
|
|
class TestReportedSSAIssues(SSABaseTest):
|
||
|
|
# Tests from issues
|
||
|
|
# https://github.com/numba/numba/issues?q=is%3Aopen+is%3Aissue+label%3ASSA
|
||
|
|
|
||
|
|
def test_issue2194(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo():
|
||
|
|
V = np.empty(1)
|
||
|
|
s = np.uint32(1)
|
||
|
|
|
||
|
|
for i in range(s):
|
||
|
|
V[i] = 1
|
||
|
|
for i in range(s, 1):
|
||
|
|
pass
|
||
|
|
|
||
|
|
self.check_func(foo, )
|
||
|
|
|
||
|
|
def test_issue3094(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def doit(x):
|
||
|
|
return x
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo(pred):
|
||
|
|
if pred:
|
||
|
|
x = True
|
||
|
|
else:
|
||
|
|
x = False
|
||
|
|
# do something with x
|
||
|
|
return doit(x)
|
||
|
|
|
||
|
|
self.check_func(foo, False)
|
||
|
|
|
||
|
|
def test_issue3931(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo(arr):
|
||
|
|
for i in range(1):
|
||
|
|
arr = arr.reshape(3 * 2)
|
||
|
|
arr = arr.reshape(3, 2)
|
||
|
|
return (arr)
|
||
|
|
|
||
|
|
np.testing.assert_allclose(foo(np.zeros((3, 2))),
|
||
|
|
foo.py_func(np.zeros((3, 2))))
|
||
|
|
|
||
|
|
def test_issue3976(self):
|
||
|
|
|
||
|
|
def overload_this(a):
|
||
|
|
return 'dummy'
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo(a):
|
||
|
|
if a:
|
||
|
|
s = 5
|
||
|
|
s = overload_this(s)
|
||
|
|
else:
|
||
|
|
s = 'b'
|
||
|
|
|
||
|
|
return s
|
||
|
|
|
||
|
|
@overload(overload_this)
|
||
|
|
def ol(a):
|
||
|
|
return overload_this
|
||
|
|
|
||
|
|
self.check_func(foo, True)
|
||
|
|
|
||
|
|
def test_issue3979(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo(A, B):
|
||
|
|
x = A[0]
|
||
|
|
y = B[0]
|
||
|
|
for i in A:
|
||
|
|
x = i
|
||
|
|
for i in B:
|
||
|
|
y = i
|
||
|
|
return x, y
|
||
|
|
|
||
|
|
self.check_func(foo, (1, 2), ('A', 'B'))
|
||
|
|
|
||
|
|
def test_issue5219(self):
|
||
|
|
|
||
|
|
def overload_this(a, b=None):
|
||
|
|
if isinstance(b, tuple):
|
||
|
|
b = b[0]
|
||
|
|
return b
|
||
|
|
|
||
|
|
@overload(overload_this)
|
||
|
|
def ol(a, b=None):
|
||
|
|
b_is_tuple = isinstance(b, (types.Tuple, types.UniTuple))
|
||
|
|
|
||
|
|
def impl(a, b=None):
|
||
|
|
if b_is_tuple is True:
|
||
|
|
b = b[0]
|
||
|
|
return b
|
||
|
|
return impl
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def test_tuple(a, b):
|
||
|
|
overload_this(a, b)
|
||
|
|
|
||
|
|
self.check_func(test_tuple, 1, (2, ))
|
||
|
|
|
||
|
|
def test_issue5223(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def bar(x):
|
||
|
|
if len(x) == 5:
|
||
|
|
return x
|
||
|
|
x = x.copy()
|
||
|
|
for i in range(len(x)):
|
||
|
|
x[i] += 1
|
||
|
|
return x
|
||
|
|
|
||
|
|
a = np.ones(5)
|
||
|
|
a.flags.writeable = False
|
||
|
|
|
||
|
|
np.testing.assert_allclose(bar(a), bar.py_func(a))
|
||
|
|
|
||
|
|
def test_issue5243(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo(q):
|
||
|
|
lin = np.array((0.1, 0.6, 0.3))
|
||
|
|
stencil = np.zeros((3, 3))
|
||
|
|
stencil[0, 0] = q[0, 0]
|
||
|
|
return lin[0]
|
||
|
|
|
||
|
|
self.check_func(foo, np.zeros((2, 2)))
|
||
|
|
|
||
|
|
def test_issue5482_missing_variable_init(self):
|
||
|
|
# Test error that lowering fails because variable is missing
|
||
|
|
# a definition before use.
|
||
|
|
@njit("(intp, intp, intp)")
|
||
|
|
def foo(x, v, n):
|
||
|
|
for i in range(n):
|
||
|
|
if i == 0:
|
||
|
|
if i == x:
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
problematic = v
|
||
|
|
else:
|
||
|
|
if i == x:
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
problematic = problematic + v
|
||
|
|
return problematic
|
||
|
|
|
||
|
|
def test_issue5482_objmode_expr_null_lowering(self):
|
||
|
|
# Existing pipelines will not have the Expr.null in objmode.
|
||
|
|
# We have to create a custom pipeline to force a SSA reconstruction
|
||
|
|
# and stripping.
|
||
|
|
from numba.core.compiler import CompilerBase, DefaultPassBuilder
|
||
|
|
from numba.core.untyped_passes import ReconstructSSA, IRProcessing
|
||
|
|
from numba.core.typed_passes import PreLowerStripPhis
|
||
|
|
|
||
|
|
class CustomPipeline(CompilerBase):
|
||
|
|
def define_pipelines(self):
|
||
|
|
pm = DefaultPassBuilder.define_objectmode_pipeline(self.state)
|
||
|
|
# Force SSA reconstruction and stripping
|
||
|
|
pm.add_pass_after(ReconstructSSA, IRProcessing)
|
||
|
|
pm.add_pass_after(PreLowerStripPhis, ReconstructSSA)
|
||
|
|
pm.finalize()
|
||
|
|
return [pm]
|
||
|
|
|
||
|
|
@jit("(intp, intp, intp)", looplift=False,
|
||
|
|
pipeline_class=CustomPipeline)
|
||
|
|
def foo(x, v, n):
|
||
|
|
for i in range(n):
|
||
|
|
if i == n:
|
||
|
|
if i == x:
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
problematic = v
|
||
|
|
else:
|
||
|
|
if i == x:
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
problematic = problematic + v
|
||
|
|
return problematic
|
||
|
|
|
||
|
|
def test_issue5493_unneeded_phi(self):
|
||
|
|
# Test error that unneeded phi is inserted because variable does not
|
||
|
|
# have a dominance definition.
|
||
|
|
data = (np.ones(2), np.ones(2))
|
||
|
|
A = np.ones(1)
|
||
|
|
B = np.ones((1,1))
|
||
|
|
|
||
|
|
def foo(m, n, data):
|
||
|
|
if len(data) == 1:
|
||
|
|
v0 = data[0]
|
||
|
|
else:
|
||
|
|
v0 = data[0]
|
||
|
|
# Unneeded PHI node for `problematic` would be placed here
|
||
|
|
for _ in range(1, len(data)):
|
||
|
|
v0 += A
|
||
|
|
|
||
|
|
for t in range(1, m):
|
||
|
|
for idx in range(n):
|
||
|
|
t = B
|
||
|
|
|
||
|
|
if idx == 0:
|
||
|
|
if idx == n - 1:
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
problematic = t
|
||
|
|
else:
|
||
|
|
if idx == n - 1:
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
problematic = problematic + t
|
||
|
|
return problematic
|
||
|
|
|
||
|
|
expect = foo(10, 10, data)
|
||
|
|
res1 = njit(foo)(10, 10, data)
|
||
|
|
res2 = jit(forceobj=True, looplift=False)(foo)(10, 10, data)
|
||
|
|
np.testing.assert_array_equal(expect, res1)
|
||
|
|
np.testing.assert_array_equal(expect, res2)
|
||
|
|
|
||
|
|
def test_issue5623_equal_statements_in_same_bb(self):
|
||
|
|
|
||
|
|
def foo(pred, stack):
|
||
|
|
i = 0
|
||
|
|
c = 1
|
||
|
|
|
||
|
|
if pred is True:
|
||
|
|
stack[i] = c
|
||
|
|
i += 1
|
||
|
|
stack[i] = c
|
||
|
|
i += 1
|
||
|
|
|
||
|
|
python = np.array([0, 666])
|
||
|
|
foo(True, python)
|
||
|
|
|
||
|
|
nb = np.array([0, 666])
|
||
|
|
njit(foo)(True, nb)
|
||
|
|
|
||
|
|
expect = np.array([1, 1])
|
||
|
|
|
||
|
|
np.testing.assert_array_equal(python, expect)
|
||
|
|
np.testing.assert_array_equal(nb, expect)
|
||
|
|
|
||
|
|
def test_issue5678_non_minimal_phi(self):
|
||
|
|
# There should be only one phi for variable "i"
|
||
|
|
|
||
|
|
from numba.core.compiler import CompilerBase, DefaultPassBuilder
|
||
|
|
from numba.core.untyped_passes import (
|
||
|
|
ReconstructSSA, FunctionPass, register_pass,
|
||
|
|
)
|
||
|
|
|
||
|
|
phi_counter = []
|
||
|
|
|
||
|
|
@register_pass(mutates_CFG=False, analysis_only=True)
|
||
|
|
class CheckSSAMinimal(FunctionPass):
|
||
|
|
# A custom pass to count the number of phis
|
||
|
|
|
||
|
|
_name = self.__class__.__qualname__ + ".CheckSSAMinimal"
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__(self)
|
||
|
|
|
||
|
|
def run_pass(self, state):
|
||
|
|
ct = 0
|
||
|
|
for blk in state.func_ir.blocks.values():
|
||
|
|
ct += len(list(blk.find_exprs('phi')))
|
||
|
|
phi_counter.append(ct)
|
||
|
|
return True
|
||
|
|
|
||
|
|
class CustomPipeline(CompilerBase):
|
||
|
|
def define_pipelines(self):
|
||
|
|
pm = DefaultPassBuilder.define_nopython_pipeline(self.state)
|
||
|
|
pm.add_pass_after(CheckSSAMinimal, ReconstructSSA)
|
||
|
|
pm.finalize()
|
||
|
|
return [pm]
|
||
|
|
|
||
|
|
@njit(pipeline_class=CustomPipeline)
|
||
|
|
def while_for(n, max_iter=1):
|
||
|
|
a = np.empty((n,n))
|
||
|
|
i = 0
|
||
|
|
while i <= max_iter:
|
||
|
|
for j in range(len(a)):
|
||
|
|
for k in range(len(a)):
|
||
|
|
a[j,k] = j + k
|
||
|
|
i += 1
|
||
|
|
return a
|
||
|
|
|
||
|
|
# Runs fine?
|
||
|
|
self.assertPreciseEqual(while_for(10), while_for.py_func(10))
|
||
|
|
# One phi?
|
||
|
|
self.assertEqual(phi_counter, [1])
|
||
|
|
|
||
|
|
def test_issue9242_use_not_dom_def(self):
|
||
|
|
from numba.core.ir import FunctionIR
|
||
|
|
from numba.core.compiler_machinery import (
|
||
|
|
AnalysisPass,
|
||
|
|
register_pass,
|
||
|
|
)
|
||
|
|
|
||
|
|
def check(fir: FunctionIR):
|
||
|
|
[blk, *_] = fir.blocks.values()
|
||
|
|
var = blk.scope.get("d")
|
||
|
|
defn = fir.get_definition(var)
|
||
|
|
self.assertEqual(defn.op, "phi")
|
||
|
|
self.assertIn(ir.UNDEFINED, defn.incoming_values)
|
||
|
|
|
||
|
|
@register_pass(mutates_CFG=False, analysis_only=True)
|
||
|
|
class SSACheck(AnalysisPass):
|
||
|
|
"""
|
||
|
|
Check SSA on variable `d`
|
||
|
|
"""
|
||
|
|
|
||
|
|
_name = "SSA_Check"
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
AnalysisPass.__init__(self)
|
||
|
|
|
||
|
|
def run_pass(self, state):
|
||
|
|
check(state.func_ir)
|
||
|
|
return False
|
||
|
|
|
||
|
|
class SSACheckPipeline(CompilerBase):
|
||
|
|
"""Inject SSACheck pass into the default pipeline following the SSA
|
||
|
|
pass
|
||
|
|
"""
|
||
|
|
|
||
|
|
def define_pipelines(self):
|
||
|
|
pipeline = DefaultPassBuilder.define_nopython_pipeline(
|
||
|
|
self.state, "ssa_check_custom_pipeline")
|
||
|
|
|
||
|
|
pipeline._finalized = False
|
||
|
|
pipeline.add_pass_after(SSACheck, ReconstructSSA)
|
||
|
|
|
||
|
|
pipeline.finalize()
|
||
|
|
return [pipeline]
|
||
|
|
|
||
|
|
@njit(pipeline_class=SSACheckPipeline)
|
||
|
|
def py_func(a):
|
||
|
|
c = a > 0
|
||
|
|
if c:
|
||
|
|
d = a + 5 # d is only defined here; undef in the else branch
|
||
|
|
|
||
|
|
return c and d > 0
|
||
|
|
|
||
|
|
py_func(10)
|
||
|
|
|
||
|
|
|
||
|
|
class TestSROAIssues(MemoryLeakMixin, TestCase):
|
||
|
|
# This tests issues related to the SROA optimization done in lowering, which
|
||
|
|
# reduces time spent in the LLVM SROA pass. The optimization is related to
|
||
|
|
# SSA and tries to reduce the number of alloca statements for variables with
|
||
|
|
# only a single assignment.
|
||
|
|
def test_issue7258_multiple_assignment_post_SSA(self):
|
||
|
|
# This test adds a pass that will duplicate assignment statements to
|
||
|
|
# variables named "foobar".
|
||
|
|
# In the reported issue, the bug will cause a memory leak.
|
||
|
|
cloned = []
|
||
|
|
|
||
|
|
@register_pass(analysis_only=False, mutates_CFG=True)
|
||
|
|
class CloneFoobarAssignments(FunctionPass):
|
||
|
|
# A pass that clones variable assignments into "foobar"
|
||
|
|
_name = "clone_foobar_assignments_pass"
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
FunctionPass.__init__(self)
|
||
|
|
|
||
|
|
def run_pass(self, state):
|
||
|
|
mutated = False
|
||
|
|
for blk in state.func_ir.blocks.values():
|
||
|
|
to_clone = []
|
||
|
|
# find assignments to "foobar"
|
||
|
|
for assign in blk.find_insts(ir.Assign):
|
||
|
|
if assign.target.name == "foobar":
|
||
|
|
to_clone.append(assign)
|
||
|
|
# clone
|
||
|
|
for assign in to_clone:
|
||
|
|
clone = copy.deepcopy(assign)
|
||
|
|
blk.insert_after(clone, assign)
|
||
|
|
mutated = True
|
||
|
|
# keep track of cloned statements
|
||
|
|
cloned.append(clone)
|
||
|
|
return mutated
|
||
|
|
|
||
|
|
class CustomCompiler(CompilerBase):
|
||
|
|
def define_pipelines(self):
|
||
|
|
pm = DefaultPassBuilder.define_nopython_pipeline(
|
||
|
|
self.state, "custom_pipeline",
|
||
|
|
)
|
||
|
|
pm._finalized = False
|
||
|
|
# Insert the cloning pass after SSA
|
||
|
|
pm.add_pass_after(CloneFoobarAssignments, ReconstructSSA)
|
||
|
|
# Capture IR post lowering
|
||
|
|
pm.add_pass_after(PreserveIR, NativeLowering)
|
||
|
|
pm.finalize()
|
||
|
|
return [pm]
|
||
|
|
|
||
|
|
@njit(pipeline_class=CustomCompiler)
|
||
|
|
def udt(arr):
|
||
|
|
foobar = arr + 1 # this assignment will be cloned
|
||
|
|
return foobar
|
||
|
|
|
||
|
|
arr = np.arange(10)
|
||
|
|
# Verify that the function works as expected
|
||
|
|
self.assertPreciseEqual(udt(arr), arr + 1)
|
||
|
|
# Verify that the expected statement is cloned
|
||
|
|
self.assertEqual(len(cloned), 1)
|
||
|
|
self.assertEqual(cloned[0].target.name, "foobar")
|
||
|
|
# Verify in the Numba IR that the expected statement is cloned
|
||
|
|
nir = udt.overloads[udt.signatures[0]].metadata['preserved_ir']
|
||
|
|
self.assertEqual(len(nir.blocks), 1,
|
||
|
|
"only one block")
|
||
|
|
[blk] = nir.blocks.values()
|
||
|
|
assigns = blk.find_insts(ir.Assign)
|
||
|
|
foobar_assigns = [stmt for stmt in assigns
|
||
|
|
if stmt.target.name == "foobar"]
|
||
|
|
self.assertEqual(
|
||
|
|
len(foobar_assigns), 2,
|
||
|
|
"expected two assignment statements into 'foobar'",
|
||
|
|
)
|
||
|
|
self.assertEqual(
|
||
|
|
foobar_assigns[0], foobar_assigns[1],
|
||
|
|
"expected the two assignment statements to be the same",
|
||
|
|
)
|