def get_func_type(state, expr):
func_ty = None
if expr.op == 'call':
# check this is a known and typed function
try:
func_ty = state.type_annotation.typemap[expr.func.name]
except KeyError:
# e.g. Calls to CUDA Intrinsic have no mapped type
# so KeyError
return None
if not hasattr(func_ty, 'get_call_type'):
return None
elif is_operator_or_getitem(expr):
func_ty = state.typingctx.resolve_value_type(expr.fn)
else:
return None
return func_ty
def check(self, test_impl, *args, **kwargs):
inline_expect = kwargs.pop('inline_expect', None)
assert inline_expect
block_count = kwargs.pop('block_count', 1)
assert not kwargs
for k, v in inline_expect.items():
assert isinstance(k, str)
assert isinstance(v, bool)
j_func = njit(pipeline_class=IRPreservingTestPipeline)(test_impl)
# check they produce the same answer first!
self.assertEqual(test_impl(*args), j_func(*args))
# make sure IR doesn't have branches
fir = j_func.overloads[j_func.signatures[0]].metadata['preserved_ir']
fir.blocks = ir_utils.simplify_CFG(fir.blocks)
if self._DEBUG:
print("FIR".center(80, "-"))
fir.dump()
if block_count != 'SKIP':
self.assertEqual(len(fir.blocks), block_count)
block = next(iter(fir.blocks.values()))
# if we don't expect the function to be inlined then make sure there is
# 'call' present still
exprs = [x for x in block.find_exprs()]
assert exprs
for k, v in inline_expect.items():
found = False
for expr in exprs:
if getattr(expr, 'op', False) == 'call':
func_defn = fir.get_definition(expr.func)
found |= func_defn.name == k
elif ir_utils.is_operator_or_getitem(expr):
found |= expr.fn.__name__ == k
self.assertFalse(found == v)
return fir # for use in further analysis