Merge pull request #1193 from IntelPython/feature/add_python_object_createion_for_unboxing_sycl_event_and_queue

Feature/add python object createion for unboxing sycl event and queue

Author: Diptorup Deb

.github/workflows/pre-commit.yml

@@ -9,6 +9,8 @@ jobs:
   pre-commit:
     runs-on: ubuntu-20.04
     steps:
-    - uses: actions/checkout@v2
-    - uses: actions/setup-python@v2
-    - uses: pre-commit/[email protected]
+    - uses: actions/checkout@v3
+    - uses: actions/setup-python@v3
+      with:
+        python-version: '3.11'
+    - uses: pre-commit/[email protected]

numba_dpex/core/datamodel/models.py

@@ -148,6 +148,10 @@ def __init__(self, dmm, fe_type):
                 "meminfo",
                 types.MemInfoPointer(types.pyobject),
             ),
+            (
+                "parent",
+                types.pyobject,
+            ),
             (
                 "queue_ref",
                 types.CPointer(types.int8),

numba_dpex/core/runtime/_dpexrt_python.c

@@ -1256,6 +1256,7 @@ static int DPEXRT_sycl_queue_from_python(NRT_api_functions *nrt,
     Py_INCREF(queue_obj);
     queue_struct->meminfo =
         nrt->manage_memory(queue_obj, NRT_MemInfo_pyobject_dtor);
+    queue_struct->parent = queue_obj;
     queue_struct->queue_ref = queue_ref;
 
     return 0;
@@ -1287,29 +1288,26 @@ static int DPEXRT_sycl_queue_from_python(NRT_api_functions *nrt,
 static PyObject *DPEXRT_sycl_queue_to_python(NRT_api_functions *nrt,
                                              queuestruct_t *queuestruct)
 {
-    PyObject *orig_queue = NULL;
-
-    orig_queue = nrt->get_data(queuestruct->meminfo);
-    // 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;
+    PyObject *queue_obj = queuestruct->parent;
+
+    if (queue_obj == NULL) {
+        // Make create copy of queue_ref so we don't need to manage nrt lifetime
+        // from python object.
+        queue_obj = SyclQueue_Make(queuestruct->queue_ref);
+    }
+    else {
+        // Unfortunately we can not optimize (nrt->release that triggers
+        // Py_DECREF() from the destructor) and Py_INCREF() because nrt may need
+        // the object even if we return it to python.
+        // We need to increase reference count because we are returning new
+        // reference to the same queue.
+        Py_INCREF(queue_obj);
     }
 
-    // TODO: is there any way to release meminfo without calling dtor so we dont
-    //  call incref, decref one after another.
-    // We need to increase reference count because we are returning new
-    // reference to the same queue.
-    Py_INCREF(orig_queue);
     // We need to release meminfo since we are taking ownership back.
     nrt->release(queuestruct->meminfo);
 
-    return orig_queue;
+    return queue_obj;
 }
 
 /*----------------------------------------------------------------------------*/
@@ -1384,33 +1382,32 @@ static int DPEXRT_sycl_event_from_python(NRT_api_functions *nrt,
 static PyObject *DPEXRT_sycl_event_to_python(NRT_api_functions *nrt,
                                              eventstruct_t *eventstruct)
 {
-    PyObject *orig_event = NULL;
+    PyObject *event_obj = NULL;
     PyGILState_STATE gstate;
 
-    orig_event = nrt->get_data(eventstruct->meminfo);
-    // FIXME: Better error checking is needed to enforce the boxing of the event
-    // object. For now, only the minimal is done as the returning of SyclEvent
-    // from a dpjit function should not be a used often and the dpctl C API for
-    // type checking etc. is not ready.
-    if (orig_event == NULL) {
-        PyErr_Format(PyExc_ValueError,
-                     "In 'box_from_eventstruct_parent', "
-                     "failed to create a new dpctl.SyclEvent object.");
-        return NULL;
-    }
+    event_obj = nrt->get_data(eventstruct->meminfo);
 
     DPEXRT_DEBUG(
         drt_debug_print("DPEXRT-DEBUG: In DPEXRT_sycl_event_to_python.\n"););
 
-    // TODO: is there any way to release meminfo without calling dtor so we dont
-    //  call incref, decref one after another.
-    // We need to increase reference count because we are returning new
-    // reference to the same event.
-    Py_INCREF(orig_event);
+    if (event_obj == NULL) {
+        // Make create copy of event_ref so we don't need to manage nrt lifetime
+        // from python object.
+        event_obj = SyclEvent_Make(eventstruct->event_ref);
+    }
+    else {
+        // Unfortunately we can not optimize (nrt->release that triggers
+        // Py_DECREF() from the destructor) and Py_INCREF() because nrt may need
+        // the object even if we return it to python.
+        // We need to increase reference count because we are returning new
+        // reference to the same event.
+        Py_INCREF(event_obj);
+    }
+
     // We need to release meminfo since we are taking ownership back.
     nrt->release(eventstruct->meminfo);
 
-    return orig_event;
+    return event_obj;
 }
 
 /*----------------------------------------------------------------------------*/

numba_dpex/core/runtime/_queuestruct.h

@@ -12,9 +12,11 @@
 #pragma once
 
 #include "numba/core/runtime/nrt_external.h"
+#include <Python.h>
 
 typedef struct
 {
     NRT_MemInfo *meminfo;
+    PyObject *parent;
     void *queue_ref;
 } queuestruct_t;

numba_dpex/dpnp_iface/_intrinsic.py

@@ -6,11 +6,12 @@
 
 from dpctl import get_device_cached_queue
 from llvmlite import ir as llvmir
-from llvmlite.ir import Constant
+from llvmlite.ir import Constant, IRBuilder
 from llvmlite.ir.types import DoubleType, FloatType
 from numba import types
 from numba.core import cgutils
 from numba.core import config as numba_config
+from numba.core import errors, imputils
 from numba.core.typing import signature
 from numba.extending import intrinsic, overload_classmethod
 from numba.np.arrayobj import (
@@ -73,8 +74,7 @@ def make_queue(context, builder, py_dpctl_sycl_queue):
         pyapi, py_dpctl_sycl_queue_addr, queue_struct_voidptr
     )
 
-    queue_struct = builder.load(queue_struct_ptr)
-    queue_ref = builder.extract_value(queue_struct, 1)
+    queue_ref = queue_struct_proxy.queue_ref
 
     return_values = namedtuple(
         "return_values", "queue_ref queue_address_ptr pyapi"
@@ -1077,3 +1077,61 @@ def codegen(context, builder, sig, args):
         return ary._getvalue()
 
     return signature, codegen
+
+
+@intrinsic
+def ol_dpnp_nd_array_sycl_queue(
+    ty_context,
+    ty_dpnp_nd_array: DpnpNdArray,
+):
+    if not isinstance(ty_dpnp_nd_array, DpnpNdArray):
+        raise errors.TypingError("Argument must be DpnpNdArray")
+
+    ty_queue: DpctlSyclQueue = ty_dpnp_nd_array.queue
+
+    sig = ty_queue(ty_dpnp_nd_array)
+
+    def codegen(context, builder: IRBuilder, sig, args: list):
+        array_proxy = cgutils.create_struct_proxy(ty_dpnp_nd_array)(
+            context,
+            builder,
+            value=args[0],
+        )
+
+        queue_ref = array_proxy.sycl_queue
+
+        queue_struct_proxy = cgutils.create_struct_proxy(ty_queue)(
+            context, builder
+        )
+
+        queue_struct_proxy.queue_ref = queue_ref
+        queue_struct_proxy.meminfo = array_proxy.meminfo
+
+        # Warning: current implementation prevents whole object from being
+        # destroyed as long as sycl_queue attribute is being used. It should be
+        # okay since anywere we use it as an argument callee creates a copy
+        # so it does not steel reference.
+        #
+        # We can avoid it by:
+        #  queue_ref_copy = sycl.dpctl_queue_copy(builder, queue_ref) #noqa E800
+        #  queue_struct_proxy.queue_ref = queue_ref_copy #noqa E800
+        #  queue_struct->meminfo =
+        #     nrt->manage_memory(queue_ref_copy, DPCTLEvent_Delete);
+        # but it will allocate new meminfo object which can negatively affect
+        # performance.
+        # Speaking philosophically attribute is a part of the object and as long
+        # as nobody can still the reference it is a part of the owner object
+        # and lifetime is tied to it.
+        # TODO: we want to have queue: queuestruct_t instead of
+        #   queue_ref: QueueRef as an attribute for DPNPNdArray.
+
+        queue_value = queue_struct_proxy._getvalue()
+
+        # We need to incref meminfo so that queue model is preventing parent
+        # ndarray from being destroyed, that can destroy queue that we are
+        # using.
+        return imputils.impl_ret_borrowed(
+            context, builder, ty_queue, queue_value
+        )
+
+    return sig, codegen

numba_dpex/dpnp_iface/arrayobj.py

@@ -11,7 +11,7 @@
 from numba.core.types.containers import UniTuple
 from numba.core.typing.npydecl import parse_dtype as _ty_parse_dtype
 from numba.core.typing.npydecl import parse_shape as _ty_parse_shape
-from numba.extending import overload
+from numba.extending import overload, overload_attribute
 from numba.np.arrayobj import getitem_arraynd_intp as np_getitem_arraynd_intp
 from numba.np.numpy_support import is_nonelike
 
@@ -27,6 +27,7 @@
     impl_dpnp_ones_like,
     impl_dpnp_zeros,
     impl_dpnp_zeros_like,
+    ol_dpnp_nd_array_sycl_queue,
 )
 
 # =========================================================================
@@ -1085,3 +1086,23 @@ def getitem_arraynd_intp(context, builder, sig, args):
         ret = builder.insert_value(ret, sycl_queue_attr, sycl_queue_attr_pos)
 
     return ret
+
+
+@overload_attribute(DpnpNdArray, "sycl_queue")
+def dpnp_nd_array_sycl_queue(arr):
+    """Returns :class:`dpctl.SyclQueue` object associated with USM data.
+
+    This is an overloaded attribute implementation for dpnp.sycl_queue.
+
+    Args:
+        arr (numba_dpex.core.types.DpnpNdArray): Input array from which to
+            take sycl_queue.
+
+    Returns:
+        function: Local function `ol_dpnp_nd_array_sycl_queue()`.
+    """
+
+    def get(arr):
+        return ol_dpnp_nd_array_sycl_queue(arr)
+
+    return get

numba_dpex/tests/core/types/DpctlSyclQueue/test_box.py

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Tests for boxing for dpnp.ndarray
+"""
+
+import dpnp
+import pytest
+from dpctl import SyclQueue
+
+from numba_dpex import dpjit
+
+
+def test_boxing_without_parent():
+    """Test unboxing of the queue that does not have parent"""
+
+    @dpjit
+    def func() -> SyclQueue:
+        arr = dpnp.empty(10)
+        queue = arr.sycl_queue
+        return queue
+
+    q: SyclQueue = func()
+
+    assert len(q.sycl_device.filter_string) > 0

numba_dpex/tests/core/types/DpctlSyclQueue/test_queue_ref_attr.py

@@ -31,16 +31,21 @@ def are_queues_equal(typingctx, ty_queue1, ty_queue2):
 
     # defines the custom code generation
     def codegen(context, builder, sig, args):
+        q1 = cgutils.create_struct_proxy(ty_queue1)(
+            context, builder, value=args[0]
+        )
+        q2 = cgutils.create_struct_proxy(ty_queue2)(
+            context, builder, value=args[1]
+        )
+
         fnty = llvmir.FunctionType(
             cgutils.bool_t, [cgutils.voidptr_t, cgutils.voidptr_t]
         )
         fn = cgutils.get_or_insert_function(
             builder.module, fnty, "DPCTLQueue_AreEq"
         )
-        qref1 = builder.extract_value(args[0], 1)
-        qref2 = builder.extract_value(args[1], 1)
 
-        ret = builder.call(fn, [qref1, qref2])
+        ret = builder.call(fn, [q1.queue_ref, q2.queue_ref])
 
         return ret
 

numba_dpex/tests/dpjit_tests/dpnp/test_dpnp_empty.py

@@ -139,3 +139,33 @@ def func(shape, queue):
     with pytest.raises(errors.TypingError):
         queue = dpctl.SyclQueue()
         func(10, queue)
+
+
+@pytest.mark.xfail(reason="dpjit allocates new dpctl.SyclQueue on boxing")
+# TODO: remove test_dpnp_empty_with_dpjit_queue_temp once pass.
+def test_dpnp_empty_with_dpjit_queue():
+    """Test if dpnp array can be created with a queue from another array"""
+
+    @dpjit
+    def func(a):
+        queue = a.sycl_queue
+        return dpnp.empty(10, sycl_queue=queue)
+
+    a = dpnp.empty(10)
+    b = func(a)
+
+    assert id(a.sycl_queue) == id(b.sycl_queue)
+
+
+def test_dpnp_empty_with_dpjit_queue_temp():
+    """Test if dpnp array can be created with a queue from another array"""
+
+    @dpjit
+    def func(a):
+        queue = a.sycl_queue
+        return dpnp.empty(10, sycl_queue=queue)
+
+    a = dpnp.empty(10)
+    b = func(a)
+
+    assert a.sycl_queue == b.sycl_queue

numba_dpex/tests/dpjit_tests/parfors/test_dpnp_logic_ops.py

@@ -78,7 +78,6 @@ def f(a, b):
     )
 
 
-@pytest.mark.xfail
 def test_unary_ops(unary_op, input_arrays):
     a = input_arrays[0]
     uop = getattr(dpnp, unary_op)