WIP for fully enable sycl queue in DPEXRT_MemInfo_alloc

Sycl queue support through DPEXRT_MemInfo_alloc working

Author: khaled

numba_dpex/_patches.py

@@ -25,6 +25,7 @@
 
 from numba_dpex.core.runtime import context as dpexrt
 from numba_dpex.core.types import DpnpNdArray
+from numba_dpex.core.types.dpctl_types import DpctlSyclQueue
 
 # Numpy array constructors
 
@@ -172,13 +173,73 @@ def _mk_alloc(
     return out
 
 
+def make_queue(context, builder, arrtype):
+    """Utility function used for allocating a new queue.
+
+    This function will allocates a new queue (e.g. SYCL queue)
+    during LLVM code generation (lowering). Given a target context,
+    builder, array type, returns a LLVM value pointing at a numba-dpex
+    runtime allocated queue.
+
+    Args:
+        context (numba.core.base.BaseContext): Any of the context
+            derived from Numba's BaseContext
+            (e.g. `numba.core.cpu.CPUContext`).
+        builder (llvmlite.ir.builder.IRBuilder): The IR builder
+            from `llvmlite` for code generation.
+        arrtype (numba_dpex.core.types.dpnp_ndarray_type.DpnpNdArray):
+            Any of the array types derived from
+            `numba.core.types.nptypes.Array`,
+            e.g. `numba_dpex.core.types.dpnp_ndarray_type.DpnpNdArray`.
+            Please refer to `numba_dpex.dpnp_iface._intrinsic.alloc_empty_arrayobj()`
+            function for details on how to construct this argument.
+
+    Returns:
+        tuple: A tuple containing `llvmlite.ir.instructions.ExtractValue`,
+            `llvmlite.ir.instructions.CastInstr` and `numba.core.pythonapi.PythonAPI`.
+    """
+
+    pyapi = context.get_python_api(builder)
+    queue_struct_proxy = cgutils.create_struct_proxy(
+        DpctlSyclQueue(arrtype.queue)
+    )(context, builder)
+    queue_struct_ptr = queue_struct_proxy._getpointer()
+    queue_struct_voidptr = builder.bitcast(queue_struct_ptr, cgutils.voidptr_t)
+
+    address = context.get_constant(types.intp, id(arrtype.queue))
+    queue_address_ptr = builder.inttoptr(address, cgutils.voidptr_t)
+
+    dpexrtCtx = dpexrt.DpexRTContext(context)
+    dpexrtCtx.queuestruct_from_python(
+        pyapi, queue_address_ptr, queue_struct_voidptr
+    )
+    # errcode = dpexrtCtx.queuestruct_from_python(
+    #     pyapi, queue_address_ptr, queue_struct_voidptr
+    # )
+    # is_error = cgutils.is_not_null(builder, errcode)
+    # # Handle error
+    # with builder.if_then(is_error, likely=False):
+    #     pyapi.err_set_string(
+    #         "_patches.make_queue(): PyExc_TypeError",
+    #         "can't unbox dpctl.SyclQueue from PyObject into a Numba "
+    #         "native value. The object maybe of a different type",
+    #     )
+
+    queue_struct = builder.load(queue_struct_ptr)
+    queue_ref = builder.extract_value(queue_struct, 1)
+
+    return (queue_ref, queue_address_ptr, pyapi)
+
+
 def _empty_nd_impl(context, builder, arrtype, shapes):
     """Utility function used for allocating a new array during LLVM code
     generation (lowering).  Given a target context, builder, array
     type, and a tuple or list of lowered dimension sizes, returns a
     LLVM value pointing at a Numba runtime allocated array.
     """
 
+    (queue, queue_ptr, pyapi) = make_queue(context, builder, arrtype)
+
     arycls = make_array(arrtype)
     ary = arycls(context, builder)
 
@@ -231,21 +292,16 @@ def _empty_nd_impl(context, builder, arrtype, shapes):
 
     if isinstance(arrtype, DpnpNdArray):
         usm_ty = arrtype.usm_type
-        usm_ty_val = 0
-        if usm_ty == "device":
-            usm_ty_val = 1
-        elif usm_ty == "shared":
-            usm_ty_val = 2
-        elif usm_ty == "host":
-            usm_ty_val = 3
-        usm_type = context.get_constant(types.uint64, usm_ty_val)
-        device = context.insert_const_string(builder.module, arrtype.device)
+        usm_ty_map = {"device": 1, "shared": 2, "host": 3}
+        usm_type = context.get_constant(
+            types.uint64, usm_ty_map[usm_ty] if usm_ty in usm_ty_map else 0
+        )
 
         args = (
             context.get_dummy_value(),
             allocsize,
             usm_type,
-            device,
+            queue,
         )
         mip = types.MemInfoPointer(types.voidptr)
         arytypeclass = types.TypeRef(type(arrtype))
@@ -265,6 +321,7 @@ def _empty_nd_impl(context, builder, arrtype, shapes):
         fnop.get_call_type(context.typing_context, sig.args, {})
         eqfn = context.get_function(fnop, sig)
         meminfo = eqfn(builder, args)
+        pyapi.decref(queue_ptr)
     else:
         dtype = arrtype.dtype
         align_val = context.get_preferred_array_alignment(dtype)
@@ -298,36 +355,36 @@ def _empty_nd_impl(context, builder, arrtype, shapes):
 
 
 @overload_classmethod(DpnpNdArray, "_usm_allocate")
-def _ol_array_allocate(cls, allocsize, usm_type, device):
+def _ol_array_allocate(cls, allocsize, usm_type, queue):
     """Implements an allocator for dpnp.ndarrays."""
 
-    def impl(cls, allocsize, usm_type, device):
-        return intrin_usm_alloc(allocsize, usm_type, device)
+    def impl(cls, allocsize, usm_type, queue):
+        return intrin_usm_alloc(allocsize, usm_type, queue)
 
     return impl
 
 
 numba_config.DISABLE_PERFORMANCE_WARNINGS = 0
 
 
-def _call_usm_allocator(arrtype, size, usm_type, device):
+def _call_usm_allocator(arrtype, size, usm_type, queue):
     """Trampoline to call the intrinsic used for allocation"""
-    return arrtype._usm_allocate(size, usm_type, device)
+    return arrtype._usm_allocate(size, usm_type, queue)
 
 
 numba_config.DISABLE_PERFORMANCE_WARNINGS = 1
 
 
 @intrinsic
-def intrin_usm_alloc(typingctx, allocsize, usm_type, device):
+def intrin_usm_alloc(typingctx, allocsize, usm_type, queue):
     """Intrinsic to call into the allocator for Array"""
 
     def codegen(context, builder, signature, args):
-        [allocsize, usm_type, device] = args
+        [allocsize, usm_type, queue] = args
         dpexrtCtx = dpexrt.DpexRTContext(context)
-        meminfo = dpexrtCtx.meminfo_alloc(builder, allocsize, usm_type, device)
+        meminfo = dpexrtCtx.meminfo_alloc(builder, allocsize, usm_type, queue)
         return meminfo
 
     mip = types.MemInfoPointer(types.voidptr)  # return untyped pointer
-    sig = signature(mip, allocsize, usm_type, device)
+    sig = signature(mip, allocsize, usm_type, queue)
     return sig, codegen

numba_dpex/core/runtime/_dpexrt_python.c

@@ -45,6 +45,9 @@ static NRT_MemInfo *NRT_MemInfo_new_from_usmndarray(PyObject *ndarrobj,
                                                     npy_intp nitems,
                                                     npy_intp itemsize,
                                                     DPCTLSyclQueueRef qref);
+static NRT_MemInfo *DPEXRT_MemInfo_alloc(npy_intp size,
+                                         size_t usm_type,
+                                         const DPCTLSyclQueueRef qref);
 static void usmndarray_meminfo_dtor(void *ptr, size_t size, void *info);
 static PyObject *box_from_arystruct_parent(arystruct_t *arystruct,
                                            int ndim,
@@ -440,33 +443,25 @@ static NRT_MemInfo *NRT_MemInfo_new_from_usmndarray(PyObject *ndarrobj,
  * @return   {return}     A new NRT_MemInfo object, NULL if no NRT_MemInfo
  *                        object could be created.
  */
-static NRT_MemInfo *
-DPEXRT_MemInfo_alloc(npy_intp size, size_t usm_type, const char *device)
+static NRT_MemInfo *DPEXRT_MemInfo_alloc(npy_intp size,
+                                         size_t usm_type,
+                                         const DPCTLSyclQueueRef qref)
 {
     NRT_MemInfo *mi = NULL;
     NRT_ExternalAllocator *ext_alloca = NULL;
     MemInfoDtorInfo *midtor_info = NULL;
-    DPCTLSyclQueueRef qref = NULL;
 
     DPEXRT_DEBUG(drt_debug_print(
         "DPEXRT-DEBUG: Inside DPEXRT_MemInfo_alloc  %s, line %d\n", __FILE__,
         __LINE__));
+
     // Allocate a new NRT_MemInfo object
     if (!(mi = (NRT_MemInfo *)malloc(sizeof(NRT_MemInfo)))) {
         DPEXRT_DEBUG(drt_debug_print(
             "DPEXRT-ERROR: Could not allocate a new NRT_MemInfo object.\n"));
         goto error;
     }
 
-    if (!(qref = (DPCTLSyclQueueRef)DPEXRTQueue_CreateFromFilterString(device)))
-    {
-        DPEXRT_DEBUG(
-            drt_debug_print("DPEXRT-ERROR: Could not create a sycl::queue from "
-                            "filter string: %s at %s %d.\n",
-                            device, __FILE__, __LINE__));
-        goto error;
-    }
-
     // Allocate a new NRT_ExternalAllocator
     if (!(ext_alloca = NRT_ExternalAllocator_new_for_usm(qref, usm_type)))
         goto error;
@@ -484,10 +479,10 @@ DPEXRT_MemInfo_alloc(npy_intp size, size_t usm_type, const char *device)
 
     mi->size = size;
     mi->external_allocator = ext_alloca;
-    DPEXRT_DEBUG(drt_debug_print(
-        "DPEXRT-DEBUG: DPEXRT_MemInfo_alloc mi=%p "
-        "external_allocator=%p for usm_type %zu on device %s, %s at %d\n",
-        mi, ext_alloca, usm_type, device, __FILE__, __LINE__));
+    // DPEXRT_DEBUG(drt_debug_print(
+    //     "DPEXRT-DEBUG: DPEXRT_MemInfo_alloc mi=%p "
+    //     "external_allocator=%p for usm_type %zu on device %s, %s at %d\n",
+    //     mi, ext_alloca, usm_type, device, __FILE__, __LINE__));
 
     return mi;
 
@@ -1157,6 +1152,11 @@ static int DPEXRT_sycl_queue_from_python(PyObject *obj,
         goto error;
     }
 
+    DPEXRT_DEBUG(DPCTLSyclDeviceRef device_ref;
+                 device_ref = DPCTLQueue_GetDevice(queue_ref);
+                 DPCTLDeviceMgr_PrintDeviceInfo(device_ref);
+                 DPCTLDevice_Delete(device_ref););
+
     queue_struct->parent = obj;
     queue_struct->queue_ref = queue_ref;
 

numba_dpex/core/runtime/context.py

@@ -20,20 +20,20 @@ def _check_null_result(func):
         @functools.wraps(func)
         def wrap(self, builder, *args, **kwargs):
             memptr = func(self, builder, *args, **kwargs)
-            msg = "USM allocation failed. Check the usm_type and filter "
-            "string values."
+            msg = "USM allocation failed. Check the usm_type and filter string values."
             cgutils.guard_memory_error(self._context, builder, memptr, msg=msg)
             return memptr
 
         return wrap
 
     @_check_null_result
-    def meminfo_alloc(self, builder, size, usm_type, device):
+    def meminfo_alloc(self, builder, size, usm_type, queue):
         """
         Wrapper to call :func:`~context.DpexRTContext.meminfo_alloc_unchecked`
         with null checking of the returned value.
         """
-        return self.meminfo_alloc_unchecked(builder, size, usm_type, device)
+
+        return self.meminfo_alloc_unchecked(builder, size, usm_type, queue)
 
     @_check_null_result
     def meminfo_fill(
@@ -60,7 +60,7 @@ def meminfo_fill(
             device,
         )
 
-    def meminfo_alloc_unchecked(self, builder, size, usm_type, device):
+    def meminfo_alloc_unchecked(self, builder, size, usm_type, queue):
         """Allocate a new MemInfo with a data payload of `size` bytes.
 
         The result of the call is checked and if it is NULL, i.e. allocation
@@ -79,6 +79,7 @@ def meminfo_alloc_unchecked(self, builder, size, usm_type, device):
 
         Returns: A pointer to the MemInfo is returned.
         """
+
         mod = builder.module
         u64 = llvmir.IntType(64)
         fnty = llvmir.FunctionType(
@@ -87,7 +88,7 @@ def meminfo_alloc_unchecked(self, builder, size, usm_type, device):
         fn = cgutils.get_or_insert_function(mod, fnty, "DPEXRT_MemInfo_alloc")
         fn.return_value.add_attribute("noalias")
 
-        ret = builder.call(fn, [size, usm_type, device])
+        ret = builder.call(fn, [size, usm_type, queue])
 
         return ret
 
@@ -164,7 +165,6 @@ def queuestruct_from_python(self, pyapi, obj, ptr):
         fn.args[1].add_attribute("nocapture")
 
         self.error = pyapi.builder.call(fn, (obj, ptr))
-
         return self.error
 
     def queuestruct_to_python(self, pyapi, val):

numba_dpex/core/types/dpctl_types.py

@@ -37,15 +37,18 @@ def unbox_sycl_queue(typ, obj, c):
     """
     Convert a SyclQueue object to a native structure.
     """
+
     qstruct = cgutils.create_struct_proxy(typ)(c.context, c.builder)
     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
     is_error = cgutils.is_not_null(c.builder, errcode)
+
     # Handle error
     with c.builder.if_then(is_error, likely=False):
         c.pyapi.err_set_string(

numba_dpex/core/types/usm_ndarray_type.py

@@ -93,7 +93,7 @@ def __init__(
             )
             name = (
                 "%s(dtype=%s, ndim=%s, layout=%s, address_space=%s, "
-                "usm_type=%s, device=%s, sycl_device=%s)" % name_parts
+                "usm_type=%s, device=%s, sycl_queue=%s)" % name_parts
             )
 
         super().__init__(

numba_dpex/dpnp_iface/_intrinsic.py

@@ -107,17 +107,17 @@ def fill_arrayobj(context, builder, ary, arrtype, fill_value):
 
 
 @intrinsic
-def intrin_usm_alloc(typingctx, allocsize, usm_type, device):
+def intrin_usm_alloc(typingctx, allocsize, usm_type, queue):
     """Intrinsic to call into the allocator for Array"""
 
     def codegen(context, builder, signature, args):
-        [allocsize, usm_type, device] = args
+        [allocsize, usm_type, queue] = args
         dpexrtCtx = dpexrt.DpexRTContext(context)
-        meminfo = dpexrtCtx.meminfo_alloc(builder, allocsize, usm_type, device)
+        meminfo = dpexrtCtx.meminfo_alloc(builder, allocsize, usm_type, queue)
         return meminfo
 
     mip = types.MemInfoPointer(types.voidptr)  # return untyped pointer
-    sig = signature(mip, allocsize, usm_type, device)
+    sig = signature(mip, allocsize, usm_type, queue)
     return sig, codegen