Make AllocOpt NewPM compatible (#44273)

JuliaLang · Feb 20, 2022 · 7fc65b4 · 7fc65b4
1 parent 3d3858f
commit 7fc65b4
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 26 deletions.
diff --git a/src/aotcompile.cpp b/src/aotcompile.cpp
@@ -891,6 +891,10 @@ static void registerCallbacks(PassBuilder &PB) {
                 PM.addPass(LateLowerGC());
                 return true;
             }
+            if (Name == "AllocOpt") {
+                PM.addPass(AllocOptPass());
+                return true;
+            }
             return false;
         });
 

diff --git a/src/llvm-alloc-opt.cpp b/src/llvm-alloc-opt.cpp
@@ -31,6 +31,7 @@
 #include "julia_internal.h"
 #include "llvm-pass-helpers.h"
 #include "llvm-alloc-helpers.h"
+#include "passes.h"
 
 #include <map>
 #include <set>
@@ -85,13 +86,7 @@ static void removeGCPreserve(CallInst *call, Instruction *val)
  * * Handle jl_box*
  */
 
-struct AllocOpt : public FunctionPass, public JuliaPassContext {
-    static char ID;
-    AllocOpt()
-        : FunctionPass(ID)
-    {
-        llvm::initializeDominatorTreeWrapperPassPass(*PassRegistry::getPassRegistry());
-    }
+struct AllocOpt : public JuliaPassContext {
 
     const DataLayout *DL;
 
@@ -100,22 +95,15 @@ struct AllocOpt : public FunctionPass, public JuliaPassContext {
 
     Type *T_int64;
 
-private:
-    bool doInitialization(Module &m) override;
-    bool runOnFunction(Function &F) override;
-    void getAnalysisUsage(AnalysisUsage &AU) const override
-    {
-        FunctionPass::getAnalysisUsage(AU);
-        AU.addRequired<DominatorTreeWrapperPass>();
-        AU.addPreserved<DominatorTreeWrapperPass>();
-        AU.setPreservesCFG();
-    }
+    bool doInitialization(Module &m);
+    bool runOnFunction(Function &F, function_ref<DominatorTree&()> GetDT);
 };
 
 struct Optimizer {
-    Optimizer(Function &F, AllocOpt &pass)
+    Optimizer(Function &F, AllocOpt &pass, function_ref<DominatorTree&()> GetDT)
         : F(F),
-          pass(pass)
+          pass(pass),
+          GetDT(std::move(GetDT))
     {}
 
     void initialize();
@@ -143,11 +131,12 @@ struct Optimizer {
     Function &F;
     AllocOpt &pass;
     DominatorTree *_DT = nullptr;
+    function_ref<DominatorTree &()> GetDT;
 
     DominatorTree &getDomTree()
     {
         if (!_DT)
-            _DT = &pass.getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+            _DT = &GetDT();
         return *_DT;
     }
     struct Lifetime {
@@ -1159,26 +1148,62 @@ bool AllocOpt::doInitialization(Module &M)
     return true;
 }
 
-bool AllocOpt::runOnFunction(Function &F)
+bool AllocOpt::runOnFunction(Function &F, function_ref<DominatorTree&()> GetDT)
 {
     if (!alloc_obj_func)
         return false;
-    Optimizer optimizer(F, *this);
+    Optimizer optimizer(F, *this, std::move(GetDT));
     optimizer.initialize();
     optimizer.optimizeAll();
     return optimizer.finalize();
 }
 
-char AllocOpt::ID = 0;
-static RegisterPass<AllocOpt> X("AllocOpt", "Promote heap allocation to stack",
+struct AllocOptLegacy : public FunctionPass {
+    static char ID;
+    AllocOpt opt;
+    AllocOptLegacy() : FunctionPass(ID) {
+        llvm::initializeDominatorTreeWrapperPassPass(*PassRegistry::getPassRegistry());
+    }
+    bool doInitialization(Module &m) override {
+        return opt.doInitialization(m);
+    }
+    bool runOnFunction(Function &F) override {
+        return opt.runOnFunction(F, [this]() -> DominatorTree & {return getAnalysis<DominatorTreeWrapperPass>().getDomTree();});
+    }
+    void getAnalysisUsage(AnalysisUsage &AU) const override
+    {
+        FunctionPass::getAnalysisUsage(AU);
+        AU.addRequired<DominatorTreeWrapperPass>();
+        AU.addPreserved<DominatorTreeWrapperPass>();
+        AU.setPreservesCFG();
+    }
+};
+
+char AllocOptLegacy::ID = 0;
+static RegisterPass<AllocOptLegacy> X("AllocOpt", "Promote heap allocation to stack",
                                 false /* Only looks at CFG */,
                                 false /* Analysis Pass */);
 
 }
 
 Pass *createAllocOptPass()
 {
-    return new AllocOpt();
+    return new AllocOptLegacy();
+}
+
+PreservedAnalyses AllocOptPass::run(Function &F, FunctionAnalysisManager &AM) {
+    AllocOpt opt;
+    bool modified = opt.doInitialization(*F.getParent());
+    if (opt.runOnFunction(F, [&]()->DominatorTree &{ return AM.getResult<DominatorTreeAnalysis>(F); })) {
+        modified = true;
+    }
+    if (modified) {
+        auto preserved = PreservedAnalyses::allInSet<CFGAnalyses>();
+        preserved.preserve<DominatorTreeAnalysis>();
+        return preserved;
+    } else {
+        return PreservedAnalyses::all();
+    }
 }
 
 extern "C" JL_DLLEXPORT void LLVMExtraAddAllocOptPass_impl(LLVMPassManagerRef PM)

diff --git a/src/llvm-julia-licm.cpp b/src/llvm-julia-licm.cpp
@@ -205,7 +205,12 @@ PreservedAnalyses JuliaLICMPass::run(Loop &L, LoopAnalysisManager &AM,
         return AR.LI;
     };
     auto juliaLICM = JuliaLICM(GetDT, GetLI);
-    juliaLICM.runOnLoop(&L);
+    if (juliaLICM.runOnLoop(&L)) {
+        auto preserved = PreservedAnalyses::allInSet<CFGAnalyses>();
+        preserved.preserve<LoopAnalysis>();
+        preserved.preserve<DominatorTreeAnalysis>();
+        return preserved;
+    }
     return PreservedAnalyses::all();
 }
 

diff --git a/src/passes.h b/src/passes.h
@@ -23,6 +23,10 @@ struct LateLowerGC : PassInfoMixin<LateLowerGC> {
     static bool isRequired() { return true; }
 };
 
+struct AllocOptPass : PassInfoMixin<AllocOptPass> {
+    PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
 // Module Passes
 struct CPUFeatures : PassInfoMixin<CPUFeatures> {
     PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);

diff --git a/test/llvmpasses/alloc-opt-gcframe.jl b/test/llvmpasses/alloc-opt-gcframe.jl
@@ -1,6 +1,7 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
 
 # RUN: julia --startup-file=no %s | opt -enable-new-pm=0 -load libjulia-codegen%shlibext -AllocOpt -LateLowerGCFrame -FinalLowerGC -S - | FileCheck %s
+# RUN: julia --startup-file=no %s | opt -enable-new-pm=1 --load-pass-plugin=libjulia-codegen%shlibext -passes='function(AllocOpt,LateLowerGCFrame),FinalLowerGC' -S - | FileCheck %s
 
 isz = sizeof(UInt) == 8 ? "i64" : "i32"
 

diff --git a/test/llvmpasses/alloc-opt-pass.jl b/test/llvmpasses/alloc-opt-pass.jl
@@ -1,6 +1,7 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
 
 # RUN: julia --startup-file=no %s | opt -enable-new-pm=0 -load libjulia-codegen%shlibext -AllocOpt -S - | FileCheck %s
+# RUN: julia --startup-file=no %s | opt -enable-new-pm=1 --load-pass-plugin=libjulia-codegen%shlibext -passes='function(AllocOpt)' -S - | FileCheck %s
 
 isz = sizeof(UInt) == 8 ? "i64" : "i32"