Support decorator wrapper (#60)

- Remove CustomFunctionCompiler and directly modify the pipeline class
in the numba.jit decorator implementation, to support decorator
wrappers

Author: ggeorgakoudis

src/numba/openmp/__init__.py

@@ -39,11 +39,6 @@
     omp_get_initial_device,
 )
 
-from .compiler import (
-    CustomCompiler,
-    CustomFunctionCompiler,
-)
-
 from .exceptions import (  # noqa: F401
     UnspecifiedVarInDefaultNone,
     ParallelForExtraCode,

src/numba/openmp/compiler.py

@@ -321,15 +321,14 @@ def _finalize_specific(self):
                 raise RuntimeError("error registering OpenMP offloading descriptor")
 
 
-class CustomFunctionCompiler(_FunctionCompiler):
-    def _customize_flags(self, flags):
-        # We need to disable SSA form for OpenMP analysis to detect variables
-        # used within regions.
+class CustomCompiler(compiler.CompilerBase):
+    def __init__(self, typingctx, targetctx, library, args, return_type, flags, locals):
+        # Ensure SSA form is disabled for OpenMP analysis to detect variables used within regions.
         flags.enable_ssa = False
-        return flags
-
+        super().__init__(
+            typingctx, targetctx, library, args, return_type, flags, locals
+        )
 
-class CustomCompiler(compiler.CompilerBase):
     @staticmethod
     def custom_untyped_pipeline(state, name="untyped-openmp"):
         """Returns an untyped part of the nopython OpenMP pipeline"""
@@ -366,14 +365,15 @@ def custom_untyped_pipeline(state, name="untyped-openmp"):
         pm.add_pass(FindLiterallyCalls, "find literally calls")
         pm.add_pass(LiteralUnroll, "handles literal_unroll")
 
-        if state.flags.enable_ssa:
-            assert False, "SSA form is not supported in OpenMP"
+        assert not state.flags.enable_ssa, (
+            "SSA form is not supported in OpenMP compilation"
+        )
 
         pm.add_pass(LiteralPropagationSubPipelinePass, "Literal propagation")
-        # Run WithLifting late to for make_implicit_explicit to work.  TODO: We
-        # should create a pass that does this instead of replicating and hacking
-        # the untyped pipeline. This handling may also negatively affect
-        # optimizations.
+        # Run WithLifting late to for make_implicit_explicit to work.
+        # TODO: We should create a pass that does this instead of replicating
+        # and hacking the untyped pipeline. This handling may also negatively
+        # affect optimizations.
         pm.add_pass(WithLifting, "Handle with contexts")
 
         pm.finalize()

src/numba/openmp/decorators.py

@@ -1,10 +1,7 @@
 import warnings
 import numba
 
-from .compiler import (
-    CustomCompiler,
-    CustomFunctionCompiler,
-)
+from .compiler import CustomCompiler
 
 
 def jit(*args, **kws):
@@ -16,10 +13,8 @@ def jit(*args, **kws):
     if "forceobj" in kws:
         warnings.warn("forceobj is set for njit and is ignored", RuntimeWarning)
         del kws["forceobj"]
-    kws.update({"nopython": True, "nogil": True})
+    kws.update({"nopython": True, "nogil": True, "pipeline_class": CustomCompiler})
     dispatcher = numba.jit(*args, **kws)
-    dispatcher._compiler.__class__ = CustomFunctionCompiler
-    dispatcher._compiler.pipeline_class = CustomCompiler
     return dispatcher