numpy · seberg · Dec 18, 2024 · Dec 14, 2024 · Dec 13, 2024 · Dec 14, 2024
diff --git a/doc/release/upcoming_changes/27998.c_api.rst b/doc/release/upcoming_changes/27998.c_api.rst
@@ -0,0 +1,10 @@
+New `NpyIter_GetTransferFlags` and ``NpyIter_IterationNeedsAPI`` change
+-----------------------------------------------------------------------
+NumPy now has the new `NpyIter_GetTransferFlags` function as a more precise
+way checking of iterator/buffering needs.  I.e. whether the Python API/GIL is
+required or floating point errors may occur.
+This function is also faster if you already know your needs without buffering.
+
+The ``NpyIter_IterationNeedsAPI`` function now performs all the checks that were
+previously performed at setup time.  While it was never necessary to call it
+multiple times, doing so will now have a larger cost.
diff --git a/doc/source/reference/c-api/iterator.rst b/doc/source/reference/c-api/iterator.rst
@@ -434,6 +434,9 @@ Construction and destruction
     is enabled, the caller must be sure to check whether
     ``NpyIter_IterationNeedsAPI(iter)`` is true, in which case
     it may not release the GIL during iteration.
+    If you are working with known dtypes `NpyIter_GetTransferFlags` is
+    a faster and more precise way to check for whether the iterator needs
+    the API due to buffering.
 
 .. c:macro:: NPY_ITER_ZEROSIZE_OK
 
@@ -823,6 +826,20 @@ Construction and destruction
 
     Returns ``NPY_SUCCEED`` or ``NPY_FAIL``.
 
+.. c:function:: NPY_ARRAYMETHOD_FLAGS NpyIter_GetTransferFlags(NpyIter *iter)
+
+    .. versionadded:: 2.3
+
+    Fetches the `NPY_METH_RUNTIME_FLAGS` which provide the information on
+    whether buffering needs the Python GIL (`NPY_METH_REQUIRES_PYAPI`) or
+    floating point errors may be set (`NPY_METH_NO_FLOATINGPOINT_ERRORS`).
+
+    Prior to NumPy 2.3, the public function available was
+    ``NpyIter_IterationNeedsAPI``, which is still available and additionally
+    checks for object (or similar) dtypes and not exclusively for
+    buffering/iteration needs itself.
+    In general, this function should be preferred.
+
 .. c:function:: int NpyIter_Reset(NpyIter* iter, char** errmsg)
 
     Resets the iterator back to its initial state, at the beginning

diff --git a/numpy/_core/code_generators/cversions.txt b/numpy/_core/code_generators/cversions.txt
@@ -77,5 +77,5 @@
 0x00000012 = 2b8f1f4da822491ff030b2b37dff07e3
 # Version 19 (NumPy 2.1.0) Only header additions
 # Version 19 (NumPy 2.2.0) No change
-# Version 19 (NumPy 2.3.0) No change
-0x00000013 = 2b8f1f4da822491ff030b2b37dff07e3
+# Version 19 (NumPy 2.3.0)
+0x00000013 = e56b74d32a934d085e7c3414cb9999b8,
diff --git a/numpy/_core/code_generators/numpy_api.py b/numpy/_core/code_generators/numpy_api.py
@@ -106,7 +106,7 @@ def get_annotations():
     '__unused_indices__': (
         [1, 4, 40, 41, 66, 67, 68, 81, 82, 83,
          103, 115, 117, 122, 163, 164, 171, 173, 197,
-         201, 202, 208, 219, 220, 221, 222, 223, 278,
+         201, 202, 208, 219, 220, 221, 222, 278,
          291, 293, 294, 295, 301]
         # range/slots reserved DType classes (see _public_dtype_api_table.h):
         + list(range(320, 361)) + [366, 367, 368]
@@ -293,8 +293,8 @@ def get_annotations():
     # Unused slot 220, was `PyArray_DatetimeToDatetimeStruct`
     # Unused slot 221, was `PyArray_TimedeltaToTimedeltaStruct`
     # Unused slot 222, was `PyArray_DatetimeStructToDatetime`
-    # Unused slot 223, was `PyArray_TimedeltaStructToTimedelta`
     # NDIter API
+    'NpyIter_GetTransferFlags':             (223, MinVersion("2.3")),
     'NpyIter_New':                          (224,),
     'NpyIter_MultiNew':                     (225,),
     'NpyIter_AdvancedNew':                  (226,),
@@ -407,6 +407,8 @@ def get_annotations():
     # `PyDataType_GetArrFuncs` checks for the NumPy runtime version.
     '_PyDataType_GetArrFuncs':                       (365,),
     # End 2.0 API
+    # NpyIterGetTransferFlags (slot 223) added.
+    # End 2.3 API
 }
 
 ufunc_types_api = {

diff --git a/numpy/_core/src/multiarray/common.h b/numpy/_core/src/multiarray/common.h
@@ -18,17 +18,6 @@ extern "C" {
 
 #define error_converting(x)  (((x) == -1) && PyErr_Occurred())
 
-#ifdef NPY_ALLOW_THREADS
-#define NPY_BEGIN_THREADS_NDITER(iter) \
-        do { \
-            if (!NpyIter_IterationNeedsAPI(iter)) { \
-                NPY_BEGIN_THREADS_THRESHOLDED(NpyIter_GetIterSize(iter)); \
-            } \
-        } while(0)
-#else
-#define NPY_BEGIN_THREADS_NDITER(iter)
-#endif
-
 
 NPY_NO_EXPORT PyArray_Descr *
 PyArray_DTypeFromObjectStringDiscovery(

diff --git a/numpy/_core/src/multiarray/ctors.c b/numpy/_core/src/multiarray/ctors.c
@@ -2696,7 +2696,6 @@ PyArray_CopyAsFlat(PyArrayObject *dst, PyArrayObject *src, NPY_ORDER order)
 
     npy_intp dst_count, src_count, count;
     npy_intp dst_size, src_size;
-    int needs_api;
 
     NPY_BEGIN_THREADS_DEF;
 
@@ -2757,13 +2756,13 @@ PyArray_CopyAsFlat(PyArrayObject *dst, PyArrayObject *src, NPY_ORDER order)
     /* Get all the values needed for the inner loop */
     dst_iternext = NpyIter_GetIterNext(dst_iter, NULL);
     dst_dataptr = NpyIter_GetDataPtrArray(dst_iter);
-    /* Since buffering is disabled, we can cache the stride */
+    /* The inner stride is also the fixed stride for the whole iteration. */
     dst_stride = NpyIter_GetInnerStrideArray(dst_iter)[0];
     dst_countptr = NpyIter_GetInnerLoopSizePtr(dst_iter);
 
     src_iternext = NpyIter_GetIterNext(src_iter, NULL);
     src_dataptr = NpyIter_GetDataPtrArray(src_iter);
-    /* Since buffering is disabled, we can cache the stride */
+    /* The inner stride is also the fixed stride for the whole iteration. */
     src_stride = NpyIter_GetInnerStrideArray(src_iter)[0];
     src_countptr = NpyIter_GetInnerLoopSizePtr(src_iter);
 
@@ -2773,15 +2772,6 @@ PyArray_CopyAsFlat(PyArrayObject *dst, PyArrayObject *src, NPY_ORDER order)
         return -1;
     }
 
-    needs_api = NpyIter_IterationNeedsAPI(dst_iter) ||
-                NpyIter_IterationNeedsAPI(src_iter);
-
-    /*
-     * Because buffering is disabled in the iterator, the inner loop
-     * strides will be the same throughout the iteration loop.  Thus,
-     * we can pass them to this function to take advantage of
-     * contiguous strides, etc.
-     */
     NPY_cast_info cast_info;
     NPY_ARRAYMETHOD_FLAGS flags;
     if (PyArray_GetDTypeTransferFunction(
@@ -2795,7 +2785,8 @@ PyArray_CopyAsFlat(PyArrayObject *dst, PyArrayObject *src, NPY_ORDER order)
         NpyIter_Deallocate(src_iter);
         return -1;
     }
-    needs_api |= (flags & NPY_METH_REQUIRES_PYAPI) != 0;
+    /* No need to worry about API use in unbuffered iterator */
+    int needs_api = (flags & NPY_METH_REQUIRES_PYAPI) != 0;
     if (!(flags & NPY_METH_NO_FLOATINGPOINT_ERRORS)) {
         npy_clear_floatstatus_barrier((char *)src_iter);
     }

diff --git a/numpy/_core/src/multiarray/einsum.c.src b/numpy/_core/src/multiarray/einsum.c.src
@@ -520,14 +520,16 @@ unbuffered_loop_nop1_ndim2(NpyIter *iter)
         return -1;
     }
 
-    /*
-     * Since the iterator wasn't tracking coordinates, the
-     * loop provided by the iterator is in Fortran-order.
-     */
+    /* IterationNeedsAPI effectively only checks for object dtype here. */
     int needs_api = NpyIter_IterationNeedsAPI(iter);
     if (!needs_api) {
         NPY_BEGIN_THREADS_THRESHOLDED(shape[1] * shape[0]);
     }
+
+    /*
+     * Since the iterator wasn't tracking coordinates, the
+     * loop provided by the iterator is in Fortran-order.
+     */
     for (coord = shape[1]; coord > 0; --coord) {
         sop(1, ptrs[0], strides[0], shape[0]);
 
@@ -581,14 +583,16 @@ unbuffered_loop_nop1_ndim3(NpyIter *iter)
         return -1;
     }
 
-    /*
-     * Since the iterator wasn't tracking coordinates, the
-     * loop provided by the iterator is in Fortran-order.
-     */
+    /* IterationNeedsAPI effectively only checks for object dtype here. */
     int needs_api = NpyIter_IterationNeedsAPI(iter);
     if (!needs_api) {
         NPY_BEGIN_THREADS_THRESHOLDED(shape[2] * shape[1] * shape[0]);
     }
+
+    /*
+     * Since the iterator wasn't tracking coordinates, the
+     * loop provided by the iterator is in Fortran-order.
+     */
     for (coords[1] = shape[2]; coords[1] > 0; --coords[1]) {
         for (coords[0] = shape[1]; coords[0] > 0; --coords[0]) {
             sop(1, ptrs[0], strides[0], shape[0]);
@@ -645,14 +649,16 @@ unbuffered_loop_nop2_ndim2(NpyIter *iter)
         return -1;
     }
 
-    /*
-     * Since the iterator wasn't tracking coordinates, the
-     * loop provided by the iterator is in Fortran-order.
-     */
+    /* IterationNeedsAPI effectively only checks for object dtype here. */
     int needs_api = NpyIter_IterationNeedsAPI(iter);
     if (!needs_api) {
         NPY_BEGIN_THREADS_THRESHOLDED(shape[1] * shape[0]);
     }
+
+    /*
+     * Since the iterator wasn't tracking coordinates, the
+     * loop provided by the iterator is in Fortran-order.
+     */
     for (coord = shape[1]; coord > 0; --coord) {
         sop(2, ptrs[0], strides[0], shape[0]);
 
@@ -708,14 +714,16 @@ unbuffered_loop_nop2_ndim3(NpyIter *iter)
         return -1;
     }
 
-    /*
-     * Since the iterator wasn't tracking coordinates, the
-     * loop provided by the iterator is in Fortran-order.
-     */
+    /* IterationNeedsAPI effectively only checks for object dtype here. */
     int needs_api = NpyIter_IterationNeedsAPI(iter);
     if (!needs_api) {
         NPY_BEGIN_THREADS_THRESHOLDED(shape[2] * shape[1] * shape[0]);
     }
+
+    /*
+     * Since the iterator wasn't tracking coordinates, the
+     * loop provided by the iterator is in Fortran-order.
+     */
     for (coords[1] = shape[2]; coords[1] > 0; --coords[1]) {
         for (coords[0] = shape[1]; coords[0] > 0; --coords[0]) {
             sop(2, ptrs[0], strides[0], shape[0]);
@@ -1120,7 +1128,6 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
         char **dataptr;
         npy_intp *stride;
         npy_intp *countptr;
-        int needs_api;
         NPY_BEGIN_THREADS_DEF;
 
         iternext = NpyIter_GetIterNext(iter, NULL);
@@ -1130,17 +1137,20 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
         dataptr = NpyIter_GetDataPtrArray(iter);
         stride = NpyIter_GetInnerStrideArray(iter);
         countptr = NpyIter_GetInnerLoopSizePtr(iter);
-        needs_api = NpyIter_IterationNeedsAPI(iter);
+        /* IterationNeedsAPI additionally checks for object dtype here. */
+        int needs_api = NpyIter_IterationNeedsAPI(iter);
+        if (!needs_api) {
+            NPY_BEGIN_THREADS_THRESHOLDED(NpyIter_GetIterSize(iter));
+        }
 
-        NPY_BEGIN_THREADS_NDITER(iter);
         NPY_EINSUM_DBG_PRINT("Einsum loop\n");
         do {
             sop(nop, dataptr, stride, *countptr);
         } while (!(needs_api && PyErr_Occurred()) && iternext(iter));
         NPY_END_THREADS;
 
         /* If the API was needed, it may have thrown an error */
-        if (NpyIter_IterationNeedsAPI(iter) && PyErr_Occurred()) {
+        if (needs_api && PyErr_Occurred()) {
             goto fail;
         }
     }

diff --git a/numpy/_core/src/multiarray/item_selection.c b/numpy/_core/src/multiarray/item_selection.c
@@ -2752,6 +2752,7 @@ PyArray_CountNonzero(PyArrayObject *self)
     if (iter == NULL) {
         return -1;
     }
+    /* IterationNeedsAPI also checks dtype for whether `nonzero` may need it */
     needs_api = NpyIter_IterationNeedsAPI(iter);
 
     /* Get the pointers for inner loop iteration */
@@ -2761,7 +2762,9 @@ PyArray_CountNonzero(PyArrayObject *self)
         return -1;
     }
 
-    NPY_BEGIN_THREADS_NDITER(iter);
+    if (!needs_api) {
+        NPY_BEGIN_THREADS_THRESHOLDED(NpyIter_GetIterSize(iter));
+    }
 
     dataptr = NpyIter_GetDataPtrArray(iter);
     strideptr = NpyIter_GetInnerStrideArray(iter);
@@ -2982,9 +2985,12 @@ PyArray_Nonzero(PyArrayObject *self)
             return NULL;
         }
 
+        /* IterationNeedsAPI also checks dtype for whether `nonzero` may need it */
         needs_api = NpyIter_IterationNeedsAPI(iter);
 
-        NPY_BEGIN_THREADS_NDITER(iter);
+        if (!needs_api) {
+            NPY_BEGIN_THREADS_THRESHOLDED(NpyIter_GetIterSize(iter));
+        }
 
         dataptr = NpyIter_GetDataPtrArray(iter);
 

diff --git a/numpy/_core/src/multiarray/mapping.c b/numpy/_core/src/multiarray/mapping.c
@@ -976,10 +976,7 @@ array_boolean_subscript(PyArrayObject *self,
         /* Get a dtype transfer function */
         NpyIter_GetInnerFixedStrideArray(iter, fixed_strides);
         NPY_cast_info cast_info;
-        /*
-         * TODO: Ignoring cast flags, since this is only ever a copy. In
-         *       principle that may not be quite right in some future?
-         */
+
         NPY_ARRAYMETHOD_FLAGS cast_flags;
         if (PyArray_GetDTypeTransferFunction(
                         IsUintAligned(self) && IsAligned(self),
@@ -992,6 +989,8 @@ array_boolean_subscript(PyArrayObject *self,
             NpyIter_Deallocate(iter);
             return NULL;
         }
+        cast_flags = PyArrayMethod_COMBINED_FLAGS(
+                cast_flags, NpyIter_GetTransferFlags(iter));
 
         /* Get the values needed for the inner loop */
         iternext = NpyIter_GetIterNext(iter, NULL);
@@ -1002,7 +1001,10 @@ array_boolean_subscript(PyArrayObject *self,
             return NULL;
         }
 
-        NPY_BEGIN_THREADS_NDITER(iter);
+        /* NOTE: Don't worry about floating point errors as this is a copy. */
+        if (!(cast_flags & NPY_METH_REQUIRES_PYAPI)) {
+            NPY_BEGIN_THREADS_THRESHOLDED(NpyIter_GetIterSize(iter));
+        }
 
         innerstrides = NpyIter_GetInnerStrideArray(iter);
         dataptrs = NpyIter_GetDataPtrArray(iter);
@@ -1195,8 +1197,11 @@ array_assign_boolean_subscript(PyArrayObject *self,
             return -1;
         }
 
+        cast_flags = PyArrayMethod_COMBINED_FLAGS(
+                cast_flags, NpyIter_GetTransferFlags(iter));
+
         if (!(cast_flags & NPY_METH_REQUIRES_PYAPI)) {
-            NPY_BEGIN_THREADS_NDITER(iter);
+            NPY_BEGIN_THREADS_THRESHOLDED(NpyIter_GetIterSize(iter));
         }
         if (!(flags & NPY_METH_NO_FLOATINGPOINT_ERRORS)) {
             npy_clear_floatstatus_barrier((char *)self);
@@ -2662,7 +2667,9 @@ PyArray_MapIterCheckIndices(PyArrayMapIterObject *mit)
             return -1;
         }
 
-        NPY_BEGIN_THREADS_NDITER(op_iter);
+        if (!(NpyIter_GetTransferFlags(op_iter) & NPY_METH_REQUIRES_PYAPI)) {
+            NPY_BEGIN_THREADS_THRESHOLDED(NpyIter_GetIterSize(op_iter));
+        }
         iterptr = NpyIter_GetDataPtrArray(op_iter);
         iterstride = NpyIter_GetInnerStrideArray(op_iter);
         do {