Create a new module to store all box/unbox routines.

Diptorup Deb · Diptorup Deb · commit 4292cecdba7c · 2024-02-02T12:36:09.000-06:00
diff --git a/numba_dpex/__init__.py b/numba_dpex/__init__.py
@@ -98,6 +98,7 @@ 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 import ranges_overloads  # noqa E402
 
diff --git a/numba_dpex/core/boxing/__init__.py b/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 *
diff --git a/numba_dpex/core/boxing/ranges.py b/numba_dpex/core/boxing/ranges.py
diff --git a/numba_dpex/core/boxing/usm_ndarray.py b/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
diff --git a/numba_dpex/core/types/dpnp_ndarray_type.py b/numba_dpex/core/types/dpnp_ndarray_type.py
@@ -3,14 +3,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import dpnp
-from numba.core import cgutils, ir, types
-from numba.core.errors import NumbaNotImplementedError
+from numba.core import ir, types
 from numba.core.ir_utils import get_np_ufunc_typ, mk_unique_var
 from numba.core.pythonapi import NativeValue, PythonAPI, 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 .usm_ndarray_type import USMNdArray
 
@@ -197,117 +192,3 @@ def __allocate__(
             out.extend([g_np_assign, attr_assign, typ_var_assign, alloc_assign])
 
         return out
-
-
-# TODO: move this section to separate file
-# --------------- Boxing/Unboxing logic for dpnp.ndarray ----------------------#
-
-
-@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
diff --git a/numba_dpex/experimental/__init__.py b/numba_dpex/experimental/__init__.py
@@ -9,7 +9,7 @@
 from numba.core.imputils import Registry
 
 # Temporary so that Range and NdRange work in experimental call_kernel
-from numba_dpex.core import boxing
+from numba_dpex.core.boxing import *
 
 from ._kernel_dpcpp_spirv_overloads import _atomic_ref_overloads
 from .decorators import device_func, kernel
