Merge pull request #1306 from IntelPython/refactor/ranges

Add range classes to kernel_api

Author: Diptorup Deb

.pre-commit-config.yaml

@@ -52,7 +52,7 @@ repos:
     - id: pylint
       name: pylint
       entry: pylint
-      files: ^numba_dpex/experimental|^numba_dpex/core/utils/kernel_launcher.py
+      files: ^numba_dpex/kernel_api|^numba_dpex/experimental|^numba_dpex/core/utils/kernel_launcher.py
       language: system
       types: [python]
       require_serial: true

numba_dpex/__init__.py

@@ -98,17 +98,19 @@ def parse_sem_version(version_string: str) -> Tuple[int, int, int]:
 
 # Re-export types itself
 import numba_dpex.core.types as types  # noqa E402
+from numba_dpex.core import boxing  # noqa E402
 from numba_dpex.core import config  # noqa E402
-from numba_dpex.core.kernel_interface.indexers import (  # noqa E402
-    NdRange,
-    Range,
-)
+from numba_dpex.core.kernel_interface import ranges_overloads  # noqa E402
 
 # Re-export all type names
 from numba_dpex.core.types import *  # noqa E402
 from numba_dpex.dpctl_iface import _intrinsic  # noqa E402
 from numba_dpex.dpnp_iface import dpnpimpl  # noqa E402
 
+# Importing NdRange and Range into numba_dpex for
+# backward compatibility
+from numba_dpex.kernel_api import NdRange, Range  # noqa E402
+
 from .core.targets import dpjit_target, kernel_target  # noqa E402
 from .decorators import dpjit, func, kernel  # noqa E402
 from .ocl.stubs import (  # noqa E402

numba_dpex/core/boxing/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Contains the ``box`` and ``unbox`` functions for numba_dpex types that are
+passable as arguments to a kernel or dpjit decorated function.
+"""
+
+from .ranges import *
+from .usm_ndarray import *

numba_dpex/core/types/range_types.py renamed to numba_dpex/core/boxing/ranges.py

@@ -4,55 +4,12 @@
 
 from contextlib import ExitStack
 
-from numba.core import cgutils, errors, types
+from numba.core import cgutils, types
 from numba.core.datamodel import default_manager
 from numba.extending import NativeValue, box, unbox
 
-from ..kernel_interface.indexers import NdRange, Range
-
-
-class RangeType(types.Type):
-    """Numba-dpex type corresponding to
-    :class:`numba_dpex.core.kernel_interface.indexers.Range`
-    """
-
-    def __init__(self, ndim: int):
-        self._ndim = ndim
-        if ndim < 1 or ndim > 3:
-            raise errors.TypingError(
-                "RangeType can only have 1,2, or 3 dimensions"
-            )
-        super(RangeType, self).__init__(name="Range<" + str(ndim) + ">")
-
-    @property
-    def ndim(self):
-        return self._ndim
-
-    @property
-    def key(self):
-        return self._ndim
-
-
-class NdRangeType(types.Type):
-    """Numba-dpex type corresponding to
-    :class:`numba_dpex.core.kernel_interface.indexers.NdRange`
-    """
-
-    def __init__(self, ndim: int):
-        self._ndim = ndim
-        if ndim < 1 or ndim > 3:
-            raise errors.TypingError(
-                "RangeType can only have 1,2, or 3 dimensions"
-            )
-        super(NdRangeType, self).__init__(name="NdRange<" + str(ndim) + ">")
-
-    @property
-    def ndim(self):
-        return self._ndim
-
-    @property
-    def key(self):
-        return self._ndim
+from numba_dpex.core.types import NdRangeType, RangeType
+from numba_dpex.kernel_api import NdRange, Range
 
 
 @unbox(RangeType)

numba_dpex/core/boxing/usm_ndarray.py

@@ -0,0 +1,126 @@
+# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from contextlib import ExitStack
+
+from numba.core import cgutils, types
+from numba.core.datamodel import default_manager
+from numba.core.errors import NumbaNotImplementedError
+from numba.extending import NativeValue, box, unbox
+from numba.np import numpy_support
+
+from numba_dpex.core.exceptions import UnreachableError
+from numba_dpex.core.runtime import context as dpexrt
+from numba_dpex.core.types import USMNdArray
+from numba_dpex.kernel_api import NdRange, Range
+
+
+@unbox(USMNdArray)
+def unbox_dpnp_nd_array(typ, obj, c):
+    """Converts a dpctl.tensor.usm_ndarray/dpnp.ndarray object to a Numba-dpex
+    internal array structure.
+
+    Args:
+        typ : The Numba type of the PyObject
+        obj : The actual PyObject to be unboxed
+        c : The unboxing context
+
+    Returns: A NativeValue object representing an unboxed
+        dpctl.tensor.usm_ndarray/dpnp.ndarray
+    """
+    # Reusing the numba.core.base.BaseContext's make_array function to get a
+    # struct allocated. The same struct is used for numpy.ndarray
+    # and dpnp.ndarray. It is possible to do so, as the extra information
+    # specific to dpnp.ndarray such as sycl_queue is inferred statically and
+    # stored as part of the DpnpNdArray type.
+
+    # --------------- Original Numba comment from @ubox(types.Array)
+    #
+    # This is necessary because unbox_buffer() does not work on some
+    # dtypes, e.g. datetime64 and timedelta64.
+    # TODO check matching dtype.
+    #      currently, mismatching dtype will still work and causes
+    #      potential memory corruption
+    #
+    # --------------- End of Numba comment from @ubox(types.Array)
+    nativearycls = c.context.make_array(typ)
+    nativeary = nativearycls(c.context, c.builder)
+    aryptr = nativeary._getpointer()
+
+    ptr = c.builder.bitcast(aryptr, c.pyapi.voidptr)
+    # FIXME : We need to check if Numba_RT as well as DPEX RT are enabled.
+    if c.context.enable_nrt:
+        dpexrtCtx = dpexrt.DpexRTContext(c.context)
+        errcode = dpexrtCtx.arraystruct_from_python(c.pyapi, obj, ptr)
+    else:
+        raise UnreachableError
+
+    # TODO: here we have minimal typechecking by the itemsize.
+    #       need to do better
+    try:
+        expected_itemsize = numpy_support.as_dtype(typ.dtype).itemsize
+    except NumbaNotImplementedError:
+        # Don't check types that can't be `as_dtype()`-ed
+        itemsize_mismatch = cgutils.false_bit
+    else:
+        expected_itemsize = nativeary.itemsize.type(expected_itemsize)
+        itemsize_mismatch = c.builder.icmp_unsigned(
+            "!=",
+            nativeary.itemsize,
+            expected_itemsize,
+        )
+
+    failed = c.builder.or_(
+        cgutils.is_not_null(c.builder, errcode),
+        itemsize_mismatch,
+    )
+    # Handle error
+    with c.builder.if_then(failed, likely=False):
+        c.pyapi.err_set_string(
+            "PyExc_TypeError",
+            "can't unbox usm array from PyObject into "
+            "native value.  The object maybe of a "
+            "different type",
+        )
+    return NativeValue(c.builder.load(aryptr), is_error=failed)
+
+
+@box(USMNdArray)
+def box_array(typ, val, c):
+    """Boxes a NativeValue representation of USMNdArray/DpnpNdArray type into a
+    dpctl.tensor.usm_ndarray/dpnp.ndarray PyObject
+
+    Args:
+        typ: The representation of the USMNdArray/DpnpNdArray type.
+        val: A native representation of a Numba USMNdArray/DpnpNdArray type
+            object.
+        c: The boxing context.
+
+    Returns: A Pyobject for a dpctl.tensor.usm_ndarray/dpnp.ndarray boxed from
+        the Numba-dpex native value.
+    """
+    if c.context.enable_nrt:
+        np_dtype = numpy_support.as_dtype(typ.dtype)
+        dtypeptr = c.env_manager.read_const(c.env_manager.add_const(np_dtype))
+        dpexrtCtx = dpexrt.DpexRTContext(c.context)
+        newary = dpexrtCtx.usm_ndarray_to_python_acqref(
+            c.pyapi, typ, val, dtypeptr
+        )
+
+        if not newary:
+            c.pyapi.err_set_string(
+                "PyExc_TypeError",
+                "could not box native array into a dpnp.ndarray PyObject.",
+            )
+
+        # Steals NRT ref
+        # Refer:
+        #   numba.core.base.nrt -> numba.core.runtime.context -> decref
+        #   The `NRT_decref` function is generated directly as LLVM IR inside
+        #   numba.core.runtime.nrtdynmod.py
+        c.context.nrt.decref(c.builder, typ, val)
+
+        return newary
+    else:
+        raise UnreachableError

numba_dpex/core/datamodel/models.py

@@ -196,9 +196,7 @@ def __init__(self, dmm, fe_type):
 
 
 class RangeModel(StructModel):
-    """The native data model for a
-    numba_dpex.core.kernel_interface.indexers.Range PyObject.
-    """
+    """The numba_dpex data model for a numba_dpex.kernel_api.Range PyObject."""
 
     def __init__(self, dmm, fe_type):
         members = [
@@ -216,8 +214,8 @@ def flattened_field_count(self):
 
 
 class NdRangeModel(StructModel):
-    """The native data model for a
-    numba_dpex.core.kernel_interface.indexers.NdRange PyObject.
+    """
+    The numba_dpex data model for a numba_dpex.kernel_api.NdRange PyObject.
     """
 
     def __init__(self, dmm, fe_type):