Don't reuse passmanagers across compilations

Author: gbaraldi

doc/src/devdocs/locks.md

@@ -30,6 +30,7 @@ The following are definitely leaf locks (level 1), and must not try to acquire a
 >   * jl_in_stackwalk (Win32)
 >   * ResourcePool<?>::mutex
 >   * RLST_mutex
+>   * llvm_printing_mutex
 >   * jl_locked_stream::mutex
 >   * debuginfo_asyncsafe
 >   * inference_timing_mutex

src/aotcompile.cpp

@@ -1032,7 +1032,7 @@ static AOTOutputs add_output_impl(Module &M, TargetMachine &SourceTM, ShardTimer
                 SourceTM.getCodeModel(),
                 SourceTM.getOptLevel()));
         fixupTM(*PMTM);
-        NewPM optimizer{std::move(PMTM), getOptLevel(jl_options.opt_level), OptimizationOptions::defaults(true, true)};
+        NewPM optimizer{std::move(PMTM), getOptLevel(jl_options.opt_level), M.getContext(), OptimizationOptions::defaults(true, true)};
         optimizer.run(M);
         assert(!verifyLLVMIR(M));
         bool inject_aliases = false;
@@ -1909,7 +1909,7 @@ void jl_get_llvmf_defn_impl(jl_llvmf_dump_t* dump, jl_method_instance_t *mi, siz
             }
             assert(!verifyLLVMIR(*m.getModuleUnlocked()));
             if (optimize) {
-                NewPM PM{jl_ExecutionEngine->cloneTargetMachine(), getOptLevel(jl_options.opt_level)};
+                NewPM PM{jl_ExecutionEngine->cloneTargetMachine(), getOptLevel(jl_options.opt_level), output.getContext()};
                 //Safe b/c context lock is held by output
                 PM.run(*m.getModuleUnlocked());
                 assert(!verifyLLVMIR(*m.getModuleUnlocked()));

src/jitlayers.cpp

@@ -1245,37 +1245,20 @@ namespace {
         }
     };
 
-    struct PMCreator {
-        orc::JITTargetMachineBuilder JTMB;
-        OptimizationLevel O;
-        SmallVector<std::function<void()>, 0> &printers;
-        PMCreator(TargetMachine &TM, int optlevel, SmallVector<std::function<void()>, 0> &printers) JL_NOTSAFEPOINT
-            : JTMB(createJTMBFromTM(TM, optlevel)), O(getOptLevel(optlevel)), printers(printers) {}
-
-        auto operator()() JL_NOTSAFEPOINT {
-            auto TM = cantFail(JTMB.createTargetMachine());
-            fixupTM(*TM);
-            auto NPM = std::make_unique<NewPM>(std::move(TM), O);
-            // TODO this needs to be locked, as different resource pools may add to the printer vector at the same time
-            printers.push_back([NPM = NPM.get()]() JL_NOTSAFEPOINT {
-                NPM->printTimers();
-            });
-            return NPM;
-        }
-    };
 
     template<size_t N>
     struct OptimizerT {
-        OptimizerT(TargetMachine &TM, SmallVector<std::function<void()>, 0> &printers) JL_NOTSAFEPOINT {
+        OptimizerT(TargetMachine &TM, SmallSet<std::function<void()>*, 4> &printers, std::mutex &print_mutex) JL_NOTSAFEPOINT
+        :TM(TM), printers(printers), print_mutex(print_mutex) {
             for (size_t i = 0; i < N; i++) {
-                PMs[i] = std::make_unique<JuliaOJIT::ResourcePool<std::unique_ptr<PassManager>>>(PMCreator(TM, i, printers));
+                JTMBs.push_back(createJTMBFromTM(TM, i));
             }
-        }
+         }
 
         OptimizerResultT operator()(orc::ThreadSafeModule TSM, orc::MaterializationResponsibility &R) JL_NOTSAFEPOINT {
             TSM.withModuleDo([&](Module &M) JL_NOTSAFEPOINT {
-                auto PoolIdx = cast<ConstantInt>(cast<ConstantAsMetadata>(M.getModuleFlag("julia.optlevel"))->getValue())->getZExtValue();
-                assert(PoolIdx < N && "Invalid optimization pool index");
+                auto optlevel = cast<ConstantInt>(cast<ConstantAsMetadata>(M.getModuleFlag("julia.optlevel"))->getValue())->getZExtValue();
+                assert(optlevel < N && "Invalid optimization pool index");
 
                 uint64_t start_time = 0;
 
@@ -1312,7 +1295,22 @@ namespace {
                 {
                     JL_TIMING(LLVM_JIT, JIT_Opt);
                     //Run the optimization
-                    (****PMs[PoolIdx]).run(M);
+                    auto TM = cantFail(JTMBs[optlevel].createTargetMachine());
+                    fixupTM(*TM);
+                    NewPM NPM(std::move(TM), getOptLevel(optlevel), M.getContext());
+                    // TODO this needs to be locked, as different resource pools may add to the printer vector at the same time
+                    std::function<void()> printfun = [&]() JL_NOTSAFEPOINT {
+                        NPM.printTimers();
+                    };
+                    {
+                    std::lock_guard<std::mutex> lock(print_mutex);
+                    printers.insert(&printfun);
+                    }
+                    NPM.run(M);
+                    {
+                    std::lock_guard<std::mutex> lock(print_mutex);
+                    printers.erase(&printfun);
+                    }
                     assert(!verifyLLVMIR(M));
                 }
 
@@ -1346,7 +1344,7 @@ namespace {
                             s.dump(stream);
                         }
                         ios_printf(stream, "  time_ns: %" PRIu64 "\n", end_time - start_time);
-                        ios_printf(stream, "  optlevel: %d\n", PoolIdx);
+                        ios_printf(stream, "  optlevel: %d\n", optlevel);
 
                         // Print LLVM function statistics _after_ optimization
                         ios_printf(stream, "  after: \n");
@@ -1359,7 +1357,7 @@ namespace {
                     }
                 }
                 ++ModulesOptimized;
-                switch (PoolIdx) {
+                switch (optlevel) {
                     case 0:
                         ++OptO0;
                         break;
@@ -1380,7 +1378,10 @@ namespace {
             return Expected<orc::ThreadSafeModule>{std::move(TSM)};
         }
     private:
-        std::array<std::unique_ptr<JuliaOJIT::ResourcePool<std::unique_ptr<PassManager>>>, N> PMs;
+        TargetMachine &TM;
+        SmallSet<std::function<void()>*, 4> &printers;
+        std::mutex &print_mutex;
+        SmallVector<orc::JITTargetMachineBuilder, 4> JTMBs;
     };
 
     template<size_t N>
@@ -1697,7 +1698,7 @@ JuliaOJIT::JuliaOJIT()
     LockLayer(ObjectLayer),
     CompileLayer(ES, LockLayer, std::make_unique<CompilerT<N_optlevels>>(orc::irManglingOptionsFromTargetOptions(TM->Options), *TM)),
     JITPointersLayer(ES, CompileLayer, orc::IRTransformLayer::TransformFunction(JITPointersT(ES))),
-    OptimizeLayer(ES, JITPointersLayer, orc::IRTransformLayer::TransformFunction(OptimizerT<N_optlevels>(*TM, PrintLLVMTimers))),
+    OptimizeLayer(ES, JITPointersLayer, orc::IRTransformLayer::TransformFunction(OptimizerT<N_optlevels>(*TM, PrintLLVMTimers, llvm_printing_mutex))),
     OptSelLayer(ES, OptimizeLayer, orc::IRTransformLayer::TransformFunction(selectOptLevel)),
     DepsVerifyLayer(ES, OptSelLayer, orc::IRTransformLayer::TransformFunction(validateExternRelocations)),
     ExternalCompileLayer(ES, LockLayer,
@@ -2043,8 +2044,11 @@ size_t JuliaOJIT::getTotalBytes() const
 
 void JuliaOJIT::printTimers()
 {
-    for (auto &printer : PrintLLVMTimers) {
-        printer();
+    {
+        std::lock_guard<std::mutex> lock(llvm_printing_mutex);
+        for (auto printer : PrintLLVMTimers) {
+            (*printer)();
+        }
     }
     reportAndResetTimings();
 }

src/jitlayers.h

@@ -1,6 +1,7 @@
 // This file is a part of Julia. License is MIT: https://julialang.org/license
 
 #include <llvm/ADT/MapVector.h>
+#include "llvm/ADT/SmallSet.h"
 
 #include <llvm/IR/LLVMContext.h>
 #include <llvm/IR/Constants.h>
@@ -99,15 +100,13 @@ struct OptimizationOptions {
 
 struct NewPM {
     std::unique_ptr<TargetMachine> TM;
-#if JL_LLVM_VERSION < 160000
     StandardInstrumentations SI;
-#endif
     std::unique_ptr<PassInstrumentationCallbacks> PIC;
     PassBuilder PB;
     ModulePassManager MPM;
     OptimizationLevel O;
 
-    NewPM(std::unique_ptr<TargetMachine> TM, OptimizationLevel O, OptimizationOptions options = OptimizationOptions::defaults()) JL_NOTSAFEPOINT;
+    NewPM(std::unique_ptr<TargetMachine> TM, OptimizationLevel O, LLVMContext &ctx, OptimizationOptions options = OptimizationOptions::defaults()) JL_NOTSAFEPOINT;
     ~NewPM() JL_NOTSAFEPOINT;
 
     void run(Module &M) JL_NOTSAFEPOINT;
@@ -537,7 +536,8 @@ class JuliaOJIT {
     jl_locked_stream dump_compiles_stream;
     jl_locked_stream dump_llvm_opt_stream;
 
-    SmallVector<std::function<void()>, 0> PrintLLVMTimers;
+    std::mutex llvm_printing_mutex{};
+    SmallSet<std::function<void()>*, 4> PrintLLVMTimers;
 
     ResourcePool<orc::ThreadSafeContext, 0, std::queue<orc::ThreadSafeContext>> ContextPool;
 

src/pipeline.cpp

@@ -701,20 +701,12 @@ PIC.addClassToPassName(decltype(CREATE_PASS)::name(), NAME);
         PIC.addClassToPassName("AfterOptimizationMarkerPass", "AfterOptimization");
     }
 
-#if JL_LLVM_VERSION >= 160000
-    auto createPIC() JL_NOTSAFEPOINT {
-        auto PIC = std::make_unique<PassInstrumentationCallbacks>();
-        adjustPIC(*PIC);
-        return PIC;
-    }
-#else
     auto createPIC(StandardInstrumentations &SI) JL_NOTSAFEPOINT {
         auto PIC = std::make_unique<PassInstrumentationCallbacks>();
         adjustPIC(*PIC);
         SI.registerCallbacks(*PIC);
         return PIC;
     }
-#endif
 
     FunctionAnalysisManager createFAM(OptimizationLevel O, TargetMachine &TM) JL_NOTSAFEPOINT {
 
@@ -743,14 +735,14 @@ PIC.addClassToPassName(decltype(CREATE_PASS)::name(), NAME);
     }
 }
 
-NewPM::NewPM(std::unique_ptr<TargetMachine> TM, OptimizationLevel O, OptimizationOptions options) :
+NewPM::NewPM(std::unique_ptr<TargetMachine> TM, OptimizationLevel O, LLVMContext &ctx, OptimizationOptions options) :
     TM(std::move(TM)),
 #if JL_LLVM_VERSION < 160000
     SI(false),
-    PIC(createPIC(SI)),
 #else
-    PIC(createPIC()),
+    SI(ctx, false),
 #endif
+    PIC(createPIC(SI)),
     PB(this->TM.get(), PipelineTuningOptions(), None, PIC.get()),
     MPM(createMPM(PB, O, options)), O(O) {}