Upgrade pybind11 API calls for 3.13t (pytorch#136370)

This is a modified version of pytorch#130341 that preserve support for older pybind version.

Pull Request resolved: pytorch#136370
Approved by: https://github.com/Skylion007, https://github.com/malfet

Author: albanD

torch/csrc/dynamo/guards.cpp

@@ -26,6 +26,7 @@
 #endif
 
 #include <sstream>
+#include <tuple>
 #include <utility>
 
 // For TupleIteratorGetItemAccessor, we need a fast way to retrieve the
@@ -2461,6 +2462,26 @@ std::unique_ptr<GuardManager> make_guard_manager(
     std::string source,
     py::handle example_value,
     py::handle guard_manager_enum) {
+#if IS_PYBIND_2_13_PLUS
+  using fourobjects =
+      std::tuple<py::object, py::object, py::object, py::object>;
+  PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<fourobjects>
+      storage;
+
+  auto& [guard_manager_enum_class, base_guard_manager_enum, dict_guard_manager_enum, dict_subclass_guard_manager_enum] =
+      storage
+          .call_once_and_store_result([]() -> fourobjects {
+            py::object guard_manager_enum_class =
+                py::module_::import("torch._dynamo.guards")
+                    .attr("GuardManagerType");
+            return {
+                guard_manager_enum_class,
+                guard_manager_enum_class.attr("GUARD_MANAGER"),
+                guard_manager_enum_class.attr("DICT_GUARD_MANAGER"),
+                guard_manager_enum_class.attr("DICT_SUBCLASS_GUARD_MANAGER")};
+          })
+          .get_stored();
+#else
   static py::object guard_manager_enum_class =
       py::module_::import("torch._dynamo.guards").attr("GuardManagerType");
   static py::object base_guard_manager_enum =
@@ -2469,6 +2490,7 @@ std::unique_ptr<GuardManager> make_guard_manager(
       guard_manager_enum_class.attr("DICT_GUARD_MANAGER");
   static py::object dict_subclass_guard_manager_enum =
       guard_manager_enum_class.attr("DICT_SUBCLASS_GUARD_MANAGER");
+#endif
   if (py::isinstance<py::dict>(example_value)) {
     // The purpose of having both DictGuardManager and DictSubclassGuardManager
     // is to handle the variability in how dictionaries and their subclasses

torch/csrc/jit/python/module_python.h

@@ -3,29 +3,58 @@
 #include <pybind11/stl.h>
 #include <torch/csrc/jit/api/module.h>
 #include <torch/csrc/utils/pybind.h>
+#include <tuple>
 
 namespace py = pybind11;
 
 namespace torch::jit {
 
 inline std::optional<Module> as_module(py::handle obj) {
+#if IS_PYBIND_2_13_PLUS
+  PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
+      storage;
+  auto& ScriptModule =
+      storage
+          .call_once_and_store_result([]() -> py::object {
+            return py::module_::import("torch.jit").attr("ScriptModule");
+          })
+          .get_stored();
+#else
   static py::handle ScriptModule =
       py::module::import("torch.jit").attr("ScriptModule");
+#endif
   if (py::isinstance(obj, ScriptModule)) {
     return py::cast<Module>(obj.attr("_c"));
   }
   return std::nullopt;
 }
 
 inline std::optional<Object> as_object(py::handle obj) {
+#if IS_PYBIND_2_13_PLUS
+  PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<
+      std::tuple<py::object, py::object>>
+      storage;
+  auto& [ScriptObject, RecursiveScriptClass] =
+      storage
+          .call_once_and_store_result(
+              []() -> std::tuple<py::object, py::object> {
+                return {
+                    py::module_::import("torch").attr("ScriptObject"),
+                    py::module_::import("torch.jit")
+                        .attr("RecursiveScriptClass")};
+              })
+          .get_stored();
+#else
   static py::handle ScriptObject =
       py::module::import("torch").attr("ScriptObject");
-  if (py::isinstance(obj, ScriptObject)) {
-    return py::cast<Object>(obj);
-  }
 
   static py::handle RecursiveScriptClass =
       py::module::import("torch.jit").attr("RecursiveScriptClass");
+#endif
+
+  if (py::isinstance(obj, ScriptObject)) {
+    return py::cast<Object>(obj);
+  }
   if (py::isinstance(obj, RecursiveScriptClass)) {
     return py::cast<Object>(obj.attr("_c"));
   }

torch/csrc/jit/python/python_ivalue.h

@@ -49,8 +49,22 @@ struct C10_EXPORT ConcretePyObjectHolder final : PyObjectHolder {
     // when using C++. The reason is unclear.
     try {
       pybind11::gil_scoped_acquire ag;
+
+#if IS_PYBIND_2_13_PLUS
+      PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
+          storage;
+      auto& extractorFn =
+          storage
+              .call_once_and_store_result([]() -> py::object {
+                return py::module_::import("torch._jit_internal")
+                    .attr("_extract_tensors");
+              })
+              .get_stored();
+#else
       static py::object& extractorFn = *new py::object(
           py::module::import("torch._jit_internal").attr("_extract_tensors"));
+#endif
+
       return extractorFn(py_obj_).cast<std::vector<at::Tensor>>();
     } catch (py::error_already_set& e) {
       auto err = std::runtime_error(

torch/csrc/utils/python_arg_parser.cpp

@@ -266,10 +266,22 @@ static py::object maybe_get_registered_torch_dispatch_rule(
   // This is a static object, so we must leak the Python object
   // "release()" is used here to preserve 1 refcount on the
   // object, preventing it from ever being de-allocated by CPython.
+#if IS_PYBIND_2_13_PLUS
+  PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
+      storage;
+  py::object find_torch_dispatch_rule =
+      storage
+          .call_once_and_store_result([]() -> py::object {
+            return py::module_::import("torch._library.simple_registry")
+                .attr("find_torch_dispatch_rule");
+          })
+          .get_stored();
+#else
   static const py::handle find_torch_dispatch_rule =
       py::object(py::module_::import("torch._library.simple_registry")
                      .attr("find_torch_dispatch_rule"))
           .release();
+#endif
   auto result = find_torch_dispatch_rule(
       py::reinterpret_borrow<py::object>(torch_api_function),
       torch_dispatch_object.get_type());

torch/csrc/utils/python_symnode.cpp

@@ -4,23 +4,53 @@ namespace torch {
 
 py::handle get_symint_class() {
   // NB: leak
+#if IS_PYBIND_2_13_PLUS
+  PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
+      storage;
+  return storage
+      .call_once_and_store_result([]() -> py::object {
+        return py::module::import("torch").attr("SymInt");
+      })
+      .get_stored();
+#else
   static py::handle symint_class =
       py::object(py::module::import("torch").attr("SymInt")).release();
   return symint_class;
+#endif
 }
 
 py::handle get_symfloat_class() {
   // NB: leak
+#if IS_PYBIND_2_13_PLUS
+  PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
+      storage;
+  return storage
+      .call_once_and_store_result([]() -> py::object {
+        return py::module::import("torch").attr("SymFloat");
+      })
+      .get_stored();
+#else
   static py::handle symfloat_class =
       py::object(py::module::import("torch").attr("SymFloat")).release();
   return symfloat_class;
+#endif
 }
 
 py::handle get_symbool_class() {
   // NB: leak
+#if IS_PYBIND_2_13_PLUS
+  PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
+      storage;
+  return storage
+      .call_once_and_store_result([]() -> py::object {
+        return py::module::import("torch").attr("SymBool");
+      })
+      .get_stored();
+#else
   static py::handle symbool_class =
       py::object(py::module::import("torch").attr("SymBool")).release();
   return symbool_class;
+#endif
 }
 
 } // namespace torch