ENH: umath: Implement the default type resolution for masked loops

This function returns a loop which calls the unmasked loop on runs
of unmasked data.

4 files changed, 132 insertions, 1 deletions

diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index a789ae683..b8cde7b3b 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -364,6 +364,7 @@ ufunc_funcs_api = {
     # End 1.6 API
     'PyUFunc_DefaultTypeResolution':            39,
     'PyUFunc_ValidateCasting':                  40,
+    'PyUFunc_DefaultTypeResolutionMasked':      41,
 }
 
 # List of all the dicts which define the C API
diff --git a/numpy/core/include/numpy/ufuncobject.h b/numpy/core/include/numpy/ufuncobject.h
index adf5c7dae..0b8b15bc5 100644
--- a/numpy/core/include/numpy/ufuncobject.h
+++ b/numpy/core/include/numpy/ufuncobject.h
@@ -20,7 +20,8 @@ typedef void (*PyUFuncGenericMaskedFunction)
              npy_bool *mask,
              npy_intp *dimensions,
              npy_intp *steps,
-             void *innerloopdata);
+             npy_intp maskstep,
+             NpyAuxData *innerloopdata);
 
 /* Forward declaration for the type resolution function */
 struct _tagPyUFuncObject;
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index b7fecefc8..8300c5d5e 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -3740,6 +3740,8 @@ PyUFunc_FromFuncAndDataAndSignature(PyUFuncGenericFunction *func, void **data,
     self->userloops=NULL;
 
     self->type_resolution_function = &PyUFunc_DefaultTypeResolution;
+    self->type_resolution_masked_function =
+                            &PyUFunc_DefaultTypeResolutionMasked;
 
     if (name == NULL) {
         self->name = "?";
diff --git a/numpy/core/src/umath/ufunc_type_resolution.c b/numpy/core/src/umath/ufunc_type_resolution.c
index 7f12fcc87..52d2e7ffc 100644
--- a/numpy/core/src/umath/ufunc_type_resolution.c
+++ b/numpy/core/src/umath/ufunc_type_resolution.c
@@ -1348,6 +1348,133 @@ type_reso_error: {
     }
 }
 
+typedef struct {
+    NpyAuxData base;
+    PyUFuncGenericFunction unmasked_innerloop;
+    void *unmasked_innerloopdata;
+    int nargs;
+} _ufunc_masker_data;
+
+static NpyAuxData *
+ufunc_masker_data_clone(NpyAuxData *data)
+{
+    _ufunc_masker_data *n;
+    
+    /* Allocate a new one */
+    n = (_ufunc_masker_data *)PyArray_malloc(sizeof(_ufunc_masker_data));
+    if (n == NULL) {
+        return NULL;
+    }
+
+    /* Copy the data (unmasked data doesn't have object semantics) */
+    memcpy(n, data, sizeof(_ufunc_masker_data));
+
+    return (NpyAuxData *)n;
+}
+
+/*
+ * This function wraps a regular unmasked ufunc inner loop as a
+ * masked ufunc inner loop, only calling the function for
+ * elements where the mask is True.
+ */
+static void
+unmasked_ufunc_loop_as_masked(
+             char **args,
+             npy_bool *mask,
+             npy_intp *dimensions,
+             npy_intp *steps,
+             npy_intp maskstep,
+             NpyAuxData *innerloopdata)
+{
+    _ufunc_masker_data *data;
+    int iargs, nargs;
+    PyUFuncGenericFunction unmasked_innerloop;
+    void *unmasked_innerloopdata;
+    npy_intp loopsize, subloopsize;
+
+    /* Put the aux data into local variables */
+    data = (_ufunc_masker_data *)innerloopdata;
+    unmasked_innerloop = data->unmasked_innerloop;
+    unmasked_innerloopdata = data->unmasked_innerloopdata;
+    nargs = data->nargs;
+    loopsize = *dimensions;
+
+    /* Process the data as runs of unmasked values */
+    do {
+        /* Skip masked values */
+        subloopsize = 0;
+        while (subloopsize < loopsize && *mask == 0) {
+            ++subloopsize;
+            mask += maskstep;
+        }
+        for (iargs = 0; iargs < nargs; ++iargs) {
+            args[iargs] += subloopsize * steps[iargs];
+        }
+        loopsize -= subloopsize;
+        /*
+         * Process unmasked values (assumes unmasked loop doesn't
+         * mess with the 'args' pointer values)
+         */
+        subloopsize = 0;
+        while (subloopsize < loopsize && *mask != 0) {
+            ++subloopsize;
+            mask += maskstep;
+        }
+        unmasked_innerloop(args, &subloopsize, steps, unmasked_innerloopdata);
+        for (iargs = 0; iargs < nargs; ++iargs) {
+            args[iargs] += subloopsize * steps[iargs];
+        }
+        loopsize -= subloopsize;
+    } while (loopsize > 0);
+}
+
+
+/*UFUNC_API
+ *
+ * This function calls the unmasked type resolution function of the
+ * ufunc, then wraps it with a function which only calls the inner
+ * loop where the mask is True.
+ *
+ * Returns 0 on success, -1 on error.
+ */
+NPY_NO_EXPORT int
+PyUFunc_DefaultTypeResolutionMasked(PyUFuncObject *ufunc,
+                                NPY_CASTING casting,
+                                PyArrayObject **operands,
+                                PyObject *type_tup,
+                                PyArray_Descr **out_dtypes,
+                                PyUFuncGenericMaskedFunction *out_innerloop,
+                                NpyAuxData **out_innerloopdata)
+{
+    int retcode;
+    _ufunc_masker_data *data;
+
+    /* Create a new NpyAuxData object for the masker data */
+    data = (_ufunc_masker_data *)PyArray_malloc(sizeof(_ufunc_masker_data));
+    if (data == NULL) {
+        PyErr_NoMemory();
+        return -1;
+    }
+    memset(data, 0, sizeof(_ufunc_masker_data));
+    data->base.free = (NpyAuxData_FreeFunc *)&PyArray_free;
+    data->base.clone = &ufunc_masker_data_clone;
+    data->nargs = ufunc->nin + ufunc->nout;
+
+    /* Get the unmasked ufunc inner loop */
+    retcode = ufunc->type_resolution_function(ufunc, casting,
+                    operands, type_tup, out_dtypes,
+                    &data->unmasked_innerloop, &data->unmasked_innerloopdata);
+    if (retcode < 0) {
+        PyArray_free(data);
+        return retcode;
+    }
+
+    /* Return the loop function + aux data */
+    *out_innerloop = &unmasked_ufunc_loop_as_masked;
+    *out_innerloopdata = (NpyAuxData *)data;
+    return 0;
+}
+
 static int
 ufunc_loop_matches(PyUFuncObject *self,
                     PyArrayObject **op,