ENH: Add Wrapping array-method to wrap existing ufunc loop

This is useful for Unit DTypes, which can cast their inputs
(if necessary), and then re-use existing ufuncs.

numpy/core/include/numpy/experimental_dtype_api.h

@@ -24,6 +24,12 @@
  *     Register a new loop for a ufunc.  This uses the `PyArrayMethod_Spec`
  *     which must be filled in (see in-line comments).
  *
+ * - PyUFunc_AddWrappingLoop:
+ *
+ *     Register a new loop which reuses an existing one, but modifies the
+ *     result dtypes.  Please search the internal NumPy docs for more info
+ *     at this point.  (Used for physical units dtype.)
+ *
  * - PyUFunc_AddPromoter:
  *
  *     Register a new promoter for a ufunc.  A promoter is a function stored
@@ -58,6 +64,16 @@
  *     also promote C; where "promotes" means implements the promotion.
  *     (There are some exceptions for abstract DTypes)
  *
+ * - PyArray_GetDefaultDescr:
+ *
+ *     Given a DType class, returns the default instance (descriptor).
+ *     This is an inline function checking for `singleton` first and only
+ *     calls the `default_descr` function if necessary.
+ *
+ * - PyArray_DoubleDType, etc.:
+ *
+ *     Aliases to the DType classes for the builtin NumPy DTypes.
+ *
  * WARNING
  * =======
  *
@@ -208,6 +224,21 @@ typedef PyObject *_ufunc_addloop_fromspec_func(
     (*(_ufunc_addloop_fromspec_func *)(__experimental_dtype_api_table[0]))
 
 
+/* Please see the NumPy definitions in `array_method.h` for details on these */
+typedef int translate_given_descrs_func(int nin, int nout,
+        PyArray_DTypeMeta *wrapped_dtypes[],
+        PyArray_Descr *given_descrs[], PyArray_Descr *new_descrs[]);
+typedef int translate_loop_descrs_func(int nin, int nout,
+        PyArray_DTypeMeta *new_dtypes[], PyArray_Descr *given_descrs[],
+        PyArray_Descr *original_descrs[], PyArray_Descr *loop_descrs[]);
+
+typedef int _ufunc_wrapping_loop_func(PyObject *ufunc_obj,
+        PyArray_DTypeMeta *new_dtypes[], PyArray_DTypeMeta *wrapped_dtypes[],
+        translate_given_descrs_func *translate_given_descrs,
+        translate_loop_descrs_func *translate_loop_descrs);
+#define PyUFunc_AddWrappingLoop \
+    (*(_ufunc_wrapping_loop_func *)(__experimental_dtype_api_table[7]))
+
 /*
  * Type of the C promoter function, which must be wrapped into a
  * PyCapsule with name "numpy._ufunc_promoter".
@@ -371,17 +402,6 @@ PyArray_GetDefaultDescr(PyArray_DTypeMeta *DType)
 }
 
 
-typedef PyArray_DTypeMeta *get_builtin_dtype_from_typenum(int);
-/*
- * Helper to fetch a builtin NumPy DType using the type number API.
- * This function cannot fail, but must only be used together with NPY_DOUBLE,
- * etc.
- * Eventually, we may expose these directly in the API
- */
-#define PyArray_DTypeFromTypeNum \
-    ((get_builtin_dtype_from_typenum *)(__experimental_dtype_api_table[7]))
-
-
 /*
  * NumPy's builtin DTypes:
  * TODO: Should these be dereferenced: `(&(PyArray_DTypeMeta *)table[10]`?

numpy/core/setup.py

@@ -1009,6 +1009,7 @@ def generate_umath_doc_header(ext, build_dir):
             join('src', 'umath', 'clip.cpp'),
             join('src', 'umath', 'dispatching.c'),
             join('src', 'umath', 'legacy_array_method.c'),
+            join('src', 'umath', 'wrapping_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

@@ -465,6 +465,15 @@ arraymethod_dealloc(PyObject *self)
 
     PyMem_Free(meth->name);
 
+    if (meth->wrapped_meth != NULL) {
+        /* Cleanup for wrapping array method (defined in umath) */
+        Py_DECREF(meth->wrapped_meth);
+        for (int i = 0; i < meth->nin + meth->nout; i++) {
+            Py_XDECREF(meth->wrapped_dtypes[i]);
+        }
+        PyMem_Free(meth->wrapped_dtypes);
+    }
+
     Py_TYPE(self)->tp_free(self);
 }
 

numpy/core/src/multiarray/array_method.h

@@ -83,6 +83,58 @@ typedef int (get_loop_function)(
         NPY_ARRAYMETHOD_FLAGS *flags);
 
 
+/*
+ * The following functions are only used be the wrapping array method defined
+ * in umath/wrapping_array_method.c
+ */
+
+/*
+ * The function to convert the given descriptors (passed in to
+ * `resolve_descriptors`) and translates them for the wrapped loop.
+ * The new descriptors MUST be viewable with the old ones, `NULL` must be
+ * supported (for outputs) and should normally be forwarded.
+ *
+ * The function must clean up on error.
+ *
+ * NOTE: We currently assume that this translation gives "viewable" results.
+ *       I.e. there is no additional casting related to the wrapping process.
+ *       In principle that could be supported, but not sure it is useful.
+ *       This currently also means that e.g. alignment must apply identically
+ *       to the new dtypes.
+ *
+ * TODO: Due to the fact that `resolve_descriptors` is also used for `can_cast`
+ *       there is no way to "pass out" the result of this function.  This means
+ *       it will be called twice for every ufunc call.
+ *       (I am considering including `auxdata` as an "optional" parameter to
+ *       `resolve_descriptors`, so that it can be filled there if not NULL.)
+ */
+typedef int translate_given_descrs_func(int nin, int nout,
+        PyArray_DTypeMeta *wrapped_dtypes[],
+        PyArray_Descr *given_descrs[], PyArray_Descr *new_descrs[]);
+
+/**
+ * The function to convert the actual loop descriptors (as returned by the
+ * original `resolve_descriptors` function) to the ones the output array
+ * should use.
+ * This function must return "viewable" types, it must not mutate them in any
+ * form that would break the inner-loop logic.  Does not need to support NULL.
+ *
+ * The function must clean up on error.
+ *
+ * @param nargs Number of arguments
+ * @param new_dtypes The DTypes of the output (usually probably not needed)
+ * @param given_descrs Original given_descrs to the resolver, necessary to
+ *        fetch any information related to the new dtypes from the original.
+ * @param original_descrs The `loop_descrs` returned by the wrapped loop.
+ * @param loop_descrs The output descriptors, compatible to `original_descrs`.
+ *
+ * @returns 0 on success, -1 on failure.
+ */
+typedef int translate_loop_descrs_func(int nin, int nout,
+        PyArray_DTypeMeta *new_dtypes[], PyArray_Descr *given_descrs[],
+        PyArray_Descr *original_descrs[], PyArray_Descr *loop_descrs[]);
+
+
 /*
  * This struct will be public and necessary for creating a new ArrayMethod
  * object (casting and ufuncs).
@@ -125,6 +177,11 @@ typedef struct PyArrayMethodObject_tag {
     PyArrayMethod_StridedLoop *contiguous_loop;
     PyArrayMethod_StridedLoop *unaligned_strided_loop;
     PyArrayMethod_StridedLoop *unaligned_contiguous_loop;
+    /* Chunk only used for wrapping array method defined in umath */
+    struct PyArrayMethodObject_tag *wrapped_meth;
+    PyArray_DTypeMeta **wrapped_dtypes;
+    translate_given_descrs_func *translate_given_descrs;
+    translate_loop_descrs_func *translate_loop_descrs;
 } PyArrayMethodObject;
 
 

numpy/core/src/multiarray/experimental_public_dtype_api.c

@@ -332,6 +332,16 @@ PyUFunc_AddLoopFromSpec(PyObject *ufunc, PyArrayMethod_Spec *spec)
     return PyUFunc_AddLoop((PyUFuncObject *)ufunc, info, 0);
 }
 
+/*
+ * Function is defined in umath/wrapping_array_method.c
+ * (same/one compilation unit)
+ */
+NPY_NO_EXPORT int
+PyUFunc_AddWrappingLoop(PyObject *ufunc_obj,
+        PyArray_DTypeMeta *new_dtypes[], PyArray_DTypeMeta *wrapped_dtypes[],
+        translate_given_descrs_func *translate_given_descrs,
+        translate_loop_descrs_func *translate_loop_descrs);
+
 
 static int
 PyUFunc_AddPromoter(
@@ -382,7 +392,7 @@ _get_experimental_dtype_api(PyObject *NPY_UNUSED(mod), PyObject *arg)
             &PyArray_CommonDType,
             &PyArray_PromoteDTypeSequence,
             &_PyArray_GetDefaultDescr,
-            NULL,
+            &PyUFunc_AddWrappingLoop,
             NULL,
             NULL,
             /* NumPy's builtin DTypes (starting at offset 10 going to 41) */