Merge pull request #1064 from IntelPython/change/DpctlSyclQueue_type

Changes to the DpctlSyclQueue and USMNdArray types.

Author: Diptorup Deb

numba_dpex/core/kernel_interface/dispatcher.py

@@ -409,9 +409,14 @@ def __call__(self, *args):
         # FIXME: For specialized and ahead of time compiled and cached kernels,
         # the CFD check was already done statically. The run-time check is
         # redundant. We should avoid these checks for the specialized case.
-        exec_queue = determine_kernel_launch_queue(
+        ty_queue = determine_kernel_launch_queue(
             args, argtypes, self.kernel_name
         )
+
+        # FIXME: We need a better way than having to create a queue every time.
+        device = ty_queue.sycl_device
+        exec_queue = dpctl.get_device_cached_queue(device)
+
         backend = exec_queue.backend
 
         if exec_queue.backend not in [

numba_dpex/core/parfors/kernel_builder.py

@@ -6,6 +6,7 @@
 import sys
 import warnings
 
+import dpctl
 import dpctl.program as dpctl_prog
 from numba.core import ir, types
 from numba.core.errors import NumbaParallelSafetyWarning
@@ -426,7 +427,10 @@ def create_kernel_for_parfor(
     for arg in parfor_args:
         obj = typemap[arg]
         if isinstance(obj, DpnpNdArray):
-            exec_queue = obj.queue
+            filter_string = obj.queue.sycl_device
+            # FIXME: A better design is required so that we do not have to
+            # create a queue every time.
+            exec_queue = dpctl.get_device_cached_queue(filter_string)
 
     if not exec_queue:
         raise AssertionError(

numba_dpex/core/parfors/reduction_kernel_builder.py

@@ -4,6 +4,7 @@
 
 import warnings
 
+import dpctl
 from numba.core import types
 from numba.core.errors import NumbaParallelSafetyWarning
 from numba.core.ir_utils import (
@@ -18,6 +19,8 @@
 )
 from numba.core.typing import signature
 
+from numba_dpex.core.types import DpctlSyclQueue
+
 from ..utils.kernel_templates.reduction_template import (
     RemainderReduceIntermediateKernelTemplate,
     TreeReduceIntermediateKernelTemplate,
@@ -134,7 +137,13 @@ def create_reduction_main_kernel_for_parfor(
         flags.noalias = True
 
     kernel_sig = signature(types.none, *kernel_param_types)
-    exec_queue = typemap[reductionKernelVar.parfor_params[0]].queue
+
+    # FIXME: A better design is required so that we do not have to create a
+    # queue every time.
+    ty_queue: DpctlSyclQueue = typemap[
+        reductionKernelVar.parfor_params[0]
+    ].queue
+    exec_queue = dpctl.get_device_cached_queue(ty_queue.sycl_device)
 
     sycl_kernel = _compile_kernel_parfor(
         exec_queue,
@@ -331,11 +340,12 @@ def create_reduction_remainder_kernel_for_parfor(
 
     kernel_sig = signature(types.none, *kernel_param_types)
 
-    # FIXME: Enable check after CFD pass has been added
-    # exec_queue = determine_kernel_launch_queue(
-    #     args=parfor_args, argtypes=kernel_param_types, kernel_name=kernel_name
-    # )
-    exec_queue = typemap[reductionKernelVar.parfor_params[0]].queue
+    # FIXME: A better design is required so that we do not have to create a
+    # queue every time.
+    ty_queue: DpctlSyclQueue = typemap[
+        reductionKernelVar.parfor_params[0]
+    ].queue
+    exec_queue = dpctl.get_device_cached_queue(ty_queue.sycl_device)
 
     sycl_kernel = _compile_kernel_parfor(
         exec_queue,

numba_dpex/core/runtime/_dpexrt_python.c

@@ -1180,7 +1180,7 @@ static int DPEXRT_sycl_queue_from_python(PyObject *obj,
     PyGILState_STATE gstate;
 
     // Increment the ref count on obj to prevent CPython from garbage
-    // collecting the array.
+    // collecting the dpctl.SyclQueue object
     Py_IncRef(obj);
 
     // We are unconditionally casting obj to a struct PySyclQueueObject*. If

numba_dpex/core/typeconv/array_conversion.py

@@ -4,8 +4,7 @@
 
 from numba.np import numpy_support
 
-from numba_dpex.core.types import USMNdArray
-from numba_dpex.core.utils import get_info_from_suai
+from numba_dpex.core.types import DpctlSyclQueue, USMNdArray
 from numba_dpex.utils.constants import address_space
 
 
@@ -37,7 +36,7 @@ def to_usm_ndarray(suai_attrs, addrspace=address_space.GLOBAL):
         ndim=suai_attrs.dimensions,
         layout=layout,
         usm_type=suai_attrs.usm_type,
-        queue=suai_attrs.queue,
+        queue=DpctlSyclQueue(suai_attrs.queue),
         readonly=not suai_attrs.is_writable,
         name=None,
         aligned=True,

numba_dpex/core/types/dpctl_types.py

@@ -14,20 +14,18 @@
 
 
 class DpctlSyclQueue(types.Type):
-    """A Numba type to represent a dpctl.SyclQueue PyObject.
-
-    For now, a dpctl.SyclQueue is represented as a Numba opaque type that allows
-    passing in and using a SyclQueue object as an opaque pointer type inside
-    Numba.
-    """
+    """A Numba type to represent a dpctl.SyclQueue PyObject."""
 
     def __init__(self, sycl_queue):
         if not isinstance(sycl_queue, SyclQueue):
             raise TypeError("The argument sycl_queue is not of type SyclQueue.")
 
-        self._sycl_queue = sycl_queue
+        # XXX: Storing the device filter string is a temporary workaround till
+        # the compute follows data inference pass is fixed to use SyclQueue
+        self._device = sycl_queue.sycl_device.filter_string
+
         try:
-            self._unique_id = hash(self._sycl_queue)
+            self._unique_id = hash(sycl_queue)
         except Exception:
             self._unique_id = self.rand_digit_str(16)
         super(DpctlSyclQueue, self).__init__(name="DpctlSyclQueue")
@@ -38,8 +36,14 @@ def rand_digit_str(self, n):
         )
 
     @property
-    def sycl_queue(self):
-        return self._sycl_queue
+    def sycl_device(self):
+        """Returns the SYCL oneAPI extension filter string associated with the
+        queue.
+
+        Returns:
+            str: A SYCL oneAPI extension filter string
+        """
+        return self._device
 
     @property
     def key(self):
@@ -69,11 +73,8 @@ def unbox_sycl_queue(typ, obj, c):
     qptr = qstruct._getpointer()
     ptr = c.builder.bitcast(qptr, c.pyapi.voidptr)
 
-    if c.context.enable_nrt:
-        dpexrtCtx = dpexrt.DpexRTContext(c.context)
-        errcode = dpexrtCtx.queuestruct_from_python(c.pyapi, obj, ptr)
-    else:
-        raise UnreachableError
+    dpexrtCtx = dpexrt.DpexRTContext(c.context)
+    errcode = dpexrtCtx.queuestruct_from_python(c.pyapi, obj, ptr)
     is_error = cgutils.is_not_null(c.builder, errcode)
 
     # Handle error

numba_dpex/core/types/usm_ndarray_type.py

@@ -12,6 +12,7 @@
 from numba.core.types.npytypes import Array
 from numba.np.numpy_support import from_dtype
 
+from numba_dpex.core.types.dpctl_types import DpctlSyclQueue
 from numba_dpex.utils import address_space
 
 
@@ -31,22 +32,28 @@ def __init__(
         aligned=True,
         addrspace=address_space.GLOBAL,
     ):
-        if queue and not isinstance(queue, types.misc.Omitted) and device:
+        if (
+            queue is not None
+            and not (
+                isinstance(queue, types.misc.Omitted)
+                or isinstance(queue, types.misc.NoneType)
+            )
+            and device is not None
+        ):
             raise TypeError(
                 "numba_dpex.core.types.usm_ndarray_type.USMNdArray.__init__(): "
-                "`device` and `sycl_queue` are exclusive keywords, i.e. use one or other."
+                "`device` and `sycl_queue` are exclusive keywords, "
+                "i.e. use one or other."
             )
 
-        self.usm_type = usm_type
-        self.addrspace = addrspace
-
-        if queue and not isinstance(queue, types.misc.Omitted):
-            if not isinstance(queue, dpctl.SyclQueue):
+        if queue is not None and not (
+            isinstance(queue, types.misc.Omitted)
+            or isinstance(queue, types.misc.NoneType)
+        ):
+            if not isinstance(queue, DpctlSyclQueue):
                 raise TypeError(
-                    "numba_dpex.core.types.usm_ndarray_type.USMNdArray.__init__(): "
-                    "The queue keyword arg should be a dpctl.SyclQueue object or None."
-                    "Found type(queue) ="
-                    + str(type(queue) + " and queue =" + queue)
+                    "The queue keyword arg should be either DpctlSyclQueue or "
+                    "NoneType. Found type(queue) = " + str(type(queue))
                 )
             self.queue = queue
         else:
@@ -55,24 +62,23 @@ def __init__(
             else:
                 if not isinstance(device, str):
                     raise TypeError(
-                        "numba_dpex.core.types.usm_ndarray_type.USMNdArray.__init__(): "
-                        "The device keyword arg should be a str object specifying "
-                        "a SYCL filter selector."
+                        "The device keyword arg should be a str object "
+                        "specifying a SYCL filter selector."
                     )
                 sycl_device = dpctl.SyclDevice(device)
 
-            self.queue = dpctl._sycl_queue_manager.get_device_cached_queue(
+            sycl_queue = dpctl._sycl_queue_manager.get_device_cached_queue(
                 sycl_device
             )
+            self.queue = DpctlSyclQueue(sycl_queue=sycl_queue)
 
-        self.device = self.queue.sycl_device.filter_string
+        self.device = self.queue.sycl_device
+        self.usm_type = usm_type
+        self.addrspace = addrspace
 
         if not dtype:
             dummy_tensor = dpctl.tensor.empty(
-                1,
-                order=layout,
-                usm_type=usm_type,
-                sycl_queue=self.queue,
+                1, order=layout, usm_type=usm_type, device=self.device
             )
             # convert dpnp type to numba/numpy type
             _dtype = dummy_tensor.dtype

numba_dpex/core/typing/typeof.py

@@ -42,15 +42,18 @@ def _typeof_helper(val, array_class_type):
             "The usm_type for the usm_ndarray could not be inferred"
         )
 
-    assert val.sycl_queue is not None
+    if not val.sycl_queue:
+        raise AssertionError
+
+    ty_queue = DpctlSyclQueue(sycl_queue=val.sycl_queue)
 
     return array_class_type(
         dtype=dtype,
         ndim=val.ndim,
         layout=layout,
         readonly=readonly,
         usm_type=usm_type,
-        queue=val.sycl_queue,
+        queue=ty_queue,
         addrspace=address_space.GLOBAL,
     )