Revert "Use global variables to register the return_types namedtuples (#107000)"

This reverts commit ae8eb7a3f9.

Reverted https://github.com/pytorch/pytorch/pull/107000 on behalf of https://github.com/huydhn due to Sorry for reverting your change, but it is failing internal build ([comment](https://github.com/pytorch/pytorch/pull/107000#issuecomment-1708862325))

build.bzl

@@ -255,7 +255,6 @@ 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 = [

tools/autograd/gen_python_functions.py

@@ -359,9 +359,6 @@ 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)
 
@@ -439,24 +436,22 @@ def create_python_return_type_bindings(
 ) -> None:
     """
     Generate function to initialize and return named tuple for native functions
-    which returns named tuple and registration invocations in `python_return_types.cpp`.
+    which returns named tuple and relevant entry for the map in `python_return_types.cpp`.
     """
     py_return_types_definition: List[str] = []
-    py_return_types_registrations: List[str] = []
+    py_return_types_map: List[str] = []
 
     grouped = group_filter_overloads(pairs, pred)
 
     for name in sorted(grouped.keys(), key=lambda x: str(x)):
         overloads = grouped[name]
-        definitions, registrations = generate_return_type_definition_and_registrations(
+        definitions, map_entries = generate_return_type_definition_and_map_entry(
             overloads
         )
         py_return_types_definition.append(
             "" if not definitions else "\n".join(definitions)
         )
-        py_return_types_registrations.append(
-            "" if not registrations else "\n".join(registrations)
-        )
+        py_return_types_map.append("" if not map_entries else "\n".join(map_entries))
 
     fm.write_with_template(
         filename,
@@ -465,39 +460,7 @@ 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_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,
+            "py_return_types_map": py_return_types_map,
         },
     )
 
@@ -720,25 +683,27 @@ def emit_namedtuple_call(
             typenames[tn_key] = typename
             typedefs.append(
                 f"""\
-static PyTypeObject* {typename} = generated::get_{name}_namedtuple();"""
+static PyTypeObject* {typename} = get_namedtuple("{name}");"""
             )
 
     return typedefs, typenames
 
 
-def generate_return_type_definition_and_registrations(
+def generate_return_type_definition_and_map_entry(
     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 registration invocations in same file.
+    and relevant entry for the map in same file.
     """
     typenames: Dict[
         str, str
     ] = {}  # map from unique name + field name lists to typedef name
-    definitions: List[str] = []  # function definition to register the typedef
-    registrations: List[str] = []  # register call for the typedef
+    definitions: List[str] = []  # function defintion to register the typedef
+    map_entries: List[
+        str
+    ] = []  # C++ map entry of <function_name, function creates it namedtuple>
 
     for overload in overloads:
         fieldnames = namedtuple_fieldnames(overload.function.func.returns)
@@ -770,42 +735,9 @@ PyTypeObject* get_{name}_namedtuple() {{
 }}
 """
             )
-            registrations.append(
-                f'addReturnType(return_types_module, "{name}", generated::get_{name}_namedtuple());'
-            )
+            map_entries.append(f'{{"{name}", get_{name}_namedtuple()}}, ')
 
-    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
+    return definitions, map_entries
 
 
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #

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/generated/python_return_types.h"
+#include "torch/csrc/autograd/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/generated/python_return_types.h"
+#include "torch/csrc/autograd/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/generated/python_return_types.h"
+#include "torch/csrc/autograd/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/generated/python_return_types.h"
+#include "torch/csrc/autograd/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,33 +4,34 @@
 #include <map>
 #include <string>
 
-#include "torch/csrc/autograd/generated/python_return_types.h"
+#include "torch/csrc/autograd/python_return_types.h"
 #include "torch/csrc/utils/structseq.h"
 #include "torch/csrc/Exceptions.h"
 
-namespace torch { namespace autograd { namespace generated {
-
+namespace {
 ${py_return_types}
-
-}}}
+}
 
 namespace torch {
 namespace autograd {
 
-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();
-  }
+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];
 }
 
 void initReturnTypes(PyObject* module) {
@@ -41,7 +42,18 @@ void initReturnTypes(PyObject* module) {
     throw python_error();
   }
 
-  ${py_return_types_registrations}
+  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();
+    }
+  }
 
   // steals a reference to return_types on success
   if (PyModule_AddObject(module, "_return_types", return_types_module) != 0) {

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/generated/python_return_types.h"
+#include "torch/csrc/autograd/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/generated/python_return_types.h"
+#include "torch/csrc/autograd/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/generated/python_return_types.h"
+#include "torch/csrc/autograd/python_return_types.h"
 
 #include <ATen/core/Tensor.h>
 #include <ATen/FuncTorchTLS.h>

torch/csrc/Module.cpp

@@ -43,7 +43,6 @@
 #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>
@@ -52,6 +51,7 @@
 #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>

torch/csrc/autograd/python_return_types.h

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