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