[mypyc] Enable free threading when compiling multiple modules (#19541)

Previously we used multi-phase initialization to enable free threading
on Python builds that support it, but only if a single module was
compiled in a group.

Implements it also for multiple modules in a group. Add support for
multi-phase initialization in module shims and the shared library. It's
still only used on free-threaded builds, and we fall back to the old
approach on other Python versions/builds.

This enables compiling mypy and mypyc on free-threaded Python builds. At
least almost all mypy and mypyc tests now pass when compiled and on
3.14.0b4 with free threading (only tested on macOS so far).

Author: JukkaL

mypyc/build.py

@@ -36,7 +36,7 @@
 from mypy.util import write_junit_xml
 from mypyc.annotate import generate_annotated_html
 from mypyc.codegen import emitmodule
-from mypyc.common import RUNTIME_C_FILES, shared_lib_name
+from mypyc.common import IS_FREE_THREADED, RUNTIME_C_FILES, shared_lib_name
 from mypyc.errors import Errors
 from mypyc.ir.pprint import format_modules
 from mypyc.namegen import exported_name
@@ -176,9 +176,15 @@ def generate_c_extension_shim(
     cname = "%s.c" % full_module_name.replace(".", os.sep)
     cpath = os.path.join(dir_name, cname)
 
+    if IS_FREE_THREADED:
+        # We use multi-phase init in free-threaded builds to enable free threading.
+        shim_name = "module_shim_no_gil_multiphase.tmpl"
+    else:
+        shim_name = "module_shim.tmpl"
+
     # We load the C extension shim template from a file.
     # (So that the file could be reused as a bazel template also.)
-    with open(os.path.join(include_dir(), "module_shim.tmpl")) as f:
+    with open(os.path.join(include_dir(), shim_name)) as f:
         shim_template = f.read()
 
     write_file(

mypyc/codegen/emitmodule.py

@@ -784,28 +784,15 @@ def generate_shared_lib_init(self, emitter: Emitter) -> None:
         assert self.group_name is not None
 
         emitter.emit_line()
+
+        short_name = shared_lib_name(self.group_name).split(".")[-1]
+
         emitter.emit_lines(
-            "PyMODINIT_FUNC PyInit_{}(void)".format(
-                shared_lib_name(self.group_name).split(".")[-1]
-            ),
+            f"static int exec_{short_name}(PyObject *module)",
             "{",
-            (
-                'static PyModuleDef def = {{ PyModuleDef_HEAD_INIT, "{}", NULL, -1, NULL, NULL }};'.format(
-                    shared_lib_name(self.group_name)
-                )
-            ),
             "int res;",
             "PyObject *capsule;",
             "PyObject *tmp;",
-            "static PyObject *module;",
-            "if (module) {",
-            "Py_INCREF(module);",
-            "return module;",
-            "}",
-            "module = PyModule_Create(&def);",
-            "if (!module) {",
-            "goto fail;",
-            "}",
             "",
         )
 
@@ -827,15 +814,26 @@ def generate_shared_lib_init(self, emitter: Emitter) -> None:
 
         for mod in self.modules:
             name = exported_name(mod)
+            if self.multi_phase_init:
+                capsule_func_prefix = "CPyExec_"
+                capsule_name_prefix = "exec_"
+                emitter.emit_line(f"extern int CPyExec_{name}(PyObject *);")
+            else:
+                capsule_func_prefix = "CPyInit_"
+                capsule_name_prefix = "init_"
+                emitter.emit_line(f"extern PyObject *CPyInit_{name}(void);")
             emitter.emit_lines(
-                f"extern PyObject *CPyInit_{name}(void);",
-                'capsule = PyCapsule_New((void *)CPyInit_{}, "{}.init_{}", NULL);'.format(
-                    name, shared_lib_name(self.group_name), name
+                'capsule = PyCapsule_New((void *){}{}, "{}.{}{}", NULL);'.format(
+                    capsule_func_prefix,
+                    name,
+                    shared_lib_name(self.group_name),
+                    capsule_name_prefix,
+                    name,
                 ),
                 "if (!capsule) {",
                 "goto fail;",
                 "}",
-                f'res = PyObject_SetAttrString(module, "init_{name}", capsule);',
+                f'res = PyObject_SetAttrString(module, "{capsule_name_prefix}{name}", capsule);',
                 "Py_DECREF(capsule);",
                 "if (res < 0) {",
                 "goto fail;",
@@ -861,7 +859,56 @@ def generate_shared_lib_init(self, emitter: Emitter) -> None:
                 "",
             )
 
-        emitter.emit_lines("return module;", "fail:", "Py_XDECREF(module);", "return NULL;", "}")
+        emitter.emit_lines("return 0;", "fail:", "return -1;", "}")
+
+        if self.multi_phase_init:
+            emitter.emit_lines(
+                f"static PyModuleDef_Slot slots_{short_name}[] = {{",
+                f"{{Py_mod_exec, exec_{short_name}}},",
+                "{Py_mod_multiple_interpreters, Py_MOD_MULTIPLE_INTERPRETERS_NOT_SUPPORTED},",
+                "{Py_mod_gil, Py_MOD_GIL_NOT_USED},",
+                "{0, NULL},",
+                "};",
+            )
+
+        size = 0 if self.multi_phase_init else -1
+        emitter.emit_lines(
+            f"static PyModuleDef module_def_{short_name} = {{",
+            "PyModuleDef_HEAD_INIT,",
+            f'.m_name = "{shared_lib_name(self.group_name)}",',
+            ".m_doc = NULL,",
+            f".m_size = {size},",
+            ".m_methods = NULL,",
+        )
+        if self.multi_phase_init:
+            emitter.emit_line(f".m_slots = slots_{short_name},")
+        emitter.emit_line("};")
+
+        if self.multi_phase_init:
+            emitter.emit_lines(
+                f"PyMODINIT_FUNC PyInit_{short_name}(void) {{",
+                f"return PyModuleDef_Init(&module_def_{short_name});",
+                "}",
+            )
+        else:
+            emitter.emit_lines(
+                f"PyMODINIT_FUNC PyInit_{short_name}(void) {{",
+                "static PyObject *module = NULL;",
+                "if (module) {",
+                "Py_INCREF(module);",
+                "return module;",
+                "}",
+                f"module = PyModule_Create(&module_def_{short_name});",
+                "if (!module) {",
+                "return NULL;",
+                "}",
+                f"if (exec_{short_name}(module) < 0) {{",
+                "Py_DECREF(module);",
+                "return NULL;",
+                "}",
+                "return module;",
+                "}",
+            )
 
     def generate_globals_init(self, emitter: Emitter) -> None:
         emitter.emit_lines(
@@ -887,16 +934,22 @@ def generate_globals_init(self, emitter: Emitter) -> None:
     def generate_module_def(self, emitter: Emitter, module_name: str, module: ModuleIR) -> None:
         """Emit the PyModuleDef struct for a module and the module init function."""
         module_prefix = emitter.names.private_name(module_name)
-        self.emit_module_exec_func(emitter, module_name, module_prefix, module)
-        if self.multi_phase_init:
-            self.emit_module_def_slots(emitter, module_prefix)
         self.emit_module_methods(emitter, module_name, module_prefix, module)
-        self.emit_module_def_struct(emitter, module_name, module_prefix)
-        self.emit_module_init_func(emitter, module_name, module_prefix)
+        self.emit_module_exec_func(emitter, module_name, module_prefix, module)
 
-    def emit_module_def_slots(self, emitter: Emitter, module_prefix: str) -> None:
+        # If using multi-phase init and a shared lib, parts of module definition
+        # will happen in the shim modules, so we skip some steps here.
+        if not (self.multi_phase_init and self.use_shared_lib):
+            if self.multi_phase_init:
+                self.emit_module_def_slots(emitter, module_prefix, module_name)
+            self.emit_module_def_struct(emitter, module_name, module_prefix)
+            self.emit_module_init_func(emitter, module_name, module_prefix)
+
+    def emit_module_def_slots(
+        self, emitter: Emitter, module_prefix: str, module_name: str
+    ) -> None:
         name = f"{module_prefix}_slots"
-        exec_name = f"{module_prefix}_exec"
+        exec_name = f"CPyExec_{exported_name(module_name)}"
 
         emitter.emit_line(f"static PyModuleDef_Slot {name}[] = {{")
         emitter.emit_line(f"{{Py_mod_exec, {exec_name}}},")
@@ -951,7 +1004,7 @@ def emit_module_def_struct(
             "0,       /* size of per-interpreter state of the module */",
             f"{module_prefix}module_methods,",
         )
-        if self.multi_phase_init:
+        if self.multi_phase_init and not self.use_shared_lib:
             slots_name = f"{module_prefix}_slots"
             emitter.emit_line(f"{slots_name}, /* m_slots */")
         else:
@@ -962,15 +1015,16 @@ def emit_module_def_struct(
     def emit_module_exec_func(
         self, emitter: Emitter, module_name: str, module_prefix: str, module: ModuleIR
     ) -> None:
-        """Emit the module init function.
+        """Emit the module exec function.
 
-        If we are compiling just one module, this will be the C API init
-        function. If we are compiling 2+ modules, we generate a shared
+        If we are compiling just one module, this will be the normal C API
+        exec function. If we are compiling 2+ modules, we generate a shared
         library for the modules and shims that call into the shared
-        library, and in this case we use an internal module initialized
-        function that will be called by the shim.
+        library, and in this case the shared module defines an internal
+        exec function for each module and these will be called by the shims
+        via Capsules.
         """
-        declaration = f"static int {module_prefix}_exec(PyObject *module)"
+        declaration = f"int CPyExec_{exported_name(module_name)}(PyObject *module)"
         module_static = self.module_internal_static_name(module_name, emitter)
         emitter.emit_lines(declaration, "{")
         emitter.emit_line("PyObject* modname = NULL;")
@@ -987,6 +1041,12 @@ def emit_module_exec_func(
             "    goto fail;",
         )
 
+        if self.multi_phase_init:
+            emitter.emit_lines(
+                f"if (PyModule_AddFunctions(module, {module_prefix}module_methods) < 0)",
+                "    goto fail;",
+            )
+
         # HACK: Manually instantiate generated classes here
         type_structs: list[str] = []
         for cl in module.classes:
@@ -1038,7 +1098,7 @@ def emit_module_init_func(
             emitter.emit_line("}")
             return
 
-        exec_func = f"{module_prefix}_exec"
+        exec_func = f"CPyExec_{exported_name(module_name)}"
 
         # Store the module reference in a static and return it when necessary.
         # This is separate from the *global* reference to the module that will

mypyc/lib-rt/module_shim_no_gil_multiphase.tmpl

@@ -0,0 +1,41 @@
+#include <Python.h>
+
+static int {modname}_exec(PyObject *module)
+{{
+    PyObject *tmp;
+    if (!(tmp = PyImport_ImportModule("{libname}"))) return -1;
+    PyObject *capsule = PyObject_GetAttrString(tmp, "exec_{full_modname}");
+    Py_DECREF(tmp);
+    if (capsule == NULL) return -1;
+    void *exec_func = PyCapsule_GetPointer(capsule, "{libname}.exec_{full_modname}");
+    Py_DECREF(capsule);
+    if (!exec_func) return -1;
+    if (((int (*)(PyObject *))exec_func)(module) != 0) return -1;
+    return 0;
+}}
+
+static PyModuleDef_Slot {modname}_slots[] = {{
+    {{Py_mod_exec, {modname}_exec}},
+    {{Py_mod_multiple_interpreters, Py_MOD_MULTIPLE_INTERPRETERS_NOT_SUPPORTED}},
+    {{Py_mod_gil, Py_MOD_GIL_NOT_USED}},
+    {{0, NULL}},
+}};
+
+static struct PyModuleDef {modname}_module = {{
+    PyModuleDef_HEAD_INIT,
+    .m_name = "{modname}",
+    .m_doc = NULL,
+    .m_methods = NULL,
+    .m_size = 0,
+    .m_slots = {modname}_slots,
+}};
+
+PyMODINIT_FUNC
+PyInit_{modname}(void)
+{{
+    return PyModuleDef_Init(&{modname}_module);
+}}
+
+// distutils sometimes spuriously tells cl to export CPyInit___init__,
+// so provide that so it chills out
+PyMODINIT_FUNC PyInit___init__(void) {{ return PyInit_{modname}(); }}