Skip to content

Commit

Permalink
ENH: Add frozen dimensions to gufunc signatures
Browse files Browse the repository at this point in the history
  • Loading branch information
@jaimefrio @mhvk
jaimefrio authored and mhvk committed Jul 17, 2018
1 parent 977431a commit 75006e9
Show file tree
Hide file tree
Showing 3 changed files with 127 additions and 6 deletions.
5 changes: 5 additions & 0 deletions numpy/core/include/numpy/ufuncobject.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,11 @@ typedef struct _tagPyUFuncObject {
* set by nditer object.
*/
npy_uint32 iter_flags;

/*
* sizes of frozen core dimensions, or -1 if unset
*/
npy_intp *core_dim_szs;
} PyUFuncObject;

#include "arrayobject.h"
Expand Down
69 changes: 68 additions & 1 deletion numpy/core/src/umath/_umath_tests.c.src
View file
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,41 @@ static void

/**end repeat**/

char *cross1d_signature = "(3),(3)->(3)";

/**begin repeat
#TYPE=LONG,DOUBLE#
#typ=npy_long, npy_double#
*/

/*
* This implements the cross product:
* out[n, 0] = in1[n, 1]*in2[n, 2] - in1[n, 2]*in2[n, 1]
* out[n, 1] = in1[n, 2]*in2[n, 0] - in1[n, 0]*in2[n, 2]
* out[n, 2] = in1[n, 0]*in2[n, 1] - in1[n, 1]*in2[n, 0]
*/
static void
@TYPE@_cross1d(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
{
INIT_OUTER_LOOP_3
npy_intp is1=steps[0], is2=steps[1], os = steps[2];
BEGIN_OUTER_LOOP_3
char *ip1=args[0], *ip2=args[1], *op=args[2];

*(@typ@ *)op = *(@typ@ *)(ip1 + is1) * *(@typ@ *)(ip2 + 2*is2) -
*(@typ@ *)(ip1 + 2*is1) * *(@typ@ *)(ip2 + is2);
op += os;
*(@typ@ *)op = *(@typ@ *)(ip1 + 2*is1) * *(@typ@ *)ip2 -
*(@typ@ *)ip1 * *(@typ@ *)(ip2 + 2*is2);
op += os;
*(@typ@ *)op = *(@typ@ *)ip1 * *(@typ@ *)(ip2 + is2) -
*(@typ@ *)(ip1 + is1) * *(@typ@ *)ip2;
END_OUTER_LOOP
}

/**end repeat**/

char *euclidean_pdist_signature = "(n,d)->(p)";

/**begin repeat
Expand Down Expand Up @@ -285,6 +320,26 @@ static void

/**end repeat**/

/* The following lines were generated using a slightly modified
version of code_generators/generate_umath.py and adding these
lines to defdict:
defdict = {
'inner1d' :
Ufunc(2, 1, None_,
r'''inner on the last dimension and broadcast on the rest \n"
" \"(i),(i)->()\" \n''',
TD('ld'),
),
'innerwt' :
Ufunc(3, 1, None_,
r'''inner1d with a weight argument \n"
" \"(i),(i),(i)->()\" \n''',
TD('ld'),
),
}
*/

static PyUFuncGenericFunction inner1d_functions[] = { LONG_inner1d, DOUBLE_inner1d };
static void * inner1d_data[] = { (void *)NULL, (void *)NULL };
Expand All @@ -295,7 +350,9 @@ static char innerwt_signatures[] = { NPY_LONG, NPY_LONG, NPY_LONG, NPY_LONG, NPY
static PyUFuncGenericFunction matrix_multiply_functions[] = { LONG_matrix_multiply, FLOAT_matrix_multiply, DOUBLE_matrix_multiply };
static void *matrix_multiply_data[] = { (void *)NULL, (void *)NULL, (void *)NULL };
static char matrix_multiply_signatures[] = { NPY_LONG, NPY_LONG, NPY_LONG, NPY_FLOAT, NPY_FLOAT, NPY_FLOAT, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE };

static PyUFuncGenericFunction cross1d_functions[] = { LONG_cross1d, DOUBLE_cross1d };
static void * cross1d_data[] = { (void *)NULL, (void *)NULL };
static char cross1d_signatures[] = { NPY_LONG, NPY_LONG, NPY_LONG, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE };
static PyUFuncGenericFunction euclidean_pdist_functions[] =
{ FLOAT_euclidean_pdist, DOUBLE_euclidean_pdist };
static void *eucldiean_pdist_data[] = { (void *)NULL, (void *)NULL };
Expand Down Expand Up @@ -376,6 +433,16 @@ addUfuncs(PyObject *dictionary) {
}
PyDict_SetItemString(dictionary, "inner1d_no_doc", f);
Py_DECREF(f);
f = PyUFunc_FromFuncAndDataAndSignature(cross1d_functions, cross1d_data,
cross1d_signatures, 2, 2, 1, PyUFunc_None, "cross1d",
"cross product on the last dimension and broadcast on the rest \n"\
" \"(3),(3)->(3)\" \n",
0, cross1d_signature);
if (f == NULL) {
return -1;
}
PyDict_SetItemString(dictionary, "cross1d", f);
Py_DECREF(f);

return 0;
}
Expand Down
59 changes: 54 additions & 5 deletions numpy/core/src/umath/ufunc_object.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -339,6 +339,28 @@ _is_alnum_underscore(char ch)
return _is_alpha_underscore(ch) || (ch >= '0' && ch <= '9');
}

/*
* Convert a string into a number
*/
static npy_intp
_get_size(const char* str)
{
char *stop;
#if defined(_MSC_VER)
#define strtoll _strtoi64
#endif
npy_intp size = (npy_intp)strtoll(str, &stop, 10);
#if defined(_MSC_VER)
#undef strtoll
#endif

if (stop == str || _is_alpha_underscore(*stop)) {
/* not a well formed number */
return -1;
}
return size;
}

/*
* Return the ending position of a variable name
*/
Expand Down Expand Up @@ -406,10 +428,13 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
ufunc->core_enabled = 1;
ufunc->core_num_dim_ix = 0;
ufunc->core_num_dims = PyArray_malloc(sizeof(int) * ufunc->nargs);
ufunc->core_dim_ixs = PyArray_malloc(sizeof(int) * len); /* shrink this later */
ufunc->core_offsets = PyArray_malloc(sizeof(int) * ufunc->nargs);
if (ufunc->core_num_dims == NULL || ufunc->core_dim_ixs == NULL
|| ufunc->core_offsets == NULL) {
/* The next two items will be shrunk later */
ufunc->core_dim_ixs = PyArray_malloc(sizeof(int) * len);
ufunc->core_dim_szs = PyArray_malloc(sizeof(npy_intp) * len);

if (ufunc->core_num_dims == NULL || ufunc->core_dim_ixs == NULL ||
ufunc->core_offsets == NULL || ufunc->core_dim_szs == NULL) {
PyErr_NoMemory();
goto fail;
}
Expand Down Expand Up @@ -438,8 +463,15 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
while (signature[i] != ')') {
/* loop over core dimensions */
int j = 0;
if (!_is_alpha_underscore(signature[i])) {
parse_error = "expect dimension name";
npy_intp frozen_size = -1;
if (signature[i] == '\0') {
parse_error = "unexpected end of signature string";
goto fail;
}

if (!_is_alpha_underscore(signature[i]) &&
(frozen_size = _get_size(signature + i)) < 0) {
parse_error = "expect dimension name or frozen size";
goto fail;
}
while (j < ufunc->core_num_dim_ix) {
Expand All @@ -450,6 +482,7 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
}
if (j >= ufunc->core_num_dim_ix) {
var_names[j] = signature+i;
ufunc->core_dim_szs[j] = frozen_size;
ufunc->core_num_dim_ix++;
}
ufunc->core_dim_ixs[cur_core_dim] = j;
Expand Down Expand Up @@ -494,6 +527,9 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
}
ufunc->core_dim_ixs = PyArray_realloc(ufunc->core_dim_ixs,
sizeof(int)*cur_core_dim);
// ufunc->core_dim_szs = PyArray_realloc(ufunc->core_dim_szs,
// sizeof(npy_intp)*ufunc->core_num_dim_ix);

/* check for trivial core-signature, e.g. "(),()->()" */
if (cur_core_dim == 0) {
ufunc->core_enabled = 0;
Expand Down Expand Up @@ -2379,6 +2415,17 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
goto fail;
}
}
/*
* Validate the core dimensions of all the operands,
* and collect all of the labeled core dimension sizes
* into the array 'core_dim_sizes'. Initialize them to
* 1, for example in the case where the operand broadcasts
* to a core dimension, it won't be visited.
*/
for (i = 0; i < ufunc->core_num_dim_ix; ++i) {
npy_intp frozen_size = ufunc->core_dim_szs[i];
core_dim_sizes[i] = frozen_size == -1 ? 1 : frozen_size;
}

/* Collect the lengths of the labelled core dimensions */
retval = _get_coredim_sizes(ufunc, op, core_dim_sizes, remap_axis);
Expand Down Expand Up @@ -4651,6 +4698,7 @@ PyUFunc_FromFuncAndDataAndSignature(PyUFuncGenericFunction *func, void **data,
ufunc->core_dim_ixs = NULL;
ufunc->core_offsets = NULL;
ufunc->core_signature = NULL;
ufunc->core_dim_szs = NULL;
if (signature != NULL) {
if (_parse_signature(ufunc, signature) != 0) {
Py_DECREF(ufunc);
Expand Down Expand Up @@ -4997,6 +5045,7 @@ ufunc_dealloc(PyUFuncObject *ufunc)
{
PyArray_free(ufunc->core_num_dims);
PyArray_free(ufunc->core_dim_ixs);
PyArray_free(ufunc->core_dim_szs);
PyArray_free(ufunc->core_offsets);
PyArray_free(ufunc->core_signature);
PyArray_free(ufunc->ptr);
Expand Down

0 comments on commit 75006e9

Please sign in to comment.