MAINT: Reorganize array dealloc and nditer buffer cleanup

numpy/core/src/multiarray/arrayobject.c

@@ -56,6 +56,7 @@ maintainer email:  oliphant.travis@ieee.org
 #include "alloc.h"
 #include "mem_overlap.h"
 #include "numpyos.h"
+#include "refcount.h"
 #include "strfuncs.h"
 
 #include "binop_override.h"
@@ -452,10 +453,6 @@ array_dealloc(PyArrayObject *self)
     }
 
     if ((fa->flags & NPY_ARRAY_OWNDATA) && fa->data) {
-        /* Free internal references if an Object array */
-        if (PyDataType_FLAGCHK(fa->descr, NPY_ITEM_REFCOUNT)) {
-            PyArray_XDECREF(self);
-        }
         if (fa->mem_handler == NULL) {
             char *env = getenv("NUMPY_WARN_IF_NO_MEM_POLICY");
             if ((env != NULL) && (strncmp(env, "1", 1) == 0)) {
@@ -464,10 +461,25 @@ array_dealloc(PyArrayObject *self)
                     "set a base owning the data (e.g. a PyCapsule).";
                 WARN_IN_DEALLOC(PyExc_RuntimeWarning, msg);
             }
+            /* Use old style decref, just in case of strange things */
+            if (PyDataType_FLAGCHK(fa->descr, NPY_ITEM_REFCOUNT)) {
+                PyArray_XDECREF(self);
+            }
             // Guess at malloc/free ???
             free(fa->data);
         }
         else {
+            /* Free internal references */
+            if (PyDataType_FLAGCHK(fa->descr, NPY_ITEM_REFCOUNT)) {
+                npy_intp itemsize = PyArray_ITEMSIZE(self);
+                npy_intp size = PyArray_SIZE(self);
+                int aligned = fa->flags & NPY_ARRAY_ALIGNED;
+                if (PyArray_ClearData(
+                        fa->descr, fa->data, itemsize, size, aligned) < 0) {
+                    PyErr_WriteUnraisable(NULL);
+                }
+            }
+            /* And actual data */
             size_t nbytes = PyArray_NBYTES(self);
             if (nbytes == 0) {
                 nbytes = 1;

numpy/core/src/multiarray/dtype_transfer.c

@@ -2690,7 +2690,7 @@ npy_clear_object_strided_loop(
     while (N > 0) {
         /* Release the reference in src and set it to NULL */
         NPY_DT_DBG_REFTRACE("dec src ref (null dst)", src_ref);
-        memcpy(&src_ref, src, sizeof(src_ref));
+        memcpy(&src_ref, src, sizeof(PyObject *));
         Py_XDECREF(src_ref);
         memset(src, 0, sizeof(PyObject *));
 

numpy/core/src/multiarray/item_selection.c

@@ -2642,7 +2642,7 @@ PyArray_Nonzero(PyArrayObject *self)
                 }
             }
             /*
-             * Fallback to a branchless strategy to avoid branch misprediction 
+             * Fallback to a branchless strategy to avoid branch misprediction
              * stalls that are very expensive on most modern processors.
              */
             else {

numpy/core/src/multiarray/nditer_api.c

@@ -17,6 +17,7 @@
 #include "nditer_impl.h"
 #include "templ_common.h"
 #include "ctors.h"
+#include "refcount.h"
 
 /* Internal helper functions private to this file */
 static npy_intp
@@ -2626,21 +2627,11 @@ npyiter_clear_buffers(NpyIter *iter)
         return;
     }
 
-    if (!(NIT_ITFLAGS(iter) & NPY_ITFLAG_NEEDSAPI)) {
-        /* Buffers do not require clearing, but should not be copied back */
-        NBF_SIZE(bufferdata) = 0;
-        return;
-    }
-
     /*
-     * The iterator may be using a dtype with references, which always
-     * requires the API. In that case, further cleanup may be necessary.
-     *
-     * TODO: At this time, we assume that a dtype having references
-     *       implies the need to hold the GIL at all times. In theory
-     *       we could broaden this definition for a new
-     *       `PyArray_Item_XDECREF` API and the assumption may become
-     *       incorrect.
+     * The iterator may be using a dtype with references (as of writing this
+     * means Python objects, but does not need to stay that way).
+     * In that case, further cleanup may be necessary and we clear buffers
+     * explicitly.
      */
     PyObject *type, *value, *traceback;
     PyErr_Fetch(&type,  &value, &traceback);
@@ -2650,13 +2641,6 @@ npyiter_clear_buffers(NpyIter *iter)
     PyArray_Descr **dtypes = NIT_DTYPES(iter);
     npyiter_opitflags *op_itflags = NIT_OPITFLAGS(iter);
     for (int iop = 0; iop < nop; ++iop, ++buffers) {
-        /*
-         * We may want to find a better way to do this, on the other hand,
-         * this cleanup seems rare and fairly special.  A dtype using
-         * references (right now only us) must always keep the buffer in
-         * a well defined state (either NULL or owning the reference).
-         * Only we implement cleanup
-         */
         if (!PyDataType_REFCHK(dtypes[iop]) ||
                 !(op_itflags[iop]&NPY_OP_ITFLAG_USINGBUFFER)) {
             continue;
@@ -2665,15 +2649,11 @@ npyiter_clear_buffers(NpyIter *iter)
             continue;
         }
         int itemsize = dtypes[iop]->elsize;
-        for (npy_intp i = 0; i < NBF_SIZE(bufferdata); i++) {
-            /*
-             * See above comment, if this API is expanded the GIL assumption
-             * could become incorrect.
-             */
-            PyArray_Item_XDECREF(*buffers + (itemsize * i), dtypes[iop]);
+        if (PyArray_ClearData(
+                dtypes[iop], *buffers, itemsize, NBF_SIZE(bufferdata), 1) < 0) {
+            /* This should never fail; if it does write it out */
+            PyErr_WriteUnraisable(NULL);
         }
-        /* Clear out the buffer just to be sure */
-        memset(*buffers, 0, NBF_SIZE(bufferdata) * itemsize);
     }
     /* Signal that the buffers are empty */
     NBF_SIZE(bufferdata) = 0;

numpy/core/src/multiarray/refcount.c

@@ -2,6 +2,10 @@
  * This module corresponds to the `Special functions for NPY_OBJECT`
  * section in the numpy reference for C-API.
  */
+#include "array_method.h"
+#include "dtype_transfer.h"
+#include "lowlevel_strided_loops.h"
+#include "pyerrors.h"
 #define NPY_NO_DEPRECATED_API NPY_API_VERSION
 #define _MULTIARRAYMODULE
 
@@ -21,6 +25,38 @@ static void
 _fillobject(char *optr, PyObject *obj, PyArray_Descr *dtype);
 
 
+/*
+ * Helper function to clear a strided memory (normally or always contiguous)
+ * from all Python (or other) references.  The function does nothing if the
+ * array dtype does not indicate holding references.
+ *
+ * It is safe to call this function more than once, failing here is usually
+ * critical (during cleanup) and should be set up to minimize the risk or
+ * avoid it fully.
+ */
+NPY_NO_EXPORT int
+PyArray_ClearData(
+        PyArray_Descr *descr, char *data,
+        npy_intp stride, npy_intp size, int aligned)
+{
+    if (!PyDataType_REFCHK(descr)) {
+        return 0;
+    }
+
+    NPY_cast_info cast_info;
+    NPY_ARRAYMETHOD_FLAGS flags;
+    if (PyArray_GetDTypeTransferFunction(
+            aligned, stride, 0, descr, NULL, 1, &cast_info, &flags) < 0) {
+        return -1;
+    }
+
+    int res = cast_info.func(
+            &cast_info.context, &data, &size, &stride, cast_info.auxdata);
+    NPY_cast_info_xfree(&cast_info);
+    return res;
+}
+
+
 /*NUMPY_API
  * XINCREF all objects in a single array item. This is complicated for
  * structured datatypes where the position of objects needs to be extracted.
@@ -204,6 +240,10 @@ PyArray_INCREF(PyArrayObject *mp)
 /*NUMPY_API
   Decrement all internal references for object arrays.
   (or arrays with object fields)
+
+  The use of this function is strongly discouraged, within NumPy
+  use PyArray_Clear, which DECREF's and sets everything to NULL and can
+  work with any dtype.
 */
 NPY_NO_EXPORT int
 PyArray_XDECREF(PyArrayObject *mp)

numpy/core/src/multiarray/refcount.h

@@ -7,6 +7,11 @@ PyArray_Item_INCREF(char *data, PyArray_Descr *descr);
 NPY_NO_EXPORT void
 PyArray_Item_XDECREF(char *data, PyArray_Descr *descr);
 
+NPY_NO_EXPORT int
+PyArray_ClearData(
+        PyArray_Descr *descr, char *data,
+        npy_intp stride, npy_intp size, int aligned);
+
 NPY_NO_EXPORT int
 PyArray_INCREF(PyArrayObject *mp);