8000 ENH: Add frozen dimensions to gufunc signatures · numpy/numpy@64267ce · GitHub
[go: up one dir, main page]
Skip to content

Commit 64267ce

Browse files
jaimefriojaimefrio
committed
ENH: Add frozen dimensions to gufunc signatures
1 parent 502cb12 commit 64267ce

File tree

4 files changed

+101
-7
lines changed

4 files changed

+101
-7
lines changed

‎numpy/core/include/numpy/ufuncobject.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,11 @@ typedef struct _tagPyUFuncObject {
227227
* set by nditer object.
228228
*/
229229
npy_uint32 iter_flags;
230+
231+
/*
232+
* sizes of frozen core dimensions, or -1 if unset
233+
*/
234+
npy_intp *core_dim_szs;
230235
} PyUFuncObject;
231236

232237
#include "arrayobject.h"

‎numpy/core/src/umath/ufunc_object.c

Lines changed: 47 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -578,6 +578,28 @@ _is_alnum_underscore(char ch)
578578
return _is_alpha_underscore(ch) || (ch >= '0' && ch <= '9');
579579
}
580580

581+
/*
582+
* Convert a string into a number
583+
*/
584+
static npy_intp
585+
_get_size(const char* str)
586+
{
587+
char *stop;
588+
#if defined(_MSC_VER)
589+
#define strtoll _strtoi64
590+
#endif
591+
npy_intp size = (npy_intp)strtoll(str, &stop, 10);
592+
#if defined(_MSC_VER)
593+
#undef strtoll
594+
#endif
595+
596+
if (stop == str || _is_alpha_underscore(*stop)) {
597+
/* not a well formed number */
598+
return -1;
599+
}
600+
return size;
601+
}
602+
581603
/*
582604
* Return the ending position of a variable name
583605
*/
@@ -645,10 +667,13 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
645667
ufunc->core_enabled = 1;
646668
ufunc->core_num_dim_ix = 0;
647669
ufunc->core_num_dims = PyArray_malloc(sizeof(int) * ufunc->nargs);
648-
ufunc->core_dim_ixs = PyArray_malloc(sizeof(int) * len); /* shrink this later */
649670
ufunc->core_offsets = PyArray_malloc(sizeof(int) * ufunc->nargs);
650-
if (ufunc->core_num_dims == NULL || ufunc->core_dim_ixs == NULL
651-
|| ufunc->core_offsets == NULL) {
671+
/* The next two items will be shrunk later */
672+
ufunc->core_dim_ixs = PyArray_malloc(sizeof(int) * len);
673+
ufunc->core_dim_szs = PyArray_malloc(sizeof(npy_intp) * len);
674+
675+
if (ufunc->core_num_dims == NULL || ufunc->core_dim_ixs == NULL ||
676+
ufunc->core_offsets == NULL || ufunc->core_dim_szs == NULL) {
652677
PyErr_NoMemory();
653678
goto fail;
654679
}
@@ -677,8 +702,15 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
677702
while (signature[i] != ')') {
678703
/* loop over core dimensions */
679704
int j = 0;
680-
if (!_is_alpha_underscore(signature[i])) {
681-
parse_error = "expect dimension name";
705+
npy_intp frozen_size = -1;
706+
if (signature[i] == '\0') {
707+
parse_error = "unexpected end of signature string";
708+
goto fail;
709+
}
710+
711+
if (!_is_alpha_underscore(signature[i]) &&
712+
(frozen_size = _get_size(signature + i)) < 0) {
713+
parse_error = "expect dimension name or frozen size";
682714
goto fail;
683715
}
684716
while (j < ufunc->core_num_dim_ix) {
@@ -689,6 +721,7 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
689721
}
690722
if (j >= ufunc->core_num_dim_ix) {
691723
var_names[j] = signature+i;
724+
ufunc->core_dim_szs[j] = frozen_size;
692725
ufunc->core_num_dim_ix++;
693726
}
694727
ufunc->core_dim_ixs[cur_core_dim] = j;
@@ -733,6 +766,9 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
733766
}
734767
ufunc->core_dim_ixs = PyArray_realloc(ufunc->core_dim_ixs,
735768
sizeof(int)*cur_core_dim);
769+
// ufunc->core_dim_szs = PyArray_realloc(ufunc->core_dim_szs,
770+
// sizeof(npy_intp)*ufunc->core_num_dim_ix);
771+
736772
/* check for trivial core-signature, e.g. "(),()->()" */
737773
if (cur_core_dim == 0) {
738774
ufunc->core_enabled = 0;
@@ -1914,7 +1950,8 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
19141950
* to a core dimension, it won't be visited.
19151951
*/
19161952
for (i = 0; i < ufunc->core_num_dim_ix; ++i) {
1917-
core_dim_sizes[i] = 1;
1953+
npy_intp frozen_size = ufunc->core_dim_szs[i];
1954+
core_dim_sizes[i] = frozen_size == -1 ? 1 : frozen_size;
19181955
}
19191956
for (i = 0; i < nop; ++i) {
19201957
if (op[i] != NULL) {
@@ -1949,9 +1986,10 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
19491986
}
19501987
for (; idim < num_dims; ++idim) {
19511988
int core_dim_index = ufunc->core_dim_ixs[dim_offset + idim];
1989+
npy_intp frozen_size = ufunc->core_dim_szs[core_dim_index];
19521990
npy_intp op_dim_size =
19531991
PyArray_SHAPE(op[i])[core_start_dim + idim];
1954-
if (core_dim_sizes[core_dim_index] == 1) {
1992+
if (frozen_size == -1 && core_dim_sizes[core_dim_index] == 1) {
19551993
core_dim_sizes[core_dim_index] = op_dim_size;
19561994
} else if ((i >= nin || op_dim_size != 1) &&
19571995
core_dim_sizes[core_dim_index] != op_dim_size) {
@@ -4370,6 +4408,7 @@ PyUFunc_FromFuncAndDataAndSignature(PyUFuncGenericFunction *func, void **data,
43704408
ufunc->core_dim_ixs = NULL;
43714409
ufunc->core_offsets = NULL;
43724410
ufunc->core_signature = NULL;
4411+
ufunc->core_dim_szs = NULL;
43734412
if (signature != NULL) {
43744413
if (_parse_signature(ufunc, signature) != 0) {
43754414
Py_DECREF(ufunc);
@@ -4716,6 +4755,7 @@ ufunc_dealloc(PyUFuncObject *ufunc)
47164755
{
47174756
PyArray_free(ufunc->core_num_dims);
47184757
PyArray_free(ufunc->core_dim_ixs);
4758+
PyArray_free(ufunc->core_dim_szs);
47194759
PyArray_free(ufunc->core_offsets);
47204760
PyArray_free(ufunc->core_signature);
47214761
PyArray_free(ufunc->ptr);

‎numpy/core/src/umath/umath_tests.c.src

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,43 @@ static void
177177

178178
/**end repeat**/
179179

180+
char *cross1d_signature = "(3),(3)->(3)";
181+
182+
/**begin repeat
183+
184+
#TYPE=LONG,DOUBLE#
185+
#typ=npy_long, npy_double#
186+
*/
187+
188+
/*
189+
* This implements the cross product:
190+
* out[n, 0] = in1[n, 1]*in2[n, 2] - in1[n, 2]*in2[n, 1]
191+
* out[n, 1] = in1[n, 2]*in2[n, 0] - in1[n, 0]*in2[n, 2]
192+
* out[n, 2] = in1[n, 0]*in2[n, 1] - in1[n, 1]*in2[n, 0]
193+
*/
194+
static void
195+
@TYPE@_cross1d(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
196+
{
197+
INIT_OUTER_LOOP_3
198+
npy_intp di = dimensions[0];
199+
npy_intp i;
200+
npy_intp is1=steps[0], is2=steps[1], os = steps[2];
201+
BEGIN_OUTER_LOOP_3
202+
char *ip1=args[0], *ip2=args[1], *op=args[2];
203+
204+
*(@typ@ *)op = *(@typ@ *)(ip1 + is1) * *(@typ@ *)(ip2 + 2*is2) -
205+
*(@typ@ *)(ip1 + 2*is1) * *(@typ@ *)(ip2 + is2);
206+
op += os;
207+
*(@typ@ *)op = *(@typ@ *)(ip1 + 2*is1) * *(@typ@ *)ip2 -
208+
*(@typ@ *)ip1 * *(@typ@ *)(ip2 + 2*is2);
209+
op += os;
210+
*(@typ@ *)op = *(@typ@ *)ip1 * *(@typ@ *)(ip2 + is2) -
211+
*(@typ@ *)(ip1 + is1) * *(@typ@ *)ip2;
212+
END_OUTER_LOOP
213+
}
214+
215+
/**end repeat**/
216+
180217
/* The following lines were generated using a slightly modified
181218
version of code_generators/generate_umath.py and adding these
182219
lines to defdict:
@@ -207,6 +244,9 @@ static char innerwt_signatures[] = { NPY_LONG, NPY_LONG, NPY_LONG, NPY_LONG, NPY
207244
static PyUFuncGenericFunction matrix_multiply_functions[] = { LONG_matrix_multiply, FLOAT_matrix_multiply, DOUBLE_matrix_multiply };
208245
static void *matrix_multiply_data[] = { (void *)NULL, (void *)NULL, (void *)NULL };
209246
static char matrix_multiply_signatures[] = { NPY_LONG, NPY_LONG, NPY_LONG, NPY_FLOAT, NPY_FLOAT, NPY_FLOAT, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE };
247+
static PyUFuncGenericFunction cross1d_functions[] = { LONG_cross1d, DOUBLE_cross1d };
248+
static void * cross1d_data[] = { (void *)NULL, (void *)NULL };
249+
static char cross1d_signatures[] = { NPY_LONG, NPY_LONG, NPY_LONG, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE };
210250

211251
static void
212252
addUfuncs(PyObject *dictionary) {
@@ -234,6 +274,14 @@ addUfuncs(PyObject *dictionary) {
234274
0, matrix_multiply_signature);
235275
PyDict_SetItemString(dictionary, "matrix_multiply", f);
236276
Py_DECREF(f);
277+
f = PyUFunc_FromFuncAndDataAndSignature(cross1d_functions, cross1d_data, cross1d_signatures, 2,
278+
2, 1, PyUFunc_None, "cross1d",
279+
"cross product on the last dimension and broadcast on the rest \n"\
280+
" \"(3),(3)->(3)\" \n",
281+
0, cross1d_signature);
282+
PyDict_SetItemString(dictionary, "cross1d", f);
283+
Py_DECREF(f);
284+
237285
}
238286

239287
/*

‎numpy/core/src/umath/umathmodule.c

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ ufunc_frompyfunc(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *NPY_UNUS
121121
self->core_num_dim_ix = 0;
122122
self->core_num_dims = NULL;
123123
self->core_dim_ixs = NULL;
124+
self->core_dim_szs = NULL;
124125
self->core_offsets = NULL;
125 42A2 126
self->core_signature = NULL;
126127
self->op_flags = PyArray_malloc(sizeof(npy_uint32)*self->nargs);

0 commit comments

Comments
 (0)
0