Rework Python Extension C++ code (again)

cmake/HalideGeneratorHelpers.cmake

@@ -100,11 +100,12 @@ function(add_halide_generator TARGET)
                "#include \"Halide.h\"\n"
                "HALIDE_GENERATOR_PYSTUB(${GEN_NAME}, ${MODULE_NAME})\n")
 
-        file(WRITE
-             "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.py_stub_generated.cpp"
-             "${stub_text}")
+        set(stub_file "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.${GEN_NAME}.${MODULE_NAME}.py_stub_generated.cpp")
+        if (NOT EXISTS "${stub_file}")
+            file(WRITE "${stub_file}" "${stub_text}")
+        endif()
 
-        Python3_add_library(${TARGET}_pystub MODULE WITH_SOABI "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.py_stub_generated.cpp" ${ARG_SOURCES})
+        Python3_add_library(${TARGET}_pystub MODULE WITH_SOABI "${stub_file}" ${ARG_SOURCES})
         set_target_properties(${TARGET}_pystub PROPERTIES
                               CXX_VISIBILITY_PRESET hidden
                               VISIBILITY_INLINES_HIDDEN ON
@@ -454,30 +455,22 @@ function(add_halide_python_extension_library TARGET)
         message(FATAL_ERROR "${TARGET} requires all libraries to use the same Halide Runtime, but saw ${len}: ${runtimes}")
     endif ()
 
-    # Module def code is the same in all of them, so arbitrarily choose the first
-    list(GET pycpps 0 module_definition_cpp)
-    add_library(${TARGET}_module_definition OBJECT ${module_definition_cpp})
-    set_target_properties(${TARGET}_module_definition PROPERTIES
-                          CXX_VISIBILITY_PRESET hidden
-                          VISIBILITY_INLINES_HIDDEN ON
-                          POSITION_INDEPENDENT_CODE ON)
-    target_link_libraries(${TARGET}_module_definition PRIVATE Halide::Runtime Python3::Module)
-
-    list(GET ARG_HALIDE_LIBRARIES 0 module_definition_lib)
-    add_dependencies(${TARGET}_module_definition ${module_definition_lib})
-
-    # Compile it with the right preprocessor definitions to provide the module defs,
-    # but not the function implementations
-    target_compile_definitions(${TARGET}_module_definition PRIVATE
-                               HALIDE_PYTHON_EXTENSION_OMIT_FUNCTION_DEFINITIONS
-                               HALIDE_PYTHON_EXTENSION_MODULE=${ARG_MODULE_NAME}
-                               "HALIDE_PYTHON_EXTENSION_FUNCTIONS=${function_names}")
+    set(pyext_runtime_name ${TARGET}_module_definition)
+    set(pyext_module_definition_src "${CMAKE_CURRENT_BINARY_DIR}/${pyext_runtime_name}.py.cpp")
+    _Halide_gengen_ensure()
+    add_custom_command(OUTPUT ${pyext_module_definition_src}
+                       COMMAND _Halide_gengen -r "${pyext_runtime_name}" -e python_extension -o "${CMAKE_CURRENT_BINARY_DIR}" target=host
+                       DEPENDS _Halide_gengen
+                       VERBATIM)
 
-    # Now compile all the pycpps to build the function implementations (but not the module def)
-    Python3_add_library(${TARGET} MODULE WITH_SOABI ${pycpps} $<TARGET_OBJECTS:${TARGET}_module_definition>)
+    Python3_add_library(${TARGET} MODULE WITH_SOABI ${pycpps} ${pyext_module_definition_src})
     target_link_libraries(${TARGET} PRIVATE ${ARG_HALIDE_LIBRARIES})
     target_compile_definitions(${TARGET} PRIVATE
-                               HALIDE_PYTHON_EXTENSION_OMIT_MODULE_DEFINITION)
+                               # Skip the default module-definition code in each file
+                               HALIDE_PYTHON_EXTENSION_OMIT_MODULE_DEFINITION
+                               # Gotta explicitly specify the module name and function(s) for this mode
+                               HALIDE_PYTHON_EXTENSION_MODULE_NAME=${ARG_MODULE_NAME}
+                               "HALIDE_PYTHON_EXTENSION_FUNCTIONS=${function_names}")
     set_target_properties(${TARGET} PROPERTIES OUTPUT_NAME ${ARG_MODULE_NAME})
     _Halide_target_export_single_symbol(${TARGET} "PyInit_${ARG_MODULE_NAME}")
 endfunction()
@@ -534,12 +527,8 @@ function(add_halide_runtime RT)
     # so that GCD calculation doesn't get confused.
     list(TRANSFORM ARG_TARGETS APPEND "-no_runtime")
 
-    # Create a Generator that is GenGen.cpp and nothing else; all it can do is generate a runtime.
-    if (NOT TARGET _Halide_gengen)
-        add_executable(_Halide_gengen )
-        target_link_libraries(_Halide_gengen PRIVATE Halide::Generator)
-        _Halide_place_dll(_Halide_gengen)
-    endif ()
+    # Ensure _Halide_gengen is defined
+    _Halide_gengen_ensure()
 
     _Halide_get_platform_details(
             is_crosscompiling
@@ -668,15 +657,38 @@ function(_Halide_fix_xcode TARGET)
 endfunction()
 
 function(_Halide_target_export_single_symbol TARGET SYMBOL)
-    file(WRITE
-         "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.ldscript.apple"
-         "_${SYMBOL}\n")
-    file(WRITE
-         "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.ldscript"
-         "{ global: ${SYMBOL}; local: *; };\n")
+    if (NOT EXISTS "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.${SYMBOL}.ldscript.apple")
+        file(WRITE
+             "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.${SYMBOL}.ldscript.apple"
+             "_${SYMBOL}\n")
+    endif()
+    if (NOT EXISTS "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.${SYMBOL}.ldscript")
+        file(WRITE
+             "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.${SYMBOL}.ldscript"
+             "{ global: ${SYMBOL}; local: *; };\n")
+    endif ()
     target_export_script(
         ${TARGET}
-        APPLE_LD "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.ldscript.apple"
-        GNU_LD "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.ldscript"
+        APPLE_LD "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.${SYMBOL}.ldscript.apple"
+        GNU_LD "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.${SYMBOL}.ldscript"
     )
 endfunction()
+
+function(_Halide_gengen_ensure)
+    # Create a Generator that is GenGen.cpp and nothing else; all it can do is generate a runtime.
+    if (NOT TARGET _Halide_gengen)
+
+        # add_executable requires at least one source file for some
+        # configs (e.g. Xcode), because, uh, reasons, so we'll create
+        # an empty one here to satisfy it
+        set(empty "${CMAKE_CURRENT_BINARY_DIR}/_Halide_gengen.empty.cpp")
+        if (NOT EXISTS "${empty}")
+            file(WRITE "${empty}" "/* nothing */\n")
+        endif ()
+
+        add_executable(_Halide_gengen "${empty}")
+        target_link_libraries(_Halide_gengen PRIVATE Halide::Generator)
+        _Halide_place_dll(_Halide_gengen)
+    endif ()
+endfunction()
+

src/CodeGen_C.cpp

@@ -1272,14 +1272,20 @@ class NativeVectorOps {
 
 void CodeGen_C::set_name_mangling_mode(NameMangling mode) {
     if (extern_c_open && mode != NameMangling::C) {
-        stream << "\n#ifdef __cplusplus\n";
-        stream << "}  // extern \"C\"\n";
-        stream << "#endif\n\n";
+        stream << R"INLINE_CODE(
+#ifdef __cplusplus
+}  // extern "C"
+#endif
+
+)INLINE_CODE";
         extern_c_open = false;
     } else if (!extern_c_open && mode == NameMangling::C) {
-        stream << "\n#ifdef __cplusplus\n";
-        stream << "extern \"C\" {\n";
-        stream << "#endif\n\n";
+        stream << R"INLINE_CODE(
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+)INLINE_CODE";
         extern_c_open = true;
     }
 }

src/Module.cpp

@@ -729,7 +729,10 @@ std::map<OutputFileType, std::string> compile_standalone_runtime(const std::map<
     // For runtime, it only makes sense to output object files or static_library, so ignore
     // everything else.
     std::map<OutputFileType, std::string> actual_outputs;
-    for (auto key : {OutputFileType::object, OutputFileType::static_library}) {
+    // If the python_extension output is specified, we'll generate just the module-registration code,
+    // with no functions at all. This is useful when gluing together multiple Halide functions
+    // into the same Python extension.
+    for (auto key : {OutputFileType::object, OutputFileType::static_library, OutputFileType::python_extension}) {
         auto it = output_files.find(key);
         if (it != output_files.end()) {
             actual_outputs[key] = it->second;