Skip to content

Commit

Permalink
Use vectorcall (where possible) when calling Python functions
Browse files Browse the repository at this point in the history
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<Py<PyTuple>> that specialize (for example, for arrays and `Vec`), but this i a good start.
  • Loading branch information
ChayimFriedman2 committed Aug 24, 2024
1 parent 07da500 commit f5aa30a
Show file tree
Hide file tree
Showing 5 changed files with 270 additions and 28 deletions.
1 change: 1 addition & 0 deletions newsfragments/4456.changed.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Improve performance of calls to Python by using the vectorcall calling convention where possible.
2 changes: 1 addition & 1 deletion pyo3-ffi/src/cpython/abstract_.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
136 changes: 135 additions & 1 deletion src/conversion.rs
Original file line number Diff line number Diff line change
@@ -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,
};
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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 () {
Expand Down
46 changes: 24 additions & 22 deletions src/types/any.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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,
Expand All @@ -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>>,
Expand All @@ -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> {
Expand Down
Loading

0 comments on commit f5aa30a

Please sign in to comment.