Merge pull request #963 from chudur-budur/github-930

Proper DpctlSyclQueue support

numba_dpex/core/datamodel/models.py

@@ -4,7 +4,7 @@
 
 from numba.core import datamodel, types
 from numba.core.datamodel.models import ArrayModel as DpnpNdArrayModel
-from numba.core.datamodel.models import OpaqueModel, PrimitiveModel, StructModel
+from numba.core.datamodel.models import PrimitiveModel, StructModel
 from numba.core.extending import register_model
 
 from numba_dpex.utils import address_space
@@ -54,6 +54,33 @@ def __init__(self, dmm, fe_type):
         super(ArrayModel, self).__init__(dmm, fe_type, members)
 
 
+class SyclQueueModel(StructModel):
+    """Represents the native data model for a dpctl.SyclQueue PyObject.
+
+    Numba-dpex uses a C struct as defined in
+    numba_dpex/core/runtime._queuestruct.h to store the required attributes for
+    a ``dpctl.SyclQueue`` Python object.
+
+        - ``queue_ref``: An opaque C pointer to an actual SYCL queue C++ object.
+        - ``parent``: A PyObject* that stores a reference back to the original
+                      ``dpctl.SyclQueue`` PyObject if the native struct is
+                      created by unboxing the PyObject.
+    """
+
+    def __init__(self, dmm, fe_type):
+        members = [
+            (
+                "parent",
+                types.CPointer(types.int8),
+            ),
+            (
+                "queue_ref",
+                types.CPointer(types.int8),
+            ),
+        ]
+        super(SyclQueueModel, self).__init__(dmm, fe_type, members)
+
+
 def _init_data_model_manager():
     dmm = datamodel.default_manager.copy()
     dmm.register(types.CPointer, GenericPointerModel)
@@ -83,6 +110,6 @@ def _init_data_model_manager():
 register_model(DpnpNdArray)(DpnpNdArrayModel)
 dpex_data_model_manager.register(DpnpNdArray, DpnpNdArrayModel)
 
-# Register the DpctlSyclQueue type with Numba's OpaqueModel
-register_model(DpctlSyclQueue)(OpaqueModel)
-dpex_data_model_manager.register(DpctlSyclQueue, OpaqueModel)
+# Register the DpctlSyclQueue type
+register_model(DpctlSyclQueue)(SyclQueueModel)
+dpex_data_model_manager.register(DpctlSyclQueue, SyclQueueModel)

numba_dpex/core/runtime/_dbg_printer.h

@@ -0,0 +1,28 @@
+// SPDX-FileCopyrightText: 2023 Intel Corporation
+//
+// SPDX-License-Identifier: Apache-2.0
+
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// A helper macro to print debug prints.
+///
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+/* Debugging facilities - enabled at compile-time */
+/* #undef NDEBUG */
+#if 0
+#include <stdio.h>
+#define DPEXRT_DEBUG(X)                                                        \
+    {                                                                          \
+        X;                                                                     \
+        fflush(stdout);                                                        \
+    }
+#else
+#define DPEXRT_DEBUG(X)                                                        \
+    if (0) {                                                                   \
+        X;                                                                     \
+    }
+#endif

numba_dpex/core/runtime/_dpexrt_python.c

@@ -19,24 +19,10 @@
 #include "_nrt_helper.h"
 #include "_nrt_python_helper.h"
 
+#include "_dbg_printer.h"
+#include "_queuestruct.h"
 #include "numba/_arraystruct.h"
 
-/* Debugging facilities - enabled at compile-time */
-/* #undef NDEBUG */
-#if 0
-#include <stdio.h>
-#define DPEXRT_DEBUG(X)                                                        \
-    {                                                                          \
-        X;                                                                     \
-        fflush(stdout);                                                        \
-    }
-#else
-#define DPEXRT_DEBUG(X)                                                        \
-    if (0) {                                                                   \
-        X;                                                                     \
-    }
-#endif
-
 // forward declarations
 static struct PyUSMArrayObject *PyUSMNdArray_ARRAYOBJ(PyObject *obj);
 static npy_intp product_of_shape(npy_intp *shape, npy_intp ndim);
@@ -66,6 +52,9 @@ DPEXRT_sycl_usm_ndarray_to_python_acqref(arystruct_t *arystruct,
                                          int ndim,
                                          int writeable,
                                          PyArray_Descr *descr);
+static int DPEXRT_sycl_queue_from_python(PyObject *obj,
+                                         queuestruct_t *queue_struct);
+static PyObject *DPEXRT_sycl_queue_to_python(queuestruct_t *queuestruct);
 
 /*
  * Debugging printf function used internally
@@ -663,7 +652,9 @@ static npy_intp product_of_shape(npy_intp *shape, npy_intp ndim)
     return nelems;
 }
 
+/*----------------------------------------------------------------------------*/
 /*----- Boxing and Unboxing implementations for a dpnp.ndarray PyObject ------*/
+/*----------------------------------------------------------------------------*/
 
 /*!
  * @brief Unboxes a PyObject that may represent a dpnp.ndarray into a Numba
@@ -1049,6 +1040,107 @@ DPEXRT_sycl_usm_ndarray_to_python_acqref(arystruct_t *arystruct,
     return (PyObject *)dpnp_ary;
 }
 
+/*----------------------------------------------------------------------------*/
+/*--------------------- Box-unbox helpers for dpctl.SyclQueue       ----------*/
+/*----------------------------------------------------------------------------*/
+
+/*!
+ * @brief Helper to unbox a Python dpctl.SyclQueue object to a Numba-native
+ * queuestruct_t instance.
+ *
+ * @param    obj            A dpctl.SyclQueue Python object
+ * @param    queue_struct   An instance of the struct numba-dpex uses to
+ *                          represent a dpctl.SyclQueue inside Numba.
+ * @return   {return}       Return code indicating success (0) or failure (-1).
+ */
+static int DPEXRT_sycl_queue_from_python(PyObject *obj,
+                                         queuestruct_t *queue_struct)
+{
+
+    struct PySyclQueueObject *queue_obj = NULL;
+    DPCTLSyclQueueRef queue_ref = NULL;
+    PyGILState_STATE gstate;
+
+    // Increment the ref count on obj to prevent CPython from garbage
+    // collecting the array.
+    Py_IncRef(obj);
+
+    // We are unconditionally casting obj to a struct PySyclQueueObject*. If
+    // the obj is not a struct PySyclQueueObject* then the SyclQueue_GetQueueRef
+    // will error out.
+    queue_obj = (struct PySyclQueueObject *)obj;
+
+    DPEXRT_DEBUG(
+        nrt_debug_print("DPEXRT-DEBUG: In DPEXRT_sycl_queue_from_python.\n"));
+
+    if (!(queue_ref = SyclQueue_GetQueueRef(queue_obj))) {
+        DPEXRT_DEBUG(nrt_debug_print(
+            "DPEXRT-ERROR: SyclQueue_GetQueueRef returned NULL at "
+            "%s, line %d.\n",
+            __FILE__, __LINE__));
+        goto error;
+    }
+
+    queue_struct->parent = obj;
+    queue_struct->queue_ref = queue_ref;
+
+    return 0;
+
+error:
+    // If the check failed then decrement the refcount and return an error
+    // code of -1.
+    // Decref the Pyobject of the array
+    // ensure the GIL
+    DPEXRT_DEBUG(nrt_debug_print(
+        "DPEXRT-ERROR: Failed to unbox dpctl SyclQueue into a Numba "
+        "queuestruct at %s, line %d\n",
+        __FILE__, __LINE__));
+    gstate = PyGILState_Ensure();
+    // decref the python object
+    Py_DECREF(obj);
+    // release the GIL
+    PyGILState_Release(gstate);
+
+    return -1;
+}
+
+/*!
+ * @brief A helper function that boxes a Numba-dpex queuestruct_t object into a
+ * dctl.SyclQueue PyObject using the queuestruct_t's parent attribute.
+ *
+ * If there is no parent pointer stored in the queuestruct, then an error will
+ * be raised.
+ *
+ * @param    queuestruct    A Numba-dpex queuestruct object.
+ * @return   {return}       A PyObject created from the queuestruct->parent, if
+ *                          the PyObject could not be created return NULL.
+ */
+static PyObject *DPEXRT_sycl_queue_to_python(queuestruct_t *queuestruct)
+{
+    PyObject *orig_queue = NULL;
+    PyGILState_STATE gstate;
+
+    orig_queue = queuestruct->parent;
+    // FIXME: Better error checking is needed to enforce the boxing of the queue
+    // object. For now, only the minimal is done as the returning of SyclQueue
+    // from a dpjit function should not be a used often and the dpctl C API for
+    // type checking etc. is not ready.
+    if (orig_queue == NULL) {
+        PyErr_Format(PyExc_ValueError,
+                     "In 'box_from_queuestruct_parent', "
+                     "failed to create a new dpctl.SyclQueue object.");
+        return NULL;
+    }
+
+    gstate = PyGILState_Ensure();
+    // decref the parent python object as we did an incref when unboxing it
+    Py_DECREF(orig_queue);
+    // release the GIL
+    PyGILState_Release(gstate);
+
+    return orig_queue;
+}
+
 /*----------------------------------------------------------------------------*/
 /*--------------------- The _dpexrt_python Python extension module  -- -------*/
 /*----------------------------------------------------------------------------*/
@@ -1082,6 +1174,9 @@ static PyObject *build_c_helpers_dict(void)
     _declpointer("DPEXRT_MemInfo_fill", &DPEXRT_MemInfo_fill);
     _declpointer("NRT_ExternalAllocator_new_for_usm",
                  &NRT_ExternalAllocator_new_for_usm);
+    _declpointer("DPEXRT_sycl_queue_from_python",
+                 &DPEXRT_sycl_queue_from_python);
+    _declpointer("DPEXRT_sycl_queue_to_python", &DPEXRT_sycl_queue_to_python);
 
 #undef _declpointer
     return dct;
@@ -1129,6 +1224,11 @@ MOD_INIT(_dpexrt_python)
         m, "DPEXRT_sycl_usm_ndarray_to_python_acqref",
         PyLong_FromVoidPtr(&DPEXRT_sycl_usm_ndarray_to_python_acqref));
 
+    PyModule_AddObject(m, "DPEXRT_sycl_queue_from_python",
+                       PyLong_FromVoidPtr(&DPEXRT_sycl_queue_from_python));
+    PyModule_AddObject(m, "DPEXRT_sycl_queue_to_python",
+                       PyLong_FromVoidPtr(&DPEXRT_sycl_queue_to_python));
+
     PyModule_AddObject(m, "DPEXRTQueue_CreateFromFilterString",
                        PyLong_FromVoidPtr(&DPEXRTQueue_CreateFromFilterString));
     PyModule_AddObject(m, "DpexrtQueue_SubmitRange",

numba_dpex/core/runtime/_queuestruct.h

@@ -0,0 +1,17 @@
+#ifndef NUMBA_DPEX_QUEUESTRUCT_H_
+#define NUMBA_DPEX_QUEUESTRUCT_H_
+/*
+ * Fill in the *queuestruct* with information from the Numpy array *obj*.
+ * *queuestruct*'s layout is defined in numba.targets.arrayobj (look
+ * for the ArrayTemplate class).
+ */
+
+#include <Python.h>
+
+typedef struct
+{
+    PyObject *parent;
+    void *queue_ref;
+} queuestruct_t;
+
+#endif /* NUMBA_DPEX_QUEUESTRUCT_H_ */

numba_dpex/core/runtime/context.py

@@ -109,13 +109,6 @@ def arraystruct_from_python(self, pyapi, obj, ptr):
         """Generates a call to DPEXRT_sycl_usm_ndarray_from_python C function
         defined in the _DPREXRT_python Python extension.
 
-        Args:
-            pyapi (_type_): _description_
-            obj (_type_): _description_
-            ptr (_type_): _description_
-
-        Returns:
-            _type_: _description_
         """
         fnty = llvmir.FunctionType(
             llvmir.IntType(32), [pyapi.pyobj, pyapi.voidptr]
@@ -128,6 +121,34 @@ def arraystruct_from_python(self, pyapi, obj, ptr):
 
         return self.error
 
+    def queuestruct_from_python(self, pyapi, obj, ptr):
+        """Calls the c function DPEXRT_sycl_queue_from_python"""
+
+        fnty = llvmir.FunctionType(
+            llvmir.IntType(32), [pyapi.pyobj, pyapi.voidptr]
+        )
+
+        fn = pyapi._get_function(fnty, "DPEXRT_sycl_queue_from_python")
+        fn.args[0].add_attribute("nocapture")
+        fn.args[1].add_attribute("nocapture")
+
+        self.error = pyapi.builder.call(fn, (obj, ptr))
+
+        return self.error
+
+    def queuestruct_to_python(self, pyapi, val):
+        """Calls the c function DPEXRT_sycl_queue_to_python"""
+
+        fnty = llvmir.FunctionType(pyapi.pyobj, [pyapi.voidptr])
+
+        fn = pyapi._get_function(fnty, "DPEXRT_sycl_queue_to_python")
+        fn.args[0].add_attribute("nocapture")
+        qptr = cgutils.alloca_once_value(pyapi.builder, val)
+        ptr = pyapi.builder.bitcast(qptr, pyapi.voidptr)
+        self.error = pyapi.builder.call(fn, [ptr])
+
+        return self.error
+
     def usm_ndarray_to_python_acqref(self, pyapi, aryty, ary, dtypeptr):
         """Boxes a DpnpNdArray native object into a Python dpnp.ndarray.
 
@@ -230,7 +251,6 @@ def submit_range(
         fn = cgutils.get_or_insert_function(
             mod, fnty, "DpexrtQueue_SubmitRange"
         )
-        # fn.return_value.add_attribute("noalias")
 
         ret = builder.call(
             fn,