Merge pull request #1030 from IntelPython/feature/support_numba057

Feature/support numba057

Author: Diptorup Deb

.github/workflows/conda-package.yml

@@ -91,7 +91,7 @@ jobs:
     strategy:
       matrix:
         python: ["3.8", "3.9", "3.10"]
-        numba: ["0.56"]
+        numba: ["0.57"]
         dpnp: ["0.11"]
 
     steps:

conda-recipe/meta.yaml

@@ -18,14 +18,14 @@ requirements:
         - python
         - setuptools >=63.*
         - cython
-        - numba 0.56*
+        - numba 0.57*
         - dpctl >=0.14*
         - dpnp >=0.11*
         - dpcpp-llvm-spirv
         - wheel
     run:
         - python
-        - numba >=0.56*
+        - numba >=0.57*
         - dpctl >=0.14*
         - spirv-tools
         - dpcpp-llvm-spirv

environment/coverage.yml

@@ -12,7 +12,7 @@ dependencies:
   - gxx_linux-64
   - dpcpp_linux-64
   - cython
-  - numba 0.56*
+  - numba 0.57*
   - dppy/label/dev:dpctl
   - dppy/label/dev:dpnp
   - spirv-tools

environment/docs.yml

@@ -11,7 +11,7 @@ dependencies:
   - gxx_linux-64
   - dpcpp_linux-64
   - cython
-  - numba 0.56*
+  - numba 0.57*
   - dpctl 0.14*
   - dpnp >=0.10.2
   - spirv-tools

numba_dpex/__init__.py

@@ -14,18 +14,23 @@
 
 import dpctl
 import llvmlite.binding as ll
-import numba
-from numba.core import ir_utils
+from numba import __version__ as numba_version
 from numba.np import arrayobj
-from numba.np.ufunc import array_exprs
 from numba.np.ufunc.decorators import Vectorize
 
-from numba_dpex._patches import _empty_nd_impl, _is_ufunc, _mk_alloc
+from numba_dpex._patches import _empty_nd_impl
 from numba_dpex.vectorizers import Vectorize as DpexVectorize
 
+from .numba_patches import (
+    patch_arrayexpr_tree_to_ir,
+    patch_is_ufunc,
+    patch_mk_alloc,
+)
+
 # Monkey patches
-array_exprs._is_ufunc = _is_ufunc
-ir_utils.mk_alloc = _mk_alloc
+patch_is_ufunc.patch()
+patch_mk_alloc.patch()
+patch_arrayexpr_tree_to_ir.patch()
 arrayobj._empty_nd_impl = _empty_nd_impl
 
 
@@ -58,8 +63,6 @@ def load_dpctl_sycl_interface():
     else:
         raise ImportError
 
-    Vectorize.target_registry.ondemand["dpex"] = lambda: DpexVectorize
-
 
 def parse_sem_version(version_string: str) -> Tuple[int, int, int]:
     """Parse sem version into tuple of three integers. If there is a suffix like
@@ -76,19 +79,19 @@ def parse_sem_version(version_string: str) -> Tuple[int, int, int]:
     )
 
 
-numba_version = parse_sem_version(numba.__version__)
-if numba_version < (0, 56, 4):
+numba_sem_version = parse_sem_version(numba_version)
+if numba_sem_version < (0, 57, 0):
     logging.warning(
-        "numba_dpex needs numba 0.56.4, using "
+        "numba_dpex needs numba 0.57.0, using "
         f"numba={numba_version} may cause unexpected behavior"
     )
 
 
-dpctl_version = tuple(map(int, dpctl.__version__.split(".")[:2]))
-if dpctl_version < (0, 14):
+dpctl_sem_version = parse_sem_version(dpctl.__version__)
+if dpctl_sem_version < (0, 14):
     logging.warning(
         "numba_dpex needs dpctl 0.14 or greater, using "
-        f"dpctl={dpctl_version} may cause unexpected behavior"
+        f"dpctl={dpctl_sem_version} may cause unexpected behavior"
     )
 
 from numba import prange  # noqa E402
@@ -139,6 +142,8 @@ def parse_sem_version(version_string: str) -> Tuple[int, int, int]:
 else:
     raise ImportError("No non-host SYCL device found to execute kernels.")
 
+Vectorize.target_registry.ondemand["dpex"] = lambda: DpexVectorize
+
 from numba_dpex._version import get_versions  # noqa E402
 
 __version__ = get_versions()["version"]

numba_dpex/_patches.py

@@ -2,17 +2,11 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import numpy
 from llvmlite import ir as llvmir
 from llvmlite.ir import Constant
 from numba.core import cgutils
 from numba.core import config as numba_config
-from numba.core import ir, types
-from numba.core.ir_utils import (
-    convert_size_to_var,
-    get_np_ufunc_typ,
-    mk_unique_var,
-)
+from numba.core import types
 from numba.core.typing import signature
 from numba.extending import intrinsic, overload_classmethod
 from numba.np.arrayobj import (
@@ -21,156 +15,10 @@
     make_array,
     populate_array,
 )
-from numba.np.ufunc.dufunc import DUFunc
 
 from numba_dpex.core.runtime import context as dpexrt
 from numba_dpex.core.types import DpnpNdArray
 
-# Numpy array constructors
-
-
-def _is_ufunc(func):
-    return isinstance(func, (numpy.ufunc, DUFunc)) or hasattr(
-        func, "is_dpnp_ufunc"
-    )
-
-
-def _mk_alloc(
-    typingctx, typemap, calltypes, lhs, size_var, dtype, scope, loc, lhs_typ
-):
-    """generate an array allocation with np.empty() and return list of nodes.
-    size_var can be an int variable or tuple of int variables.
-    lhs_typ is the type of the array being allocated.
-    """
-    out = []
-    ndims = 1
-    size_typ = types.intp
-    if isinstance(size_var, tuple):
-        if len(size_var) == 1:
-            size_var = size_var[0]
-            size_var = convert_size_to_var(size_var, typemap, scope, loc, out)
-        else:
-            # tuple_var = build_tuple([size_var...])
-            ndims = len(size_var)
-            tuple_var = ir.Var(scope, mk_unique_var("$tuple_var"), loc)
-            if typemap:
-                typemap[tuple_var.name] = types.containers.UniTuple(
-                    types.intp, ndims
-                )
-            # constant sizes need to be assigned to vars
-            new_sizes = [
-                convert_size_to_var(s, typemap, scope, loc, out)
-                for s in size_var
-            ]
-            tuple_call = ir.Expr.build_tuple(new_sizes, loc)
-            tuple_assign = ir.Assign(tuple_call, tuple_var, loc)
-            out.append(tuple_assign)
-            size_var = tuple_var
-            size_typ = types.containers.UniTuple(types.intp, ndims)
-
-    if hasattr(lhs_typ, "__allocate__"):
-        return lhs_typ.__allocate__(
-            typingctx,
-            typemap,
-            calltypes,
-            lhs,
-            size_var,
-            dtype,
-            scope,
-            loc,
-            lhs_typ,
-            size_typ,
-            out,
-        )
-
-    # g_np_var = Global(numpy)
-    g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)
-    if typemap:
-        typemap[g_np_var.name] = types.misc.Module(numpy)
-    g_np = ir.Global("np", numpy, loc)
-    g_np_assign = ir.Assign(g_np, g_np_var, loc)
-    # attr call: empty_attr = getattr(g_np_var, empty)
-    empty_attr_call = ir.Expr.getattr(g_np_var, "empty", loc)
-    attr_var = ir.Var(scope, mk_unique_var("$empty_attr_attr"), loc)
-    if typemap:
-        typemap[attr_var.name] = get_np_ufunc_typ(numpy.empty)
-    attr_assign = ir.Assign(empty_attr_call, attr_var, loc)
-    # Assume str(dtype) returns a valid type
-    dtype_str = str(dtype)
-    # alloc call: lhs = empty_attr(size_var, typ_var)
-    typ_var = ir.Var(scope, mk_unique_var("$np_typ_var"), loc)
-    if typemap:
-        typemap[typ_var.name] = types.functions.NumberClass(dtype)
-    # If dtype is a datetime/timedelta with a unit,
-    # then it won't return a valid type and instead can be created
-    # with a string. i.e. "datetime64[ns]")
-    if (
-        isinstance(dtype, (types.NPDatetime, types.NPTimedelta))
-        and dtype.unit != ""
-    ):
-        typename_const = ir.Const(dtype_str, loc)
-        typ_var_assign = ir.Assign(typename_const, typ_var, loc)
-    else:
-        if dtype_str == "bool":
-            # empty doesn't like 'bool' sometimes (e.g. kmeans example)
-            dtype_str = "bool_"
-        np_typ_getattr = ir.Expr.getattr(g_np_var, dtype_str, loc)
-        typ_var_assign = ir.Assign(np_typ_getattr, typ_var, loc)
-    alloc_call = ir.Expr.call(attr_var, [size_var, typ_var], (), loc)
-
-    if calltypes:
-        cac = typemap[attr_var.name].get_call_type(
-            typingctx, [size_typ, types.functions.NumberClass(dtype)], {}
-        )
-        # By default, all calls to "empty" are typed as returning a standard
-        # NumPy ndarray.  If we are allocating a ndarray subclass here then
-        # just change the return type to be that of the subclass.
-        cac._return_type = (
-            lhs_typ.copy(layout="C") if lhs_typ.layout == "F" else lhs_typ
-        )
-        calltypes[alloc_call] = cac
-    if lhs_typ.layout == "F":
-        empty_c_typ = lhs_typ.copy(layout="C")
-        empty_c_var = ir.Var(scope, mk_unique_var("$empty_c_var"), loc)
-        if typemap:
-            typemap[empty_c_var.name] = lhs_typ.copy(layout="C")
-        empty_c_assign = ir.Assign(alloc_call, empty_c_var, loc)
-
-        # attr call: asfortranarray = getattr(g_np_var, asfortranarray)
-        asfortranarray_attr_call = ir.Expr.getattr(
-            g_np_var, "asfortranarray", loc
-        )
-        afa_attr_var = ir.Var(
-            scope, mk_unique_var("$asfortran_array_attr"), loc
-        )
-        if typemap:
-            typemap[afa_attr_var.name] = get_np_ufunc_typ(numpy.asfortranarray)
-        afa_attr_assign = ir.Assign(asfortranarray_attr_call, afa_attr_var, loc)
-        # call asfortranarray
-        asfortranarray_call = ir.Expr.call(afa_attr_var, [empty_c_var], (), loc)
-        if calltypes:
-            calltypes[asfortranarray_call] = typemap[
-                afa_attr_var.name
-            ].get_call_type(typingctx, [empty_c_typ], {})
-
-        asfortranarray_assign = ir.Assign(asfortranarray_call, lhs, loc)
-
-        out.extend(
-            [
-                g_np_assign,
-                attr_assign,
-                typ_var_assign,
-                empty_c_assign,
-                afa_attr_assign,
-                asfortranarray_assign,
-            ]
-        )
-    else:
-        alloc_assign = ir.Assign(alloc_call, lhs, loc)
-        out.extend([g_np_assign, attr_assign, typ_var_assign, alloc_assign])
-
-    return out
-
 
 def _empty_nd_impl(context, builder, arrtype, shapes):
     """Utility function used for allocating a new array during LLVM code

numba_dpex/config.py

@@ -53,8 +53,6 @@ def __getattr__(name):
 # Dump offload diagnostics
 OFFLOAD_DIAGNOSTICS = _readenv("NUMBA_DPEX_OFFLOAD_DIAGNOSTICS", int, 0)
 
-FALLBACK_ON_CPU = _readenv("NUMBA_DPEX_FALLBACK_ON_CPU", int, 1)
-
 # Activate Native floating point atomcis support for supported devices.
 # Requires llvm-spirv supporting the FP atomics extension
 NATIVE_FP_ATOMICS = _readenv("NUMBA_DPEX_ACTIVATE_ATOMICS_FP_NATIVE", int, 0)

numba_dpex/core/parfors/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0

numba_dpex/core/utils/kernel_builder.py renamed to numba_dpex/core/parfors/kernel_builder.py

@@ -21,14 +21,14 @@
     replace_var_names,
 )
 from numba.core.typing import signature
+from numba.parfors import parfor
 
 import numba_dpex as dpex
 from numba_dpex import config
 
 from ..descriptor import dpex_kernel_target
-from ..passes import parfor
 from ..types.dpnp_ndarray_type import DpnpNdArray
-from .kernel_templates import RangeKernelTemplate
+from ..utils.kernel_templates import RangeKernelTemplate
 
 
 class ParforKernel:
@@ -338,7 +338,7 @@ def create_kernel_for_parfor(
         print("legal parfor_params = ", parfor_params, type(parfor_params))
 
     # Determine the unique names of the kernel functions.
-    kernel_name = "__numba_parfor_kernel_%s" % (parfor_node.id)
+    kernel_name = "__dpex_parfor_kernel_%s" % (parfor_node.id)
 
     kernel_template = RangeKernelTemplate(
         kernel_name=kernel_name,
@@ -493,8 +493,7 @@ def update_sentinel(kernel_ir, sentinel_name, kernel_body, new_label):
                 # parfor body.
                 prev_block.append(ir.Jump(body_first_label, loc))
 
-                # Add all the parfor loop body blocks to the gufunc function's
-                # IR.
+                # Add all the parfor loop body blocks to the kernel IR
                 for loop, b in kernel_body.items():
                     kernel_ir.blocks[loop] = copy.copy(b)
                     kernel_ir.blocks[loop].body = copy.copy(