Skip to content

Commit

Permalink
Merge pull request #1454 from IntelPython/correct-refcount-handling-f…
Browse files Browse the repository at this point in the history
…ixing-memory-leak

Fixed ref-counting of Python object temporaries in unboxing code
  • Loading branch information
Diptorup Deb authored May 10, 2024
2 parents e36c979 + 5dcf8af commit 574daab
Showing 1 changed file with 17 additions and 8 deletions.
25 changes: 17 additions & 8 deletions numba_dpex/core/runtime/_dpexrt_python.c
Original file line number Diff line number Diff line change
Expand Up @@ -744,14 +744,18 @@ static struct PyUSMArrayObject *PyUSMNdArray_ARRAYOBJ(PyObject *obj)
DPEXRT_DEBUG(
drt_debug_print("DPEXRT-DEBUG: usm array was passed directly\n"));
arrayobj = obj;
Py_INCREF(arrayobj);
}
else if (PyObject_HasAttrString(obj, "_array_obj")) {
// PyObject_GetAttrString gives reference
arrayobj = PyObject_GetAttrString(obj, "_array_obj");

if (!arrayobj)
return NULL;
if (!PyObject_TypeCheck(arrayobj, &PyUSMArrayType))
if (!PyObject_TypeCheck(arrayobj, &PyUSMArrayType)) {
Py_DECREF(arrayobj);
return NULL;
}
}

struct PyUSMArrayObject *pyusmarrayobj =
Expand Down Expand Up @@ -803,17 +807,13 @@ static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj,
PyGILState_STATE gstate;
npy_intp itemsize = 0;

// Increment the ref count on obj to prevent CPython from garbage
// collecting the array.
// TODO: add extra description why do we need this
Py_IncRef(obj);

DPEXRT_DEBUG(drt_debug_print(
"DPEXRT-DEBUG: In DPEXRT_sycl_usm_ndarray_from_python at %s, line %d\n",
__FILE__, __LINE__));

// Check if the PyObject obj has an _array_obj attribute that is of
// dpctl.tensor.usm_ndarray type.
// arrayobj is a new reference, reference of obj is borrowed
if (!(arrayobj = PyUSMNdArray_ARRAYOBJ(obj))) {
DPEXRT_DEBUG(drt_debug_print(
"DPEXRT-ERROR: PyUSMNdArray_ARRAYOBJ check failed at %s, line %d\n",
Expand All @@ -832,6 +832,7 @@ static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj,
data = (void *)UsmNDArray_GetData(arrayobj);
nitems = product_of_shape(shape, ndim);
itemsize = (npy_intp)UsmNDArray_GetElementSize(arrayobj);

if (!(qref = UsmNDArray_GetQueueRef(arrayobj))) {
DPEXRT_DEBUG(drt_debug_print(
"DPEXRT-ERROR: UsmNDArray_GetQueueRef returned NULL at "
Expand All @@ -850,6 +851,9 @@ 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;
Expand Down Expand Up @@ -906,7 +910,7 @@ static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj,
__FILE__, __LINE__));
gstate = PyGILState_Ensure();
// decref the python object
Py_DECREF(obj);
Py_XDECREF((PyObject *)arrayobj);
// release the GIL
PyGILState_Release(gstate);

Expand Down Expand Up @@ -938,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 574daab

Please sign in to comment.