PERF: Skip probing __array_ufunc__ for NumPy builtin scalars (#21470)

eendebakpt · web-flow · commit 79875572916a · 2022-05-11T10:02:54.000+02:00
The method `PyUFuncOverride_GetNonDefaultArrayUfunc` is expensive on numpy scalars because these objects do not have a `__array_ufunc__` set and for a missing attribute lookup cpython generates an exception that is later cleared by numpy. This is a performance bottleneck, see #21455. An issue has been submitted to cpython (python/cpython#92216). But even if this is addressed in cpython, it will take untill python 3.12+ before this will be useable by numpy. As an alternative solution, this PR adds a fast path to `PyUFuncOverride_GetNonDefaultArrayUfunc` to determine whether an object is a numpy scalar.
diff --git a/numpy/core/src/common/get_attr_string.h b/numpy/core/src/common/get_attr_string.h
@@ -42,7 +42,7 @@ _is_basic_python_type(PyTypeObject *tp)
  * on the type object, rather than on the instance itself.
  *
  * Assumes that the special method is a numpy-specific one, so does not look
- * at builtin types, nor does it look at a base ndarray.
+ * at builtin types. It does check base ndarray and numpy scalar types.
  *
  * In future, could be made more like _Py_LookupSpecial
  */
diff --git a/numpy/core/src/common/ufunc_override.c b/numpy/core/src/common/ufunc_override.c
@@ -5,6 +5,7 @@
 #include "get_attr_string.h"
 #include "npy_import.h"
 #include "ufunc_override.h"
+#include "scalartypes.h"
 
 /*
  * Check whether an object has __array_ufunc__ defined on its class and it
@@ -30,6 +31,11 @@ PyUFuncOverride_GetNonDefaultArrayUfunc(PyObject *obj)
     if (PyArray_CheckExact(obj)) {
         return NULL;
     }
+   /* Fast return for numpy scalar types */
+    if (is_anyscalar_exact(obj)) {
+        return NULL;
+    }
+
     /*
      * Does the class define __array_ufunc__? (Note that LookupSpecial has fast
      * return for basic python types, so no need to worry about those here)

Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ _is_basic_python_type(PyTypeObject *tp)`
`42`	`42`	`* on the type object, rather than on the instance itself.`
`43`	`43`	`*`
`44`	`44`	`* Assumes that the special method is a numpy-specific one, so does not look`
`45`		`- * at builtin types, nor does it look at a base ndarray.`
	`45`	`+ * at builtin types. It does check base ndarray and numpy scalar types.`
`46`	`46`	`*`
`47`	`47`	`* In future, could be made more like _Py_LookupSpecial`
`48`	`48`	`*/`