diff --git a/numba_dpex/__init__.py b/numba_dpex/__init__.py index 661aeb218b..a9e57fdffa 100644 --- a/numba_dpex/__init__.py +++ b/numba_dpex/__init__.py @@ -6,30 +6,36 @@ The numba-dpex extension module adds data-parallel offload support to Numba. """ -import numba_dpex.core.dpjit_dispatcher -import numba_dpex.core.offload_dispatcher +from numba.core import ir_utils -# Initialize the _dpexrt_python extension -import numba_dpex.core.runtime -import numba_dpex.core.targets.dpjit_target +from .numba_patches.patch_mk_alloc import _mk_alloc -# Re-export types itself -import numba_dpex.core.types as types -from numba_dpex.core.kernel_interface.indexers import NdRange, Range +ir_utils.mk_alloc = _mk_alloc -# Re-export all type names -from numba_dpex.core.types import * -from numba_dpex.retarget import offload_to_sycl_device +if True: # noqa: E402 + import numba_dpex.core.dpjit_dispatcher + import numba_dpex.core.offload_dispatcher -from . import config + # Initialize the _dpexrt_python extension + import numba_dpex.core.runtime + import numba_dpex.core.targets.dpjit_target -if config.HAS_NON_HOST_DEVICE: - from .device_init import * -else: - raise ImportError("No non-host SYCL device found to execute kernels.") + # Re-export types itself + import numba_dpex.core.types as types + from numba_dpex.core.kernel_interface.indexers import NdRange, Range + # Re-export all type names + from numba_dpex.core.types import * + from numba_dpex.retarget import offload_to_sycl_device -from ._version import get_versions + from . import config + + if config.HAS_NON_HOST_DEVICE: + from .device_init import * + else: + raise ImportError("No non-host SYCL device found to execute kernels.") + + from ._version import get_versions __version__ = get_versions()["version"] del get_versions diff --git a/numba_dpex/numba_patches/__init__.py b/numba_dpex/numba_patches/__init__.py new file mode 100644 index 0000000000..b2ed211648 --- /dev/null +++ b/numba_dpex/numba_patches/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2020 - 2022 Intel Corporation +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/numba_dpex/numba_patches/patch_mk_alloc.py b/numba_dpex/numba_patches/patch_mk_alloc.py new file mode 100644 index 0000000000..4f8341a6e7 --- /dev/null +++ b/numba_dpex/numba_patches/patch_mk_alloc.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: 2020 - 2022 Intel Corporation +# +# SPDX-License-Identifier: Apache-2.0 + +import numpy +from numba.core import ir, types +from numba.core.ir_utils import ( + convert_size_to_var, + get_np_ufunc_typ, + mk_unique_var, +) + + +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