Use global variables to register the return_types namedtuples (#108832)

Fixes #69221. Builds on top of #107000, fixing the buck build issue linked [here](https://github.com/pytorch/pytorch/pull/107000#issuecomment-1708857375).

Pull Request resolved: https://github.com/pytorch/pytorch/pull/108832
Approved by: https://github.com/zou3519

build.bzl

@@ -255,6 +255,7 @@ GENERATED_CPP_CORE = [
 
 _GENERATED_AUTOGRAD_PYTHON_HEADERS = [
     "torch/csrc/autograd/generated/python_functions.h",
+    "torch/csrc/autograd/generated/python_return_types.h",
 ]
 
 _GENERATED_AUTOGRAD_CPP_HEADERS = [

caffe2/CMakeLists.txt

@@ -400,6 +400,7 @@ set(GENERATED_CXX_PYTHON
 
 set(GENERATED_H_PYTHON
   "${TORCH_SRC_DIR}/csrc/autograd/generated/python_functions.h"
+  "${TORCH_SRC_DIR}/csrc/autograd/generated/python_return_types.h"
   )
 
 set(GENERATED_TESTING_PYTHON

pt_template_srcs.bzl

@@ -146,6 +146,7 @@ def get_generate_code_bin_outs():
             "autograd/generated/python_linalg_functions.cpp": ["autograd/generated/python_linalg_functions.cpp"],
             "autograd/generated/python_nested_functions.cpp": ["autograd/generated/python_nested_functions.cpp"],
             "autograd/generated/python_nn_functions.cpp": ["autograd/generated/python_nn_functions.cpp"],
+            "autograd/generated/python_return_types.h": ["autograd/generated/python_return_types.h"],
             "autograd/generated/python_return_types.cpp": ["autograd/generated/python_return_types.cpp"],
             "autograd/generated/python_sparse_functions.cpp": ["autograd/generated/python_sparse_functions.cpp"],
             "autograd/generated/python_special_functions.cpp": ["autograd/generated/python_special_functions.cpp"],

tools/BUCK.bzl

@@ -132,6 +132,7 @@ def define_tools_targets(
             "autograd/templates/python_linalg_functions.cpp",
             "autograd/templates/python_nested_functions.cpp",
             "autograd/templates/python_nn_functions.cpp",
+            "autograd/templates/python_return_types.h",
             "autograd/templates/python_return_types.cpp",
             "autograd/templates/python_sparse_functions.cpp",
             "autograd/templates/python_special_functions.cpp",

tools/autograd/gen_python_functions.py

@@ -359,6 +359,9 @@ def gen(
     create_python_return_type_bindings(
         fm, functions, lambda fn: True, "python_return_types.cpp"
     )
+    create_python_return_type_bindings_header(
+        fm, functions, lambda fn: True, "python_return_types.h"
+    )
 
     valid_tags = parse_tags_yaml(tags_yaml_path)
 
@@ -436,22 +439,24 @@ def create_python_return_type_bindings(
 ) -> None:
     """
     Generate function to initialize and return named tuple for native functions
-    which returns named tuple and relevant entry for the map in `python_return_types.cpp`.
+    which returns named tuple and registration invocations in `python_return_types.cpp`.
     """
     py_return_types_definition: List[str] = []
-    py_return_types_map: List[str] = []
+    py_return_types_registrations: List[str] = []
 
     grouped = group_filter_overloads(pairs, pred)
 
     for name in sorted(grouped.keys(), key=lambda x: str(x)):
         overloads = grouped[name]
-        definitions, map_entries = generate_return_type_definition_and_map_entry(
+        definitions, registrations = generate_return_type_definition_and_registrations(
             overloads
         )
         py_return_types_definition.append(
             "" if not definitions else "\n".join(definitions)
         )
-        py_return_types_map.append("" if not map_entries else "\n".join(map_entries))
+        py_return_types_registrations.append(
+            "" if not registrations else "\n".join(registrations)
+        )
 
     fm.write_with_template(
         filename,
@@ -460,7 +465,39 @@ def create_python_return_type_bindings(
             "generated_comment": "@"
             + f"generated from {fm.template_dir_for_comments()}/{filename}",
             "py_return_types": py_return_types_definition,
-            "py_return_types_map": py_return_types_map,
+            "py_return_types_registrations": py_return_types_registrations,
+        },
+    )
+
+
+def create_python_return_type_bindings_header(
+    fm: FileManager,
+    pairs: Sequence[PythonSignatureNativeFunctionPair],
+    pred: Callable[[NativeFunction], bool],
+    filename: str,
+) -> None:
+    """
+    Generate function to initialize and return named tuple for native functions
+    which returns named tuple and relevant entry for the map in `python_return_types.cpp`.
+    """
+    py_return_types_declarations: List[str] = []
+
+    grouped = group_filter_overloads(pairs, pred)
+
+    for name in sorted(grouped.keys(), key=lambda x: str(x)):
+        overloads = grouped[name]
+        declarations = generate_return_type_declarations(overloads)
+        py_return_types_declarations.append(
+            "" if not declarations else "\n".join(declarations)
+        )
+
+    fm.write_with_template(
+        filename,
+        filename,
+        lambda: {
+            "generated_comment": "@"
+            + f"generated from {fm.template_dir_for_comments()}/{filename}",
+            "py_return_types_declarations": py_return_types_declarations,
         },
     )
 
@@ -683,27 +720,25 @@ def emit_namedtuple_call(
             typenames[tn_key] = typename
             typedefs.append(
                 f"""\
-static PyTypeObject* {typename} = get_namedtuple("{name}");"""
+static PyTypeObject* {typename} = generated::get_{name}_namedtuple();"""
             )
 
     return typedefs, typenames
 
 
-def generate_return_type_definition_and_map_entry(
+def generate_return_type_definition_and_registrations(
     overloads: Sequence[PythonSignatureNativeFunctionPair],
 ) -> Tuple[List[str], List[str]]:
     """
     Generate block of function in `python_return_types.cpp` to initialize
     and return named tuple for a native function which returns named tuple
-    and relevant entry for the map in same file.
+    and registration invocations in same file.
     """
     typenames: Dict[
         str, str
     ] = {}  # map from unique name + field name lists to typedef name
-    definitions: List[str] = []  # function defintion to register the typedef
-    map_entries: List[
-        str
-    ] = []  # C++ map entry of <function_name, function creates it namedtuple>
+    definitions: List[str] = []  # function definition to register the typedef
+    registrations: List[str] = []  # register call for the typedef
 
     for overload in overloads:
         fieldnames = namedtuple_fieldnames(overload.function.func.returns)
@@ -735,9 +770,42 @@ PyTypeObject* get_{name}_namedtuple() {{
 }}
 """
             )
-            map_entries.append(f'{{"{name}", get_{name}_namedtuple()}}, ')
+            registrations.append(
+                f'addReturnType(return_types_module, "{name}", generated::get_{name}_namedtuple());'
+            )
 
-    return definitions, map_entries
+    return definitions, registrations
+
+
+def generate_return_type_declarations(
+    overloads: Sequence[PythonSignatureNativeFunctionPair],
+) -> List[str]:
+    """
+    Generate block of function declarations in `python_return_types.h` to initialize
+    and return named tuple for a native function.
+    """
+    typenames: Dict[
+        str, str
+    ] = {}  # map from unique name + field name lists to typedef name
+    declarations: List[str] = []  # function declaration to register the typedef
+
+    for overload in overloads:
+        fieldnames = namedtuple_fieldnames(overload.function.func.returns)
+        if not fieldnames:
+            continue
+
+        name = cpp.name(overload.function.func)  # use @with_native_function?
+        tn_key = gen_namedtuple_typename_key(overload.function)
+        typename = typenames.get(tn_key)
+
+        if typename is None:
+            typename = (
+                f'{name}NamedTuple{"" if not declarations else len(declarations)}'
+            )
+            typenames[tn_key] = typename
+            declarations.append(f"PyTypeObject* get_{name}_namedtuple();")
+
+    return declarations
 
 
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #

tools/autograd/templates/python_fft_functions.cpp

@@ -5,7 +5,7 @@
 #include "torch/csrc/DynamicTypes.h"
 #include "torch/csrc/Exceptions.h"
 #include "torch/csrc/autograd/python_fft_functions.h"
-#include "torch/csrc/autograd/python_return_types.h"
+#include "torch/csrc/autograd/generated/python_return_types.h"
 #include "torch/csrc/autograd/python_variable.h"
 #include "torch/csrc/autograd/utils/wrap_outputs.h"
 #include "torch/csrc/autograd/utils/python_arg_parsing.h"

tools/autograd/templates/python_linalg_functions.cpp

@@ -5,7 +5,7 @@
 #include "torch/csrc/DynamicTypes.h"
 #include "torch/csrc/Exceptions.h"
 #include "torch/csrc/autograd/python_linalg_functions.h"
-#include "torch/csrc/autograd/python_return_types.h"
+#include "torch/csrc/autograd/generated/python_return_types.h"
 #include "torch/csrc/autograd/python_variable.h"
 #include "torch/csrc/autograd/utils/wrap_outputs.h"
 #include "torch/csrc/autograd/utils/python_arg_parsing.h"

tools/autograd/templates/python_nested_functions.cpp

@@ -5,7 +5,7 @@
 #include "torch/csrc/DynamicTypes.h"
 #include "torch/csrc/Exceptions.h"
 #include "torch/csrc/autograd/python_nested_functions.h"
-#include "torch/csrc/autograd/python_return_types.h"
+#include "torch/csrc/autograd/generated/python_return_types.h"
 #include "torch/csrc/autograd/python_variable.h"
 #include "torch/csrc/autograd/utils/wrap_outputs.h"
 #include "torch/csrc/autograd/utils/python_arg_parsing.h"

tools/autograd/templates/python_nn_functions.cpp

@@ -5,7 +5,7 @@
 #include "torch/csrc/DynamicTypes.h"
 #include "torch/csrc/Exceptions.h"
 #include "torch/csrc/autograd/python_nn_functions.h"
-#include "torch/csrc/autograd/python_return_types.h"
+#include "torch/csrc/autograd/generated/python_return_types.h"
 #include "torch/csrc/autograd/python_variable.h"
 #include "torch/csrc/autograd/utils/wrap_outputs.h"
 #include "torch/csrc/autograd/utils/python_arg_parsing.h"

tools/autograd/templates/python_return_types.cpp

@@ -4,34 +4,33 @@
 #include <map>
 #include <string>
 
-#include "torch/csrc/autograd/python_return_types.h"
+#include "torch/csrc/autograd/generated/python_return_types.h"
 #include "torch/csrc/utils/structseq.h"
 #include "torch/csrc/Exceptions.h"
 
-namespace {
+namespace torch { namespace autograd { namespace generated {
+
 ${py_return_types}
-}
+
+}}}
 
 namespace torch {
 namespace autograd {
 
-std::map<std::string, PyTypeObject*>& get_namedtuple_types_map() {
-  // [NOTE] Non-global map
-  // This map calls Python functions during its initialization.
-  // If it is a global static variable and in case it is loaded
-  // before Python interpreter is ready, then the calls it makes during
-  // initialization will SEGFAULT.
-  // To avoid this we make it function static variable so that it is
-  // initialized only after the Python interpreter is ready.
-  static std::map<std::string, PyTypeObject*> namedtuple_types_map = {
-    ${py_return_types_map}
-  };
-  return namedtuple_types_map;
-}
-
-PyTypeObject* get_namedtuple(std::string name) {
-  static auto& namedtuple_types_map = get_namedtuple_types_map();
-  return namedtuple_types_map[name];
+static void addReturnType(
+    PyObject* module,
+    const char* name,
+    PyTypeObject* type) {
+  // hold onto the TypeObject for the unlikely case of user
+  // deleting or overriding it.
+  Py_INCREF(type);
+  if (PyModule_AddObject(
+          module,
+          name,
+          (PyObject*)type) != 0) {
+    Py_DECREF(type);
+    throw python_error();
+  }
 }
 
 void initReturnTypes(PyObject* module) {
@@ -42,18 +41,7 @@ void initReturnTypes(PyObject* module) {
     throw python_error();
   }
 
-  for (const auto& return_type_pair : get_namedtuple_types_map()) {
-    // hold onto the TypeObject for the unlikely case of user
-    // deleting or overriding it.
-    Py_INCREF(return_type_pair.second);
-    if (PyModule_AddObject(
-            return_types_module,
-            return_type_pair.first.c_str(),
-            (PyObject*)return_type_pair.second) != 0) {
-      Py_DECREF((PyObject*)return_type_pair.second);
-      throw python_error();
-    }
-  }
+  ${py_return_types_registrations}
 
   // steals a reference to return_types on success
   if (PyModule_AddObject(module, "_return_types", return_types_module) != 0) {

tools/autograd/templates/python_return_types.h

@@ -2,8 +2,12 @@
 
 namespace torch {
 namespace autograd {
+namespace generated {
+
+${py_return_types_declarations}
+
+}
 
-PyTypeObject* get_namedtuple(std::string name);
 void initReturnTypes(PyObject* module);
 
 } // namespace autograd

tools/autograd/templates/python_special_functions.cpp

@@ -5,7 +5,7 @@
 #include "torch/csrc/DynamicTypes.h"
 #include "torch/csrc/Exceptions.h"
 #include "torch/csrc/autograd/python_special_functions.h"
-#include "torch/csrc/autograd/python_return_types.h"
+#include "torch/csrc/autograd/generated/python_return_types.h"
 #include "torch/csrc/autograd/python_variable.h"
 #include "torch/csrc/autograd/utils/wrap_outputs.h"
 #include "torch/csrc/autograd/utils/python_arg_parsing.h"

tools/autograd/templates/python_torch_functions.cpp

@@ -32,7 +32,7 @@
 #include "torch/csrc/autograd/generated/variable_factories.h"
 #include "torch/csrc/utils/structseq.h"
 #include "torch/csrc/utils/cuda_lazy_init.h"
-#include "torch/csrc/autograd/python_return_types.h"
+#include "torch/csrc/autograd/generated/python_return_types.h"
 
 #include <ATen/core/Tensor.h>
 

tools/autograd/templates/python_variable_methods.cpp

@@ -34,7 +34,7 @@
 #include "torch/csrc/utils/tensor_numpy.h"
 #include "torch/csrc/utils/tensor_types.h"
 #include "torch/csrc/utils/structseq.h"
-#include "torch/csrc/autograd/python_return_types.h"
+#include "torch/csrc/autograd/generated/python_return_types.h"
 
 #include <ATen/core/Tensor.h>
 #include <ATen/FuncTorchTLS.h>

torch/csrc/Module.cpp

@@ -43,6 +43,7 @@
 #include <torch/csrc/THP.h>
 #include <torch/csrc/TypeInfo.h>
 #include <torch/csrc/api/include/torch/python/init.h>
+#include <torch/csrc/autograd/generated/python_return_types.h>
 #include <torch/csrc/autograd/python_cpp_function.h>
 #include <torch/csrc/autograd/python_enum_tag.h>
 #include <torch/csrc/autograd/python_fft_functions.h>
@@ -51,7 +52,6 @@
 #include <torch/csrc/autograd/python_linalg_functions.h>
 #include <torch/csrc/autograd/python_nested_functions.h>
 #include <torch/csrc/autograd/python_nn_functions.h>
-#include <torch/csrc/autograd/python_return_types.h>
 #include <torch/csrc/autograd/python_sparse_functions.h>
 #include <torch/csrc/autograd/python_special_functions.h>
 #include <torch/csrc/autograd/python_variable.h>