diff --git a/numba/ir_utils.py b/numba/ir_utils.py index 635107e16f1..6eee8cc9642 100644 --- a/numba/ir_utils.py +++ b/numba/ir_utils.py @@ -549,7 +549,8 @@ def remove_dead_block(block, lives, call_table, arg_aliases, alias_map, # let external calls handle stmt if type matches if type(stmt) in remove_dead_extensions: f = remove_dead_extensions[type(stmt)] - stmt = f(stmt, lives, arg_aliases, alias_map, func_ir, typemap) + stmt = f(stmt, lives_n_aliases, arg_aliases, alias_map, func_ir, + typemap) if stmt is None: removed = True continue diff --git a/numba/tests/test_remove_dead.py b/numba/tests/test_remove_dead.py index 8c0d58d4649..40b2dd61a92 100644 --- a/numba/tests/test_remove_dead.py +++ b/numba/tests/test_remove_dead.py @@ -213,5 +213,59 @@ def func(A, i): # recover global state numba.ir_utils.alias_func_extensions = old_ext_handlers + @skip_parfors_unsupported + def test_alias_parfor_extension(self): + """Make sure aliases are considered in remove dead extension for + parfors. + """ + def func(): + n = 11 + numba.parfor.init_prange() + A = np.empty(n) + B = A # create alias to A + for i in numba.prange(n): + A[i] = i + + return B + + class TestPipeline(numba.compiler.Pipeline): + """Test pipeline that just converts prange() to parfor and calls + remove_dead(). Copy propagation can replace B in the example code + which this pipeline avoids. + """ + def define_pipelines(self, pm): + pm.create_pipeline("test parfor aliasing") + self.add_preprocessing_stage(pm) + self.add_pre_typing_stage(pm) + self.add_typing_stage(pm) + pm.add_stage( + self.stage_limited_parfor, "just prange and rm dead") + self.add_lowering_stage(pm) + self.add_cleanup_stage(pm) + + def stage_limited_parfor(self): + parfor_pass = numba.parfor.ParforPass( + self.func_ir, + self.type_annotation.typemap, + self.type_annotation.calltypes, + self.return_type, + self.typingctx, + self.flags.auto_parallel, + self.flags, + self.parfor_diagnostics + ) + remove_dels(self.func_ir.blocks) + parfor_pass.array_analysis.run(self.func_ir.blocks) + parfor_pass._convert_loop(self.func_ir.blocks) + remove_dead(self.func_ir.blocks, self.func_ir.arg_names, + self.func_ir, self.type_annotation.typemap) + numba.parfor.get_parfor_params(self.func_ir.blocks, + parfor_pass.options.fusion, parfor_pass.nested_fusion_info) + + test_res = numba.jit(pipeline_class=TestPipeline)(func)() + py_res = func() + np.testing.assert_array_equal(test_res, py_res) + + if __name__ == "__main__": unittest.main()