Merge pull request #23262 from ngoldbaum/public-array-clear-on-dtype

API: expose traversal typedefs and `get_clear_loop` slot

5 files changed, 117 insertions, 71 deletions

diff --git a/numpy/core/include/numpy/_dtype_api.h b/numpy/core/include/numpy/_dtype_api.h
index d657928a7..2f801eace 100644
--- a/numpy/core/include/numpy/_dtype_api.h
+++ b/numpy/core/include/numpy/_dtype_api.h
@@ -5,7 +5,7 @@
 #ifndef NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_
 #define NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_
 
-#define __EXPERIMENTAL_DTYPE_API_VERSION 8
+#define __EXPERIMENTAL_DTYPE_API_VERSION 9
 
 struct PyArrayMethodObject_tag;
 
@@ -140,10 +140,6 @@ typedef struct {
 #define NPY_METH_unaligned_contiguous_loop 7
 #define NPY_METH_contiguous_indexed_loop 8
 
-/* other slots are in order, so note last one (internal use!) */
-#define _NPY_NUM_DTYPE_SLOTS 8
-#define _NPY_NUM_DTYPE_PYARRAY_ARRFUNC_SLOTS 22 + (1 << 10)
-
 /*
  * The resolve descriptors function, must be able to handle NULL values for
  * all output (but not input) `given_descrs` and fill `loop_descrs`.
@@ -257,6 +253,46 @@ typedef int translate_loop_descrs_func(int nin, int nout,
 
 
 /*
+ * A traverse loop working on a single array. This is similar to the general
+ * strided-loop function. This is designed for loops that need to visit every
+ * element of a single array.
+ *
+ * Currently this is only used for array clearing, via the
+ * NPY_DT_get_clear_loop, api hook, particularly for arrays storing embedded
+ * references to python objects or heap-allocated data. If you define a dtype
+ * that uses embedded references, the NPY_ITEM_REFCOUNT flag must be set on the
+ * dtype instance.
+ *
+ * The `void *traverse_context` is passed in because we may need to pass in
+ * Intepreter state or similar in the futurem, but we don't want to pass in
+ * a full context (with pointers to dtypes, method, caller which all make
+ * no sense for a traverse function).
+ *
+ * We assume for now that this context can be just passed through in the
+ * the future (for structured dtypes).
+ *
+ */
+typedef int (traverse_loop_function)(
+        void *traverse_context, PyArray_Descr *descr, char *data,
+        npy_intp size, npy_intp stride, NpyAuxData *auxdata);
+
+
+/*
+ * Simplified get_loop function specific to dtype traversal
+ *
+ * Currently this is only used for clearing arrays. It should set the flags
+ * needed for the traversal loop and set out_loop to the loop function, which
+ * must be a valid traverse_loop_function pointer.
+ *
+ */
+typedef int (get_traverse_loop_function)(
+        void *traverse_context, PyArray_Descr *descr,
+        int aligned, npy_intp fixed_stride,
+        traverse_loop_function **out_loop, NpyAuxData **out_auxdata,
+        NPY_ARRAYMETHOD_FLAGS *flags);
+
+
+/*
  * ****************************
  *          DTYPE API
  * ****************************
@@ -266,7 +302,15 @@ typedef int translate_loop_descrs_func(int nin, int nout,
 #define NPY_DT_PARAMETRIC 1 << 2
 #define NPY_DT_NUMERIC 1 << 3
 
+/*
+ * These correspond to slots in the NPY_DType_Slots struct and must
+ * be in the same order as the members of that struct. If new slots
+ * get added or old slots get removed NPY_NUM_DTYPE_SLOTS must also
+ * be updated
+ */
+
 #define NPY_DT_discover_descr_from_pyobject 1
+// this slot is considered private because its API hasn't beed decided
 #define _NPY_DT_is_known_scalar_type 2
 #define NPY_DT_default_descr 3
 #define NPY_DT_common_dtype 4
@@ -274,6 +318,7 @@ typedef int translate_loop_descrs_func(int nin, int nout,
 #define NPY_DT_ensure_canonical 6
 #define NPY_DT_setitem 7
 #define NPY_DT_getitem 8
+#define NPY_DT_get_clear_loop 9
 
 // These PyArray_ArrFunc slots will be deprecated and replaced eventually
 // getitem and setitem can be defined as a performance optimization;
@@ -281,38 +326,42 @@ typedef int translate_loop_descrs_func(int nin, int nout,
 // `legacy_setitem_using_DType`, respectively. This functionality is
 // only supported for basic NumPy DTypes.
 
+
+// used to separate dtype slots from arrfuncs slots
+// intended only for internal use but defined here for clarity
+#define _NPY_DT_ARRFUNCS_OFFSET (1 << 10)
+
 // Cast is disabled
-// #define NPY_DT_PyArray_ArrFuncs_cast 0 + (1 << 10)
-
-#define NPY_DT_PyArray_ArrFuncs_getitem 1 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_setitem 2 + (1 << 10)
-
-#define NPY_DT_PyArray_ArrFuncs_copyswapn 3 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_copyswap 4 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_compare 5 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_argmax 6 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_dotfunc 7 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_scanfunc 8 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_fromstr 9 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_nonzero 10 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_fill 11 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_fillwithscalar 12 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_sort 13 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_argsort 14 + (1 << 10)
+// #define NPY_DT_PyArray_ArrFuncs_cast 0 + _NPY_DT_ARRFUNCS_OFFSET
+
+#define NPY_DT_PyArray_ArrFuncs_getitem 1 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_setitem 2 + _NPY_DT_ARRFUNCS_OFFSET
+
+#define NPY_DT_PyArray_ArrFuncs_copyswapn 3 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_copyswap 4 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_compare 5 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_argmax 6 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_dotfunc 7 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_scanfunc 8 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_fromstr 9 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_nonzero 10 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_fill 11 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_fillwithscalar 12 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_sort 13 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_argsort 14 + _NPY_DT_ARRFUNCS_OFFSET
 
 // Casting related slots are disabled. See
 // https://github.com/numpy/numpy/pull/23173#discussion_r1101098163
-// #define NPY_DT_PyArray_ArrFuncs_castdict 15 + (1 << 10)
-// #define NPY_DT_PyArray_ArrFuncs_scalarkind 16 + (1 << 10)
-// #define NPY_DT_PyArray_ArrFuncs_cancastscalarkindto 17 + (1 << 10)
-// #define NPY_DT_PyArray_ArrFuncs_cancastto 18 + (1 << 10)
+// #define NPY_DT_PyArray_ArrFuncs_castdict 15 + _NPY_DT_ARRFUNCS_OFFSET
+// #define NPY_DT_PyArray_ArrFuncs_scalarkind 16 + _NPY_DT_ARRFUNCS_OFFSET
+// #define NPY_DT_PyArray_ArrFuncs_cancastscalarkindto 17 + _NPY_DT_ARRFUNCS_OFFSET
+// #define NPY_DT_PyArray_ArrFuncs_cancastto 18 + _NPY_DT_ARRFUNCS_OFFSET
 
 // These are deprecated in NumPy 1.19, so are disabled here.
-// #define NPY_DT_PyArray_ArrFuncs_fastclip 19 + (1 << 10)
-// #define NPY_DT_PyArray_ArrFuncs_fastputmask 20 + (1 << 10)
-// #define NPY_DT_PyArray_ArrFuncs_fasttake 21 + (1 << 10)
-#define NPY_DT_PyArray_ArrFuncs_argmin 22 + (1 << 10)
-
+// #define NPY_DT_PyArray_ArrFuncs_fastclip 19 + _NPY_DT_ARRFUNCS_OFFSET
+// #define NPY_DT_PyArray_ArrFuncs_fastputmask 20 + _NPY_DT_ARRFUNCS_OFFSET
+// #define NPY_DT_PyArray_ArrFuncs_fasttake 21 + _NPY_DT_ARRFUNCS_OFFSET
+#define NPY_DT_PyArray_ArrFuncs_argmin 22 + _NPY_DT_ARRFUNCS_OFFSET
 
 // TODO: These slots probably still need some thought, and/or a way to "grow"?
 typedef struct {
diff --git a/numpy/core/src/multiarray/dtype_traversal.h b/numpy/core/src/multiarray/dtype_traversal.h
index 5bf983a78..fd060a0f0 100644
--- a/numpy/core/src/multiarray/dtype_traversal.h
+++ b/numpy/core/src/multiarray/dtype_traversal.h
@@ -3,30 +3,6 @@
 
 #include "array_method.h"
 
-/*
- * A traverse loop working on a single array. This is similar to the general
- * strided-loop function.
- * Using a `void *traverse_context`, because I we may need to pass in
- * Intepreter state or similar in the future.  But I don't want to pass in
- * a full context (with pointers to dtypes, method, caller which all make
- * no sense for a traverse function).
- *
- * We assume for now that this context can be just passed through in the
- * the future (for structured dtypes).
- */
-typedef int (traverse_loop_function)(
-        void *traverse_context, PyArray_Descr *descr, char *data,
-        npy_intp size, npy_intp stride, NpyAuxData *auxdata);
-
-
-/* Simplified get_loop function specific to dtype traversal */
-typedef int (get_traverse_loop_function)(
-        void *traverse_context, PyArray_Descr *descr,
-        int aligned, npy_intp fixed_stride,
-        traverse_loop_function **out_loop, NpyAuxData **out_auxdata,
-        NPY_ARRAYMETHOD_FLAGS *flags);
-
-
 /* NumPy DType clear (object DECREF + NULLing) implementations */
 
 NPY_NO_EXPORT int
diff --git a/numpy/core/src/multiarray/dtypemeta.h b/numpy/core/src/multiarray/dtypemeta.h
index 75baf611c..3b4dbad24 100644
--- a/numpy/core/src/multiarray/dtypemeta.h
+++ b/numpy/core/src/multiarray/dtypemeta.h
@@ -29,11 +29,6 @@ typedef struct {
     setitemfunction *setitem;
     getitemfunction *getitem;
     /*
-     * The casting implementation (ArrayMethod) to convert between two
-     * instances of this DType, stored explicitly for fast access:
-     */
-    PyArrayMethodObject *within_dtype_castingimpl;
-    /*
      * Either NULL or fetches a clearing function.  Clearing means deallocating
      * any referenced data and setting it to a safe state.  For Python objects
      * this means using `Py_CLEAR` which is equivalent to `Py_DECREF` and
@@ -47,6 +42,11 @@ typedef struct {
      */
     get_traverse_loop_function *get_clear_loop;
     /*
+     * The casting implementation (ArrayMethod) to convert between two
+     * instances of this DType, stored explicitly for fast access:
+     */
+    PyArrayMethodObject *within_dtype_castingimpl;
+    /*
      * Dictionary of ArrayMethods representing most possible casts
      * (structured and object are exceptions).
      * This should potentially become a weak mapping in the future.
@@ -61,6 +61,13 @@ typedef struct {
     PyArray_ArrFuncs f;
 } NPY_DType_Slots;
 
+// This must be updated if new slots before within_dtype_castingimpl
+// are added
+#define NPY_NUM_DTYPE_SLOTS 9
+#define NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS 22
+#define NPY_DT_MAX_ARRFUNCS_SLOT \
+  NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS + _NPY_DT_ARRFUNCS_OFFSET
+
 
 #define NPY_DTYPE(descr) ((PyArray_DTypeMeta *)Py_TYPE(descr))
 #define NPY_DT_SLOTS(dtype) ((NPY_DType_Slots *)(dtype)->dt_slots)
diff --git a/numpy/core/src/multiarray/experimental_public_dtype_api.c b/numpy/core/src/multiarray/experimental_public_dtype_api.c
index 63a9aa0a8..70a6f1995 100644
--- a/numpy/core/src/multiarray/experimental_public_dtype_api.c
+++ b/numpy/core/src/multiarray/experimental_public_dtype_api.c
@@ -161,6 +161,7 @@ PyArrayInitDTypeMeta_FromSpec(
     NPY_DT_SLOTS(DType)->common_instance = NULL;
     NPY_DT_SLOTS(DType)->setitem = NULL;
     NPY_DT_SLOTS(DType)->getitem = NULL;
+    NPY_DT_SLOTS(DType)->get_clear_loop = NULL;
     NPY_DT_SLOTS(DType)->f = default_funcs;
 
     PyType_Slot *spec_slot = spec->slots;
@@ -171,26 +172,29 @@ PyArrayInitDTypeMeta_FromSpec(
         if (slot == 0) {
             break;
         }
-        if (slot > _NPY_NUM_DTYPE_PYARRAY_ARRFUNC_SLOTS || slot < 0) {
+        if ((slot < 0) ||
+            ((slot > NPY_NUM_DTYPE_SLOTS) &&
+             (slot <= _NPY_DT_ARRFUNCS_OFFSET)) ||
+            (slot > NPY_DT_MAX_ARRFUNCS_SLOT)) {
             PyErr_Format(PyExc_RuntimeError,
                     "Invalid slot with value %d passed in.", slot);
             return -1;
         }
         /*
-         * It is up to the user to get this right, and slots are sorted
-         * exactly like they are stored right now:
+         * It is up to the user to get this right, the slots in the public API
+         * are sorted exactly like they are stored in the NPY_DT_Slots struct
+         * right now:
          */
-        if (slot <= _NPY_NUM_DTYPE_SLOTS) {
-            // slot > 8 are PyArray_ArrFuncs
+        if (slot <= NPY_NUM_DTYPE_SLOTS) {
+            // slot > NPY_NUM_DTYPE_SLOTS are PyArray_ArrFuncs
             void **current = (void **)(&(
                     NPY_DT_SLOTS(DType)->discover_descr_from_pyobject));
             current += slot - 1;
             *current = pfunc;
         }
         else {
-            // Remove PyArray_ArrFuncs offset
-            int f_slot = slot - (1 << 10);
-            if (1 <= f_slot && f_slot <= 22) {
+            int f_slot = slot - _NPY_DT_ARRFUNCS_OFFSET;
+            if (1 <= f_slot && f_slot <= NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS) {
                 switch (f_slot) {
                     case 1:
                         NPY_DT_SLOTS(DType)->f.getitem = pfunc;
@@ -290,6 +294,7 @@ PyArrayInitDTypeMeta_FromSpec(
             return -1;
         }
     }
+
     /* invalid type num. Ideally, we get away with it! */
     DType->type_num = -1;
 
diff --git a/numpy/core/src/multiarray/refcount.c b/numpy/core/src/multiarray/refcount.c
index e735e3b8f..20527f7af 100644
--- a/numpy/core/src/multiarray/refcount.c
+++ b/numpy/core/src/multiarray/refcount.c
@@ -16,6 +16,7 @@
 #include "numpy/arrayobject.h"
 #include "numpy/arrayscalars.h"
 #include "iterators.h"
+#include "dtypemeta.h"
 
 #include "npy_config.h"
 
@@ -358,9 +359,17 @@ PyArray_XDECREF(PyArrayObject *mp)
 NPY_NO_EXPORT void
 PyArray_FillObjectArray(PyArrayObject *arr, PyObject *obj)
 {
+    PyArray_Descr* descr = PyArray_DESCR(arr);
+
+    // non-legacy dtypes are responsible for initializing
+    // their own internal references
+    if (!NPY_DT_is_legacy(NPY_DTYPE(descr))) {
+        return;
+    }
+
     npy_intp i,n;
     n = PyArray_SIZE(arr);
-    if (PyArray_DESCR(arr)->type_num == NPY_OBJECT) {
+    if (descr->type_num == NPY_OBJECT) {
         PyObject **optr;
         optr = (PyObject **)(PyArray_DATA(arr));
         n = PyArray_SIZE(arr);
@@ -380,8 +389,8 @@ PyArray_FillObjectArray(PyArrayObject *arr, PyObject *obj)
         char *optr;
         optr = PyArray_DATA(arr);
         for (i = 0; i < n; i++) {
-            _fillobject(optr, obj, PyArray_DESCR(arr));
-            optr += PyArray_DESCR(arr)->elsize;
+            _fillobject(optr, obj, descr);
+            optr += descr->elsize;
         }
     }
 }