[LLVM][PM] Make MultiVersioning NewPM compatible

Author: vchuravy

src/aotcompile.cpp

@@ -934,6 +934,10 @@ static void registerCallbacks(PassBuilder &PB) {
                 PM.addPass(RemoveJuliaAddrspacesPass());
                 return true;
             }
+            if (Name == "MultiVersioning") {
+                PM.addPass(MultiVersioning());
+                return true;
+            }
             return false;
         });
 

src/llvm-multiversioning.cpp

@@ -7,6 +7,7 @@
 // LLVM pass to clone function for different archs
 
 #include "llvm-version.h"
+#include "passes.h"
 
 #include <llvm-c/Core.h>
 #include <llvm-c/Types.h>
@@ -46,8 +47,6 @@ namespace {
 constexpr uint32_t clone_mask =
     JL_TARGET_CLONE_LOOP | JL_TARGET_CLONE_SIMD | JL_TARGET_CLONE_MATH | JL_TARGET_CLONE_CPU;
 
-struct MultiVersioning;
-
 // Treat identical mapping as missing and return `def` in that case.
 // We mainly need this to identify cloned function using value map after LLVM cloning
 // functions fills the map with identity entries.
@@ -243,7 +242,7 @@ struct CloneCtx {
             return cast<Function>(vmap->lookup(orig_f));
         }
     };
-    CloneCtx(MultiVersioning *pass, Module &M);
+    CloneCtx(Module &M, function_ref<LoopInfo&()> GetLI, function_ref<CallGraph&()> GetCG);
     void clone_bases();
     void collect_func_infos();
     void clone_all_partials();
@@ -277,12 +276,14 @@ struct CloneCtx {
     Type *T_void;
     PointerType *T_psize;
     MDNode *tbaa_const;
-    MultiVersioning *pass;
     std::vector<jl_target_spec_t> specs;
     std::vector<Group> groups{};
     std::vector<Function*> fvars;
     std::vector<Constant*> gvars;
     Module &M;
+    function_ref<LoopInfo&()> GetLI;
+    function_ref<CallGraph&()> GetCG;
+
     // Map from original functiton to one based index in `fvars`
     std::map<const Function*,uint32_t> func_ids{};
     std::vector<Function*> orig_funcs{};
@@ -298,23 +299,6 @@ struct CloneCtx {
     bool has_cloneall{false};
 };
 
-struct MultiVersioning: public ModulePass {
-    static char ID;
-    MultiVersioning()
-        : ModulePass(ID)
-    {}
-
-private:
-    bool runOnModule(Module &M) override;
-    void getAnalysisUsage(AnalysisUsage &AU) const override
-    {
-        AU.addRequired<LoopInfoWrapperPass>();
-        AU.addRequired<CallGraphWrapperPass>();
-        AU.addPreserved<LoopInfoWrapperPass>();
-    }
-    friend struct CloneCtx;
-};
-
 template<typename T>
 static inline std::vector<T*> consume_gv(Module &M, const char *name)
 {
@@ -335,18 +319,19 @@ static inline std::vector<T*> consume_gv(Module &M, const char *name)
 }
 
 // Collect basic information about targets and functions.
-CloneCtx::CloneCtx(MultiVersioning *pass, Module &M)
+CloneCtx::CloneCtx(Module &M, function_ref<LoopInfo&()> GetLI, function_ref<CallGraph&()> GetCG)
     : ctx(M.getContext()),
       T_size(M.getDataLayout().getIntPtrType(ctx, 0)),
       T_int32(Type::getInt32Ty(ctx)),
       T_void(Type::getVoidTy(ctx)),
       T_psize(PointerType::get(T_size, 0)),
       tbaa_const(tbaa_make_child_with_context(ctx, "jtbaa_const", nullptr, true).first),
-      pass(pass),
       specs(jl_get_llvm_clone_targets()),
       fvars(consume_gv<Function>(M, "jl_sysimg_fvars")),
       gvars(consume_gv<Constant>(M, "jl_sysimg_gvars")),
-      M(M)
+      M(M),
+      GetLI(GetLI),
+      GetCG(GetCG)
 {
     groups.emplace_back(0, specs[0]);
     uint32_t ntargets = specs.size();
@@ -449,7 +434,7 @@ bool CloneCtx::is_vector(FunctionType *ty) const
 uint32_t CloneCtx::collect_func_info(Function &F)
 {
     uint32_t flag = 0;
-    if (!pass->getAnalysis<LoopInfoWrapperPass>(F).getLoopInfo().empty())
+    if (!GetLI().empty())
         flag |= JL_TARGET_CLONE_LOOP;
     if (is_vector(F.getFunctionType())) {
         flag |= JL_TARGET_CLONE_SIMD;
@@ -563,7 +548,7 @@ void CloneCtx::check_partial(Group &grp, Target &tgt)
     auto *next_set = &sets[1];
     // Reduce dispatch by expand the cloning set to functions that are directly called by
     // and calling cloned functions.
-    auto &graph = pass->getAnalysis<CallGraphWrapperPass>().getCallGraph();
+    auto &graph = GetCG();
     while (!cur_set->empty()) {
         for (auto orig_f: *cur_set) {
             // Use the uncloned function since it's already in the call graph
@@ -1079,7 +1064,7 @@ void CloneCtx::emit_metadata()
     }
 }
 
-bool MultiVersioning::runOnModule(Module &M)
+static bool runMultiVersioning(Module &M, function_ref<LoopInfo&()> GetLI, function_ref<CallGraph&()> GetCG)
 {
     // Group targets and identify cloning bases.
     // Also initialize function info maps (we'll update these maps as we go)
@@ -1092,7 +1077,7 @@ bool MultiVersioning::runOnModule(Module &M)
     if (M.getName() == "sysimage")
         return false;
 
-    CloneCtx clone(this, M);
+    CloneCtx clone(M, GetLI, GetCG);
 
     // Collect a list of original functions and clone base functions
     clone.clone_bases();
@@ -1130,16 +1115,60 @@ bool MultiVersioning::runOnModule(Module &M)
     return true;
 }
 
-char MultiVersioning::ID = 0;
-static RegisterPass<MultiVersioning> X("JuliaMultiVersioning", "JuliaMultiVersioning Pass",
+struct MultiVersioningLegacy: public ModulePass {
+    static char ID;
+    MultiVersioningLegacy()
+        : ModulePass(ID)
+    {}
+
+private:
+    bool runOnModule(Module &M) override;
+    void getAnalysisUsage(AnalysisUsage &AU) const override
+    {
+        AU.addRequired<LoopInfoWrapperPass>();
+        AU.addRequired<CallGraphWrapperPass>();
+        AU.addPreserved<LoopInfoWrapperPass>();
+    }
+};
+
+bool MultiVersioningLegacy::runOnModule(Module &M)
+{
+    auto GetLI = [this]() -> LoopInfo & {
+        return getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
+    };
+    auto GetCG = [this]() -> CallGraph & {
+        return getAnalysis<CallGraphWrapperPass>().getCallGraph();
+    };
+    return runMultiVersioning(M, GetLI, GetCG);
+}
+
+
+char MultiVersioningLegacy::ID = 0;
+static RegisterPass<MultiVersioningLegacy> X("JuliaMultiVersioning", "JuliaMultiVersioning Pass",
                                        false /* Only looks at CFG */,
                                        false /* Analysis Pass */);
 
+} // anonymous namespace
+
+PreservedAnalyses MultiVersioning::run(Module &M, ModuleAnalysisManager &AM)
+{
+    auto GetLI = [&]() -> LoopInfo & {
+        return AM.getResult<LoopAnalysis>(M);;
+    };
+    auto GetCG = [&]() -> CallGraph & {
+        return AM.getResult<CallGraphAnalysis>(M);
+    };
+    if (runMultiVersioning(M, GetLI, GetCG)) {
+        auto preserved = PreservedAnalyses::allInSet<CFGAnalyses>();
+        preserved.preserve<LoopAnalysis>();
+        return preserved;
+    }
+    return PreservedAnalyses::all();
 }
 
 Pass *createMultiVersioningPass()
 {
-    return new MultiVersioning();
+    return new MultiVersioningLegacy();
 }
 
 extern "C" JL_DLLEXPORT void LLVMExtraAddMultiVersioningPass_impl(LLVMPassManagerRef PM)

src/passes.h

@@ -64,6 +64,11 @@ struct FinalLowerGCPass : PassInfoMixin<LateLowerGC> {
     static bool isRequired() { return true; }
 };
 
+struct MultiVersioning : PassInfoMixin<MultiVersioning> {
+    PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+    static bool isRequired() { return true; }
+};
+
 struct RemoveJuliaAddrspacesPass : PassInfoMixin<RemoveJuliaAddrspacesPass> {
     PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
     static bool isRequired() { return true; }