Skip to content

Commit

Permalink
Use obj as usm array parent and clean up arrayobj individually
Browse files Browse the repository at this point in the history
  • Loading branch information
ZzEeKkAa committed May 9, 2024
1 parent 96e95f6 commit 5dcf8af
Showing 1 changed file with 11 additions and 3 deletions.
14 changes: 11 additions & 3 deletions numba_dpex/core/runtime/_dpexrt_python.c
Original file line number Diff line number Diff line change
Expand Up @@ -842,7 +842,7 @@ static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj,
}

if (!(arystruct->meminfo = NRT_MemInfo_new_from_usmndarray(
(PyObject *)arrayobj, data, nitems, itemsize, qref)))
obj, data, nitems, itemsize, qref)))
{
DPEXRT_DEBUG(drt_debug_print(
"DPEXRT-ERROR: NRT_MemInfo_new_from_usmndarray failed "
Expand All @@ -851,11 +851,14 @@ static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj,
goto error;
}

Py_XDECREF(arrayobj);
Py_IncRef(obj);

arystruct->data = data;
arystruct->sycl_queue = qref;
arystruct->nitems = nitems;
arystruct->itemsize = itemsize;
arystruct->parent = (PyObject *)arrayobj;
arystruct->parent = obj;

p = arystruct->shape_and_strides;

Expand Down Expand Up @@ -939,26 +942,31 @@ static PyObject *box_from_arystruct_parent(usmarystruct_t *arystruct,
drt_debug_print("DPEXRT-DEBUG: In box_from_arystruct_parent.\n"));

if (!(arrayobj = PyUSMNdArray_ARRAYOBJ(arystruct->parent))) {
Py_XDECREF(arrayobj);
DPEXRT_DEBUG(
drt_debug_print("DPEXRT-DEBUG: Arrayobj cannot be boxed from "
"parent as parent pointer is NULL.\n"));
return NULL;
}

if ((void *)UsmNDArray_GetData(arrayobj) != arystruct->data) {
Py_XDECREF(arrayobj);
DPEXRT_DEBUG(drt_debug_print(
"DPEXRT-DEBUG: Arrayobj cannot be boxed "
"from parent as data pointer in the arystruct is not the same as "
"the data pointer in the parent object.\n"));
return NULL;
}

if (UsmNDArray_GetNDim(arrayobj) != ndim)
if (UsmNDArray_GetNDim(arrayobj) != ndim) {
Py_XDECREF(arrayobj);
return NULL;
}

p = arystruct->shape_and_strides;
shape = UsmNDArray_GetShape(arrayobj);
strides = UsmNDArray_GetStrides(arrayobj);
Py_XDECREF(arrayobj);

// Ensure the shape of the array to be boxed matches the shape of the
// original parent.
Expand Down

0 comments on commit 5dcf8af

Please sign in to comment.