1141 lines
42 KiB
Python
1141 lines
42 KiB
Python
import itertools
|
|
|
|
import numpy as np
|
|
import sys
|
|
from collections import namedtuple
|
|
from io import StringIO
|
|
|
|
from numba import njit, typeof, prange
|
|
from numba.core import (
|
|
types,
|
|
typing,
|
|
ir,
|
|
bytecode,
|
|
postproc,
|
|
cpu,
|
|
registry,
|
|
utils,
|
|
)
|
|
from numba.tests.support import (TestCase, tag, skip_parfors_unsupported,
|
|
skip_unless_scipy)
|
|
from numba.parfors.array_analysis import EquivSet, ArrayAnalysis
|
|
from numba.core.compiler import Compiler, Flags, PassManager
|
|
from numba.core.ir_utils import remove_dead
|
|
from numba.core.untyped_passes import (ExtractByteCode, TranslateByteCode, FixupArgs,
|
|
IRProcessing, DeadBranchPrune,
|
|
RewriteSemanticConstants, GenericRewrites,
|
|
WithLifting, PreserveIR, InlineClosureLikes)
|
|
|
|
from numba.core.typed_passes import (NopythonTypeInference, AnnotateTypes,
|
|
NopythonRewrites, IRLegalization)
|
|
|
|
from numba.core.compiler_machinery import FunctionPass, PassManager, register_pass
|
|
from numba.experimental import jitclass
|
|
import unittest
|
|
|
|
|
|
skip_unsupported = skip_parfors_unsupported
|
|
|
|
|
|
# test class for #3700
|
|
@jitclass([('L', types.int32), ('T', types.int32)])
|
|
class ExampleClass3700(object):
|
|
def __init__(self, n):
|
|
self.L = n
|
|
self.T = n + 1
|
|
|
|
|
|
# test value for test_global_tuple
|
|
GVAL = (1.2,)
|
|
GVAL2 = (3, 4)
|
|
|
|
|
|
class TestEquivSet(TestCase):
|
|
|
|
"""
|
|
Test array_analysis.EquivSet.
|
|
"""
|
|
def test_insert_equiv(self):
|
|
s1 = EquivSet()
|
|
s1.insert_equiv('a', 'b')
|
|
self.assertTrue(s1.is_equiv('a', 'b'))
|
|
self.assertTrue(s1.is_equiv('b', 'a'))
|
|
s1.insert_equiv('c', 'd')
|
|
self.assertTrue(s1.is_equiv('c', 'd'))
|
|
self.assertFalse(s1.is_equiv('c', 'a'))
|
|
s1.insert_equiv('a', 'c')
|
|
self.assertTrue(s1.is_equiv('a', 'b', 'c', 'd'))
|
|
self.assertFalse(s1.is_equiv('a', 'e'))
|
|
|
|
def test_intersect(self):
|
|
s1 = EquivSet()
|
|
s2 = EquivSet()
|
|
r = s1.intersect(s2)
|
|
self.assertTrue(r.is_empty())
|
|
s1.insert_equiv('a', 'b')
|
|
r = s1.intersect(s2)
|
|
self.assertTrue(r.is_empty())
|
|
s2.insert_equiv('b', 'c')
|
|
r = s1.intersect(s2)
|
|
self.assertTrue(r.is_empty())
|
|
s2.insert_equiv('d', 'a')
|
|
r = s1.intersect(s2)
|
|
self.assertTrue(r.is_empty())
|
|
s1.insert_equiv('a', 'e')
|
|
s2.insert_equiv('c', 'd')
|
|
r = s1.intersect(s2)
|
|
self.assertTrue(r.is_equiv('a', 'b'))
|
|
self.assertFalse(r.is_equiv('a', 'e'))
|
|
self.assertFalse(r.is_equiv('c', 'd'))
|
|
|
|
|
|
@register_pass(analysis_only=False, mutates_CFG=True)
|
|
class ArrayAnalysisPass(FunctionPass):
|
|
_name = "array_analysis_pass"
|
|
|
|
def __init__(self):
|
|
FunctionPass.__init__(self)
|
|
|
|
def run_pass(self, state):
|
|
state.array_analysis = ArrayAnalysis(state.typingctx, state.func_ir,
|
|
state.typemap, state.calltypes)
|
|
state.array_analysis.run(state.func_ir.blocks)
|
|
post_proc = postproc.PostProcessor(state.func_ir)
|
|
post_proc.run()
|
|
state.func_ir_copies.append(state.func_ir.copy())
|
|
if state.test_idempotence and len(state.func_ir_copies) > 1:
|
|
state.test_idempotence(state.func_ir_copies)
|
|
return False
|
|
|
|
|
|
class ArrayAnalysisTester(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 = registry.cpu_target.typing_context
|
|
if target_context is None:
|
|
target_context = registry.cpu_target.target_context
|
|
return cls(typing_context, target_context, library, args, return_type,
|
|
flags, locals)
|
|
|
|
def compile_to_ir(self, func, test_idempotence=None):
|
|
"""
|
|
Populate and run compiler pipeline
|
|
"""
|
|
self.state.func_id = bytecode.FunctionIdentity.from_function(func)
|
|
ExtractByteCode().run_pass(self.state)
|
|
|
|
self.state.lifted = ()
|
|
self.state.lifted_from = None
|
|
state = self.state
|
|
state.func_ir_copies = []
|
|
state.test_idempotence = test_idempotence
|
|
|
|
name = 'array_analysis_testing'
|
|
pm = PassManager(name)
|
|
pm.add_pass(TranslateByteCode, "analyzing bytecode")
|
|
pm.add_pass(FixupArgs, "fix up args")
|
|
pm.add_pass(IRProcessing, "processing IR")
|
|
# pre typing
|
|
if not state.flags.no_rewrites:
|
|
pm.add_pass(GenericRewrites, "nopython rewrites")
|
|
pm.add_pass(RewriteSemanticConstants, "rewrite semantic constants")
|
|
pm.add_pass(DeadBranchPrune, "dead branch pruning")
|
|
pm.add_pass(InlineClosureLikes,
|
|
"inline calls to locally defined closures")
|
|
# typing
|
|
pm.add_pass(NopythonTypeInference, "nopython frontend")
|
|
|
|
if not state.flags.no_rewrites:
|
|
pm.add_pass(NopythonRewrites, "nopython rewrites")
|
|
|
|
# Array Analysis pass
|
|
pm.add_pass(ArrayAnalysisPass, "array analysis")
|
|
if test_idempotence:
|
|
# Do another pass of array analysis to test idempotence
|
|
pm.add_pass(ArrayAnalysisPass, "idempotence array analysis")
|
|
# legalise
|
|
pm.add_pass(IRLegalization, "ensure IR is legal prior to lowering")
|
|
pm.add_pass(AnnotateTypes, "annotate types")
|
|
|
|
# partial compile
|
|
pm.finalize()
|
|
pm.run(state)
|
|
return state.array_analysis
|
|
|
|
|
|
class TestArrayAnalysis(TestCase):
|
|
|
|
def compare_ir(self, ir_list):
|
|
outputs = []
|
|
for func_ir in ir_list:
|
|
remove_dead(func_ir.blocks, func_ir.arg_names, func_ir)
|
|
output = StringIO()
|
|
func_ir.dump(file=output)
|
|
outputs.append(output.getvalue())
|
|
self.assertTrue(len(set(outputs)) == 1) # assert all outputs are equal
|
|
|
|
def _compile_and_test(self, fn, arg_tys, asserts=[], equivs=[], idempotent=True):
|
|
"""
|
|
Compile the given function and get its IR.
|
|
"""
|
|
test_pipeline = ArrayAnalysisTester.mk_pipeline(arg_tys)
|
|
test_idempotence = self.compare_ir if idempotent else lambda x:()
|
|
analysis = test_pipeline.compile_to_ir(fn, test_idempotence)
|
|
if equivs:
|
|
for func in equivs:
|
|
# only test the equiv_set of the first block
|
|
func(analysis.equiv_sets[0])
|
|
if asserts is None:
|
|
self.assertTrue(self._has_no_assertcall(analysis.func_ir))
|
|
else:
|
|
for func in asserts:
|
|
func(analysis.func_ir, analysis.typemap)
|
|
|
|
def _has_assertcall(self, func_ir, typemap, args):
|
|
msg = "Sizes of {} do not match".format(', '.join(args))
|
|
for label, block in func_ir.blocks.items():
|
|
for expr in block.find_exprs(op='call'):
|
|
fn = func_ir.get_definition(expr.func.name)
|
|
if isinstance(fn, ir.Global) and fn.name == 'assert_equiv':
|
|
typ = typemap[expr.args[0].name]
|
|
if typ.literal_value.startswith(msg):
|
|
return True
|
|
return False
|
|
|
|
def _has_shapecall(self, func_ir, x):
|
|
for label, block in func_ir.blocks.items():
|
|
for expr in block.find_exprs(op='getattr'):
|
|
if expr.attr == 'shape':
|
|
y = func_ir.get_definition(expr.value, lhs_only=True)
|
|
z = func_ir.get_definition(x, lhs_only=True)
|
|
y = y.name if isinstance(y, ir.Var) else y
|
|
z = z.name if isinstance(z, ir.Var) else z
|
|
if y == z:
|
|
return True
|
|
return False
|
|
|
|
def _has_no_assertcall(self, func_ir):
|
|
for label, block in func_ir.blocks.items():
|
|
for expr in block.find_exprs(op='call'):
|
|
fn = func_ir.get_definition(expr.func.name)
|
|
if isinstance(fn, ir.Global) and fn.name == 'assert_equiv':
|
|
return False
|
|
return True
|
|
|
|
def with_assert(self, *args):
|
|
return lambda func_ir, typemap: self.assertTrue(
|
|
self._has_assertcall(func_ir, typemap, args))
|
|
|
|
def without_assert(self, *args):
|
|
return lambda func_ir, typemap: self.assertFalse(
|
|
self._has_assertcall(func_ir, typemap, args))
|
|
|
|
def with_equiv(self, *args):
|
|
def check(equiv_set):
|
|
n = len(args)
|
|
for i in range(n - 1):
|
|
if not equiv_set.is_equiv(args[i], args[n - 1]):
|
|
return False
|
|
return True
|
|
return lambda equiv_set: self.assertTrue(check(equiv_set))
|
|
|
|
def without_equiv(self, *args):
|
|
def check(equiv_set):
|
|
n = len(args)
|
|
for i in range(n - 1):
|
|
if equiv_set.is_equiv(args[i], args[n - 1]):
|
|
return False
|
|
return True
|
|
return lambda equiv_set: self.assertTrue(check(equiv_set))
|
|
|
|
def with_shapecall(self, x):
|
|
return lambda func_ir, s: self.assertTrue(self._has_shapecall(func_ir, x))
|
|
|
|
def without_shapecall(self, x):
|
|
return lambda func_ir, s: self.assertFalse(self._has_shapecall(func_ir, x))
|
|
|
|
def test_base_cases(self):
|
|
def test_0():
|
|
a = np.zeros(0)
|
|
b = np.zeros(1)
|
|
m = 0
|
|
n = 1
|
|
c = np.zeros((m, n))
|
|
return
|
|
self._compile_and_test(test_0, (),
|
|
equivs=[self.with_equiv('a', (0,)),
|
|
self.with_equiv('b', (1,)),
|
|
self.with_equiv('c', (0, 1))])
|
|
|
|
def test_1(n):
|
|
a = np.zeros(n)
|
|
b = np.zeros(n)
|
|
return a + b
|
|
self._compile_and_test(test_1, (types.intp,), asserts=None)
|
|
|
|
def test_2(m, n):
|
|
a = np.zeros(n)
|
|
b = np.zeros(m)
|
|
return a + b
|
|
self._compile_and_test(test_2, (types.intp, types.intp),
|
|
asserts=[self.with_assert('a', 'b')])
|
|
|
|
def test_3(n):
|
|
a = np.zeros(n)
|
|
return a + n
|
|
self._compile_and_test(test_3, (types.intp,), asserts=None)
|
|
|
|
def test_4(n):
|
|
a = np.zeros(n)
|
|
b = a + 1
|
|
c = a + 2
|
|
return a + c
|
|
self._compile_and_test(test_4, (types.intp,), asserts=None)
|
|
|
|
def test_5(n):
|
|
a = np.zeros((n, n))
|
|
m = n
|
|
b = np.zeros((m, n))
|
|
return a + b
|
|
self._compile_and_test(test_5, (types.intp,), asserts=None)
|
|
|
|
def test_6(m, n):
|
|
a = np.zeros(n)
|
|
b = np.zeros(m)
|
|
d = a + b
|
|
e = a - b
|
|
return d + e
|
|
self._compile_and_test(test_6, (types.intp, types.intp),
|
|
asserts=[self.with_assert('a', 'b'),
|
|
self.without_assert('d', 'e')])
|
|
|
|
def test_7(m, n):
|
|
a = np.zeros(n)
|
|
b = np.zeros(m)
|
|
if m == 10:
|
|
d = a + b
|
|
else:
|
|
d = a - b
|
|
return d + a
|
|
self._compile_and_test(test_7, (types.intp, types.intp),
|
|
asserts=[self.with_assert('a', 'b'),
|
|
self.without_assert('d', 'a')])
|
|
|
|
def test_8(m, n):
|
|
a = np.zeros(n)
|
|
b = np.zeros(m)
|
|
if m == 10:
|
|
d = b + a
|
|
else:
|
|
d = a + a
|
|
return b + d
|
|
self._compile_and_test(test_8, (types.intp, types.intp),
|
|
asserts=[self.with_assert('b', 'a'),
|
|
self.with_assert('b', 'd')])
|
|
|
|
def test_9(m):
|
|
A = np.ones(m)
|
|
s = 0
|
|
while m < 2:
|
|
m += 1
|
|
B = np.ones(m)
|
|
s += np.sum(A + B)
|
|
return s
|
|
self._compile_and_test(test_9, (types.intp,),
|
|
asserts=[self.with_assert('A', 'B')])
|
|
|
|
def test_10(m, n):
|
|
p = m - 1
|
|
q = n + 1
|
|
r = q + 1
|
|
A = np.zeros(p)
|
|
B = np.zeros(q)
|
|
C = np.zeros(r)
|
|
D = np.zeros(m)
|
|
s = np.sum(A + B)
|
|
t = np.sum(C + D)
|
|
return s + t
|
|
self._compile_and_test(test_10, (types.intp,types.intp,),
|
|
asserts=[self.with_assert('A', 'B'),
|
|
self.without_assert('C', 'D')])
|
|
|
|
def test_11():
|
|
a = np.ones(5)
|
|
b = np.ones(5)
|
|
c = a[1:]
|
|
d = b[:-1]
|
|
e = len(c)
|
|
f = len(d)
|
|
return e == f
|
|
self._compile_and_test(test_11, (),
|
|
equivs=[self.with_equiv('e', 'f')])
|
|
|
|
def test_12():
|
|
a = np.ones(25).reshape((5,5))
|
|
b = np.ones(25).reshape((5,5))
|
|
c = a[1:,:]
|
|
d = b[:-1,:]
|
|
e = c.shape[0]
|
|
f = d.shape[0]
|
|
g = len(d)
|
|
return e == f
|
|
self._compile_and_test(test_12, (),
|
|
equivs=[self.with_equiv('e', 'f', 'g')])
|
|
|
|
def test_tup_arg(T):
|
|
T2 = T
|
|
return T2[0]
|
|
|
|
int_arr_typ = types.Array(types.intp, 1, 'C')
|
|
self._compile_and_test(test_tup_arg,
|
|
(types.Tuple((int_arr_typ, int_arr_typ)),), asserts=None)
|
|
|
|
def test_arr_in_tup(m):
|
|
A = np.ones(m)
|
|
S = (A,)
|
|
B = np.ones(len(S[0]))
|
|
return B
|
|
|
|
self._compile_and_test(test_arr_in_tup, (types.intp,),
|
|
equivs=[self.with_equiv("A", "B")])
|
|
|
|
T = namedtuple("T", ['a','b'])
|
|
def test_namedtuple(n):
|
|
r = T(n, n)
|
|
return r[0]
|
|
self._compile_and_test(test_namedtuple, (types.intp,),
|
|
equivs=[self.with_equiv('r', ('n', 'n'))],)
|
|
|
|
# np.where is tricky since it returns tuple of arrays
|
|
def test_np_where_tup_return(A):
|
|
c = np.where(A)
|
|
return len(c[0])
|
|
|
|
self._compile_and_test(test_np_where_tup_return,
|
|
(types.Array(types.intp, 1, 'C'),), asserts=None)
|
|
|
|
def test_shape(A):
|
|
(m, n) = A.shape
|
|
B = np.ones((m, n))
|
|
return A + B
|
|
self._compile_and_test(test_shape, (types.Array(types.intp, 2, 'C'),),
|
|
asserts=None)
|
|
|
|
def test_cond(l, m, n):
|
|
A = np.ones(l)
|
|
B = np.ones(m)
|
|
C = np.ones(n)
|
|
if l == m:
|
|
r = np.sum(A + B)
|
|
else:
|
|
r = 0
|
|
if m != n:
|
|
s = 0
|
|
else:
|
|
s = np.sum(B + C)
|
|
t = 0
|
|
if l == m:
|
|
if m == n:
|
|
t = np.sum(A + B + C)
|
|
return r + s + t
|
|
self._compile_and_test(test_cond, (types.intp, types.intp, types.intp),
|
|
asserts=None)
|
|
|
|
def test_assert_1(m, n):
|
|
assert(m == n)
|
|
A = np.ones(m)
|
|
B = np.ones(n)
|
|
return np.sum(A + B)
|
|
self._compile_and_test(test_assert_1, (types.intp, types.intp),
|
|
asserts=None)
|
|
|
|
def test_assert_2(A, B):
|
|
assert(A.shape == B.shape)
|
|
return np.sum(A + B)
|
|
|
|
self._compile_and_test(test_assert_2, (types.Array(types.intp, 1, 'C'),
|
|
types.Array(types.intp, 1, 'C'),),
|
|
asserts=None)
|
|
self._compile_and_test(test_assert_2, (types.Array(types.intp, 2, 'C'),
|
|
types.Array(types.intp, 2, 'C'),),
|
|
asserts=None)
|
|
# expected failure
|
|
with self.assertRaises(AssertionError) as raises:
|
|
self._compile_and_test(test_assert_2, (types.Array(types.intp, 1, 'C'),
|
|
types.Array(types.intp, 2, 'C'),),
|
|
asserts=None)
|
|
msg = "Dimension mismatch"
|
|
self.assertIn(msg, str(raises.exception))
|
|
|
|
|
|
def test_stencilcall(self):
|
|
from numba.stencils.stencil import stencil
|
|
@stencil
|
|
def kernel_1(a):
|
|
return 0.25 * (a[0,1] + a[1,0] + a[0,-1] + a[-1,0])
|
|
|
|
def test_1(n):
|
|
a = np.ones((n,n))
|
|
b = kernel_1(a)
|
|
return a + b
|
|
|
|
self._compile_and_test(test_1, (types.intp,),
|
|
equivs=[self.with_equiv('a', 'b')],
|
|
asserts=[self.without_assert('a', 'b')])
|
|
|
|
def test_2(n):
|
|
a = np.ones((n,n))
|
|
b = np.ones((n+1,n+1))
|
|
kernel_1(a, out=b)
|
|
return a
|
|
|
|
self._compile_and_test(test_2, (types.intp,),
|
|
equivs=[self.without_equiv('a', 'b')])
|
|
|
|
@stencil(standard_indexing=('c',))
|
|
def kernel_2(a, b, c):
|
|
return a[0,1,0] + b[0,-1,0] + c[0]
|
|
|
|
def test_3(n):
|
|
a = np.arange(64).reshape(4,8,2)
|
|
b = np.arange(64).reshape(n,8,2)
|
|
u = np.zeros(1)
|
|
v = kernel_2(a, b, u)
|
|
return v
|
|
|
|
# standard indexed arrays are not considered in size equivalence
|
|
self._compile_and_test(test_3, (types.intp,),
|
|
equivs=[self.with_equiv('a', 'b', 'v'),
|
|
self.without_equiv('a', 'u')],
|
|
asserts=[self.with_assert('a', 'b')])
|
|
|
|
def test_slice(self):
|
|
def test_1(m, n):
|
|
A = np.zeros(m)
|
|
B = np.zeros(n)
|
|
s = np.sum(A + B)
|
|
C = A[1:m-1]
|
|
D = B[1:n-1]
|
|
t = np.sum(C + D)
|
|
return s + t
|
|
self._compile_and_test(test_1, (types.intp,types.intp,),
|
|
asserts=[self.with_assert('A', 'B'),
|
|
self.without_assert('C', 'D')],
|
|
idempotent=False)
|
|
|
|
def test_2(m):
|
|
A = np.zeros(m)
|
|
B = A[0:m-3]
|
|
C = A[1:m-2]
|
|
D = A[2:m-1]
|
|
E = B + C
|
|
return D + E
|
|
self._compile_and_test(test_2, (types.intp,),
|
|
asserts=[self.without_assert('B', 'C'),
|
|
self.without_assert('D', 'E')],
|
|
idempotent=False)
|
|
|
|
def test_3(m):
|
|
A = np.zeros((m,m))
|
|
B = A[0:m-2,0:m-2]
|
|
C = A[1:m-1,1:m-1]
|
|
E = B + C
|
|
return E
|
|
self._compile_and_test(test_3, (types.intp,),
|
|
asserts=[self.without_assert('B', 'C')],
|
|
idempotent=False)
|
|
|
|
def test_4(m):
|
|
A = np.zeros((m,m))
|
|
B = A[0:m-2,:]
|
|
C = A[1:m-1,:]
|
|
E = B + C
|
|
return E
|
|
self._compile_and_test(test_4, (types.intp,),
|
|
asserts=[self.without_assert('B', 'C')],
|
|
idempotent=False)
|
|
|
|
def test_5(m,n):
|
|
A = np.zeros(m)
|
|
B = np.zeros(m)
|
|
B[0:m-2] = A[1:m-1]
|
|
C = np.zeros(n)
|
|
D = A[1:m-1]
|
|
C[0:n-2] = D
|
|
# B and C are not necessarily of the same size because we can't
|
|
# derive m == n from (m-2) % m == (n-2) % n
|
|
return B + C
|
|
self._compile_and_test(test_5, (types.intp,types.intp),
|
|
asserts=[self.without_assert('B', 'A'),
|
|
self.with_assert('C', 'D'),
|
|
self.with_assert('B', 'C')],
|
|
idempotent=False)
|
|
|
|
def test_6(m):
|
|
A = np.zeros((m,m))
|
|
B = A[0:m-2,:-1]
|
|
C = A[1:m-1,:-1]
|
|
E = B + C
|
|
return E
|
|
self._compile_and_test(test_6, (types.intp,),
|
|
asserts=[self.without_assert('B', 'C')],
|
|
idempotent=False)
|
|
|
|
def test_7(m):
|
|
A = np.zeros((m,m))
|
|
B = A[0:m-2,-3:-1]
|
|
C = A[1:m-1,-4:-2]
|
|
E = B + C
|
|
return E
|
|
self._compile_and_test(test_7, (types.intp,),
|
|
asserts=[self.with_assert('B', 'C')],
|
|
idempotent=False)
|
|
|
|
def test_8(m):
|
|
A = np.zeros((m,m))
|
|
B = A[:m-2,0:]
|
|
C = A[1:-1,:]
|
|
E = B + C
|
|
return E
|
|
self._compile_and_test(test_8, (types.intp,),
|
|
asserts=[self.without_assert('B', 'C')],
|
|
idempotent=False)
|
|
|
|
def test_9(m):
|
|
# issues #3461 and #3554, checks equivalence on empty slices
|
|
# and across binop
|
|
A = np.zeros((m))
|
|
B = A[:0] # B = array([], dtype=int64)
|
|
C = A[1:]
|
|
D = A[:-1:-1] # D = array([], dtype=int64)
|
|
E = B + D
|
|
F = E
|
|
F += 1 # F = array([], dtype=int64)
|
|
return A, C, F
|
|
self._compile_and_test(test_9, (types.intp,),
|
|
equivs=[self.without_equiv('B', 'C'),
|
|
self.with_equiv('A', 'm'),
|
|
self.with_equiv('B', 'D'),
|
|
self.with_equiv('F', 'D'),],
|
|
idempotent=False)
|
|
|
|
@skip_unless_scipy
|
|
def test_numpy_calls(self):
|
|
def test_zeros(n):
|
|
a = np.zeros(n)
|
|
b = np.zeros((n, n))
|
|
c = np.zeros(shape=(n, n))
|
|
self._compile_and_test(test_zeros, (types.intp,),
|
|
equivs=[self.with_equiv('a', 'n'),
|
|
self.with_equiv('b', ('n', 'n')),
|
|
self.with_equiv('b', 'c')])
|
|
|
|
def test_0d_array(n):
|
|
a = np.array(1)
|
|
b = np.ones(2)
|
|
return a + b
|
|
self._compile_and_test(test_0d_array, (types.intp,),
|
|
equivs=[self.without_equiv('a', 'b')],
|
|
asserts=[self.without_shapecall('a')])
|
|
|
|
def test_ones(n):
|
|
a = np.ones(n)
|
|
b = np.ones((n, n))
|
|
c = np.ones(shape=(n, n))
|
|
self._compile_and_test(test_ones, (types.intp,),
|
|
equivs=[self.with_equiv('a', 'n'),
|
|
self.with_equiv('b', ('n', 'n')),
|
|
self.with_equiv('b', 'c')])
|
|
|
|
def test_empty(n):
|
|
a = np.empty(n)
|
|
b = np.empty((n, n))
|
|
c = np.empty(shape=(n, n))
|
|
self._compile_and_test(test_empty, (types.intp,),
|
|
equivs=[self.with_equiv('a', 'n'),
|
|
self.with_equiv('b', ('n', 'n')),
|
|
self.with_equiv('b', 'c')])
|
|
|
|
def test_eye(n):
|
|
a = np.eye(n)
|
|
b = np.eye(N=n)
|
|
c = np.eye(N=n, M=n)
|
|
d = np.eye(N=n, M=n + 1)
|
|
self._compile_and_test(test_eye, (types.intp,),
|
|
equivs=[self.with_equiv('a', ('n', 'n')),
|
|
self.with_equiv('b', ('n', 'n')),
|
|
self.with_equiv('b', 'c'),
|
|
self.without_equiv('b', 'd')])
|
|
|
|
def test_identity(n):
|
|
a = np.identity(n)
|
|
self._compile_and_test(test_identity, (types.intp,),
|
|
equivs=[self.with_equiv('a', ('n', 'n'))])
|
|
|
|
def test_diag(n):
|
|
a = np.identity(n)
|
|
b = np.diag(a)
|
|
c = np.diag(b)
|
|
d = np.diag(a, k=1)
|
|
self._compile_and_test(test_diag, (types.intp,),
|
|
equivs=[self.with_equiv('b', ('n',)),
|
|
self.with_equiv('c', ('n', 'n'))],
|
|
asserts=[self.with_shapecall('d'),
|
|
self.without_shapecall('c')])
|
|
|
|
def test_array_like(a):
|
|
b = np.empty_like(a)
|
|
c = np.zeros_like(a)
|
|
d = np.ones_like(a)
|
|
e = np.full_like(a, 1)
|
|
f = np.asfortranarray(a)
|
|
|
|
self._compile_and_test(test_array_like, (types.Array(types.intp, 2, 'C'),),
|
|
equivs=[
|
|
self.with_equiv('a', 'b', 'd', 'e', 'f')],
|
|
asserts=[self.with_shapecall('a'),
|
|
self.without_shapecall('b')])
|
|
|
|
def test_reshape(n):
|
|
a = np.ones(n * n)
|
|
b = a.reshape((n, n))
|
|
return a.sum() + b.sum()
|
|
self._compile_and_test(test_reshape, (types.intp,),
|
|
equivs=[self.with_equiv('b', ('n', 'n'))],
|
|
asserts=[self.without_shapecall('b')])
|
|
|
|
|
|
def test_transpose(m, n):
|
|
a = np.ones((m, n))
|
|
b = a.T
|
|
c = a.transpose()
|
|
# Numba njit cannot compile explicit transpose call!
|
|
# c = np.transpose(b)
|
|
self._compile_and_test(test_transpose, (types.intp, types.intp),
|
|
equivs=[self.with_equiv('a', ('m', 'n')),
|
|
self.with_equiv('b', ('n', 'm')),
|
|
self.with_equiv('c', ('n', 'm'))])
|
|
|
|
|
|
def test_transpose_3d(m, n, k):
|
|
a = np.ones((m, n, k))
|
|
b = a.T
|
|
c = a.transpose()
|
|
d = a.transpose(2,0,1)
|
|
dt = a.transpose((2,0,1))
|
|
e = a.transpose(0,2,1)
|
|
et = a.transpose((0,2,1))
|
|
# Numba njit cannot compile explicit transpose call!
|
|
# c = np.transpose(b)
|
|
self._compile_and_test(test_transpose_3d, (types.intp, types.intp, types.intp),
|
|
equivs=[self.with_equiv('a', ('m', 'n', 'k')),
|
|
self.with_equiv('b', ('k', 'n', 'm')),
|
|
self.with_equiv('c', ('k', 'n', 'm')),
|
|
self.with_equiv('d', ('k', 'm', 'n')),
|
|
self.with_equiv('dt', ('k', 'm', 'n')),
|
|
self.with_equiv('e', ('m', 'k', 'n')),
|
|
self.with_equiv('et', ('m', 'k', 'n'))])
|
|
|
|
def test_real_imag_attr(m, n):
|
|
a = np.ones((m, n))
|
|
b = a.real
|
|
c = a.imag
|
|
|
|
self._compile_and_test(test_real_imag_attr, (types.intp, types.intp),
|
|
equivs=[self.with_equiv('a', ('m', 'n')),
|
|
self.with_equiv('b', ('m', 'n')),
|
|
self.with_equiv('c', ('m', 'n')),
|
|
])
|
|
|
|
def test_random(n):
|
|
a0 = np.random.rand(n)
|
|
a1 = np.random.rand(n, n)
|
|
b0 = np.random.randn(n)
|
|
b1 = np.random.randn(n, n)
|
|
c0 = np.random.ranf(n)
|
|
c1 = np.random.ranf((n, n))
|
|
c2 = np.random.ranf(size=(n, n))
|
|
d0 = np.random.random_sample(n)
|
|
d1 = np.random.random_sample((n, n))
|
|
d2 = np.random.random_sample(size=(n, n))
|
|
e0 = np.random.sample(n)
|
|
e1 = np.random.sample((n, n))
|
|
e2 = np.random.sample(size=(n, n))
|
|
f0 = np.random.random(n)
|
|
f1 = np.random.random((n, n))
|
|
f2 = np.random.random(size=(n, n))
|
|
g0 = np.random.standard_normal(n)
|
|
g1 = np.random.standard_normal((n, n))
|
|
g2 = np.random.standard_normal(size=(n, n))
|
|
h0 = np.random.chisquare(10, n)
|
|
h1 = np.random.chisquare(10, (n, n))
|
|
h2 = np.random.chisquare(10, size=(n, n))
|
|
i0 = np.random.weibull(10, n)
|
|
i1 = np.random.weibull(10, (n, n))
|
|
i2 = np.random.weibull(10, size=(n, n))
|
|
j0 = np.random.power(10, n)
|
|
j1 = np.random.power(10, (n, n))
|
|
j2 = np.random.power(10, size=(n, n))
|
|
k0 = np.random.geometric(0.1, n)
|
|
k1 = np.random.geometric(0.1, (n, n))
|
|
k2 = np.random.geometric(0.1, size=(n, n))
|
|
l0 = np.random.exponential(10, n)
|
|
l1 = np.random.exponential(10, (n, n))
|
|
l2 = np.random.exponential(10, size=(n, n))
|
|
m0 = np.random.poisson(10, n)
|
|
m1 = np.random.poisson(10, (n, n))
|
|
m2 = np.random.poisson(10, size=(n, n))
|
|
n0 = np.random.rayleigh(10, n)
|
|
n1 = np.random.rayleigh(10, (n, n))
|
|
n2 = np.random.rayleigh(10, size=(n, n))
|
|
o0 = np.random.normal(0, 1, n)
|
|
o1 = np.random.normal(0, 1, (n, n))
|
|
o2 = np.random.normal(0, 1, size=(n, n))
|
|
p0 = np.random.uniform(0, 1, n)
|
|
p1 = np.random.uniform(0, 1, (n, n))
|
|
p2 = np.random.uniform(0, 1, size=(n, n))
|
|
q0 = np.random.beta(0.1, 1, n)
|
|
q1 = np.random.beta(0.1, 1, (n, n))
|
|
q2 = np.random.beta(0.1, 1, size=(n, n))
|
|
r0 = np.random.binomial(0, 1, n)
|
|
r1 = np.random.binomial(0, 1, (n, n))
|
|
r2 = np.random.binomial(0, 1, size=(n, n))
|
|
s0 = np.random.f(0.1, 1, n)
|
|
s1 = np.random.f(0.1, 1, (n, n))
|
|
s2 = np.random.f(0.1, 1, size=(n, n))
|
|
t0 = np.random.gamma(0.1, 1, n)
|
|
t1 = np.random.gamma(0.1, 1, (n, n))
|
|
t2 = np.random.gamma(0.1, 1, size=(n, n))
|
|
u0 = np.random.lognormal(0, 1, n)
|
|
u1 = np.random.lognormal(0, 1, (n, n))
|
|
u2 = np.random.lognormal(0, 1, size=(n, n))
|
|
v0 = np.random.laplace(0, 1, n)
|
|
v1 = np.random.laplace(0, 1, (n, n))
|
|
v2 = np.random.laplace(0, 1, size=(n, n))
|
|
w0 = np.random.randint(0, 10, n)
|
|
w1 = np.random.randint(0, 10, (n, n))
|
|
w2 = np.random.randint(0, 10, size=(n, n))
|
|
x0 = np.random.triangular(-3, 0, 10, n)
|
|
x1 = np.random.triangular(-3, 0, 10, (n, n))
|
|
x2 = np.random.triangular(-3, 0, 10, size=(n, n))
|
|
|
|
last = ord('x') + 1
|
|
vars1d = [('n',)] + [chr(x) + '0' for x in range(ord('a'), last)]
|
|
vars2d = [('n', 'n')] + [chr(x) + '1' for x in range(ord('a'), last)]
|
|
vars2d += [chr(x) + '1' for x in range(ord('c'), last)]
|
|
self._compile_and_test(test_random, (types.intp,),
|
|
equivs=[self.with_equiv(*vars1d),
|
|
self.with_equiv(*vars2d)])
|
|
|
|
def test_concatenate(m, n):
|
|
a = np.ones(m)
|
|
b = np.ones(n)
|
|
c = np.concatenate((a, b))
|
|
d = np.ones((2, n))
|
|
e = np.ones((3, n))
|
|
f = np.concatenate((d, e))
|
|
# Numba njit cannot compile concatenate with single array!
|
|
# g = np.ones((3,4,5))
|
|
# h = np.concatenate(g)
|
|
i = np.ones((m, 2))
|
|
j = np.ones((m, 3))
|
|
k = np.concatenate((i, j), axis=1)
|
|
l = np.ones((m, n))
|
|
o = np.ones((m, n))
|
|
p = np.concatenate((l, o))
|
|
# Numba njit cannot support list argument!
|
|
# q = np.concatenate([d, e])
|
|
self._compile_and_test(test_concatenate, (types.intp, types.intp),
|
|
equivs=[self.with_equiv('f', (5, 'n')),
|
|
#self.with_equiv('h', (3 + 4 + 5, )),
|
|
self.with_equiv('k', ('m', 5))],
|
|
asserts=[self.with_shapecall('c'),
|
|
self.without_shapecall('f'),
|
|
self.without_shapecall('k'),
|
|
self.with_shapecall('p')])
|
|
|
|
def test_vsd_stack():
|
|
k = np.ones((2,))
|
|
l = np.ones((2, 3))
|
|
o = np.ones((2, 3, 4))
|
|
p = np.vstack((k, k))
|
|
q = np.vstack((l, l))
|
|
r = np.hstack((k, k))
|
|
s = np.hstack((l, l))
|
|
t = np.dstack((k, k))
|
|
u = np.dstack((l, l))
|
|
v = np.dstack((o, o))
|
|
|
|
self._compile_and_test(test_vsd_stack, (),
|
|
equivs=[self.with_equiv('p', (2, 2)),
|
|
self.with_equiv('q', (4, 3)),
|
|
self.with_equiv('r', (4,)),
|
|
self.with_equiv('s', (2, 6)),
|
|
self.with_equiv('t', (1, 2, 2)),
|
|
self.with_equiv('u', (2, 3, 2)),
|
|
self.with_equiv('v', (2, 3, 8)),
|
|
])
|
|
|
|
def test_stack(m, n):
|
|
a = np.ones(m)
|
|
b = np.ones(n)
|
|
c = np.stack((a, b))
|
|
d = np.ones((m, n))
|
|
e = np.ones((m, n))
|
|
f = np.stack((d, e))
|
|
g = np.stack((d, e), axis=0)
|
|
h = np.stack((d, e), axis=1)
|
|
i = np.stack((d, e), axis=2)
|
|
j = np.stack((d, e), axis=-1)
|
|
|
|
self._compile_and_test(test_stack, (types.intp, types.intp),
|
|
equivs=[self.with_equiv('m', 'n'),
|
|
self.with_equiv('c', (2, 'm')),
|
|
self.with_equiv(
|
|
'f', 'g', (2, 'm', 'n')),
|
|
self.with_equiv(
|
|
'h', ('m', 2, 'n')),
|
|
self.with_equiv(
|
|
'i', 'j', ('m', 'n', 2)),
|
|
])
|
|
|
|
def test_linspace(m, n):
|
|
a = np.linspace(m, n)
|
|
b = np.linspace(m, n, 10)
|
|
# Numba njit does not support num keyword to linspace call!
|
|
# c = np.linspace(m,n,num=10)
|
|
self._compile_and_test(test_linspace, (types.float64, types.float64),
|
|
equivs=[self.with_equiv('a', (50,)),
|
|
self.with_equiv('b', (10,))])
|
|
|
|
def test_dot(l, m, n):
|
|
a = np.dot(np.ones(1), np.ones(1))
|
|
b = np.dot(np.ones(2), np.ones((2, 3)))
|
|
# Numba njit does not support higher dimensional inputs
|
|
#c = np.dot(np.ones(2),np.ones((3,2,4)))
|
|
#d = np.dot(np.ones(2),np.ones((3,5,2,4)))
|
|
e = np.dot(np.ones((1, 2)), np.ones(2,))
|
|
#f = np.dot(np.ones((1,2,3)),np.ones(3,))
|
|
#g = np.dot(np.ones((1,2,3,4)),np.ones(4,))
|
|
h = np.dot(np.ones((2, 3)), np.ones((3, 4)))
|
|
i = np.dot(np.ones((m, n)), np.ones((n, m)))
|
|
j = np.dot(np.ones((m, m)), np.ones((l, l)))
|
|
|
|
self._compile_and_test(test_dot, (types.intp, types.intp, types.intp),
|
|
equivs=[self.without_equiv('a', (1,)), # not array
|
|
self.with_equiv('b', (3,)),
|
|
self.with_equiv('e', (1,)),
|
|
self.with_equiv('h', (2, 4)),
|
|
self.with_equiv('i', ('m', 'm')),
|
|
self.with_equiv('j', ('m', 'm')),
|
|
],
|
|
asserts=[self.with_assert('m', 'l')])
|
|
|
|
def test_broadcast(m, n):
|
|
a = np.ones((m, n))
|
|
b = np.ones(n)
|
|
c = a + b
|
|
d = np.ones((1, n))
|
|
e = a + c - d
|
|
self._compile_and_test(test_broadcast, (types.intp, types.intp),
|
|
equivs=[self.with_equiv('a', 'c', 'e')],
|
|
asserts=None)
|
|
|
|
# make sure shape of a global tuple of ints is handled properly
|
|
def test_global_tuple():
|
|
a = np.ones(GVAL2)
|
|
b = np.ones(GVAL2)
|
|
|
|
self._compile_and_test(test_global_tuple, (),
|
|
equivs=[self.with_equiv('a', 'b')],
|
|
asserts=None)
|
|
|
|
|
|
class TestArrayAnalysisParallelRequired(TestCase):
|
|
"""This is to just split out tests that need the parallel backend and
|
|
therefore serialised execution.
|
|
"""
|
|
|
|
_numba_parallel_test_ = False
|
|
|
|
@skip_unsupported
|
|
def test_misc(self):
|
|
|
|
@njit
|
|
def swap(x, y):
|
|
return(y, x)
|
|
|
|
def test_bug2537(m):
|
|
a = np.ones(m)
|
|
b = np.ones(m)
|
|
for i in range(m):
|
|
a[i], b[i] = swap(a[i], b[i])
|
|
|
|
try:
|
|
njit(test_bug2537, parallel=True)(10)
|
|
except IndexError:
|
|
self.fail("test_bug2537 raised IndexError!")
|
|
|
|
@skip_unsupported
|
|
def test_global_namedtuple(self):
|
|
Row = namedtuple('Row', ['A'])
|
|
row = Row(3)
|
|
|
|
def test_impl():
|
|
rr = row
|
|
res = rr.A
|
|
if res == 2:
|
|
res = 3
|
|
return res
|
|
|
|
self.assertEqual(njit(test_impl, parallel=True)(), test_impl())
|
|
|
|
@skip_unsupported
|
|
def test_array_T_issue_3700(self):
|
|
|
|
def test_impl(t_obj, X):
|
|
for i in prange(t_obj.T):
|
|
X[i] = i
|
|
return X.sum()
|
|
|
|
n = 5
|
|
t_obj = ExampleClass3700(n)
|
|
X1 = np.zeros(t_obj.T)
|
|
X2 = np.zeros(t_obj.T)
|
|
self.assertEqual(
|
|
njit(test_impl, parallel=True)(t_obj, X1), test_impl(t_obj, X2))
|
|
|
|
@skip_unsupported
|
|
def test_slice_shape_issue_3380(self):
|
|
# these tests shouldn't throw error in array analysis
|
|
def test_impl1():
|
|
a = slice(None, None)
|
|
return True
|
|
|
|
self.assertEqual(njit(test_impl1, parallel=True)(), test_impl1())
|
|
|
|
def test_impl2(A, a):
|
|
b = a
|
|
return A[b]
|
|
|
|
A = np.arange(10)
|
|
a = slice(None)
|
|
np.testing.assert_array_equal(
|
|
njit(test_impl2, parallel=True)(A, a), test_impl2(A, a))
|
|
|
|
@skip_unsupported
|
|
def test_slice_dtype_issue_5056(self):
|
|
# see issue 5056
|
|
|
|
@njit(parallel=True)
|
|
def test_impl(data):
|
|
N = data.shape[0]
|
|
sums = np.zeros(N)
|
|
for i in prange(N):
|
|
sums[i] = np.sum(data[np.int32(0):np.int32(1)])
|
|
return sums
|
|
|
|
data = np.arange(10.)
|
|
np.testing.assert_array_equal(test_impl(data), test_impl.py_func(data))
|
|
|
|
@skip_unsupported
|
|
def test_global_tuple(self):
|
|
"""make sure a global tuple with non-integer values does not cause errors
|
|
(test for #6726).
|
|
"""
|
|
|
|
def test_impl():
|
|
d = GVAL[0]
|
|
return d
|
|
|
|
self.assertEqual(njit(test_impl, parallel=True)(), test_impl())
|
|
|
|
|
|
class TestArrayAnalysisInterface(TestCase):
|
|
def test_analyze_op_call_interface(self):
|
|
# gather _analyze_op_call_*
|
|
aoc = {}
|
|
for fname in dir(ArrayAnalysis):
|
|
if fname.startswith('_analyze_op_call_'):
|
|
aoc[fname] = getattr(ArrayAnalysis, fname)
|
|
# check interface
|
|
def iface_stub(self, scope, equiv_set, loc, args, kws):
|
|
pass
|
|
expected = utils.pysignature(iface_stub)
|
|
for k, v in aoc.items():
|
|
got = utils.pysignature(v)
|
|
with self.subTest(fname=k, sig=got):
|
|
self.assertEqual(got, expected)
|
|
|
|
@skip_unsupported
|
|
def test_array_analysis_extensions(self):
|
|
# Test that the `array_analysis` object in `array_analysis_extensions`
|
|
# can perform analysis on the scope using `equiv_sets`.
|
|
from numba.parfors.parfor import Parfor
|
|
from numba.parfors import array_analysis
|
|
|
|
orig_parfor = array_analysis.array_analysis_extensions[Parfor]
|
|
|
|
shared = {'counter': 0}
|
|
|
|
def testcode(array_analysis):
|
|
# Find call node corresponding to the ``A = empty(n)``
|
|
func_ir = array_analysis.func_ir
|
|
for call in func_ir.blocks[0].find_exprs('call'):
|
|
callee = func_ir.get_definition(call.func)
|
|
if getattr(callee, "value", None) is empty:
|
|
if getattr(call.args[0], 'name', None) == 'n':
|
|
break
|
|
else:
|
|
return
|
|
|
|
variable_A = func_ir.get_assignee(call)
|
|
# n must be equiv to
|
|
es = array_analysis.equiv_sets[0]
|
|
self.assertTrue(es.is_equiv('n', variable_A.name))
|
|
shared['counter'] += 1
|
|
|
|
def new_parfor(parfor, equiv_set, typemap, array_analysis):
|
|
"""Recursive array analysis for parfor nodes.
|
|
"""
|
|
testcode(array_analysis)
|
|
# Call original
|
|
return orig_parfor(
|
|
parfor, equiv_set, typemap, array_analysis,
|
|
)
|
|
|
|
try:
|
|
# Replace the array-analysis extension for Parfor node
|
|
array_analysis.array_analysis_extensions[Parfor] = new_parfor
|
|
|
|
empty = np.empty # avoid scanning a getattr in the IR
|
|
def f(n):
|
|
A = empty(n)
|
|
for i in prange(n):
|
|
S = np.arange(i)
|
|
A[i] = S.sum()
|
|
return A + 1
|
|
|
|
got = njit(parallel=True)(f)(10)
|
|
executed_count = shared['counter']
|
|
self.assertGreater(executed_count, 0)
|
|
finally:
|
|
# Re-install the original handler
|
|
array_analysis.array_analysis_extensions[Parfor] = orig_parfor
|
|
|
|
# Check normal execution
|
|
expected = njit(parallel=True)(f)(10)
|
|
self.assertPreciseEqual(got, expected)
|
|
# Make sure we have uninstalled the handler
|
|
self.assertEqual(executed_count, shared['counter'])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|