Merge pull request numpy#19571 from seberg/ufunc-refactor-2021-normal

MAINT: Refactor UFunc core to use NEP 43 style dispatching

Author: mattip

doc/source/reference/c-api/types-and-structures.rst

@@ -818,6 +818,7 @@ PyUFunc_Type and PyUFuncObject
           npy_intp *core_dim_sizes;
           npy_uint32 *core_dim_flags;
           PyObject *identity_value;
+          /* Further private slots (size depends on the NumPy version) */
       } PyUFuncObject;
 
    .. c:macro: PyObject_HEAD
@@ -957,9 +958,12 @@ PyUFunc_Type and PyUFuncObject
 
    .. c:member:: PyUFunc_LegacyInnerLoopSelectionFunc *legacy_inner_loop_selector
 
-       A function which returns an inner loop. The ``legacy`` in the name arises
-       because for NumPy 1.6 a better variant had been planned. This variant
-       has not yet come about.
+       .. deprecated:: 1.22
+
+            Some fallback support for this slot exists, but will be removed
+            eventually.  A univiersal function which relied on this will have
+            eventually have to be ported.
+            See ref:`NEP 41 <NEP41>` and ref:`NEP 43 <NEP43>`
 
    .. c:member:: void *reserved2
 

numpy/core/include/numpy/ufuncobject.h

@@ -211,6 +211,12 @@ typedef struct _tagPyUFuncObject {
         /* Identity for reduction, when identity == PyUFunc_IdentityValue */
         PyObject *identity_value;
 
+        /* New in NPY_API_VERSION 0x0000000F and above */
+
+        /* New private fields related to dispatching */
+        void *_dispatch_cache;
+        /* A PyListObject of `(tuple of DTypes, ArrayMethod/Promoter)` */
+        PyObject *_loops;
 } PyUFuncObject;
 
 #include "arrayobject.h"

numpy/core/setup.py

@@ -928,6 +928,8 @@ def generate_umath_c(ext, build_dir):
             join('src', 'umath', 'matmul.c.src'),
             join('src', 'umath', 'clip.h.src'),
             join('src', 'umath', 'clip.c.src'),
+            join('src', 'umath', 'dispatching.c'),
+            join('src', 'umath', 'legacy_array_method.c'),
             join('src', 'umath', 'ufunc_object.c'),
             join('src', 'umath', 'extobj.c'),
             join('src', 'umath', 'scalarmath.c.src'),

numpy/core/src/multiarray/array_method.c

@@ -757,9 +757,6 @@ boundarraymethod__simple_strided_call(
 
 
 /*
- * TODO: Currently still based on the old ufunc system and not ArrayMethod!
- *       This requires fixing the ufunc code first.
- *
  * Support for masked inner-strided loops.  Masked inner-strided loops are
  * only used in the ufunc machinery.  So this special cases them.
  * In the future it probably makes sense to create an::
@@ -770,8 +767,8 @@ boundarraymethod__simple_strided_call(
  */
 typedef struct {
     NpyAuxData base;
-    PyUFuncGenericFunction unmasked_stridedloop;
-    void *innerloopdata;
+    PyArrayMethod_StridedLoop *unmasked_stridedloop;
+    NpyAuxData *unmasked_auxdata;
     int nargs;
     char *dataptrs[];
 } _masked_stridedloop_data;
@@ -781,6 +778,7 @@ static void
 _masked_stridedloop_data_free(NpyAuxData *auxdata)
 {
     _masked_stridedloop_data *data = (_masked_stridedloop_data *)auxdata;
+    NPY_AUXDATA_FREE(data->unmasked_auxdata);
     PyMem_Free(data);
 }
 
@@ -790,15 +788,15 @@ _masked_stridedloop_data_free(NpyAuxData *auxdata)
  * masked strided-loop, only calling the function for elements
  * where the mask is True.
  */
-static void
-unmasked_ufunc_loop_as_masked(
-        char **data, const npy_intp *dimensions,
-        const npy_intp *strides, void *_auxdata)
+static int
+generic_masked_strided_loop(PyArrayMethod_Context *context,
+        char *const *data, const npy_intp *dimensions,
+        const npy_intp *strides, NpyAuxData *_auxdata)
 {
     _masked_stridedloop_data *auxdata = (_masked_stridedloop_data *)_auxdata;
     int nargs = auxdata->nargs;
-    PyUFuncGenericFunction strided_loop = auxdata->unmasked_stridedloop;
-    void *innerloopdata = auxdata->innerloopdata;
+    PyArrayMethod_StridedLoop *strided_loop = auxdata->unmasked_stridedloop;
+    NpyAuxData *strided_loop_auxdata = auxdata->unmasked_auxdata;
 
     char **dataptrs = auxdata->dataptrs;
     memcpy(dataptrs, data, nargs * sizeof(char *));
@@ -819,39 +817,37 @@ unmasked_ufunc_loop_as_masked(
 
         /* Process unmasked values */
         mask = npy_memchr(mask, 0, mask_stride, N, &subloopsize, 0);
-        strided_loop(dataptrs, &subloopsize, strides, innerloopdata);
+        int res = strided_loop(context,
+                dataptrs, &subloopsize, strides, strided_loop_auxdata);
+        if (res != 0) {
+            return res;
+        }
         for (int i = 0; i < nargs; i++) {
             dataptrs[i] += subloopsize * strides[i];
         }
         N -= subloopsize;
     } while (N > 0);
+
+    return 0;
 }
 
 
 /*
- * TODO: This function will be the masked equivalent to `get_loop`.
- * This function wraps a legacy inner loop so it becomes masked.
- *
- * Returns 0 on success, -1 on error.
+ * Identical to the `get_loop` functions and wraps it.  This adds support
+ * to a boolean mask being passed in as a last, additional, operand.
+ * The wrapped loop will only be called for unmasked elements.
+ * (Does not support `move_references` or inner dimensions!)
  */
 NPY_NO_EXPORT int
-PyUFunc_DefaultMaskedInnerLoopSelector(PyUFuncObject *ufunc,
-        PyArray_Descr **dtypes,
-        PyUFuncGenericFunction *out_innerloop,
-        NpyAuxData **out_innerloopdata,
-        int *out_needs_api)
+PyArrayMethod_GetMaskedStridedLoop(
+        PyArrayMethod_Context *context,
+        int aligned, npy_intp *fixed_strides,
+        PyArrayMethod_StridedLoop **out_loop,
+        NpyAuxData **out_transferdata,
+        NPY_ARRAYMETHOD_FLAGS *flags)
 {
-    int retcode;
     _masked_stridedloop_data *data;
-    int nargs = ufunc->nin + ufunc->nout;
-
-    if (ufunc->legacy_inner_loop_selector == NULL) {
-        PyErr_SetString(PyExc_RuntimeError,
-                "the ufunc default masked inner loop selector doesn't "
-                "yet support wrapping the new inner loop selector, it "
-                "still only wraps the legacy inner loop selector");
-        return -1;
-    }
+    int nargs = context->method->nin + context->method->nout;
 
     /* Add working memory for the data pointers, to modify them in-place */
     data = PyMem_Malloc(sizeof(_masked_stridedloop_data) +
@@ -865,18 +861,14 @@ PyUFunc_DefaultMaskedInnerLoopSelector(PyUFuncObject *ufunc,
     data->unmasked_stridedloop = NULL;
     data->nargs = nargs;
 
-    /* Get the unmasked ufunc inner loop */
-    retcode = ufunc->legacy_inner_loop_selector(ufunc, dtypes,
-            &data->unmasked_stridedloop, &data->innerloopdata,
-            out_needs_api);
-    if (retcode < 0) {
-        PyArray_free(data);
-        return retcode;
+    if (context->method->get_strided_loop(context,
+            aligned, 0, fixed_strides,
+            &data->unmasked_stridedloop, &data->unmasked_auxdata, flags) < 0) {
+        PyMem_Free(data);
+        return -1;
     }
-
-    /* Return the loop function + aux data */
-    *out_innerloop = &unmasked_ufunc_loop_as_masked;
-    *out_innerloopdata = (NpyAuxData *)data;
+    *out_transferdata = (NpyAuxData *)data;
+    *out_loop = generic_masked_strided_loop;
     return 0;
 }
 

numpy/core/src/multiarray/array_method.h

@@ -17,6 +17,7 @@ typedef enum {
      * setup/check. No function should set error flags and ignore them
      * since it would interfere with chaining operations (e.g. casting).
      */
+    /* TODO: Change this into a positive flag */
     NPY_METH_NO_FLOATINGPOINT_ERRORS = 1 << 2,
     /* Whether the method supports unaligned access (not runtime) */
     NPY_METH_SUPPORTS_UNALIGNED = 1 << 3,
@@ -158,17 +159,16 @@ npy_default_get_strided_loop(
         PyArrayMethod_StridedLoop **out_loop, NpyAuxData **out_transferdata,
         NPY_ARRAYMETHOD_FLAGS *flags);
 
-/*
- * TODO: This function will not rely on the current ufunc code after the
- *       ufunc refactor.
- */
-#include "numpy/ufuncobject.h"
+
 NPY_NO_EXPORT int
-PyUFunc_DefaultMaskedInnerLoopSelector(PyUFuncObject *ufunc,
-        PyArray_Descr **dtypes,
-        PyUFuncGenericFunction *out_innerloop,
-        NpyAuxData **out_innerloopdata,
-        int *out_needs_api);
+PyArrayMethod_GetMaskedStridedLoop(
+        PyArrayMethod_Context *context,
+        int aligned,
+        npy_intp *fixed_strides,
+        PyArrayMethod_StridedLoop **out_loop,
+        NpyAuxData **out_transferdata,
+        NPY_ARRAYMETHOD_FLAGS *flags);
+
 
 /*
  * TODO: This function is the internal version, and its error paths may

numpy/core/src/multiarray/nditer_constr.c

@@ -449,6 +449,11 @@ NpyIter_AdvancedNew(int nop, PyArrayObject **op_in, npy_uint32 flags,
     /*
      * If REFS_OK was specified, check whether there are any
      * reference arrays and flag it if so.
+     *
+     * NOTE: This really should be unnecessary, but chances are someone relies
+     *       on it.  The iterator itself does not require the API here
+     *       as it only does so for casting/buffering.  But in almost all
+     *       use-cases the API will be required for whatever operation is done.
      */
     if (flags & NPY_ITER_REFS_OK) {
         for (iop = 0; iop < nop; ++iop) {