Add fast_type_map, use it authoritatively for local types and as a hint for global types

nanobind has a similar two-level lookup strategy, added and explained
by
wjakob/nanobind@b515b1f

In this PR I've ported this approach to pybind11. To avoid an ABI
break, I've kept the fast maps to the `local_internals`. I think this
should be safe because any particular module should see its
`local_internals` reset at least as often as the global `internals`,
and misses in the fast "hint" map for global types fall back to the
global `internals`.

Performance seems to have improved. Using my patched fork of
pybind11_benchmark
(https://github.com/swolchok/pybind11_benchmark/tree/benchmark-updates,
specifically commit hash b6613d12607104d547b1c10a8145d1b3e9937266), I
run bench.py and observe the MyInt case. Each time, I do 3 runs and
just report all 3.

master, Mac: 75.9, 76.9, 75.3 nsec/loop
this PR, Mac: 73.8, 73.8, 73.6 nsec/loop
master, Linux box: 188, 187, 188 nsec/loop
this PR, Linux box: 164, 165, 164 nsec/loop

Note that the "real" percentage improvement is larger than implied by the
above because master does not yet include pybind#5824.

Author: swolchok

include/pybind11/detail/class.h

@@ -221,10 +221,12 @@ extern "C" inline void pybind11_meta_dealloc(PyObject *obj) {
             auto tindex = std::type_index(*tinfo->cpptype);
             internals.direct_conversions.erase(tindex);
 
+            auto &local_internals = get_local_internals();
             if (tinfo->module_local) {
-                get_local_internals().registered_types_cpp.erase(tindex);
+                local_internals.registered_types_cpp.erase(tinfo->cpptype);
             } else {
                 internals.registered_types_cpp.erase(tindex);
+                local_internals.global_registered_types_cpp_fast.erase(tinfo->cpptype);
             }
             internals.registered_types_py.erase(tinfo->type);
 

include/pybind11/detail/internals.h

@@ -177,6 +177,12 @@ struct type_equal_to {
 };
 #endif
 
+template <typename value_type>
+// REVIEW: do we need to add a fancy hash for pointers or is the
+// possible identity hash function from the standard library (e.g.,
+// libstdc++) sufficient?
+using fast_type_map = std::unordered_map<const std::type_info *, value_type>;
+
 template <typename value_type>
 using type_map = std::unordered_map<std::type_index, value_type, type_hash, type_equal_to>;
 
@@ -302,7 +308,15 @@ struct internals {
 // impact any other modules, because the only things accessing the local internals is the
 // module that contains them.
 struct local_internals {
-    type_map<type_info *> registered_types_cpp;
+    // It should be safe to use fast_type_map here because this entire
+    // data structure is scoped to our single module, and thus a single
+    // DSO and single instance of type_info for any particular type.
+    fast_type_map<type_info *> registered_types_cpp;
+
+    // fast hint for the *global* internals registered_types_cpp
+    // map. If we lookup successfully, that's the right answer;
+    // otherwise we go to the global map and then backfill this one.
+    fast_type_map<type_info *> global_registered_types_cpp_fast;
     std::forward_list<ExceptionTranslator> registered_exception_translators;
     PyTypeObject *function_record_py_type = nullptr;
 };

include/pybind11/detail/type_caster_base.h

@@ -205,35 +205,57 @@ PYBIND11_NOINLINE detail::type_info *get_type_info(PyTypeObject *type) {
     return bases.front();
 }
 
-inline detail::type_info *get_local_type_info(const std::type_index &tp) {
-    auto &locals = get_local_internals().registered_types_cpp;
-    auto it = locals.find(tp);
+inline detail::type_info *get_local_type_info(const std::type_info &tp,
+                                              const local_internals &local_internals) {
+    auto &locals = local_internals.registered_types_cpp;
+    auto it = locals.find(&tp);
     if (it != locals.end()) {
         return it->second;
     }
     return nullptr;
 }
 
-inline detail::type_info *get_global_type_info(const std::type_index &tp) {
+inline detail::type_info *get_local_type_info(const std::type_info &tp) {
+    return get_local_type_info(tp, get_local_internals());
+}
+
+inline detail::type_info *get_global_type_info(const std::type_info &tp,
+                                               local_internals &local_internals) {
     return with_internals([&](internals &internals) {
         detail::type_info *type_info = nullptr;
+        auto &fast_types = local_internals.global_registered_types_cpp_fast;
         auto &types = internals.registered_types_cpp;
-        auto it = types.find(tp);
+        auto fast_it = fast_types.find(&tp);
+        if (fast_it != fast_types.end()) {
+#ifndef NDEBUG
+            auto types_it = types.find(std::type_index(tp));
+            assert(types_it != types.end());
+            assert(types_it->second == fast_it->second);
+#endif
+            return fast_it->second;
+        }
+        auto it = types.find(std::type_index(tp));
         if (it != types.end()) {
+            fast_types.emplace(&tp, it->second);
             type_info = it->second;
         }
         return type_info;
     });
 }
 
+inline detail::type_info *get_global_type_info(const std::type_info &tp) {
+    return get_global_type_info(tp, get_local_internals());
+}
+
 /// Return the type info for a given C++ type; on lookup failure can either throw or return
 /// nullptr.
-PYBIND11_NOINLINE detail::type_info *get_type_info(const std::type_index &tp,
+PYBIND11_NOINLINE detail::type_info *get_type_info(const std::type_info &tp,
                                                    bool throw_if_missing = false) {
-    if (auto *ltype = get_local_type_info(tp)) {
+    auto &local_internals = get_local_internals();
+    if (auto *ltype = get_local_type_info(tp, local_internals)) {
         return ltype;
     }
-    if (auto *gtype = get_global_type_info(tp)) {
+    if (auto *gtype = get_global_type_info(tp, local_internals)) {
         return gtype;
     }
 

include/pybind11/pybind11.h

@@ -1636,10 +1636,12 @@ class generic_type : public object {
         with_internals([&](internals &internals) {
             auto tindex = std::type_index(*rec.type);
             tinfo->direct_conversions = &internals.direct_conversions[tindex];
+            auto &local_internals = get_local_internals();
             if (rec.module_local) {
-                get_local_internals().registered_types_cpp[tindex] = tinfo;
+                local_internals.registered_types_cpp[rec.type] = tinfo;
             } else {
                 internals.registered_types_cpp[tindex] = tinfo;
+                local_internals.global_registered_types_cpp_fast[rec.type] = tinfo;
             }
 
             PYBIND11_WARNING_PUSH
@@ -2137,10 +2139,16 @@ class class_ : public detail::generic_type {
 
         if (has_alias) {
             with_internals([&](internals &internals) {
-                auto &instances = record.module_local ? get_local_internals().registered_types_cpp
-                                                      : internals.registered_types_cpp;
-                instances[std::type_index(typeid(type_alias))]
-                    = instances[std::type_index(typeid(type))];
+                auto &local_internals = get_local_internals();
+                if (record.module_local) {
+                    local_internals.registered_types_cpp[&typeid(type_alias)]
+                        = local_internals.registered_types_cpp[&typeid(type)];
+                } else {
+                    type_info *const val
+                        = internals.registered_types_cpp[std::type_index(typeid(type))];
+                    internals.registered_types_cpp[std::type_index(typeid(type_alias))] = val;
+                    local_internals.global_registered_types_cpp_fast[&typeid(type_alias)] = val;
+                }
             });
         }
         def("_pybind11_conduit_v1_", cpp_conduit_method);

tests/test_embed/external_module.cpp

@@ -6,11 +6,32 @@ namespace py = pybind11;
  * modules aren't preserved over a finalize/initialize.
  */
 
+namespace {
+// Compare unsafe_reset_internals_for_single_interpreter in
+// test_subinterpreter.cpp.
+void unsafe_reset_local_internals() {
+    // NOTE: This code is NOT SAFE unless the caller guarantees no other threads are alive
+    // NOTE: This code is tied to the precise implementation of the internals holder
+
+    // first, unref the thread local internals
+    py::detail::get_local_internals_pp_manager().unref();
+
+    // now we unref the static global singleton internals
+    py::detail::get_local_internals_pp_manager().unref();
+
+    // finally, we reload the static global singleton
+    py::detail::get_local_internals();
+}
+} // namespace
+
 PYBIND11_MODULE(external_module,
                 m,
                 py::mod_gil_not_used(),
                 py::multiple_interpreters::per_interpreter_gil()) {
-
+    // At least one test ("Single Subinterpreter") wants to reset
+    // internals. We have separate local internals because we are a
+    // separate DSO, so ours need to be reset too!
+    unsafe_reset_local_internals();
     class A {
     public:
         explicit A(int value) : v{value} {};