Use vectorcall (where possible) when calling Python functions

ChayimFriedman2 · ChayimFriedman2 · commit f5aa30a292da · 2024-08-24T21:42:27.000+03:00
This works without any changes to user code.

The way it works is by creating a methods on `IntoPy` to call functions, and specializing them for tuples.

This currently supports only non-kwargs for methods, and kwargs with somewhat slow approach (converting from PyDict) for functions. This can be improved, but that will require additional API.

We may consider adding more impls IntoPy&lt;Py&lt;PyTuple&gt;&gt; that specialize (for example, for arrays and `Vec`), but this i a good start.
diff --git a/newsfragments/4456.changed.md b/newsfragments/4456.changed.md
@@ -0,0 +1 @@
+Improve performance of calls to Python by using the vectorcall calling convention where possible.
diff --git a/pyo3-ffi/src/cpython/abstract_.rs b/pyo3-ffi/src/cpython/abstract_.rs
@@ -40,7 +40,7 @@ extern "C" {
 }
 
 #[cfg(Py_3_8)]
-const PY_VECTORCALL_ARGUMENTS_OFFSET: size_t =
+pub const PY_VECTORCALL_ARGUMENTS_OFFSET: size_t =
     1 << (8 * std::mem::size_of::<size_t>() as size_t - 1);
 
 #[cfg(Py_3_8)]
diff --git a/src/conversion.rs b/src/conversion.rs
@@ -1,10 +1,11 @@
 //! Defines conversions between Rust and Python types.
 use crate::err::PyResult;
+use crate::ffi_ptr_ext::FfiPtrExt;
 #[cfg(feature = "experimental-inspect")]
 use crate::inspect::types::TypeInfo;
 use crate::pyclass::boolean_struct::False;
 use crate::types::any::PyAnyMethods;
-use crate::types::PyTuple;
+use crate::types::{PyDict, PyString, PyTuple};
 use crate::{
     ffi, Borrowed, Bound, BoundObject, Py, PyAny, PyClass, PyErr, PyObject, PyRef, PyRefMut, Python,
 };
@@ -172,6 +173,93 @@ pub trait IntoPy<T>: Sized {
     fn type_output() -> TypeInfo {
         TypeInfo::Any
     }
+
+    // The following methods are helpers to use the vectorcall API where possible.
+    // They are overridden on tuples to perform a vectorcall.
+    // Be careful when you're implementing these: they can never refer to `Bound` call methods,
+    // as those refer to these methods, so this will create an infinite recursion.
+    #[doc(hidden)]
+    #[inline]
+    fn __py_call_vectorcall1<'py>(
+        self,
+        py: Python<'py>,
+        function: Borrowed<'_, 'py, PyAny>,
+        _: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>>
+    where
+        Self: IntoPy<Py<PyTuple>>,
+    {
+        #[inline]
+        fn inner<'py>(
+            py: Python<'py>,
+            function: Borrowed<'_, 'py, PyAny>,
+            args: Bound<'py, PyTuple>,
+        ) -> PyResult<Bound<'py, PyAny>> {
+            unsafe {
+                ffi::PyObject_Call(function.as_ptr(), args.as_ptr(), std::ptr::null_mut())
+                    .assume_owned_or_err(py)
+            }
+        }
+        inner(
+            py,
+            function,
+            <Self as IntoPy<Py<PyTuple>>>::into_py(self, py).into_bound(py),
+        )
+    }
+
+    #[doc(hidden)]
+    #[inline]
+    fn __py_call_vectorcall<'py>(
+        self,
+        py: Python<'py>,
+        function: Borrowed<'_, 'py, PyAny>,
+        kwargs: Option<Borrowed<'_, '_, PyDict>>,
+        _: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>>
+    where
+        Self: IntoPy<Py<PyTuple>>,
+    {
+        #[inline]
+        fn inner<'py>(
+            py: Python<'py>,
+            function: Borrowed<'_, 'py, PyAny>,
+            args: Bound<'py, PyTuple>,
+            kwargs: Option<Borrowed<'_, '_, PyDict>>,
+        ) -> PyResult<Bound<'py, PyAny>> {
+            unsafe {
+                ffi::PyObject_Call(
+                    function.as_ptr(),
+                    args.as_ptr(),
+                    kwargs.map_or_else(std::ptr::null_mut, |kwargs| kwargs.as_ptr()),
+                )
+                .assume_owned_or_err(py)
+            }
+        }
+        inner(
+            py,
+            function,
+            <Self as IntoPy<Py<PyTuple>>>::into_py(self, py).into_bound(py),
+            kwargs,
+        )
+    }
+
+    #[doc(hidden)]
+    #[inline]
+    fn __py_call_method_vectorcall1<'py>(
+        self,
+        _py: Python<'py>,
+        object: Borrowed<'_, 'py, PyAny>,
+        method_name: Bound<'py, PyString>,
+        _: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>>
+    where
+        Self: IntoPy<Py<PyTuple>>,
+    {
+        // Don't `self.into_py()`! This will lose the optimization of vectorcall.
+        object
+            .getattr(method_name)
+            .and_then(|method| method.call1(self))
+    }
 }
 
 /// Defines a conversion from a Rust type to a Python object, which may fail.
@@ -502,6 +590,52 @@ impl IntoPy<Py<PyTuple>> for () {
     fn into_py(self, py: Python<'_>) -> Py<PyTuple> {
         PyTuple::empty(py).unbind()
     }
+
+    #[inline]
+    fn __py_call_vectorcall1<'py>(
+        self,
+        py: Python<'py>,
+        function: Borrowed<'_, 'py, PyAny>,
+        _: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        unsafe { ffi::compat::PyObject_CallNoArgs(function.as_ptr()).assume_owned_or_err(py) }
+    }
+
+    #[inline]
+    fn __py_call_vectorcall<'py>(
+        self,
+        py: Python<'py>,
+        function: Borrowed<'_, 'py, PyAny>,
+        kwargs: Option<Borrowed<'_, '_, PyDict>>,
+        _: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        unsafe {
+            match kwargs {
+                Some(kwargs) => ffi::PyObject_Call(
+                    function.as_ptr(),
+                    PyTuple::empty(py).as_ptr(),
+                    kwargs.as_ptr(),
+                )
+                .assume_owned_or_err(py),
+                None => ffi::compat::PyObject_CallNoArgs(function.as_ptr()).assume_owned_or_err(py),
+            }
+        }
+    }
+
+    #[inline]
+    #[allow(clippy::used_underscore_binding)]
+    fn __py_call_method_vectorcall1<'py>(
+        self,
+        py: Python<'py>,
+        object: Borrowed<'_, 'py, PyAny>,
+        method_name: Bound<'py, PyString>,
+        _: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        unsafe {
+            ffi::compat::PyObject_CallMethodNoArgs(object.as_ptr(), method_name.as_ptr())
+                .assume_owned_or_err(py)
+        }
+    }
 }
 
 impl<'py> IntoPyObject<'py> for () {
diff --git a/src/types/any.rs b/src/types/any.rs
@@ -1,5 +1,5 @@
 use crate::class::basic::CompareOp;
-use crate::conversion::{AsPyPointer, FromPyObjectBound, IntoPy, ToPyObject};
+use crate::conversion::{private, AsPyPointer, FromPyObjectBound, IntoPy, ToPyObject};
 use crate::err::{DowncastError, DowncastIntoError, PyErr, PyResult};
 use crate::exceptions::{PyAttributeError, PyTypeError};
 use crate::ffi_ptr_ext::FfiPtrExt;
@@ -1160,33 +1160,24 @@ impl<'py> PyAnyMethods<'py> for Bound<'py, PyAny> {
         args: impl IntoPy<Py<PyTuple>>,
         kwargs: Option<&Bound<'_, PyDict>>,
     ) -> PyResult<Bound<'py, PyAny>> {
-        fn inner<'py>(
-            any: &Bound<'py, PyAny>,
-            args: Bound<'_, PyTuple>,
-            kwargs: Option<&Bound<'_, PyDict>>,
-        ) -> PyResult<Bound<'py, PyAny>> {
-            unsafe {
-                ffi::PyObject_Call(
-                    any.as_ptr(),
-                    args.as_ptr(),
-                    kwargs.map_or(std::ptr::null_mut(), |dict| dict.as_ptr()),
-                )
-                .assume_owned_or_err(any.py())
-            }
-        }
-
-        let py = self.py();
-        inner(self, args.into_py(py).into_bound(py), kwargs)
+        args.__py_call_vectorcall(
+            self.py(),
+            self.as_borrowed(),
+            kwargs.map(Bound::as_borrowed),
+            private::Token,
+        )
     }
 
+    #[inline]
     fn call0(&self) -> PyResult<Bound<'py, PyAny>> {
         unsafe { ffi::compat::PyObject_CallNoArgs(self.as_ptr()).assume_owned_or_err(self.py()) }
     }
 
     fn call1(&self, args: impl IntoPy<Py<PyTuple>>) -> PyResult<Bound<'py, PyAny>> {
-        self.call(args, None)
+        args.__py_call_vectorcall1(self.py(), self.as_borrowed(), private::Token)
     }
 
+    #[inline]
     fn call_method<N, A>(
         &self,
         name: N,
@@ -1197,10 +1188,16 @@ impl<'py> PyAnyMethods<'py> for Bound<'py, PyAny> {
         N: IntoPy<Py<PyString>>,
         A: IntoPy<Py<PyTuple>>,
     {
-        self.getattr(name)
-            .and_then(|method| method.call(args, kwargs))
+        // Don't `args.into_py()`! This will lose the optimization of vectorcall.
+        match kwargs {
+            Some(_) => self
+                .getattr(name)
+                .and_then(|method| method.call(args, kwargs)),
+            None => self.call_method1(name, args),
+        }
     }
 
+    #[inline]
     fn call_method0<N>(&self, name: N) -> PyResult<Bound<'py, PyAny>>
     where
         N: IntoPy<Py<PyString>>,
@@ -1218,7 +1215,12 @@ impl<'py> PyAnyMethods<'py> for Bound<'py, PyAny> {
         N: IntoPy<Py<PyString>>,
         A: IntoPy<Py<PyTuple>>,
     {
-        self.call_method(name, args, None)
+        args.__py_call_method_vectorcall1(
+            self.py(),
+            self.as_borrowed(),
+            name.into_py(self.py()).into_bound(self.py()),
+            private::Token,
+        )
     }
 
     fn is_truthy(&self) -> PyResult<bool> {
diff --git a/src/types/tuple.rs b/src/types/tuple.rs

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+Improve performance of calls to Python by using the vectorcall calling convention where possible.`
Original file line number	Diff line number	Diff line change
`@@ -40,7 +40,7 @@ extern "C" {`
`40`	`40`	`}`
`41`	`41`
`42`	`42`	`#[cfg(Py_3_8)]`
`43`		`-const PY_VECTORCALL_ARGUMENTS_OFFSET: size_t =`
	`43`	`+pub const PY_VECTORCALL_ARGUMENTS_OFFSET: size_t =`
`44`	`44`	`1 << (8 * std::mem::size_of::<size_t>() as size_t - 1);`
`45`	`45`
`46`	`46`	`#[cfg(Py_3_8)]`