import gc from io import StringIO import numpy as np from numba import njit, vectorize from numba import typeof from numba.core import utils, types, typing, ir, compiler, cpu, cgutils from numba.core.compiler import Compiler, Flags from numba.core.registry import cpu_target from numba.tests.support import (MemoryLeakMixin, TestCase, temp_directory, create_temp_module) from numba.extending import ( overload, models, lower_builtin, register_model, make_attribute_wrapper, type_callable, typeof_impl ) import operator import textwrap import unittest class Namespace(dict): def __getattr__(s, k): return s[k] if k in s else super(Namespace, s).__getattr__(k) def axy(a, x, y): return a * x + y def ax2(a, x, y): return a * x + y def pos_root(As, Bs, Cs): return (-Bs + (((Bs ** 2.) - (4. * As * Cs)) ** 0.5)) / (2. * As) def neg_root_common_subexpr(As, Bs, Cs): _2As = 2. * As _4AsCs = 2. * _2As * Cs _Bs2_4AsCs = (Bs ** 2. - _4AsCs) return (-Bs - (_Bs2_4AsCs ** 0.5)) / _2As def neg_root_complex_subexpr(As, Bs, Cs): _2As = 2. * As _4AsCs = 2. * _2As * Cs _Bs2_4AsCs = (Bs ** 2. - _4AsCs) + 0j # Force into the complex domain. return (-Bs - (_Bs2_4AsCs ** 0.5)) / _2As vaxy = vectorize(axy) def call_stuff(a0, a1): return np.cos(vaxy(a0, np.sin(a1) - 1., 1.)) def are_roots_imaginary(As, Bs, Cs): return (Bs ** 2 - 4 * As * Cs) < 0 def div_add(As, Bs, Cs): return As / Bs + Cs def cube(As): return As ** 3 def explicit_output(a, b, out): np.cos(a, out) return np.add(out, b, out) def variable_name_reuse(a, b, c, d): u = a + b u = u - a * b u = u * c + d return u # From issue #1264 def distance_matrix(vectors): n_vectors = vectors.shape[0] result = np.empty((n_vectors, n_vectors), dtype=np.float64) for i in range(n_vectors): for j in range(i, n_vectors): result[i,j] = result[j,i] = np.sum( (vectors[i] - vectors[j]) ** 2) ** 0.5 return result class RewritesTester(Compiler): @classmethod def mk_pipeline(cls, args, return_type=None, flags=None, locals={}, library=None, typing_context=None, target_context=None): if not flags: flags = Flags() flags.nrt = True if typing_context is None: typing_context = cpu_target.typing_context if target_context is None: target_context = cpu_target.target_context return cls(typing_context, target_context, library, args, return_type, flags, locals) @classmethod def mk_no_rw_pipeline(cls, args, return_type=None, flags=None, locals={}, library=None, **kws): if not flags: flags = Flags() flags.no_rewrites = True return cls.mk_pipeline(args, return_type, flags, locals, library, **kws) class TestArrayExpressions(MemoryLeakMixin, TestCase): def _compile_function(self, fn, arg_tys): """ Compile the given function both without and with rewrites enabled. """ control_pipeline = RewritesTester.mk_no_rw_pipeline(arg_tys) cres_0 = control_pipeline.compile_extra(fn) control_cfunc = cres_0.entry_point test_pipeline = RewritesTester.mk_pipeline(arg_tys) cres_1 = test_pipeline.compile_extra(fn) test_cfunc = cres_1.entry_point return control_pipeline, control_cfunc, test_pipeline, test_cfunc def test_simple_expr(self): ''' Using a simple array expression, verify that rewriting is taking place, and is fusing loops. ''' A = np.linspace(0,1,10) X = np.linspace(2,1,10) Y = np.linspace(1,2,10) arg_tys = [typeof(arg) for arg in (A, X, Y)] control_pipeline, nb_axy_0, test_pipeline, nb_axy_1 = \ self._compile_function(axy, arg_tys) control_pipeline2 = RewritesTester.mk_no_rw_pipeline(arg_tys) cres_2 = control_pipeline2.compile_extra(ax2) nb_ctl = cres_2.entry_point expected = nb_axy_0(A, X, Y) actual = nb_axy_1(A, X, Y) control = nb_ctl(A, X, Y) np.testing.assert_array_equal(expected, actual) np.testing.assert_array_equal(control, actual) ir0 = control_pipeline.state.func_ir.blocks ir1 = test_pipeline.state.func_ir.blocks ir2 = control_pipeline2.state.func_ir.blocks self.assertEqual(len(ir0), len(ir1)) self.assertEqual(len(ir0), len(ir2)) # The rewritten IR should be smaller than the original. self.assertGreater(len(ir0[0].body), len(ir1[0].body)) self.assertEqual(len(ir0[0].body), len(ir2[0].body)) def _get_array_exprs(self, block): for instr in block: if isinstance(instr, ir.Assign): if isinstance(instr.value, ir.Expr): if instr.value.op == 'arrayexpr': yield instr def _array_expr_to_set(self, expr, out=None): ''' Convert an array expression tree into a set of operators. ''' if out is None: out = set() if not isinstance(expr, tuple): raise ValueError("{0} not a tuple".format(expr)) operation, operands = expr processed_operands = [] for operand in operands: if isinstance(operand, tuple): operand, _ = self._array_expr_to_set(operand, out) processed_operands.append(operand) processed_expr = operation, tuple(processed_operands) out.add(processed_expr) return processed_expr, out def _test_root_function(self, fn=pos_root): A = np.random.random(10) B = np.random.random(10) + 1. # Increase likelihood of real # root (could add 2 to force all # roots to be real). C = np.random.random(10) arg_tys = [typeof(arg) for arg in (A, B, C)] control_pipeline = RewritesTester.mk_no_rw_pipeline(arg_tys) control_cres = control_pipeline.compile_extra(fn) nb_fn_0 = control_cres.entry_point test_pipeline = RewritesTester.mk_pipeline(arg_tys) test_cres = test_pipeline.compile_extra(fn) nb_fn_1 = test_cres.entry_point np_result = fn(A, B, C) nb_result_0 = nb_fn_0(A, B, C) nb_result_1 = nb_fn_1(A, B, C) np.testing.assert_array_almost_equal(np_result, nb_result_0) np.testing.assert_array_almost_equal(nb_result_0, nb_result_1) return Namespace(locals()) def _test_cube_function(self, fn=cube): A = np.arange(10, dtype=np.float64) arg_tys = (typeof(A),) control_pipeline = RewritesTester.mk_no_rw_pipeline(arg_tys) control_cres = control_pipeline.compile_extra(fn) nb_fn_0 = control_cres.entry_point test_pipeline = RewritesTester.mk_pipeline(arg_tys) test_cres = test_pipeline.compile_extra(fn) nb_fn_1 = test_cres.entry_point expected = A ** 3 self.assertPreciseEqual(expected, nb_fn_0(A)) self.assertPreciseEqual(expected, nb_fn_1(A)) return Namespace(locals()) def _test_explicit_output_function(self, fn): """ Test function having a (a, b, out) signature where *out* is an output array the function writes into. """ A = np.arange(10, dtype=np.float64) B = A + 1 arg_tys = (typeof(A),) * 3 control_pipeline, control_cfunc, test_pipeline, test_cfunc = \ self._compile_function(fn, arg_tys) def run_func(fn): out = np.zeros_like(A) fn(A, B, out) return out expected = run_func(fn) self.assertPreciseEqual(expected, run_func(control_cfunc)) self.assertPreciseEqual(expected, run_func(test_cfunc)) return Namespace(locals()) def _assert_array_exprs(self, block, expected_count): """ Assert the *block* has the expected number of array expressions in it. """ rewrite_count = len(list(self._get_array_exprs(block))) self.assertEqual(rewrite_count, expected_count) def _assert_total_rewrite(self, control_ir, test_ir, trivial=False): """ Given two dictionaries of Numba IR blocks, check to make sure the control IR has no array expressions, while the test IR contains one and only one. """ # Both IRs have the same number of blocks (presumably 1) self.assertEqual(len(control_ir), len(test_ir)) control_block = control_ir[0].body test_block = test_ir[0].body self._assert_array_exprs(control_block, 0) self._assert_array_exprs(test_block, 1) if not trivial: # If the expression wasn't trivial, the block length should # have decreased (since a sequence of exprs was replaced # with a single nested array expr). self.assertGreater(len(control_block), len(test_block)) def _assert_no_rewrite(self, control_ir, test_ir): """ Given two dictionaries of Numba IR blocks, check to make sure the control IR and the test IR both have no array expressions. """ self.assertEqual(len(control_ir), len(test_ir)) # All blocks should be identical, and not rewritten for k, v in control_ir.items(): control_block = v.body test_block = test_ir[k].body self.assertEqual(len(control_block), len(test_block)) self._assert_array_exprs(control_block, 0) self._assert_array_exprs(test_block, 0) def test_trivial_expr(self): """ Ensure even a non-nested expression is rewritten, as it can enable scalar optimizations such as rewriting `x ** 2`. """ ns = self._test_cube_function() self._assert_total_rewrite(ns.control_pipeline.state.func_ir.blocks, ns.test_pipeline.state.func_ir.blocks, trivial=True) def test_complicated_expr(self): ''' Using the polynomial root function, ensure the full expression is being put in the same kernel with no remnants of intermediate array expressions. ''' ns = self._test_root_function() self._assert_total_rewrite(ns.control_pipeline.state.func_ir.blocks, ns.test_pipeline.state.func_ir.blocks) def test_common_subexpressions(self, fn=neg_root_common_subexpr): ''' Attempt to verify that rewriting will incorporate user common subexpressions properly. ''' ns = self._test_root_function(fn) ir0 = ns.control_pipeline.state.func_ir.blocks ir1 = ns.test_pipeline.state.func_ir.blocks self.assertEqual(len(ir0), len(ir1)) self.assertGreater(len(ir0[0].body), len(ir1[0].body)) self.assertEqual(len(list(self._get_array_exprs(ir0[0].body))), 0) # Verify that we didn't rewrite everything into a monolithic # array expression since we stored temporary values in # variables that might be used later (from the optimization's # point of view). array_expr_instrs = list(self._get_array_exprs(ir1[0].body)) self.assertGreater(len(array_expr_instrs), 1) # Now check that we haven't duplicated any subexpressions in # the rewritten code. array_sets = list(self._array_expr_to_set(instr.value.expr)[1] for instr in array_expr_instrs) for expr_set_0, expr_set_1 in zip(array_sets[:-1], array_sets[1:]): intersections = expr_set_0.intersection(expr_set_1) if intersections: self.fail("Common subexpressions detected in array " "expressions ({0})".format(intersections)) def test_complex_subexpression(self): return self.test_common_subexpressions(neg_root_complex_subexpr) def test_ufunc_and_dufunc_calls(self): ''' Verify that ufunc and DUFunc calls are being properly included in array expressions. ''' A = np.random.random(10) B = np.random.random(10) arg_tys = [typeof(arg) for arg in (A, B)] vaxy_descr = vaxy._dispatcher.targetdescr control_pipeline = RewritesTester.mk_no_rw_pipeline( arg_tys, typing_context=vaxy_descr.typing_context, target_context=vaxy_descr.target_context) cres_0 = control_pipeline.compile_extra(call_stuff) nb_call_stuff_0 = cres_0.entry_point test_pipeline = RewritesTester.mk_pipeline( arg_tys, typing_context=vaxy_descr.typing_context, target_context=vaxy_descr.target_context) cres_1 = test_pipeline.compile_extra(call_stuff) nb_call_stuff_1 = cres_1.entry_point expected = call_stuff(A, B) control = nb_call_stuff_0(A, B) actual = nb_call_stuff_1(A, B) np.testing.assert_array_almost_equal(expected, control) np.testing.assert_array_almost_equal(expected, actual) self._assert_total_rewrite(control_pipeline.state.func_ir.blocks, test_pipeline.state.func_ir.blocks) def test_cmp_op(self): ''' Verify that comparison operators are supported by the rewriter. ''' ns = self._test_root_function(are_roots_imaginary) self._assert_total_rewrite(ns.control_pipeline.state.func_ir.blocks, ns.test_pipeline.state.func_ir.blocks) def test_explicit_output(self): """ Check that ufunc calls with explicit outputs are not rewritten. """ ns = self._test_explicit_output_function(explicit_output) self._assert_no_rewrite(ns.control_pipeline.state.func_ir.blocks, ns.test_pipeline.state.func_ir.blocks) class TestRewriteIssues(MemoryLeakMixin, TestCase): def test_issue_1184(self): from numba import jit import numpy as np @jit(nopython=True) def foo(arr): return arr @jit(nopython=True) def bar(arr): c = foo(arr) d = foo(arr) # two calls to trigger rewrite return c, d arr = np.arange(10) out_c, out_d = bar(arr) self.assertIs(out_c, out_d) self.assertIs(out_c, arr) def test_issue_1264(self): n = 100 x = np.random.uniform(size=n*3).reshape((n,3)) expected = distance_matrix(x) actual = njit(distance_matrix)(x) np.testing.assert_array_almost_equal(expected, actual) # Avoid sporadic failures in MemoryLeakMixin.tearDown() gc.collect() def test_issue_1372(self): """Test array expression with duplicated term""" from numba import njit @njit def foo(a, b): b = np.sin(b) return b + b + a a = np.random.uniform(10) b = np.random.uniform(10) expect = foo.py_func(a, b) got = foo(a, b) np.testing.assert_allclose(got, expect) def test_unary_arrayexpr(self): """ Typing of unary array expression (np.negate) can be incorrect. """ @njit def foo(a, b): return b - a + -a b = 1.5 a = np.arange(10, dtype=np.int32) expect = foo.py_func(a, b) got = foo(a, b) self.assertPreciseEqual(got, expect) def test_bitwise_arrayexpr(self): """ Typing of bitwise boolean array expression can be incorrect (issue #1813). """ @njit def foo(a, b): return ~(a & (~b)) a = np.array([True, True, False, False]) b = np.array([False, True, False, True]) expect = foo.py_func(a, b) got = foo(a, b) self.assertPreciseEqual(got, expect) def test_annotations(self): """ Type annotation of array expressions with disambiguated variable names (issue #1466). """ cfunc = njit(variable_name_reuse) a = np.linspace(0, 1, 10) cfunc(a, a, a, a) buf = StringIO() cfunc.inspect_types(buf) res = buf.getvalue() self.assertIn("# u.1 = ", res) self.assertIn("# u.2 = ", res) def test_issue_5599_name_collision(self): # The original error will fail in lowering of the array_expr @njit def f(x): arr = np.ones(x) for _ in range(2): val = arr * arr arr = arr.copy() return arr got = f(5) expect = f.py_func(5) np.testing.assert_array_equal(got, expect) class TestSemantics(MemoryLeakMixin, unittest.TestCase): def test_division_by_zero(self): # Array expressions should follow the Numpy error model # i.e. 1./0. returns +inf instead of raising ZeroDivisionError pyfunc = div_add cfunc = njit(pyfunc) a = np.float64([0.0, 1.0, float('inf')]) b = np.float64([0.0, 0.0, 1.0]) c = np.ones_like(a) expect = pyfunc(a, b, c) got = cfunc(a, b, c) np.testing.assert_array_equal(expect, got) class TestOptionals(MemoryLeakMixin, unittest.TestCase): """ Tests the arrival and correct lowering of Optional types at a arrayexpr derived ufunc, see #3972""" def test_optional_scalar_type(self): @njit def arr_expr(x, y): return x + y @njit def do_call(x, y): if y > 0: z = None else: z = y return arr_expr(x, z) args = (np.arange(5), -1.2) # check result res = do_call(*args) expected = do_call.py_func(*args) np.testing.assert_allclose(res, expected) # check type s = arr_expr.signatures oty = s[0][1] self.assertTrue(isinstance(oty, types.Optional)) self.assertTrue(isinstance(oty.type, types.Float)) def test_optional_array_type(self): @njit def arr_expr(x, y): return x + y @njit def do_call(x, y): if y[0] > 0: z = None else: z = y return arr_expr(x, z) args = (np.arange(5), np.arange(5.)) # check result res = do_call(*args) expected = do_call.py_func(*args) np.testing.assert_allclose(res, expected) # check type s = arr_expr.signatures oty = s[0][1] self.assertTrue(isinstance(oty, types.Optional)) self.assertTrue(isinstance(oty.type, types.Array)) self.assertTrue(isinstance(oty.type.dtype, types.Float)) class TestOptionalsExceptions(MemoryLeakMixin, unittest.TestCase): # same as above, but the Optional resolves to None and TypeError's def test_optional_scalar_type_exception_on_none(self): self.disable_leak_check() @njit def arr_expr(x, y): return x + y @njit def do_call(x, y): if y > 0: z = None else: z = y return arr_expr(x, z) args = (np.arange(5), 1.0) # check result with self.assertRaises(TypeError) as raises: do_call(*args) self.assertIn("expected float64, got None", str(raises.exception)) # check type s = arr_expr.signatures oty = s[0][1] self.assertTrue(isinstance(oty, types.Optional)) self.assertTrue(isinstance(oty.type, types.Float)) def test_optional_array_type_exception_on_none(self): self.disable_leak_check() @njit def arr_expr(x, y): return x + y @njit def do_call(x, y): if y[0] > 0: z = None else: z = y return arr_expr(x, z) args = (np.arange(5), np.arange(1., 5.)) # check result with self.assertRaises(TypeError) as raises: do_call(*args) excstr = str(raises.exception) self.assertIn("expected array(float64,", excstr) self.assertIn("got None", excstr) # check type s = arr_expr.signatures oty = s[0][1] self.assertTrue(isinstance(oty, types.Optional)) self.assertTrue(isinstance(oty.type, types.Array)) self.assertTrue(isinstance(oty.type.dtype, types.Float)) class TestExternalTypes(MemoryLeakMixin, unittest.TestCase): """ Tests RewriteArrayExprs with external (user defined) types, see #5157""" source_lines = textwrap.dedent(""" from numba.core import types class FooType(types.Type): def __init__(self): super(FooType, self).__init__(name='Foo') """) def make_foo_type(self, FooType): class Foo(object): def __init__(self, value): self.value = value @register_model(FooType) class FooModel(models.StructModel): def __init__(self, dmm, fe_type): members = [("value", types.intp)] models.StructModel.__init__(self, dmm, fe_type, members) make_attribute_wrapper(FooType, "value", "value") @type_callable(Foo) def type_foo(context): def typer(value): return FooType() return typer @lower_builtin(Foo, types.intp) def impl_foo(context, builder, sig, args): typ = sig.return_type [value] = args foo = cgutils.create_struct_proxy(typ)(context, builder) foo.value = value return foo._getvalue() @typeof_impl.register(Foo) def typeof_foo(val, c): return FooType() return Foo, FooType def test_external_type(self): with create_temp_module(self.source_lines) as test_module: Foo, FooType = self.make_foo_type(test_module.FooType) # sum of foo class instance and array return an array # binary operation with foo class instance as one of args @overload(operator.add) def overload_foo_add(lhs, rhs): if isinstance(lhs, FooType) and isinstance(rhs, types.Array): def imp(lhs, rhs): return np.array([lhs.value, rhs[0]]) return imp # sum of 2 foo class instances return an array # binary operation with 2 foo class instances as args @overload(operator.add) def overload_foo_add(lhs, rhs): if isinstance(lhs, FooType) and isinstance(rhs, FooType): def imp(lhs, rhs): return np.array([lhs.value, rhs.value]) return imp # neg of foo class instance return an array # unary operation with foo class instance arg @overload(operator.neg) def overload_foo_neg(x): if isinstance(x, FooType): def imp(x): return np.array([-x.value]) return imp @njit def arr_expr_sum1(x, y): return Foo(x) + np.array([y]) @njit def arr_expr_sum2(x, y): return Foo(x) + Foo(y) @njit def arr_expr_neg(x): return -Foo(x) np.testing.assert_array_equal(arr_expr_sum1(0, 1), np.array([0, 1])) np.testing.assert_array_equal(arr_expr_sum2(2, 3), np.array([2, 3])) np.testing.assert_array_equal(arr_expr_neg(4), np.array([-4])) if __name__ == "__main__": unittest.main()