Skip to content

Commit 5dcf8af

Browse files
committed
Use obj as usm array parent and clean up arrayobj individually
1 parent 96e95f6 commit 5dcf8af

File tree

1 file changed

+11
-3
lines changed

1 file changed

+11
-3
lines changed

numba_dpex/core/runtime/_dpexrt_python.c

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -842,7 +842,7 @@ static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj,
842842
}
843843

844844
if (!(arystruct->meminfo = NRT_MemInfo_new_from_usmndarray(
845-
(PyObject *)arrayobj, data, nitems, itemsize, qref)))
845+
obj, data, nitems, itemsize, qref)))
846846
{
847847
DPEXRT_DEBUG(drt_debug_print(
848848
"DPEXRT-ERROR: NRT_MemInfo_new_from_usmndarray failed "
@@ -851,11 +851,14 @@ static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj,
851851
goto error;
852852
}
853853

854+
Py_XDECREF(arrayobj);
855+
Py_IncRef(obj);
856+
854857
arystruct->data = data;
855858
arystruct->sycl_queue = qref;
856859
arystruct->nitems = nitems;
857860
arystruct->itemsize = itemsize;
858-
arystruct->parent = (PyObject *)arrayobj;
861+
arystruct->parent = obj;
859862

860863
p = arystruct->shape_and_strides;
861864

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

941944
if (!(arrayobj = PyUSMNdArray_ARRAYOBJ(arystruct->parent))) {
945+
Py_XDECREF(arrayobj);
942946
DPEXRT_DEBUG(
943947
drt_debug_print("DPEXRT-DEBUG: Arrayobj cannot be boxed from "
944948
"parent as parent pointer is NULL.\n"));
945949
return NULL;
946950
}
947951

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

956-
if (UsmNDArray_GetNDim(arrayobj) != ndim)
961+
if (UsmNDArray_GetNDim(arrayobj) != ndim) {
962+
Py_XDECREF(arrayobj);
957963
return NULL;
964+
}
958965

959966
p = arystruct->shape_and_strides;
960967
shape = UsmNDArray_GetShape(arrayobj);
961968
strides = UsmNDArray_GetStrides(arrayobj);
969+
Py_XDECREF(arrayobj);
962970

963971
// Ensure the shape of the array to be boxed matches the shape of the
964972
// original parent.

0 commit comments

Comments
 (0)