Skip to content

Commit df8442f

Browse files
committed
BUG: Support the fastcall mechanism for NumPy's ufunc methods
This is a start to support NumPy 1.21 use of FASTCALL methods in ufuncs. There are probably better ways to achieve this support. NumPy further uses `tp_vectorcall_offset` on newer versions. This PR does not add support for this (meaning that a DUfunc will be unnecessarily slow on NumPy >=1.21 when kwargs are used).
1 parent e617b39 commit df8442f

1 file changed

Lines changed: 97 additions & 2 deletions

File tree

‎numba/np/ufunc/_internal.c‎

Lines changed: 97 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,7 @@ static PyMemberDef dufunc_members[] = {
276276
*/
277277

278278
static struct _ufunc_dispatch {
279+
/* Note that the following may also hold `_PyCFunctionFastWithKeywords` */
279280
PyCFunctionWithKeywords ufunc_reduce;
280281
PyCFunctionWithKeywords ufunc_accumulate;
281282
PyCFunctionWithKeywords ufunc_reduceat;
@@ -286,7 +287,7 @@ static struct _ufunc_dispatch {
286287
} ufunc_dispatch;
287288

288289
static int
289-
init_ufunc_dispatch(void)
290+
init_ufunc_dispatch(int *numpy_uses_fastcall)
290291
{
291292
int result = 0;
292293
PyMethodDef * crnt = PyUFunc_Type.tp_methods;
@@ -329,6 +330,16 @@ init_ufunc_dispatch(void)
329330
result = -1; /* Unknown method */
330331
}
331332
if (result < 0) break;
333+
334+
/* Check whether NumPy uses fastcall (ufunc.at never uses it) */
335+
if (!strncmp(crnt_name, "at", 3)) {
336+
if (*numpy_uses_fastcall == -1) {
337+
*numpy_uses_fastcall = crnt->ml_flags & METH_FASTCALL;
338+
}
339+
else if (*numpy_uses_fastcall != (crnt->ml_flags & METH_FASTCALL)) {
340+
return -1;
341+
}
342+
}
332343
}
333344
if (result == 0) {
334345
/* Sanity check. */
@@ -344,6 +355,7 @@ init_ufunc_dispatch(void)
344355
return result;
345356
}
346357

358+
347359
static PyObject *
348360
dufunc_reduce(PyDUFuncObject * self, PyObject * args, PyObject *kws)
349361
{
@@ -368,6 +380,47 @@ dufunc_outer(PyDUFuncObject * self, PyObject * args, PyObject *kws)
368380
return ufunc_dispatch.ufunc_outer((PyObject*)self->ufunc, args, kws);
369381
}
370382

383+
384+
/*
385+
* The following are the vectorcall versions of the above, since NumPy
386+
* uses the FASTCALL/Vectorcall protocol starting with version 1.21.
387+
* The only NumPy versions supporting vectorcall use Python 3.7 or higher.
388+
*/
389+
static PyObject *
390+
dufunc_reduce_fast(PyDUFuncObject * self,
391+
PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames)
392+
{
393+
return ((_PyCFunctionFastWithKeywords)ufunc_dispatch.ufunc_reduce)(
394+
(PyObject*)self->ufunc, args, len_args, kwnames);
395+
}
396+
397+
static PyObject *
398+
dufunc_reduceat_fast(PyDUFuncObject * self,
399+
PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames)
400+
{
401+
return ((_PyCFunctionFastWithKeywords)ufunc_dispatch.ufunc_reduceat)(
402+
(PyObject*)self->ufunc, args, len_args, kwnames);
403+
}
404+
405+
406+
static PyObject *
407+
dufunc_accumulate_fast(PyDUFuncObject * self,
408+
PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames)
409+
{
410+
return ((_PyCFunctionFastWithKeywords)ufunc_dispatch.ufunc_accumulate)(
411+
(PyObject*)self->ufunc, args, len_args, kwnames);
412+
}
413+
414+
415+
static PyObject *
416+
dufunc_outer_fast(PyDUFuncObject * self,
417+
PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames)
418+
{
419+
return ((_PyCFunctionFastWithKeywords)ufunc_dispatch.ufunc_outer)(
420+
(PyObject*)self->ufunc, args, len_args, kwnames);
421+
}
422+
423+
371424
#if NPY_API_VERSION >= 0x00000008
372425
static PyObject *
373426
dufunc_at(PyDUFuncObject * self, PyObject * args)
@@ -568,6 +621,41 @@ static struct PyMethodDef dufunc_methods[] = {
568621
{NULL, NULL, 0, NULL} /* sentinel */
569622
};
570623

624+
625+
/*
626+
* If Python is new enough, NumPy may use fastcall. In that case we have to
627+
* also use fastcall for simplicity and speed.
628+
*/
629+
static struct PyMethodDef dufunc_methods_fast[] = {
630+
{"reduce",
631+
(PyCFunction)dufunc_reduce_fast,
632+
METH_FASTCALL | METH_KEYWORDS, NULL },
633+
{"accumulate",
634+
(PyCFunction)dufunc_accumulate_fast,
635+
METH_FASTCALL | METH_KEYWORDS, NULL },
636+
{"reduceat",
637+
(PyCFunction)dufunc_reduceat_fast,
638+
METH_FASTCALL | METH_KEYWORDS, NULL },
639+
{"outer",
640+
(PyCFunction)dufunc_outer_fast,
641+
METH_FASTCALL | METH_KEYWORDS, NULL},
642+
#if NPY_API_VERSION >= 0x00000008
643+
{"at",
644+
(PyCFunction)dufunc_at,
645+
METH_VARARGS, NULL},
646+
#endif
647+
{"_compile_for_args",
648+
(PyCFunction)dufunc__compile_for_args,
649+
METH_VARARGS | METH_KEYWORDS,
650+
"Abstract method: subclasses should overload _compile_for_args() to compile the ufunc at the given arguments' types."},
651+
{"_add_loop",
652+
(PyCFunction)dufunc__add_loop,
653+
METH_VARARGS,
654+
NULL},
655+
{NULL, NULL, 0, NULL} /* sentinel */
656+
};
657+
658+
571659
static PyObject *
572660
dufunc_getfrozen(PyDUFuncObject * self, void * closure)
573661
{
@@ -681,8 +769,15 @@ MOD_INIT(_internal)
681769
return MOD_ERROR_VAL;
682770

683771
PyDUFunc_Type.tp_new = PyType_GenericNew;
684-
if (init_ufunc_dispatch() <= 0)
772+
773+
int numpy_uses_fastcall = -1;
774+
if (init_ufunc_dispatch(&numpy_uses_fastcall) <= 0)
685775
return MOD_ERROR_VAL;
776+
777+
if (numpy_uses_fastcall) {
778+
PyDUFunc_Type.tp_methods = dufunc_methods_fast;
779+
}
780+
686781
if (PyType_Ready(&PyDUFunc_Type) < 0)
687782
return MOD_ERROR_VAL;
688783
Py_INCREF(&PyDUFunc_Type);

0 commit comments

Comments
 (0)