Avoid heap allocation for function calls with a small number of arguments

We don't have access to llvm::SmallVector or similar, but given the
limited subset of the `std::vector` API that
`function_call::args{,_convert}` need and the "reserve-then-fill"
usage pattern, it is relatively straightforward to implement custom
containers that get the job done.

Seems to improves time to call the collatz function in
pybind/pybind11_benchmark significantly; numbers are a little noisy
but there's a clear improvement from "about 60 ns per call" to "about
45 ns per call" on my machine (M4 Max Mac), as measured with
`timeit.repeat('collatz(4)', 'from pybind11_benchmark import
collatz')`.

Author: swolchok

include/pybind11/cast.h

@@ -27,6 +27,7 @@
 #include <tuple>
 #include <type_traits>
 #include <utility>
+#include <variant>
 #include <vector>
 
 PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
@@ -2037,6 +2038,226 @@ using is_pos_only = std::is_same<intrinsic_t<T>, pos_only>;
 // forward declaration (definition in attr.h)
 struct function_record;
 
+// small_vector-like container to avoid heap allocation for N or fewer
+// arguments.
+template <std::size_t N>
+struct argument_vector {
+    struct handle_array {
+        std::array<handle, N> arr;
+        std::size_t size = 0;
+    };
+
+    using handle_vector = std::vector<handle>;
+
+public:
+    argument_vector() = default;
+
+    argument_vector(const argument_vector &) = delete;
+    argument_vector &operator=(const argument_vector &) = delete;
+    argument_vector(argument_vector &&) noexcept = default;
+    argument_vector &operator=(argument_vector &&) noexcept = default;
+
+    std::size_t size() const {
+        if (std::holds_alternative<handle_array>(repr_)) {
+            return std::get<handle_array>(repr_).size;
+        } else if (std::holds_alternative<handle_vector>(repr_)) {
+            return std::get<handle_vector>(repr_).size();
+        } else {
+            // NOTE: we know that our variant is never going to be
+            // valueless_by_exception() because the relevant move operations
+            // don't throw. Marking this case unreachable allows better code
+            // generation.
+            PYBIND11_UNREACHABLE();
+        }
+    }
+
+    handle &operator[](std::size_t idx) {
+        assert(idx < size());
+        if (std::holds_alternative<handle_array>(repr_)) {
+            return std::get<handle_array>(repr_).arr[idx];
+        } else if (std::holds_alternative<handle_vector>(repr_)) {
+            return std::get<handle_vector>(repr_)[idx];
+        } else {
+            PYBIND11_UNREACHABLE();
+        }
+    }
+
+    handle operator[](std::size_t idx) const {
+        assert(idx < size());
+        if (std::holds_alternative<handle_array>(repr_)) {
+            return std::get<handle_array>(repr_).arr[idx];
+        } else if (std::holds_alternative<handle_vector>(repr_)) {
+            return std::get<handle_vector>(repr_)[idx];
+        } else {
+            PYBIND11_UNREACHABLE();
+        }
+    }
+
+    void push_back(handle x) {
+        if (std::holds_alternative<handle_array>(repr_)) {
+            auto &ha = std::get<handle_array>(repr_);
+            if (ha.size == N) {
+                move_to_vector_with_reserved_size(N + 1);
+                std::get<handle_vector>(repr_).push_back(x);
+            } else {
+                ha.arr[ha.size++] = x;
+            }
+        } else if (std::holds_alternative<handle_vector>(repr_)) {
+            std::get<handle_vector>(repr_).push_back(x);
+        } else {
+            PYBIND11_UNREACHABLE();
+        }
+    }
+
+    template <typename Arg>
+    void emplace_back(Arg &&x) {
+        push_back(handle(x));
+    }
+
+    void reserve(std::size_t sz) {
+        if (std::holds_alternative<handle_array>(repr_)) {
+            if (sz > N) {
+                move_to_vector_with_reserved_size(sz);
+            }
+        } else if (std::holds_alternative<handle_vector>(repr_)) {
+            std::get<handle_vector>(repr_).reserve(sz);
+        } else {
+            PYBIND11_UNREACHABLE();
+        }
+    }
+
+private:
+    void move_to_vector_with_reserved_size(std::size_t reserved_size) {
+        auto &ha = std::get<handle_array>(repr_);
+        handle_vector hv;
+        hv.reserve(reserved_size);
+        std::copy(ha.arr.begin(), ha.arr.begin() + ha.size, std::back_inserter(hv));
+        repr_ = std::move(hv);
+    }
+    std::variant<handle_array, handle_vector> repr_;
+};
+
+// small_vector-like container to avoid heap allocation for N or fewer
+// arguments.
+template <std::size_t kRequestedInlineSize>
+struct args_convert_vector {
+private:
+    static constexpr auto kBitsPerWord = 8 * sizeof(std::size_t);
+    static constexpr auto kWords = (kRequestedInlineSize + kBitsPerWord - 1) / kBitsPerWord;
+    static constexpr auto kInlineSize = kWords * kBitsPerWord;
+    struct inline_array {
+        std::array<std::size_t, kWords> arr;
+        std::size_t size = 0;
+    };
+    using heap_array = std::vector<bool>;
+
+public:
+    args_convert_vector() = default;
+
+    args_convert_vector(const args_convert_vector &) = delete;
+    args_convert_vector &operator=(const args_convert_vector &) = delete;
+    args_convert_vector(args_convert_vector &&) noexcept = default;
+    args_convert_vector &operator=(args_convert_vector &&) noexcept = default;
+
+    args_convert_vector(std::size_t count, bool value) {
+        if (count > kInlineSize) {
+            repr_ = heap_array(count, value);
+        } else {
+            auto &inline_arr = std::get<inline_array>(repr_);
+            inline_arr.arr.fill(value ? std::size_t(-1) : 0);
+            inline_arr.size = count;
+        }
+    }
+
+    std::size_t size() const {
+        if (std::holds_alternative<inline_array>(repr_)) {
+            return std::get<inline_array>(repr_).size;
+        } else if (std::holds_alternative<heap_array>(repr_)) {
+            return std::get<heap_array>(repr_).size();
+        } else {
+            PYBIND11_UNREACHABLE();
+        }
+    }
+
+    void reserve(std::size_t sz) {
+        if (std::holds_alternative<inline_array>(repr_)) {
+            if (sz > kInlineSize) {
+                move_to_vector_with_reserved_size(sz);
+            }
+        } else if (std::holds_alternative<heap_array>(repr_)) {
+            std::get<heap_array>(repr_).reserve(sz);
+        } else {
+            PYBIND11_UNREACHABLE();
+        }
+    }
+
+    bool operator[](std::size_t idx) const {
+        if (std::holds_alternative<inline_array>(repr_)) {
+            return inline_index(idx);
+        } else if (std::holds_alternative<heap_array>(repr_)) {
+            assert(idx < std::get<heap_array>(repr_).size());
+            return std::get<heap_array>(repr_)[idx];
+        } else {
+            PYBIND11_UNREACHABLE();
+        }
+    }
+
+    void push_back(bool b) {
+        if (std::holds_alternative<inline_array>(repr_)) {
+            auto &ha = std::get<inline_array>(repr_);
+            if (ha.size == kInlineSize) {
+                move_to_vector_with_reserved_size(kInlineSize + 1);
+                std::get<heap_array>(repr_).push_back(b);
+            } else {
+                assert(ha.size < kInlineSize);
+                const auto wbi = word_and_bit_index(ha.size++);
+                assert(wbi.word < kWords);
+                assert(wbi.bit < kBitsPerWord);
+                if (b) {
+                    ha.arr[wbi.word] |= (std::size_t(1) << wbi.bit);
+                } else {
+                    ha.arr[wbi.word] &= ~(std::size_t(1) << wbi.bit);
+                }
+                assert(operator[](ha.size - 1) == b);
+            }
+        } else if (std::holds_alternative<heap_array>(repr_)) {
+            std::get<heap_array>(repr_).push_back(b);
+        } else {
+            PYBIND11_UNREACHABLE();
+        }
+    }
+
+    void swap(args_convert_vector &rhs) { std::swap(repr_, rhs.repr_); }
+
+private:
+    struct WordAndBitIndex {
+        std::size_t word;
+        std::size_t bit;
+    };
+
+    static auto word_and_bit_index(std::size_t idx) {
+        return WordAndBitIndex{idx / kBitsPerWord, idx % kBitsPerWord};
+    }
+
+    bool inline_index(std::size_t idx) const {
+        const auto wbi = word_and_bit_index(idx);
+        assert(wbi.word < kWords);
+        assert(wbi.bit < kBitsPerWord);
+        return std::get<inline_array>(repr_).arr[wbi.word] & (std::size_t(1) << wbi.bit);
+    }
+
+    void move_to_vector_with_reserved_size(std::size_t reserved_size) {
+        auto &inline_arr = std::get<inline_array>(repr_);
+        heap_array hp;
+        hp.reserve(reserved_size);
+        for (std::size_t ii = 0; ii < inline_arr.size; ++ii) {
+            hp.push_back(inline_index(ii));
+        }
+        repr_ = std::move(hp);
+    }
+    std::variant<inline_array, heap_array> repr_;
+};
+
 /// Internal data associated with a single function call
 struct function_call {
     function_call(const function_record &f, handle p); // Implementation in attr.h
@@ -2045,10 +2266,12 @@ struct function_call {
     const function_record &func;
 
     /// Arguments passed to the function:
-    std::vector<handle> args;
+    /// (Inline size chosen mostly arbitrarily; 5 should pad function_call out to two cache lines
+    /// (16 pointers) in size.)
+    argument_vector<5> args;
 
     /// The `convert` value the arguments should be loaded with
-    std::vector<bool> args_convert;
+    args_convert_vector<5> args_convert;
 
     /// Extra references for the optional `py::args` and/or `py::kwargs` arguments (which, if
     /// present, are also in `args` but without a reference).

include/pybind11/detail/common.h

@@ -181,6 +181,18 @@
 #    define PYBIND11_MAYBE_UNUSED __attribute__((__unused__))
 #endif
 
+#if defined(__GNUC__)
+#    define PYBIND11_UNREACHABLE()                                                                \
+        assert(false);                                                                            \
+        __builtin_unreachable()
+#elif defined(_MSC_VER)
+#    define PYBIND11_UNREACHABLE()                                                                \
+        assert(false);                                                                            \
+        __assume(0)
+#else
+#    define PYBIND11_UNREACHABLE() assert(false);
+#endif
+
 // https://en.cppreference.com/w/c/chrono/localtime
 #if defined(__STDC_LIB_EXT1__) && !defined(__STDC_WANT_LIB_EXT1__)
 #    define __STDC_WANT_LIB_EXT1__

include/pybind11/pybind11.h

@@ -1048,12 +1048,12 @@ class cpp_function : public function {
                 }
 #endif
 
-                std::vector<bool> second_pass_convert;
+                args_convert_vector<5> second_pass_convert;
                 if (overloaded) {
                     // We're in the first no-convert pass, so swap out the conversion flags for a
                     // set of all-false flags.  If the call fails, we'll swap the flags back in for
                     // the conversion-allowed call below.
-                    second_pass_convert.resize(func.nargs, false);
+                    second_pass_convert = args_convert_vector<5>(func.nargs, false);
                     call.args_convert.swap(second_pass_convert);
                 }
 

tests/test_embed/CMakeLists.txt

@@ -33,7 +33,8 @@ if(PYBIND11_TEST_SMART_HOLDER)
     -DPYBIND11_RUN_TESTING_WITH_SMART_HOLDER_AS_DEFAULT_BUT_NEVER_USE_IN_PRODUCTION_PLEASE)
 endif()
 
-add_executable(test_embed catch.cpp test_interpreter.cpp test_subinterpreter.cpp)
+add_executable(test_embed catch.cpp test_args_convert_vector.cpp test_argument_vector.cpp
+                          test_interpreter.cpp test_subinterpreter.cpp)
 pybind11_enable_warnings(test_embed)
 
 target_link_libraries(test_embed PRIVATE pybind11::embed Catch2::Catch2 Threads::Threads)

tests/test_embed/test_args_convert_vector.cpp

@@ -0,0 +1,72 @@
+#include "pybind11/pybind11.h"
+#include "catch.hpp"
+
+namespace py = pybind11;
+
+// 5 is chosen to match the production value. It doesn't really matter
+// what we pick because the actual inline size is going to be 8 *
+// sizeof(size_t) given the implementation at the time of writing.
+using args_convert_vector = py::detail::args_convert_vector<5>;
+
+namespace {
+template <typename Container>
+std::vector<Container> get_sample_vectors() {
+    std::vector<Container> result;
+    result.emplace_back();
+    for (const auto sz : {0, 4, 5, 6, 31, 32, 33, 63, 64, 65}) {
+        for (const bool b : {false, true}) {
+            result.emplace_back(static_cast<std::size_t>(sz), b);
+        }
+    }
+    return result;
+}
+
+void require_vector_matches_sample(const args_convert_vector &actual,
+                                   const std::vector<bool> &expected) {
+    REQUIRE(actual.size() == expected.size());
+    for (size_t ii = 0; ii < actual.size(); ++ii) {
+        REQUIRE(actual[ii] == expected[ii]);
+    }
+}
+
+template <typename MutationFunc>
+void mutation_test_with_samples(MutationFunc mutation_func) {
+    auto sample_contents = get_sample_vectors<std::vector<bool>>();
+    auto samples = get_sample_vectors<args_convert_vector>();
+    for (size_t ii = 0; ii < samples.size(); ++ii) {
+        auto &actual = samples[ii];
+        auto &expected = sample_contents[ii];
+
+        mutation_func(actual);
+        mutation_func(expected);
+        require_vector_matches_sample(actual, expected);
+    }
+}
+} // namespace
+
+TEST_CASE("check sample args_convert_vector contents") {
+    mutation_test_with_samples([](auto &) {});
+}
+
+TEST_CASE("args_convert_vector push_back") {
+    for (const bool b : {false, true}) {
+        mutation_test_with_samples([b](auto &vec) { vec.push_back(b); });
+    }
+}
+
+TEST_CASE("args_convert_vector reserve") {
+    for (std::size_t ii = 0; ii < 4; ++ii) {
+        mutation_test_with_samples([ii](auto &vec) { vec.reserve(ii); });
+    }
+}
+
+TEST_CASE("args_convert_vector reserve then push_back") {
+    for (std::size_t ii = 0; ii < 4; ++ii) {
+        for (const bool b : {false, true}) {
+            mutation_test_with_samples([ii, b](auto &vec) {
+                vec.reserve(ii);
+                vec.push_back(b);
+            });
+        }
+    }
+}