Merge pull request #957 from IntelPython/feature/dpnp_parfor_v2

Parfor lowering as kernels and dpnp ufunc compilation to kernels

Author: Diptorup Deb

numba_dpex/__init__.py

@@ -74,6 +74,7 @@ def load_dpctl_sycl_interface():
         f"dpctl={dpctl_version} may cause unexpected behavior"
     )
 
+from numba import prange  # noqa E402
 
 import numba_dpex.core.dpjit_dispatcher  # noqa E402
 import numba_dpex.core.offload_dispatcher  # noqa E402
@@ -92,6 +93,7 @@ def load_dpctl_sycl_interface():
 
 # Re-export all type names
 from numba_dpex.core.types import *  # noqa E402
+from numba_dpex.dpnp_iface import dpnpimpl  # noqa E402
 from numba_dpex.retarget import offload_to_sycl_device  # noqa E402
 
 if config.HAS_NON_HOST_DEVICE:

numba_dpex/_patches.py

@@ -258,13 +258,11 @@ def _empty_nd_impl(context, builder, arrtype, shapes):
         )
         from numba_dpex.decorators import dpjit
 
-        numba_config.DISABLE_PERFORMANCE_WARNINGS = 0
         op = dpjit(_call_usm_allocator)
         fnop = context.typing_context.resolve_value_type(op)
         # The _call_usm_allocator function will be compiled and added to registry
         # when the get_call_type function is invoked.
         fnop.get_call_type(context.typing_context, sig.args, {})
-        numba_config.DISABLE_PERFORMANCE_WARNINGS = 1
         eqfn = context.get_function(fnop, sig)
         meminfo = eqfn(builder, args)
     else:
@@ -309,11 +307,17 @@ def impl(cls, allocsize, usm_type, device):
     return impl
 
 
+numba_config.DISABLE_PERFORMANCE_WARNINGS = 0
+
+
 def _call_usm_allocator(arrtype, size, usm_type, device):
     """Trampoline to call the intrinsic used for allocation"""
     return arrtype._usm_allocate(size, usm_type, device)
 
 
+numba_config.DISABLE_PERFORMANCE_WARNINGS = 1
+
+
 @intrinsic
 def intrin_usm_alloc(typingctx, allocsize, usm_type, device):
     """Intrinsic to call into the allocator for Array"""

numba_dpex/core/exceptions.py

@@ -221,7 +221,7 @@ def __init__(
                 f"Arguments {ndarray_args} are non-usm arrays, "
                 f"and arguments {usmarray_args} are usm arrays."
             )
-        elif usmarray_argnum_list:
+        elif usmarray_argnum_list is not None:
             usmarray_args = ",".join([str(i) for i in usmarray_argnum_list])
             self.message = (
                 f'Execution queue for kernel "{kernel_name}" could '
@@ -433,3 +433,13 @@ def __init__(self, kernel_name, argtypes) -> None:
         )
 
         super().__init__(self.message)
+
+
+class UnsupportedParforError(Exception):
+    """Exception raised when a parfor node could not be lowered by Numba-dpex"""
+
+    def __init__(self, extra_msg=None) -> None:
+        self.message = "Expression cannot be offloaded"
+        if extra_msg:
+            self.message += " due to " + extra_msg
+        super().__init__(self.message)

numba_dpex/core/passes/__init__.py

@@ -1,3 +1,27 @@
 # SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
 #
 # SPDX-License-Identifier: Apache-2.0
+
+from .parfor_legalize_cfd_pass import ParforLegalizeCFDPass
+from .parfor_lowering_pass import ParforLoweringPass
+from .passes import (
+    DumpParforDiagnostics,
+    NoPythonBackend,
+    ParforFusionPass,
+    ParforPass,
+    ParforPreLoweringPass,
+    PreParforPass,
+    SplitParforPass,
+)
+
+__all__ = [
+    "DumpParforDiagnostics",
+    "ParforLoweringPass",
+    "ParforLegalizeCFDPass",
+    "ParforFusionPass",
+    "ParforPreLoweringPass",
+    "ParforPass",
+    "PreParforPass",
+    "SplitParforPass",
+    "NoPythonBackend",
+]

numba_dpex/core/passes/parfor.py

@@ -2421,7 +2421,6 @@ def _arrayexpr_to_parfor(self, equiv_set, lhs, arrayexpr, avail_vars):
         expr = arrayexpr.expr
         arr_typ = pass_states.typemap[lhs.name]
         el_typ = arr_typ.dtype
-
         # generate loopnests and size variables from lhs correlations
         size_vars = equiv_set.get_shape(lhs)
         index_vars, loopnests = _mk_parfor_loops(
@@ -3788,6 +3787,46 @@ def _get_call_arg_types(expr, typemap):
     return tuple(new_arg_typs), new_kw_types
 
 
+def _ufunc_to_parfor_instr(
+    typemap,
+    op,
+    avail_vars,
+    loc,
+    scope,
+    func_ir,
+    out_ir,
+    arg_vars,
+    typingctx,
+    calltypes,
+    expr_out_var,
+):
+    func_var_name = _find_func_var(typemap, op, avail_vars, loc=loc)
+    func_var = ir.Var(scope, mk_unique_var(func_var_name), loc)
+    typemap[func_var.name] = typemap[func_var_name]
+    func_var_def = copy.deepcopy(func_ir.get_definition(func_var_name))
+    if (
+        isinstance(func_var_def, ir.Expr)
+        and func_var_def.op == "getattr"
+        and func_var_def.attr == "sqrt"
+    ):
+        g_math_var = ir.Var(scope, mk_unique_var("$math_g_var"), loc)
+        typemap[g_math_var.name] = types.misc.Module(math)
+        g_math = ir.Global("math", math, loc)
+        g_math_assign = ir.Assign(g_math, g_math_var, loc)
+        func_var_def = ir.Expr.getattr(g_math_var, "sqrt", loc)
+        out_ir.append(g_math_assign)
+    #                     out_ir.append(func_var_def)
+    ir_expr = ir.Expr.call(func_var, arg_vars, (), loc)
+    call_typ = typemap[func_var.name].get_call_type(
+        typingctx, tuple(typemap[a.name] for a in arg_vars), {}
+    )
+    calltypes[ir_expr] = call_typ
+    el_typ = call_typ.return_type
+    # signature(el_typ, el_typ)
+    out_ir.append(ir.Assign(func_var_def, func_var, loc))
+    out_ir.append(ir.Assign(ir_expr, expr_out_var, loc))
+
+
 def _arrayexpr_tree_to_ir(
     func_ir,
     typingctx,
@@ -3852,35 +3891,33 @@ def _arrayexpr_tree_to_ir(
                 # elif isinstance(op, (np.ufunc, DUFunc)):
                 # function calls are stored in variables which are not removed
                 # op is typing_key to the variables type
-                func_var_name = _find_func_var(typemap, op, avail_vars, loc=loc)
-                func_var = ir.Var(scope, mk_unique_var(func_var_name), loc)
-                typemap[func_var.name] = typemap[func_var_name]
-                func_var_def = copy.deepcopy(
-                    func_ir.get_definition(func_var_name)
-                )
-                if (
-                    isinstance(func_var_def, ir.Expr)
-                    and func_var_def.op == "getattr"
-                    and func_var_def.attr == "sqrt"
-                ):
-                    g_math_var = ir.Var(
-                        scope, mk_unique_var("$math_g_var"), loc
-                    )
-                    typemap[g_math_var.name] = types.misc.Module(math)
-                    g_math = ir.Global("math", math, loc)
-                    g_math_assign = ir.Assign(g_math, g_math_var, loc)
-                    func_var_def = ir.Expr.getattr(g_math_var, "sqrt", loc)
-                    out_ir.append(g_math_assign)
-                #                     out_ir.append(func_var_def)
-                ir_expr = ir.Expr.call(func_var, arg_vars, (), loc)
-                call_typ = typemap[func_var.name].get_call_type(
-                    typingctx, tuple(typemap[a.name] for a in arg_vars), {}
+                _ufunc_to_parfor_instr(
+                    typemap,
+                    op,
+                    avail_vars,
+                    loc,
+                    scope,
+                    func_ir,
+                    out_ir,
+                    arg_vars,
+                    typingctx,
+                    calltypes,
+                    expr_out_var,
                 )
-                calltypes[ir_expr] = call_typ
-                el_typ = call_typ.return_type
-                # signature(el_typ, el_typ)
-                out_ir.append(ir.Assign(func_var_def, func_var, loc))
-                out_ir.append(ir.Assign(ir_expr, expr_out_var, loc))
+        if hasattr(op, "is_dpnp_ufunc"):
+            _ufunc_to_parfor_instr(
+                typemap,
+                op,
+                avail_vars,
+                loc,
+                scope,
+                func_ir,
+                out_ir,
+                arg_vars,
+                typingctx,
+                calltypes,
+                expr_out_var,
+            )
     elif isinstance(expr, ir.Var):
         var_typ = typemap[expr.name]
         if isinstance(var_typ, types.Array):