Merge pull request #1454 from IntelPython/correct-refcount-handling-fixing-memory-leak

Fixed ref-counting of Python object temporaries in unboxing code

Author: Diptorup Deb

numba_dpex/core/runtime/_dpexrt_python.c

@@ -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 =
@@ -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",
@@ -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 "
@@ -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;
@@ -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);
 
@@ -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.