@@ -276,6 +276,7 @@ static PyMemberDef dufunc_members[] = {
276276 */
277277
278278static 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
288289static 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+
347359static PyObject *
348360dufunc_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
372425static PyObject *
373426dufunc_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+
571659static PyObject *
572660dufunc_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