Make more passes NewPM compatible (#44277)

* Make PropagateAddrspaces and RemoveAddrspaces NewPM compatible

* Fix LateGCLower analysis invalidation

* Fix missing return

Author: pchintalapudi

src/aotcompile.cpp

@@ -895,6 +895,10 @@ static void registerCallbacks(PassBuilder &PB) {
                 PM.addPass(AllocOptPass());
                 return true;
             }
+            if (Name == "PropagateJuliaAddrspaces") {
+                PM.addPass(PropagateJuliaAddrspacesPass());
+                return true;
+            }
             return false;
         });
 
@@ -917,6 +921,10 @@ static void registerCallbacks(PassBuilder &PB) {
                 PM.addPass(FinalLowerGCPass());
                 return true;
             }
+            if (Name == "RemoveJuliaAddrspaces") {
+                PM.addPass(RemoveJuliaAddrspacesPass());
+                return true;
+            }
             return false;
         });
 

src/llvm-late-gc-lowering.cpp

@@ -326,7 +326,7 @@ struct LateLowerGCFrame:  private JuliaPassContext {
     LateLowerGCFrame(function_ref<DominatorTree &()> GetDT) : GetDT(GetDT) {}
 
 public:
-    bool runOnFunction(Function &F);
+    bool runOnFunction(Function &F, bool *CFGModified = nullptr);
 
 private:
     CallInst *pgcstack;
@@ -358,7 +358,7 @@ struct LateLowerGCFrame:  private JuliaPassContext {
     void PlaceGCFrameStore(State &S, unsigned R, unsigned MinColorRoot, const std::vector<int> &Colors, Value *GCFrame, Instruction *InsertBefore);
     void PlaceGCFrameStores(State &S, unsigned MinColorRoot, const std::vector<int> &Colors, Value *GCFrame);
     void PlaceRootsAndUpdateCalls(std::vector<int> &Colors, State &S, std::map<Value *, std::pair<int, int>>);
-    bool CleanupIR(Function &F, State *S=nullptr);
+    bool CleanupIR(Function &F, State *S, bool *CFGModified);
     void NoteUseChain(State &S, BBState &BBS, User *TheUser);
     SmallVector<int, 1> GetPHIRefinements(PHINode *phi, State &S);
     void FixUpRefinements(ArrayRef<int> PHINumbers, State &S);
@@ -2234,7 +2234,7 @@ MDNode *createMutableTBAAAccessTag(MDNode *Tag) {
 }
 
 
-bool LateLowerGCFrame::CleanupIR(Function &F, State *S) {
+bool LateLowerGCFrame::CleanupIR(Function &F, State *S, bool *CFGModified) {
     bool ChangesMade = false;
     // We create one alloca for all the jlcall frames that haven't been processed
     // yet. LLVM would merge them anyway later, so might as well save it a bit
@@ -2460,6 +2460,9 @@ bool LateLowerGCFrame::CleanupIR(Function &F, State *S) {
             CI->eraseFromParent();
             continue;
         }
+        if (CFGModified) {
+            *CFGModified = true;
+        }
         IRBuilder<> builder(CI);
         builder.SetCurrentDebugLocation(CI->getDebugLoc());
         auto parBits = builder.CreateAnd(EmitLoadTag(builder, parent), 3);
@@ -2675,22 +2678,22 @@ void LateLowerGCFrame::PlaceRootsAndUpdateCalls(std::vector<int> &Colors, State
     }
 }
 
-bool LateLowerGCFrame::runOnFunction(Function &F) {
+bool LateLowerGCFrame::runOnFunction(Function &F, bool *CFGModified) {
     initAll(*F.getParent());
     LLVM_DEBUG(dbgs() << "GC ROOT PLACEMENT: Processing function " << F.getName() << "\n");
     if (!pgcstack_getter)
-        return CleanupIR(F);
+        return CleanupIR(F, nullptr, CFGModified);
 
     pgcstack = getPGCstack(F);
     if (!pgcstack)
-        return CleanupIR(F);
+        return CleanupIR(F, nullptr, CFGModified);
 
     State S = LocalScan(F);
     ComputeLiveness(S);
     std::vector<int> Colors = ColorRoots(S);
     std::map<Value *, std::pair<int, int>> CallFrames; // = OptimizeCallFrames(S, Ordering);
     PlaceRootsAndUpdateCalls(Colors, S, CallFrames);
-    CleanupIR(F, &S);
+    CleanupIR(F, &S, CFGModified);
     return true;
 }
 
@@ -2708,7 +2711,14 @@ PreservedAnalyses LateLowerGC::run(Function &F, FunctionAnalysisManager &AM)
         return AM.getResult<DominatorTreeAnalysis>(F);
     };
     auto lateLowerGCFrame = LateLowerGCFrame(GetDT);
-    lateLowerGCFrame.runOnFunction(F);
+    bool CFGModified = false;
+    if (lateLowerGCFrame.runOnFunction(F, &CFGModified)) {
+        if (CFGModified) {
+            return PreservedAnalyses::none();
+        } else {
+            return PreservedAnalyses::allInSet<CFGAnalyses>();
+        }
+    }
     return PreservedAnalyses::all();
 }
 

src/llvm-propagate-addrspaces.cpp

@@ -24,6 +24,7 @@
 
 #include "codegen_shared.h"
 #include "julia.h"
+#include "passes.h"
 
 #define DEBUG_TYPE "propagate_julia_addrspaces"
 
@@ -40,16 +41,13 @@ using namespace llvm;
     optimizations.
 */
 
-struct PropagateJuliaAddrspaces : public FunctionPass, public InstVisitor<PropagateJuliaAddrspaces> {
-    static char ID;
+struct PropagateJuliaAddrspacesVisitor : public InstVisitor<PropagateJuliaAddrspacesVisitor> {
     DenseMap<Value *, Value *> LiftingMap;
     SmallPtrSet<Value *, 4> Visited;
     std::vector<Instruction *> ToDelete;
     std::vector<std::pair<Instruction *, Instruction *>> ToInsert;
-    PropagateJuliaAddrspaces() : FunctionPass(ID) {};
 
 public:
-    bool runOnFunction(Function &F) override;
     Value *LiftPointer(Value *V, Type *LocTy = nullptr, Instruction *InsertPt=nullptr);
     void visitMemop(Instruction &I, Type *T, unsigned OpIndex);
     void visitLoadInst(LoadInst &LI);
@@ -63,19 +61,6 @@ struct PropagateJuliaAddrspaces : public FunctionPass, public InstVisitor<Propag
     void PoisonValues(std::vector<Value *> &Worklist);
 };
 
-bool PropagateJuliaAddrspaces::runOnFunction(Function &F) {
-    visit(F);
-    for (auto it : ToInsert)
-        it.first->insertBefore(it.second);
-    for (Instruction *I : ToDelete)
-        I->eraseFromParent();
-    ToInsert.clear();
-    ToDelete.clear();
-    LiftingMap.clear();
-    Visited.clear();
-    return true;
-}
-
 static unsigned getValueAddrSpace(Value *V) {
     return cast<PointerType>(V->getType())->getAddressSpace();
 }
@@ -84,7 +69,7 @@ static bool isSpecialAS(unsigned AS) {
     return AddressSpace::FirstSpecial <= AS && AS <= AddressSpace::LastSpecial;
 }
 
-void PropagateJuliaAddrspaces::PoisonValues(std::vector<Value *> &Worklist) {
+void PropagateJuliaAddrspacesVisitor::PoisonValues(std::vector<Value *> &Worklist) {
     while (!Worklist.empty()) {
         Value *CurrentV = Worklist.back();
         Worklist.pop_back();
@@ -97,7 +82,7 @@ void PropagateJuliaAddrspaces::PoisonValues(std::vector<Value *> &Worklist) {
     }
 }
 
-Value *PropagateJuliaAddrspaces::LiftPointer(Value *V, Type *LocTy, Instruction *InsertPt) {
+Value *PropagateJuliaAddrspacesVisitor::LiftPointer(Value *V, Type *LocTy, Instruction *InsertPt) {
     SmallVector<Value *, 4> Stack;
     std::vector<Value *> Worklist;
     std::set<Value *> LocalVisited;
@@ -232,7 +217,7 @@ Value *PropagateJuliaAddrspaces::LiftPointer(Value *V, Type *LocTy, Instruction
     return CollapseCastsAndLift(V, InsertPt);
 }
 
-void PropagateJuliaAddrspaces::visitMemop(Instruction &I, Type *T, unsigned OpIndex) {
+void PropagateJuliaAddrspacesVisitor::visitMemop(Instruction &I, Type *T, unsigned OpIndex) {
     Value *Original = I.getOperand(OpIndex);
     unsigned AS = Original->getType()->getPointerAddressSpace();
     if (!isSpecialAS(AS))
@@ -243,23 +228,23 @@ void PropagateJuliaAddrspaces::visitMemop(Instruction &I, Type *T, unsigned OpIn
     I.setOperand(OpIndex, Replacement);
 }
 
-void PropagateJuliaAddrspaces::visitLoadInst(LoadInst &LI) {
+void PropagateJuliaAddrspacesVisitor::visitLoadInst(LoadInst &LI) {
     visitMemop(LI, LI.getType(), LoadInst::getPointerOperandIndex());
 }
 
-void PropagateJuliaAddrspaces::visitStoreInst(StoreInst &SI) {
+void PropagateJuliaAddrspacesVisitor::visitStoreInst(StoreInst &SI) {
     visitMemop(SI, SI.getValueOperand()->getType(), StoreInst::getPointerOperandIndex());
 }
 
-void PropagateJuliaAddrspaces::visitAtomicCmpXchgInst(AtomicCmpXchgInst &SI) {
+void PropagateJuliaAddrspacesVisitor::visitAtomicCmpXchgInst(AtomicCmpXchgInst &SI) {
     visitMemop(SI, SI.getNewValOperand()->getType(), AtomicCmpXchgInst::getPointerOperandIndex());
 }
 
-void PropagateJuliaAddrspaces::visitAtomicRMWInst(AtomicRMWInst &SI) {
+void PropagateJuliaAddrspacesVisitor::visitAtomicRMWInst(AtomicRMWInst &SI) {
     visitMemop(SI, SI.getType(), AtomicRMWInst::getPointerOperandIndex());
 }
 
-void PropagateJuliaAddrspaces::visitMemSetInst(MemSetInst &MI) {
+void PropagateJuliaAddrspacesVisitor::visitMemSetInst(MemSetInst &MI) {
     unsigned AS = MI.getDestAddressSpace();
     if (!isSpecialAS(AS))
         return;
@@ -272,7 +257,7 @@ void PropagateJuliaAddrspaces::visitMemSetInst(MemSetInst &MI) {
     MI.setArgOperand(0, Replacement);
 }
 
-void PropagateJuliaAddrspaces::visitMemTransferInst(MemTransferInst &MTI) {
+void PropagateJuliaAddrspacesVisitor::visitMemTransferInst(MemTransferInst &MTI) {
     unsigned DestAS = MTI.getDestAddressSpace();
     unsigned SrcAS = MTI.getSourceAddressSpace();
     if (!isSpecialAS(DestAS) && !isSpecialAS(SrcAS))
@@ -299,11 +284,42 @@ void PropagateJuliaAddrspaces::visitMemTransferInst(MemTransferInst &MTI) {
     MTI.setArgOperand(1, Src);
 }
 
-char PropagateJuliaAddrspaces::ID = 0;
-static RegisterPass<PropagateJuliaAddrspaces> X("PropagateJuliaAddrspaces", "Propagate (non-)rootedness information", false, false);
+bool propagateJuliaAddrspaces(Function &F) {
+    PropagateJuliaAddrspacesVisitor visitor;
+    visitor.visit(F);
+    for (auto it : visitor.ToInsert)
+        it.first->insertBefore(it.second);
+    for (Instruction *I : visitor.ToDelete)
+        I->eraseFromParent();
+    visitor.ToInsert.clear();
+    visitor.ToDelete.clear();
+    visitor.LiftingMap.clear();
+    visitor.Visited.clear();
+    return true;
+}
+
+struct PropagateJuliaAddrspacesLegacy : FunctionPass {
+    static char ID;
+
+    PropagateJuliaAddrspacesLegacy() : FunctionPass(ID) {}
+    bool runOnFunction(Function &F) override {
+        return propagateJuliaAddrspaces(F);
+    }
+};
+
+char PropagateJuliaAddrspacesLegacy::ID = 0;
+static RegisterPass<PropagateJuliaAddrspacesLegacy> X("PropagateJuliaAddrspaces", "Propagate (non-)rootedness information", false, false);
 
 Pass *createPropagateJuliaAddrspaces() {
-    return new PropagateJuliaAddrspaces();
+    return new PropagateJuliaAddrspacesLegacy();
+}
+
+PreservedAnalyses PropagateJuliaAddrspacesPass::run(Function &F, FunctionAnalysisManager &AM) {
+    if (propagateJuliaAddrspaces(F)) {
+        return PreservedAnalyses::allInSet<CFGAnalyses>();
+    } else {
+        return PreservedAnalyses::all();
+    }
 }
 
 extern "C" JL_DLLEXPORT void LLVMExtraAddPropagateJuliaAddrspaces_impl(LLVMPassManagerRef PM)

src/llvm-remove-addrspaces.cpp

@@ -14,6 +14,7 @@
 
 #include "codegen_shared.h"
 #include "julia.h"
+#include "passes.h"
 
 #define DEBUG_TYPE "remove_addrspaces"
 
@@ -231,18 +232,7 @@ unsigned removeAllAddrspaces(unsigned AS)
     return AddressSpace::Generic;
 }
 
-struct RemoveAddrspacesPass : public ModulePass {
-    static char ID;
-    AddrspaceRemapFunction ASRemapper;
-    RemoveAddrspacesPass(
-            AddrspaceRemapFunction ASRemapper = removeAllAddrspaces)
-        : ModulePass(ID), ASRemapper(ASRemapper){};
-
-public:
-    bool runOnModule(Module &M) override;
-};
-
-bool RemoveAddrspacesPass::runOnModule(Module &M)
+bool removeAddrspaces(Module &M, AddrspaceRemapFunction ASRemapper)
 {
     ValueToValueMapTy VMap;
     AddrspaceRemoveTypeRemapper TypeRemapper(ASRemapper);
@@ -457,8 +447,22 @@ bool RemoveAddrspacesPass::runOnModule(Module &M)
     return true;
 }
 
-char RemoveAddrspacesPass::ID = 0;
-static RegisterPass<RemoveAddrspacesPass>
+
+struct RemoveAddrspacesPassLegacy : public ModulePass {
+    static char ID;
+    AddrspaceRemapFunction ASRemapper;
+    RemoveAddrspacesPassLegacy(
+            AddrspaceRemapFunction ASRemapper = removeAllAddrspaces)
+        : ModulePass(ID), ASRemapper(ASRemapper){};
+
+public:
+    bool runOnModule(Module &M) override {
+        return removeAddrspaces(M, ASRemapper);
+    }
+};
+
+char RemoveAddrspacesPassLegacy::ID = 0;
+static RegisterPass<RemoveAddrspacesPassLegacy>
         X("RemoveAddrspaces",
           "Remove IR address space information.",
           false,
@@ -467,7 +471,17 @@ static RegisterPass<RemoveAddrspacesPass>
 Pass *createRemoveAddrspacesPass(
         AddrspaceRemapFunction ASRemapper = removeAllAddrspaces)
 {
-    return new RemoveAddrspacesPass(ASRemapper);
+    return new RemoveAddrspacesPassLegacy(ASRemapper);
+}
+
+RemoveAddrspacesPass::RemoveAddrspacesPass() : RemoveAddrspacesPass(removeAllAddrspaces) {}
+
+PreservedAnalyses RemoveAddrspacesPass::run(Module &M, ModuleAnalysisManager &AM) {
+    if (removeAddrspaces(M, ASRemapper)) {
+        return PreservedAnalyses::allInSet<CFGAnalyses>();
+    } else {
+        return PreservedAnalyses::all();
+    }
 }
 
 
@@ -483,24 +497,28 @@ unsigned removeJuliaAddrspaces(unsigned AS)
         return AS;
 }
 
-struct RemoveJuliaAddrspacesPass : public ModulePass {
+struct RemoveJuliaAddrspacesPassLegacy : public ModulePass {
     static char ID;
-    RemoveAddrspacesPass Pass;
-    RemoveJuliaAddrspacesPass() : ModulePass(ID), Pass(removeJuliaAddrspaces){};
+    RemoveAddrspacesPassLegacy Pass;
+    RemoveJuliaAddrspacesPassLegacy() : ModulePass(ID), Pass(removeJuliaAddrspaces){};
 
-    bool runOnModule(Module &M) { return Pass.runOnModule(M); }
+    bool runOnModule(Module &M) override { return Pass.runOnModule(M); }
 };
 
-char RemoveJuliaAddrspacesPass::ID = 0;
-static RegisterPass<RemoveJuliaAddrspacesPass>
+char RemoveJuliaAddrspacesPassLegacy::ID = 0;
+static RegisterPass<RemoveJuliaAddrspacesPassLegacy>
         Y("RemoveJuliaAddrspaces",
           "Remove IR address space information.",
           false,
           false);
 
 Pass *createRemoveJuliaAddrspacesPass()
 {
-    return new RemoveJuliaAddrspacesPass();
+    return new RemoveJuliaAddrspacesPassLegacy();
+}
+
+PreservedAnalyses RemoveJuliaAddrspacesPass::run(Module &M, ModuleAnalysisManager &AM) {
+    return RemoveAddrspacesPass(removeJuliaAddrspaces).run(M, AM);
 }
 
 extern "C" JL_DLLEXPORT void LLVMExtraAddRemoveJuliaAddrspacesPass_impl(LLVMPassManagerRef PM)

src/passes.h

@@ -26,6 +26,10 @@ struct LateLowerGC : PassInfoMixin<LateLowerGC> {
 struct AllocOptPass : PassInfoMixin<AllocOptPass> {
     PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
+struct PropagateJuliaAddrspacesPass : PassInfoMixin<PropagateJuliaAddrspacesPass> {
+    PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+    static bool isRequired() { return true; }
+};
 
 // Module Passes
 struct CPUFeatures : PassInfoMixin<CPUFeatures> {
@@ -46,6 +50,21 @@ struct FinalLowerGCPass : PassInfoMixin<LateLowerGC> {
     static bool isRequired() { return true; }
 };
 
+struct RemoveJuliaAddrspacesPass : PassInfoMixin<RemoveJuliaAddrspacesPass> {
+    PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+    static bool isRequired() { return true; }
+};
+
+struct RemoveAddrspacesPass : PassInfoMixin<RemoveAddrspacesPass> {
+
+    std::function<unsigned(unsigned)> ASRemapper;
+    RemoveAddrspacesPass();
+    RemoveAddrspacesPass(std::function<unsigned(unsigned)> ASRemapper) : ASRemapper(std::move(ASRemapper)) {}
+
+    PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+    static bool isRequired() { return true; }
+};
+
 // Loop Passes
 struct JuliaLICMPass : PassInfoMixin<JuliaLICMPass> {
     PreservedAnalyses run(Loop &L, LoopAnalysisManager &AM,