def test1(self): typingctx = typing.Context() targetctx = cpu.CPUContext(typingctx) test_ir = compiler.run_frontend(test_will_propagate) #print("Num blocks = ", len(test_ir.blocks)) #print(test_ir.dump()) with cpu_target.nested_context(typingctx, targetctx): typingctx.refresh() targetctx.refresh() args = (types.int64, types.int64, types.int64) typemap, return_type, calltypes = compiler.type_inference_stage(typingctx, test_ir, args, None) #print("typemap = ", typemap) #print("return_type = ", return_type) type_annotation = type_annotations.TypeAnnotation( func_ir=test_ir, typemap=typemap, calltypes=calltypes, lifted=(), lifted_from=None, args=args, return_type=return_type, html_output=config.HTML) remove_dels(test_ir.blocks) in_cps, out_cps = copy_propagate(test_ir.blocks, typemap) apply_copy_propagate(test_ir.blocks, in_cps, get_name_var_table(test_ir.blocks), typemap, calltypes) remove_dead(test_ir.blocks, test_ir.arg_names, test_ir) self.assertFalse(findLhsAssign(test_ir, "x"))
def test1(self): typingctx = typing.Context() targetctx = cpu.CPUContext(typingctx) test_ir = compiler.run_frontend(test_will_propagate) #print("Num blocks = ", len(test_ir.blocks)) #print(test_ir.dump()) with cpu_target.nested_context(typingctx, targetctx): typingctx.refresh() targetctx.refresh() args = (types.int64, types.int64, types.int64) typemap, return_type, calltypes = compiler.type_inference_stage(typingctx, test_ir, args, None) #print("typemap = ", typemap) #print("return_type = ", return_type) type_annotation = type_annotations.TypeAnnotation( func_ir=test_ir, typemap=typemap, calltypes=calltypes, lifted=(), lifted_from=None, args=args, return_type=return_type, html_output=config.HTML) remove_dels(test_ir.blocks) in_cps, out_cps = copy_propagate(test_ir.blocks, typemap) apply_copy_propagate(test_ir.blocks, in_cps, get_name_var_table(test_ir.blocks), typemap, calltypes) remove_dead(test_ir.blocks, test_ir.arg_names, test_ir) self.assertFalse(findLhsAssign(test_ir, "x"))
def test_test1(self): typingctx = typing.Context() targetctx = cpu.CPUContext(typingctx) test_ir = compiler.run_frontend(test1) with cpu_target.nested_context(typingctx, targetctx): one_arg = numba.types.npytypes.Array( numba.types.scalars.Float(name="float64"), 1, 'C') args = (one_arg, one_arg, one_arg, one_arg, one_arg) tp = TestPipeline(typingctx, targetctx, args, test_ir) numba.rewrites.rewrite_registry.apply('before-inference', tp, tp.func_ir) tp.typemap, tp.return_type, tp.calltypes = compiler.type_inference_stage( tp.typingctx, tp.func_ir, tp.args, None) type_annotation = type_annotations.TypeAnnotation( func_ir=tp.func_ir, typemap=tp.typemap, calltypes=tp.calltypes, lifted=(), lifted_from=None, args=tp.args, return_type=tp.return_type, html_output=config.HTML) numba.rewrites.rewrite_registry.apply('after-inference', tp, tp.func_ir) parfor_pass = numba.parfor.ParforPass(tp.func_ir, tp.typemap, tp.calltypes, tp.return_type, tp.typingctx) parfor_pass.run() self.assertTrue(countParfors(test_ir) == 1)
def compile_ir(self, the_ir, args=(), return_type=None): typingctx = self.typingctx targetctx = self.targetctx flags = self.flags # Register the contexts in case for nested @jit or @overload calls with cpu_target.nested_context(typingctx, targetctx): return compile_ir(typingctx, targetctx, the_ir, args, return_type, flags, locals={})
def countParfors(test_func, args, **kws): typingctx = typing.Context() targetctx = cpu.CPUContext(typingctx) test_ir = compiler.run_frontend(test_func) if kws: options = cpu.ParallelOptions(kws) else: options = cpu.ParallelOptions(True) with cpu_target.nested_context(typingctx, targetctx): tp = TestPipeline(typingctx, targetctx, args, test_ir) inline_pass = inline_closurecall.InlineClosureCallPass( tp.func_ir, options) inline_pass.run() numba.rewrites.rewrite_registry.apply('before-inference', tp, tp.func_ir) tp.typemap, tp.return_type, tp.calltypes = compiler.type_inference_stage( tp.typingctx, tp.func_ir, tp.args, None) type_annotations.TypeAnnotation(func_ir=tp.func_ir, typemap=tp.typemap, calltypes=tp.calltypes, lifted=(), lifted_from=None, args=tp.args, return_type=tp.return_type, html_output=config.HTML) preparfor_pass = numba.parfor.PreParforPass(tp.func_ir, tp.typemap, tp.calltypes, tp.typingctx, options) preparfor_pass.run() numba.rewrites.rewrite_registry.apply('after-inference', tp, tp.func_ir) parfor_pass = numba.parfor.ParforPass(tp.func_ir, tp.typemap, tp.calltypes, tp.return_type, tp.typingctx, options) parfor_pass.run() ret_count = 0 for label, block in test_ir.blocks.items(): for i, inst in enumerate(block.body): if isinstance(inst, numba.parfor.Parfor): ret_count += 1 return ret_count
def compile(self, func, args, return_type=None, flags=DEFAULT_FLAGS): """ Compile the function or retrieve an already compiled result from the cache. """ from numba.targets.registry import cpu_target cache_key = (func, args, return_type, flags) try: cr = self.cr_cache[cache_key] except KeyError: # Register the contexts in case for nested @jit or @overload calls # (same as compile_isolated()) with cpu_target.nested_context(self.typingctx, self.targetctx): cr = compile_extra(self.typingctx, self.targetctx, func, args, return_type, flags, locals={}) self.cr_cache[cache_key] = cr return cr
def compile(self, func, args, return_type=None, flags=DEFAULT_FLAGS): """ Compile the function or retrieve an already compiled result from the cache. """ from numba.targets.registry import cpu_target cache_key = (func, args, return_type, flags) try: cr = self.cr_cache[cache_key] except KeyError: # Register the contexts in case for nested @jit or @overload calls # (same as compile_isolated()) with cpu_target.nested_context(self.typingctx, self.targetctx): cr = compile_extra(self.typingctx, self.targetctx, func, args, return_type, flags, locals={}) self.cr_cache[cache_key] = cr return cr
def get_stencil_ir(sf, typingctx, args, scope, loc, input_dict, typemap, calltypes): """get typed IR from stencil bytecode """ from numba.targets.cpu import CPUContext from numba.targets.registry import cpu_target from numba.annotations import type_annotations from numba.compiler import type_inference_stage # get untyped IR stencil_func_ir = sf.kernel_ir.copy() # copy the IR nodes to avoid changing IR in the StencilFunc object stencil_blocks = copy.deepcopy(stencil_func_ir.blocks) stencil_func_ir.blocks = stencil_blocks name_var_table = ir_utils.get_name_var_table(stencil_func_ir.blocks) if "out" in name_var_table: raise ValueError("Cannot use the reserved word 'out' in stencil kernels.") # get typed IR with a dummy pipeline (similar to test_parfors.py) targetctx = CPUContext(typingctx) with cpu_target.nested_context(typingctx, targetctx): tp = DummyPipeline(typingctx, targetctx, args, stencil_func_ir) numba.rewrites.rewrite_registry.apply( 'before-inference', tp, tp.func_ir) tp.typemap, tp.return_type, tp.calltypes = type_inference_stage( tp.typingctx, tp.func_ir, tp.args, None) type_annotations.TypeAnnotation( func_ir=tp.func_ir, typemap=tp.typemap, calltypes=tp.calltypes, lifted=(), lifted_from=None, args=tp.args, return_type=tp.return_type, html_output=numba.config.HTML) # make block labels unique stencil_blocks = ir_utils.add_offset_to_labels(stencil_blocks, ir_utils.next_label()) min_label = min(stencil_blocks.keys()) max_label = max(stencil_blocks.keys()) ir_utils._max_label = max_label if config.DEBUG_ARRAY_OPT == 1: print("Initial stencil_blocks") ir_utils.dump_blocks(stencil_blocks) # rename variables, var_dict = {} for v, typ in tp.typemap.items(): new_var = ir.Var(scope, mk_unique_var(v), loc) var_dict[v] = new_var typemap[new_var.name] = typ # add new var type for overall function ir_utils.replace_vars(stencil_blocks, var_dict) if config.DEBUG_ARRAY_OPT == 1: print("After replace_vars") ir_utils.dump_blocks(stencil_blocks) # add call types to overall function for call, call_typ in tp.calltypes.items(): calltypes[call] = call_typ arg_to_arr_dict = {} # replace arg with arr for block in stencil_blocks.values(): for stmt in block.body: if isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Arg): if config.DEBUG_ARRAY_OPT == 1: print("input_dict", input_dict, stmt.value.index, stmt.value.name, stmt.value.index in input_dict) arg_to_arr_dict[stmt.value.name] = input_dict[stmt.value.index].name stmt.value = input_dict[stmt.value.index] if config.DEBUG_ARRAY_OPT == 1: print("arg_to_arr_dict", arg_to_arr_dict) print("After replace arg with arr") ir_utils.dump_blocks(stencil_blocks) ir_utils.remove_dels(stencil_blocks) stencil_func_ir.blocks = stencil_blocks return stencil_func_ir, sf.get_return_type(args)[0], arg_to_arr_dict