diff options
| -rw-r--r-- | numpy/core/code_generators/numpy_api.py | 1 | ||||
| -rw-r--r-- | numpy/core/include/numpy/ufuncobject.h | 3 | ||||
| -rw-r--r-- | numpy/core/src/umath/ufunc_object.c | 2 | ||||
| -rw-r--r-- | numpy/core/src/umath/ufunc_type_resolution.c | 127 |
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, |