Merge pull request #1219 from IntelPython/feature/async_kernel_submition

ZzEeKkAa · web-flow · commit 04c18bf6886b · 2023-11-30T09:29:17.000-05:00
Feature/async kernel submition
diff --git a/numba_dpex/core/runtime/CMakeLists.txt b/numba_dpex/core/runtime/CMakeLists.txt
@@ -109,7 +109,7 @@ python_add_library(${PROJECT_NAME} MODULE ${SOURCES})
 
 # Add SYCL to target, this must come after python_add_library()
 # FIXME: sources incompatible with sycl include?
-# add_sycl_to_target(TARGET ${PROJECT_NAME} SOURCES ${KERNEL_SOURCES})
+add_sycl_to_target(TARGET ${PROJECT_NAME} SOURCES ${KERNEL_SOURCES})
 
 # Link the DPCTLSyclInterface library to target
 target_link_libraries(${PROJECT_NAME} PRIVATE DPCTLSyclInterface)
diff --git a/numba_dpex/core/runtime/_dpexrt_python.c b/numba_dpex/core/runtime/_dpexrt_python.c
@@ -24,6 +24,7 @@
 #include "_queuestruct.h"
 #include "_usmarraystruct.h"
 
+#include "experimental/nrt_reserve_meminfo.h"
 #include "numba/core/runtime/nrt_external.h"
 
 // forward declarations
@@ -1490,6 +1491,8 @@ static PyObject *build_c_helpers_dict(void)
                  &DPEXRT_sycl_event_from_python);
     _declpointer("DPEXRT_sycl_event_to_python", &DPEXRT_sycl_event_to_python);
     _declpointer("DPEXRT_sycl_event_init", &DPEXRT_sycl_event_init);
+    _declpointer("DPEXRT_nrt_acquire_meminfo_and_schedule_release",
+                 &DPEXRT_nrt_acquire_meminfo_and_schedule_release);
 
 #undef _declpointer
     return dct;
@@ -1557,6 +1560,9 @@ MOD_INIT(_dpexrt_python)
                        PyLong_FromVoidPtr(&DPEXRT_MemInfo_alloc));
     PyModule_AddObject(m, "DPEXRT_MemInfo_fill",
                        PyLong_FromVoidPtr(&DPEXRT_MemInfo_fill));
+    PyModule_AddObject(
+        m, "DPEXRT_nrt_acquire_meminfo_and_schedule_release",
+        PyLong_FromVoidPtr(&DPEXRT_nrt_acquire_meminfo_and_schedule_release));
     PyModule_AddObject(m, "c_helpers", build_c_helpers_dict());
     return MOD_SUCCESS_VAL(m);
 }
diff --git a/numba_dpex/core/runtime/context.py b/numba_dpex/core/runtime/context.py
@@ -433,3 +433,41 @@ def submit_ndrange(
         )
 
         return ret
+
+    def acquire_meminfo_and_schedule_release(
+        self, builder: llvmir.IRBuilder, args
+    ):
+        """Inserts LLVM IR to call nrt_acquire_meminfo_and_schedule_release.
+
+        DPCTLSyclEventRef
+        DPEXRT_nrt_acquire_meminfo_and_schedule_release(
+            NRT_api_functions *nrt,
+            DPCTLSyclQueueRef QRef,
+            NRT_MemInfo **meminfo_array,
+            size_t meminfo_array_size,
+            DPCTLSyclEventRef *depERefs,
+            size_t nDepERefs,
+            int *status,
+        );
+
+        """
+        mod = builder.module
+
+        func_ty = llvmir.FunctionType(
+            cgutils.voidptr_t,
+            [
+                cgutils.voidptr_t,
+                cgutils.voidptr_t,
+                cgutils.voidptr_t.as_pointer(),
+                llvmir.IntType(64),
+                cgutils.voidptr_t.as_pointer(),
+                llvmir.IntType(64),
+                llvmir.IntType(64).as_pointer(),
+            ],
+        )
+        fn = cgutils.get_or_insert_function(
+            mod, func_ty, "DPEXRT_nrt_acquire_meminfo_and_schedule_release"
+        )
+        ret = builder.call(fn, args)
+
+        return ret
diff --git a/numba_dpex/core/runtime/experimental/nrt_reserve_meminfo.cpp b/numba_dpex/core/runtime/experimental/nrt_reserve_meminfo.cpp
@@ -0,0 +1,71 @@
+// SPDX-FileCopyrightText: 2023 Intel Corporation
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "nrt_reserve_meminfo.h"
+
+#include "_dbg_printer.h"
+#include "syclinterface/dpctl_sycl_type_casters.hpp"
+#include <CL/sycl.hpp>
+
+extern "C"
+{
+    DPCTLSyclEventRef
+    DPEXRT_nrt_acquire_meminfo_and_schedule_release(NRT_api_functions *nrt,
+                                                    DPCTLSyclQueueRef QRef,
+                                                    NRT_MemInfo **meminfo_array,
+                                                    size_t meminfo_array_size,
+                                                    DPCTLSyclEventRef *depERefs,
+                                                    size_t nDepERefs,
+                                                    int *status)
+    {
+        DPEXRT_DEBUG(drt_debug_print(
+                         "DPEXRT-DEBUG: scheduling nrt meminfo release.\n"););
+
+        using dpctl::syclinterface::unwrap;
+        using dpctl::syclinterface::wrap;
+
+        sycl::queue *q = unwrap<sycl::queue>(QRef);
+
+        std::vector<NRT_MemInfo *> meminfo_vec(
+            meminfo_array, meminfo_array + meminfo_array_size);
+
+        for (size_t i = 0; i < meminfo_array_size; ++i) {
+            nrt->acquire(meminfo_vec[i]);
+        }
+
+        DPEXRT_DEBUG(drt_debug_print("DPEXRT-DEBUG: acquired meminfo.\n"););
+
+        try {
+            sycl::event ht_ev = q->submit([&](sycl::handler &cgh) {
+                for (size_t ev_id = 0; ev_id < nDepERefs; ++ev_id) {
+                    cgh.depends_on(*(unwrap<sycl::event>(depERefs[ev_id])));
+                }
+                cgh.host_task([meminfo_array_size, meminfo_vec, nrt]() {
+                    for (size_t i = 0; i < meminfo_array_size; ++i) {
+                        nrt->release(meminfo_vec[i]);
+                        DPEXRT_DEBUG(
+                            drt_debug_print("DPEXRT-DEBUG: released meminfo "
+                                            "from host_task.\n"););
+                    }
+                });
+            });
+
+            constexpr int result_ok = 0;
+
+            *status = result_ok;
+            auto e_ptr = new sycl::event(ht_ev);
+            return wrap<sycl::event>(e_ptr);
+        } catch (const std::exception &e) {
+            constexpr int result_std_exception = 1;
+
+            *status = result_std_exception;
+            return nullptr;
+        }
+
+        constexpr int result_other_abnormal = 2;
+
+        *status = result_other_abnormal;
+        return nullptr;
+    }
+}
diff --git a/numba_dpex/core/runtime/experimental/nrt_reserve_meminfo.h b/numba_dpex/core/runtime/experimental/nrt_reserve_meminfo.h
@@ -0,0 +1,48 @@
+// SPDX-FileCopyrightText: 2023 Intel Corporation
+//
+// SPDX-License-Identifier: Apache-2.0
+
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Defines dpctl style function(s) that interruct with nrt meminfo and sycl.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef _EXPERIMENTAL_H_
+#define _EXPERIMENTAL_H_
+
+#include "dpctl_capi.h"
+#include "numba/core/runtime/nrt_external.h"
+
+#ifdef __cplusplus
+extern "C"
+{
+#endif
+
+    /*!
+     * @brief Acquires meminfos and schedules a host task to release them.
+     *
+     * @param    nrt            NRT public API functions,
+     * @param    QRef           Queue reference,
+     * @param    meminfo_array  Array of meminfo pointers to perform actions on,
+     * @param    meminfo_array_size Length of meminfo_array,
+     * @param    depERefs       Array of dependant events for the host task,
+     * @param    nDepERefs      Length of depERefs,
+     * @param    status         Variable to write status to. Same style as
+     * dpctl,
+     * @return   {return}       Event reference to the host task.
+     */
+    DPCTLSyclEventRef
+    DPEXRT_nrt_acquire_meminfo_and_schedule_release(NRT_api_functions *nrt,
+                                                    DPCTLSyclQueueRef QRef,
+                                                    NRT_MemInfo **meminfo_array,
+                                                    size_t meminfo_array_size,
+                                                    DPCTLSyclEventRef *depERefs,
+                                                    size_t nDepERefs,
+                                                    int *status);
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* _EXPERIMENTAL_H_ */
diff --git a/numba_dpex/dpctl_iface/_intrinsic.py b/numba_dpex/dpctl_iface/_intrinsic.py
@@ -5,13 +5,12 @@
 import dpctl
 from llvmlite.ir import IRBuilder
 from numba import types
-from numba.core import cgutils, imputils
 from numba.core.datamodel import default_manager
-from numba.extending import intrinsic, overload, overload_method, type_callable
+from numba.extending import intrinsic, overload, overload_method
 
 import numba_dpex.dpctl_iface.libsyclinterface_bindings as sycl
 from numba_dpex.core import types as dpex_types
-from numba_dpex.core.runtime import context as dpexrt
+from numba_dpex.dpctl_iface.wrappers import wrap_event_reference
 
 
 @intrinsic
@@ -33,23 +32,8 @@ def sycl_event_create(
     sig = ty_event(types.void)
 
     def codegen(context, builder: IRBuilder, sig, args: list):
-        pyapi = context.get_python_api(builder)
-
-        event_struct_proxy = cgutils.create_struct_proxy(ty_event)(
-            context, builder
-        )
-
         event = sycl.dpctl_event_create(builder)
-        dpexrtCtx = dpexrt.DpexRTContext(context)
-
-        # Ref count after the call is equal to 1.
-        dpexrtCtx.eventstruct_init(
-            pyapi, event, event_struct_proxy._getpointer()
-        )
-
-        event_value = event_struct_proxy._getvalue()
-
-        return event_value
+        return wrap_event_reference(context, builder, event)
 
     return sig, codegen
 
diff --git a/numba_dpex/dpctl_iface/wrappers.py b/numba_dpex/dpctl_iface/wrappers.py
@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from numba.core import cgutils
+
+from numba_dpex.core.runtime import context as dpexrt
+from numba_dpex.core.types import DpctlSyclEvent
+
+
+def wrap_event_reference(ctx, builder, eref):
+    """Wrap dpctl event reference into datamodel so it can be boxed to
+    Python."""
+
+    ty_event = DpctlSyclEvent()
+
+    pyapi = ctx.get_python_api(builder)
+
+    event_struct_proxy = cgutils.create_struct_proxy(ty_event)(ctx, builder)
+
+    # Ref count after the call is equal to 1.
+    # TODO: get dpex RT from cached property once the PR is merged
+    # https://github.com/IntelPython/numba-dpex/pull/1027
+    # ctx.dpexrt.eventstruct_init( # noqa: W0621
+    dpexrt.DpexRTContext(ctx).eventstruct_init(
+        pyapi,
+        eref,
+        # calling _<method>() is by numba's design
+        event_struct_proxy._getpointer(),  # pylint: disable=W0212
+    )
+
+    # calling _<method>() is by numba's design
+    event_value = event_struct_proxy._getvalue()  # pylint: disable=W0212
+
+    return event_value
diff --git a/numba_dpex/experimental/__init__.py b/numba_dpex/experimental/__init__.py
@@ -10,7 +10,7 @@
 
 from .decorators import kernel
 from .kernel_dispatcher import KernelDispatcher
-from .launcher import call_kernel
+from .launcher import call_kernel, call_kernel_async
 from .models import *
 from .types import KernelDispatcherType
 
@@ -26,4 +26,4 @@ def dpex_dispatcher_const(context):
     return context.get_dummy_value()
 
 
-__all__ = ["kernel", "KernelDispatcher", "call_kernel"]
+__all__ = ["kernel", "KernelDispatcher", "call_kernel", "call_kernel_async"]
diff --git a/numba_dpex/experimental/kernel_dispatcher.py b/numba_dpex/experimental/kernel_dispatcher.py
@@ -254,7 +254,7 @@ def get_overload_device_ir(self, sig):
         args, _ = sigutils.normalize_signature(sig)
         return self.overloads[tuple(args)].kernel_device_ir_module
 
-    def compile(self, sig) -> _KernelCompileResult:
+    def compile(self, sig) -> any:
         disp = self._get_dispatcher_for_current_target()
         if disp is not self:
             return disp.compile(sig)
diff --git a/numba_dpex/experimental/launcher.py b/numba_dpex/experimental/launcher.py
diff --git a/numba_dpex/tests/experimental/test_async_kernel.py b/numba_dpex/tests/experimental/test_async_kernel.py
