From f7d62077a6f5e379a1e35abcaad0001a2c6fda66 Mon Sep 17 00:00:00 2001 From: "reazul.hoque" Date: Tue, 30 Mar 2021 16:38:11 -0500 Subject: [PATCH 1/2] Numba patch for lower_extensions --- numba/core/cpu.py | 3 +++ numba/core/lowering.py | 20 +++++--------------- numba/core/typed_passes.py | 7 +++++++ numba/parfors/parfor_lowering.py | 3 --- 4 files changed, 15 insertions(+), 18 deletions(-) diff --git a/numba/core/cpu.py b/numba/core/cpu.py index a54ae7ee339..a72681eb807 100644 --- a/numba/core/cpu.py +++ b/numba/core/cpu.py @@ -62,6 +62,9 @@ def init(self): import numba.typed.dictimpl import numba.experimental.function_type + # Add lower_extension attribute + self.lower_extensions = {} + def load_additional_registries(self): # Add target specific implementations from numba.np import npyimpl diff --git a/numba/core/lowering.py b/numba/core/lowering.py index 44009bde487..57efe28f768 100644 --- a/numba/core/lowering.py +++ b/numba/core/lowering.py @@ -267,20 +267,9 @@ def debug_print(self, msg): self.context.debug_print(self.builder, "DEBUGJIT: {0}".format(msg)) -# Dictionary mapping instruction class to its lowering function. -lower_extensions = {} - - class Lower(BaseLower): GeneratorLower = generators.GeneratorLower - def __init__(self, context, library, fndesc, func_ir, metadata=None): - BaseLower.__init__(self, context, library, fndesc, func_ir, metadata) - from numba.parfors.parfor_lowering import _lower_parfor_parallel - from numba.parfors import parfor - if parfor.Parfor not in lower_extensions: - lower_extensions[parfor.Parfor] = [_lower_parfor_parallel] - def pre_block(self, block): from numba.core.unsafe import eh @@ -445,10 +434,11 @@ def lower_inst(self, inst): self.lower_static_try_raise(inst) else: - for _class, func in lower_extensions.items(): - if isinstance(inst, _class): - func[-1](self, inst) - return + if hasattr(self.context, "lower_extensions"): + for _class, func in self.context.lower_extensions.items(): + if isinstance(inst, _class): + func(self, inst) + return raise NotImplementedError(type(inst)) def lower_setitem(self, target_var, index_var, value_var, signature): diff --git a/numba/core/typed_passes.py b/numba/core/typed_passes.py index e2cc3bf0131..6ad5747b993 100644 --- a/numba/core/typed_passes.py +++ b/numba/core/typed_passes.py @@ -277,6 +277,13 @@ def run_pass(self, state): """ Convert data-parallel computations into Parfor nodes """ + # Register lowerer for Parfor Node + from numba.parfors.parfor_lowering import _lower_parfor_parallel + if hasattr(state.targetctx, "lower_extensions"): + state.targetctx.lower_extensions[Parfor] = _lower_parfor_parallel + else: + raise AttributeError("target_context has no attribute 'lower_extensions'") + # Ensure we have an IR and type information. assert state.func_ir parfor_pass = _parfor_ParforPass(state.func_ir, diff --git a/numba/parfors/parfor_lowering.py b/numba/parfors/parfor_lowering.py index 791b9b79403..765bdb52314 100644 --- a/numba/parfors/parfor_lowering.py +++ b/numba/parfors/parfor_lowering.py @@ -479,9 +479,6 @@ def _lower_parfor_parallel(lowerer, parfor): if config.DEBUG_ARRAY_OPT: print("_lower_parfor_parallel done") -# A work-around to prevent circular imports -#lowering.lower_extensions[parfor.Parfor] = _lower_parfor_parallel - def _create_shape_signature( get_shape_classes, From 873cfa445ee995de599faf740e1b8fc01ee9749f Mon Sep 17 00:00:00 2001 From: "reazul.hoque" Date: Fri, 2 Apr 2021 11:16:43 -0500 Subject: [PATCH 2/2] Move registration of Parfor lowering to CPUContext.init --- numba/core/cpu.py | 4 ++++ numba/core/typed_passes.py | 7 ------- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/numba/core/cpu.py b/numba/core/cpu.py index a72681eb807..62eba99638c 100644 --- a/numba/core/cpu.py +++ b/numba/core/cpu.py @@ -64,6 +64,10 @@ def init(self): # Add lower_extension attribute self.lower_extensions = {} + from numba.parfors.parfor_lowering import _lower_parfor_parallel + from numba.parfors.parfor import Parfor + # Specify how to lower Parfor nodes using the lower_extensions + self.lower_extensions[Parfor] = _lower_parfor_parallel def load_additional_registries(self): # Add target specific implementations diff --git a/numba/core/typed_passes.py b/numba/core/typed_passes.py index 6ad5747b993..e2cc3bf0131 100644 --- a/numba/core/typed_passes.py +++ b/numba/core/typed_passes.py @@ -277,13 +277,6 @@ def run_pass(self, state): """ Convert data-parallel computations into Parfor nodes """ - # Register lowerer for Parfor Node - from numba.parfors.parfor_lowering import _lower_parfor_parallel - if hasattr(state.targetctx, "lower_extensions"): - state.targetctx.lower_extensions[Parfor] = _lower_parfor_parallel - else: - raise AttributeError("target_context has no attribute 'lower_extensions'") - # Ensure we have an IR and type information. assert state.func_ir parfor_pass = _parfor_ParforPass(state.func_ir,