Restore early release of threadsafe context for parallel compilation

by creating a manual co-routine to get out of the caller scope.

Author: yuyichao

src/aotcompile.cpp

@@ -1852,63 +1852,142 @@ static inline void schedule_uv_thread(uv_thread_t *worker, CB &&cb)
 
 // Entrypoint to optionally-multithreaded image compilation. This handles global coordination of the threading,
 // as well as partitioning, serialization, and deserialization.
-template<typename ModuleReleasedFunc>
-static SmallVector<AOTOutputs, 16> add_output(Module &M, TargetMachine &TM, StringRef name, unsigned threads,
-                bool unopt_out, bool opt_out, bool obj_out, bool asm_out, ModuleReleasedFunc module_released) {
-    SmallVector<AOTOutputs, 16> outputs(threads);
-    assert(threads);
-    assert(unopt_out || opt_out || obj_out || asm_out);
-    // Timers for timing purposes
-    TimerGroup timer_group("add_output", ("Time to optimize and emit LLVM module " + name).str());
-    SmallVector<ShardTimers, 1> timers(threads);
-    for (unsigned i = 0; i < threads; ++i) {
-        auto idx = std::to_string(i);
-        timers[i].name = "shard_" + idx;
-        timers[i].desc = ("Timings for " + name + " module shard " + idx).str();
-        timers[i].deserialize.init("deserialize_" + idx, "Deserialize module");
-        timers[i].materialize.init("materialize_" + idx, "Materialize declarations");
-        timers[i].construct.init("construct_" + idx, "Construct partitioned definitions");
-        timers[i].unopt.init("unopt_" + idx, "Emit unoptimized bitcode");
-        timers[i].optimize.init("optimize_" + idx, "Optimize shard");
-        timers[i].opt.init("opt_" + idx, "Emit optimized bitcode");
-        timers[i].obj.init("obj_" + idx, "Emit object file");
-        timers[i].asm_.init("asm_" + idx, "Emit assembly file");
-    }
-    Timer partition_timer("partition", "Partition module", timer_group);
-    Timer serialize_timer("serialize", "Serialize module", timer_group);
-    Timer output_timer("output", "Add outputs", timer_group);
-    bool report_timings = false;
-    if (auto env = getenv("JULIA_IMAGE_TIMINGS")) {
-        char *endptr;
-        unsigned long val = strtoul(env, &endptr, 10);
-        if (endptr != env && !*endptr && val <= 1) {
-            report_timings = val;
-        } else {
-            if (StringRef("true").compare_insensitive(env) == 0)
-                report_timings = true;
-            else if (StringRef("false").compare_insensitive(env) == 0)
-                report_timings = false;
-            else
-                errs() << "WARNING: Invalid value for JULIA_IMAGE_TIMINGS: " << env << "\n";
+
+// This is more or less a manual co-routine version of add_output
+// which allows exiting calling scope when the module is released.
+struct OutputAdder {
+    OutputAdder(Module &M, TargetMachine &TM, StringRef name, unsigned threads,
+                bool unopt_out, bool opt_out, bool obj_out, bool asm_out)
+        : TM(TM), threads(threads), unopt_out(unopt_out),
+          opt_out(opt_out), obj_out(obj_out), asm_out(asm_out),
+          timer_group("add_output", ("Time to optimize and emit LLVM module " + name).str())
+    {
+        assert(threads);
+        assert(unopt_out || opt_out || obj_out || asm_out);
+        for (unsigned i = 0; i < threads; ++i) {
+            auto idx = std::to_string(i);
+            timers[i].name = "shard_" + idx;
+            timers[i].desc = ("Timings for " + name + " module shard " + idx).str();
+            timers[i].deserialize.init("deserialize_" + idx, "Deserialize module");
+            timers[i].materialize.init("materialize_" + idx, "Materialize declarations");
+            timers[i].construct.init("construct_" + idx, "Construct partitioned definitions");
+            timers[i].unopt.init("unopt_" + idx, "Emit unoptimized bitcode");
+            timers[i].optimize.init("optimize_" + idx, "Optimize shard");
+            timers[i].opt.init("opt_" + idx, "Emit optimized bitcode");
+            timers[i].obj.init("obj_" + idx, "Emit object file");
+            timers[i].asm_.init("asm_" + idx, "Emit assembly file");
+        }
+        if (auto env = getenv("JULIA_IMAGE_TIMINGS")) {
+            char *endptr;
+            unsigned long val = strtoul(env, &endptr, 10);
+            if (endptr != env && !*endptr && val <= 1) {
+                report_timings = val;
+            } else {
+                if (StringRef("true").compare_insensitive(env) == 0)
+                    report_timings = true;
+                else if (StringRef("false").compare_insensitive(env) == 0)
+                    report_timings = false;
+                else
+                    errs() << "WARNING: Invalid value for JULIA_IMAGE_TIMINGS: " << env << "\n";
+            }
+        }
+        // Single-threaded case
+        if (threads == 1) {
+            output_timer.startTimer();
+            {
+                JL_TIMING(NATIVE_AOT, NATIVE_Opt);
+                // convert gvars to the expected offset table format for shard 0
+                if (M.getGlobalVariable("jl_gvars")) {
+                    auto gvars = consume_gv<Constant>(M, "jl_gvars", false);
+                    Type *T_size = M.getDataLayout().getIntPtrType(M.getContext());
+                    emit_offset_table(M, T_size, gvars, "jl_gvar", "_0"); // module flag "julia.mv.suffix"
+                    M.getGlobalVariable("jl_gvar_idxs")->setName("jl_gvar_idxs_0");
+                }
+                output0 = add_output_impl(M, TM, timers[0], unopt_out, opt_out, obj_out, asm_out);
+            }
+            output_timer.stopTimer();
+            return;
         }
+
+        partition_timer.startTimer();
+        uint64_t counter = 0;
+        // Partitioning requires all globals to have names.
+        // We use a prefix to avoid name conflicts with user code.
+        for (auto &G : M.global_values()) {
+            if (!G.isDeclaration() && !G.hasName()) {
+                G.setName("jl_ext_" + Twine(counter++));
+            }
+        }
+        partitions = partitionModule(M, threads);
+        partition_timer.stopTimer();
+
+        serialize_timer.startTimer();
+        serialized = serializeModule(M);
+        serialize_timer.stopTimer();
     }
-    // Single-threaded case
-    if (threads == 1) {
+
+    auto finish()
+    {
+        SmallVector<AOTOutputs, 16> outputs(threads);
+        if (threads == 1) {
+            outputs[0] = std::move(output0);
+            if (!report_timings) {
+                timer_group.clear();
+            } else {
+                timer_group.print(dbgs(), true);
+                for (auto &t : timers) {
+                    t.print(dbgs(), true);
+                }
+            }
+            return outputs;
+        }
         output_timer.startTimer();
+
+        // Start all of the worker threads
         {
             JL_TIMING(NATIVE_AOT, NATIVE_Opt);
-            // convert gvars to the expected offset table format for shard 0
-            if (M.getGlobalVariable("jl_gvars")) {
-                auto gvars = consume_gv<Constant>(M, "jl_gvars", false);
-                Type *T_size = M.getDataLayout().getIntPtrType(M.getContext());
-                emit_offset_table(M, T_size, gvars, "jl_gvar", "_0"); // module flag "julia.mv.suffix"
-                M.getGlobalVariable("jl_gvar_idxs")->setName("jl_gvar_idxs_0");
+            std::vector<uv_thread_t> workers(threads);
+            for (unsigned i = 0; i < threads; i++) {
+                schedule_uv_thread(&workers[i], [&, i]() {
+                    LLVMContext ctx;
+                    ctx.setDiscardValueNames(true);
+                    // Lazily deserialize the entire module
+                    timers[i].deserialize.startTimer();
+                    auto EM = getLazyBitcodeModule(MemoryBufferRef(StringRef(serialized.data(), serialized.size()), "Optimized"), ctx);
+                    // Make sure this also fails with only julia, but not LLVM assertions enabled,
+                    // otherwise, the first error we hit is the LLVM module verification failure,
+                    // which will look very confusing, because the module was partially deserialized.
+                    bool deser_succeeded = (bool)EM;
+                    auto M = cantFail(std::move(EM), "Error loading module");
+                    assert(deser_succeeded); (void)deser_succeeded;
+                    timers[i].deserialize.stopTimer();
+
+                    timers[i].materialize.startTimer();
+                    materializePreserved(*M, partitions[i]);
+                    timers[i].materialize.stopTimer();
+
+                    timers[i].construct.startTimer();
+                    std::string suffix = "_" + std::to_string(i);
+                    construct_vars(*M, partitions[i], suffix);
+                    M->setModuleFlag(Module::Error, "julia.mv.suffix", MDString::get(M->getContext(), suffix));
+                    // The DICompileUnit file is not used for anything, but ld64 requires it be a unique string per object file
+                    // or it may skip emitting debug info for that file. Here set it to ./julia#N
+                    DIFile *topfile = DIFile::get(M->getContext(), "julia#" + std::to_string(i), ".");
+                    if (M->getNamedMetadata("llvm.dbg.cu"))
+                        for (auto CU: M->getNamedMetadata("llvm.dbg.cu")->operands())
+                            CU->replaceOperandWith(0, topfile);
+                    timers[i].construct.stopTimer();
+
+                    outputs[i] = add_output_impl(*M, TM, timers[i], unopt_out, opt_out, obj_out, asm_out);
+                });
             }
-            outputs[0] = add_output_impl(M, TM, timers[0], unopt_out, opt_out, obj_out, asm_out);
+
+            // Wait for all of the worker threads to finish
+            for (unsigned i = 0; i < threads; i++)
+                uv_thread_join(&workers[i]);
         }
+
         output_timer.stopTimer();
-        // Don't need M anymore
-        module_released(M);
 
         if (!report_timings) {
             timer_group.clear();
@@ -1917,97 +1996,37 @@ static SmallVector<AOTOutputs, 16> add_output(Module &M, TargetMachine &TM, Stri
             for (auto &t : timers) {
                 t.print(dbgs(), true);
             }
+            dbgs() << "Partition weights: [";
+            bool comma = false;
+            for (auto &p : partitions) {
+                if (comma)
+                    dbgs() << ", ";
+                else
+                    comma = true;
+                dbgs() << p.weight;
+            }
+            dbgs() << "]\n";
         }
         return outputs;
     }
 
-    partition_timer.startTimer();
-    uint64_t counter = 0;
-    // Partitioning requires all globals to have names.
-    // We use a prefix to avoid name conflicts with user code.
-    for (auto &G : M.global_values()) {
-        if (!G.isDeclaration() && !G.hasName()) {
-            G.setName("jl_ext_" + Twine(counter++));
-        }
-    }
-    auto partitions = partitionModule(M, threads);
-    partition_timer.stopTimer();
-
-    serialize_timer.startTimer();
-    auto serialized = serializeModule(M);
-    serialize_timer.stopTimer();
-
-    // Don't need M anymore, since we'll only read from serialized from now on
-    module_released(M);
-
-    output_timer.startTimer();
-
-    // Start all of the worker threads
-    {
-        JL_TIMING(NATIVE_AOT, NATIVE_Opt);
-        std::vector<uv_thread_t> workers(threads);
-        for (unsigned i = 0; i < threads; i++) {
-            schedule_uv_thread(&workers[i], [&, i]() {
-                LLVMContext ctx;
-                ctx.setDiscardValueNames(true);
-                // Lazily deserialize the entire module
-                timers[i].deserialize.startTimer();
-                auto EM = getLazyBitcodeModule(MemoryBufferRef(StringRef(serialized.data(), serialized.size()), "Optimized"), ctx);
-                // Make sure this also fails with only julia, but not LLVM assertions enabled,
-                // otherwise, the first error we hit is the LLVM module verification failure,
-                // which will look very confusing, because the module was partially deserialized.
-                bool deser_succeeded = (bool)EM;
-                auto M = cantFail(std::move(EM), "Error loading module");
-                assert(deser_succeeded); (void)deser_succeeded;
-                timers[i].deserialize.stopTimer();
-
-                timers[i].materialize.startTimer();
-                materializePreserved(*M, partitions[i]);
-                timers[i].materialize.stopTimer();
-
-                timers[i].construct.startTimer();
-                std::string suffix = "_" + std::to_string(i);
-                construct_vars(*M, partitions[i], suffix);
-                M->setModuleFlag(Module::Error, "julia.mv.suffix", MDString::get(M->getContext(), suffix));
-                // The DICompileUnit file is not used for anything, but ld64 requires it be a unique string per object file
-                // or it may skip emitting debug info for that file. Here set it to ./julia#N
-                DIFile *topfile = DIFile::get(M->getContext(), "julia#" + std::to_string(i), ".");
-                if (M->getNamedMetadata("llvm.dbg.cu"))
-                    for (auto CU: M->getNamedMetadata("llvm.dbg.cu")->operands())
-                        CU->replaceOperandWith(0, topfile);
-                timers[i].construct.stopTimer();
-
-                outputs[i] = add_output_impl(*M, TM, timers[i], unopt_out, opt_out, obj_out, asm_out);
-            });
-        }
-
-        // Wait for all of the worker threads to finish
-        for (unsigned i = 0; i < threads; i++)
-            uv_thread_join(&workers[i]);
-    }
-
-    output_timer.stopTimer();
-
-    if (!report_timings) {
-        timer_group.clear();
-    } else {
-        timer_group.print(dbgs(), true);
-        for (auto &t : timers) {
-            t.print(dbgs(), true);
-        }
-        dbgs() << "Partition weights: [";
-        bool comma = false;
-        for (auto &p : partitions) {
-            if (comma)
-                dbgs() << ", ";
-            else
-                comma = true;
-            dbgs() << p.weight;
-        }
-        dbgs() << "]\n";
-    }
-    return outputs;
-}
+    TargetMachine &TM;
+    unsigned threads;
+    bool unopt_out;
+    bool opt_out;
+    bool obj_out;
+    bool asm_out;
+    AOTOutputs output0;
+    // Timers for timing purposes
+    TimerGroup timer_group;
+    SmallVector<ShardTimers, 1> timers{threads};
+    Timer partition_timer{"partition", "Partition module", timer_group};
+    Timer serialize_timer{"serialize", "Serialize module", timer_group};
+    Timer output_timer{"output", "Add outputs", timer_group};
+    bool report_timings{false};
+    SmallVector<Partition, 32> partitions;
+    SmallVector<char, 0> serialized;
+};
 
 extern int jl_is_timing_passes;
 static unsigned compute_image_thread_count(const ModuleInfo &info) {
@@ -2145,8 +2164,8 @@ void jl_dump_native_impl(void *native_code,
         OverrideStackAlignment = M.getOverrideStackAlignment();
     });
 
-    auto compile = [&](Module &M, StringRef name, unsigned threads, auto module_released) {
-        return add_output(M, *SourceTM, name, threads, !!unopt_bc_fname, !!bc_fname, !!obj_fname, !!asm_fname, module_released);
+    auto start_compile = [&](Module &M, StringRef name, unsigned threads) {
+        return OutputAdder(M, *SourceTM, name, threads, !!unopt_bc_fname, !!bc_fname, !!obj_fname, !!asm_fname);
     };
 
     SmallVector<AOTOutputs, 16> sysimg_outputs;
@@ -2215,7 +2234,7 @@ void jl_dump_native_impl(void *native_code,
         // Note that we don't set z to null, this allows the check in WRITE_ARCHIVE
         // to function as expected
         // no need to free the module/context, destructor handles that
-        sysimg_outputs = compile(sysimgM, "sysimg", 1, [](Module &) {});
+        sysimg_outputs = start_compile(sysimgM, "sysimg", 1).finish();
     }
 
     const bool imaging_mode = true;
@@ -2314,34 +2333,13 @@ void jl_dump_native_impl(void *native_code,
     });
 
     {
-        // Don't use withModuleDo here since we delete the TSM midway through
-        auto TSCtx = data->M.getContext();
-#if JL_LLVM_VERSION < 210000
-        auto lock = TSCtx.getLock();
-        auto dataM = data->M.getModuleUnlocked();
-
-        data_outputs = compile(*dataM, "text", threads, [data, &lock, &TSCtx](Module &) {
-            // Delete data when add_output thinks it's done with it
-            // Saves memory for use when multithreading
-            auto lock2 = std::move(lock);
-            delete data;
-            // Drop last reference to shared LLVM::Context
-            auto TSCtx2 = std::move(TSCtx);
-        });
-#else
-        TSCtx.withContextDo([&] (LLVMContext*) {
-            auto dataM = data->M.getModuleUnlocked();
-            data_outputs = compile(*dataM, "text", threads, [data, &TSCtx](Module &) {
-                // Delete data when add_output thinks it's done with it
-                // Saves memory for use when multithreading
-                // Unfortunately, for LLVM 21, this does not release the lock
-                // or release the LLVM context.
-                delete data;
-                // Drop our reference to shared LLVM::Context
-                auto TSCtx2 = std::move(TSCtx);
-            });
+        auto adder = data->M.withModuleDo([&] (auto &dataM) {
+            return start_compile(dataM, "text", threads);
         });
-#endif
+        // Delete data when add_output thinks it's done with it
+        // Saves memory for use when multithreading
+        delete data;
+        data_outputs = adder.finish();
     }
 
     if (params->emit_metadata) {
@@ -2444,7 +2442,7 @@ void jl_dump_native_impl(void *native_code,
         }
 
         // no need to free module/context, destructor handles that
-        metadata_outputs = compile(metadataM, "data", 1, [](Module &) {});
+        metadata_outputs = start_compile(metadataM, "data", 1).finish();
     }
 
     {