From 0364e76397328e09386c725d20104f338e0aeb0f Mon Sep 17 00:00:00 2001 From: Yevhenii Havrylko Date: Tue, 14 May 2024 17:43:59 -0400 Subject: [PATCH] Use nrt api to allocate meminfo object --- numba_dpex/core/runtime/_dpexrt_python.c | 32 ++++++++++++++++-------- numba_dpex/core/runtime/context.py | 12 ++++++--- 2 files changed, 30 insertions(+), 14 deletions(-) diff --git a/numba_dpex/core/runtime/_dpexrt_python.c b/numba_dpex/core/runtime/_dpexrt_python.c index 80f66e4fd2..5647a09241 100644 --- a/numba_dpex/core/runtime/_dpexrt_python.c +++ b/numba_dpex/core/runtime/_dpexrt_python.c @@ -45,12 +45,14 @@ static NRT_MemInfo *DPEXRT_MemInfo_fill(NRT_MemInfo *mi, bool value_is_float, int64_t value, const DPCTLSyclQueueRef qref); -static NRT_MemInfo *NRT_MemInfo_new_from_usmndarray(PyObject *ndarrobj, +static NRT_MemInfo *NRT_MemInfo_new_from_usmndarray(NRT_api_functions *nrt, + PyObject *ndarrobj, void *data, npy_intp nitems, npy_intp itemsize, DPCTLSyclQueueRef qref); -static NRT_MemInfo *DPEXRT_MemInfo_alloc(npy_intp size, +static NRT_MemInfo *DPEXRT_MemInfo_alloc(NRT_api_functions *nrt, + npy_intp size, size_t usm_type, const DPCTLSyclQueueRef qref); static void usmndarray_meminfo_dtor(void *ptr, size_t size, void *info); @@ -58,7 +60,8 @@ static PyObject *box_from_arystruct_parent(usmarystruct_t *arystruct, int ndim, PyArray_Descr *descr); -static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj, +static int DPEXRT_sycl_usm_ndarray_from_python(NRT_api_functions *nrt, + PyObject *obj, usmarystruct_t *arystruct); static PyObject * DPEXRT_sycl_usm_ndarray_to_python_acqref(usmarystruct_t *arystruct, @@ -336,6 +339,11 @@ NRT_ExternalAllocator_new_for_usm(DPCTLSyclQueueRef qref, size_t usm_type) static void usmndarray_meminfo_dtor(void *ptr, size_t size, void *info) { MemInfoDtorInfo *mi_dtor_info = NULL; + // Warning: we are destructing sycl memory. MI destructor is called + // separately by numba. + DPEXRT_DEBUG(drt_debug_print("DPEXRT-DEBUG: Call to " + "usmndarray_meminfo_dtor at %s, line %d\n", + __FILE__, __LINE__)); // Sanity-check to make sure the mi_dtor_info is an actual pointer. if (!(mi_dtor_info = (MemInfoDtorInfo *)info)) { @@ -416,7 +424,8 @@ static MemInfoDtorInfo *MemInfoDtorInfo_new(NRT_MemInfo *mi, PyObject *owner) * of the dpnp.ndarray was allocated. * @return {return} A new NRT_MemInfo object */ -static NRT_MemInfo *NRT_MemInfo_new_from_usmndarray(PyObject *ndarrobj, +static NRT_MemInfo *NRT_MemInfo_new_from_usmndarray(NRT_api_functions *nrt, + PyObject *ndarrobj, void *data, npy_intp nitems, npy_intp itemsize, @@ -427,8 +436,9 @@ static NRT_MemInfo *NRT_MemInfo_new_from_usmndarray(PyObject *ndarrobj, MemInfoDtorInfo *midtor_info = NULL; DPCTLSyclContextRef cref = NULL; - // Allocate a new NRT_MemInfo object - if (!(mi = (NRT_MemInfo *)malloc(sizeof(NRT_MemInfo)))) { + // Allocate a new NRT_MemInfo object. By passing 0 we are just allocating + // MemInfo and not the `data` that the MemInfo object manages. + if (!(mi = (NRT_MemInfo *)nrt->allocate(0))) { DPEXRT_DEBUG(drt_debug_print( "DPEXRT-ERROR: Could not allocate a new NRT_MemInfo " "object at %s, line %d\n", @@ -505,7 +515,8 @@ 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, +static NRT_MemInfo *DPEXRT_MemInfo_alloc(NRT_api_functions *nrt, + npy_intp size, size_t usm_type, const DPCTLSyclQueueRef qref) { @@ -517,7 +528,7 @@ static NRT_MemInfo *DPEXRT_MemInfo_alloc(npy_intp size, "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)))) { + if (!(mi = (NRT_MemInfo *)nrt->allocate(0))) { DPEXRT_DEBUG(drt_debug_print( "DPEXRT-ERROR: Could not allocate a new NRT_MemInfo object.\n")); goto error; @@ -795,7 +806,8 @@ static npy_intp product_of_shape(npy_intp *shape, npy_intp ndim) * instance of a dpnp.ndarray * @return {return} Error code representing success (0) or failure (-1). */ -static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj, +static int DPEXRT_sycl_usm_ndarray_from_python(NRT_api_functions *nrt, + PyObject *obj, usmarystruct_t *arystruct) { struct PyUSMArrayObject *arrayobj = NULL; @@ -842,7 +854,7 @@ static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj, } if (!(arystruct->meminfo = NRT_MemInfo_new_from_usmndarray( - obj, data, nitems, itemsize, qref))) + nrt, obj, data, nitems, itemsize, qref))) { DPEXRT_DEBUG(drt_debug_print( "DPEXRT-ERROR: NRT_MemInfo_new_from_usmndarray failed " diff --git a/numba_dpex/core/runtime/context.py b/numba_dpex/core/runtime/context.py index 80f0253101..4fe89abd7a 100644 --- a/numba_dpex/core/runtime/context.py +++ b/numba_dpex/core/runtime/context.py @@ -89,12 +89,14 @@ def meminfo_alloc_unchecked(self, builder, size, usm_type, queue_ref): mod = builder.module u64 = llvmir.IntType(64) fnty = llvmir.FunctionType( - cgutils.voidptr_t, [cgutils.intp_t, u64, cgutils.voidptr_t] + cgutils.voidptr_t, + [cgutils.voidptr_t, cgutils.intp_t, u64, cgutils.voidptr_t], ) fn = cgutils.get_or_insert_function(mod, fnty, "DPEXRT_MemInfo_alloc") fn.return_value.add_attribute("noalias") + nrt_api = self._context.nrt.get_nrt_api(builder) - ret = builder.call(fn, [size, usm_type, queue_ref]) + ret = builder.call(fn, [nrt_api, size, usm_type, queue_ref]) return ret @@ -168,13 +170,15 @@ def arraystruct_from_python(self, pyapi, obj, ptr): """ fnty = llvmir.FunctionType( - llvmir.IntType(32), [pyapi.pyobj, pyapi.voidptr] + llvmir.IntType(32), [pyapi.voidptr, pyapi.pyobj, pyapi.voidptr] ) + nrt_api = self._context.nrt.get_nrt_api(pyapi.builder) fn = pyapi._get_function(fnty, "DPEXRT_sycl_usm_ndarray_from_python") fn.args[0].add_attribute("nocapture") fn.args[1].add_attribute("nocapture") + fn.args[2].add_attribute("nocapture") - self.error = pyapi.builder.call(fn, (obj, ptr)) + self.error = pyapi.builder.call(fn, (nrt_api, obj, ptr)) return self.error