Numba patch for lower_extensions

Author: reazulhoque

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

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):

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,

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,