Support eraseIf in pass manager

paperchalice · paperchalice · commit 5364e1c27404 · 2024-09-14T20:29:09.000+08:00
diff --git a/llvm/include/llvm/Analysis/CGSCCPassManager.h b/llvm/include/llvm/Analysis/CGSCCPassManager.h
@@ -346,6 +346,10 @@ class ModuleToPostOrderCGSCCPassAdaptor
 
   static bool isRequired() { return true; }
 
+  bool isEmpty() const { return Pass == nullptr; }
+
+  void eraseIf(function_ref<bool(StringRef)> Pred);
+
 private:
   std::unique_ptr<PassConceptT> Pass;
 };
@@ -488,6 +492,8 @@ class CGSCCToFunctionPassAdaptor
 
   static bool isRequired() { return true; }
 
+  void eraseIf(function_ref<bool(StringRef)> Pred);
+
 private:
   std::unique_ptr<PassConceptT> Pass;
   bool EagerlyInvalidate;
diff --git a/llvm/include/llvm/CodeGen/MachinePassManager.h b/llvm/include/llvm/CodeGen/MachinePassManager.h
@@ -207,6 +207,10 @@ class FunctionToMachineFunctionPassAdaptor
 
   static bool isRequired() { return true; }
 
+  bool isEmpty() const { return Pass == nullptr; }
+
+  void eraseIf(function_ref<bool(StringRef)> Pred);
+
 private:
   std::unique_ptr<PassConceptT> Pass;
 };
diff --git a/llvm/include/llvm/IR/PassManager.h b/llvm/include/llvm/IR/PassManager.h
@@ -218,6 +218,21 @@ class PassManager : public PassInfoMixin<
 
   static bool isRequired() { return true; }
 
+  /// Erase all passes that satisfy the predicate \p Pred.
+  /// For internal use only!
+  void eraseIf(function_ref<bool(StringRef)> Pred) {
+    for (auto I = Passes.begin(); I != Passes.end();) {
+      (*I)->eraseIf(Pred);
+      bool IsSpecial = (*I)->name().ends_with("PassAdaptor") ||
+                       (*I)->name().contains("PassManager");
+      bool PredResult = Pred((*I)->name());
+      if ((!IsSpecial && PredResult) || (IsSpecial && (*I)->isEmpty()))
+        I = Passes.erase(I);
+      else
+        ++I;
+    }
+  }
+
 protected:
   using PassConceptT =
       detail::PassConcept<IRUnitT, AnalysisManagerT, ExtraArgTs...>;
@@ -836,6 +851,10 @@ class ModuleToFunctionPassAdaptor
 
   static bool isRequired() { return true; }
 
+  bool isEmpty() const { return Pass == nullptr; }
+
+  void eraseIf(function_ref<bool(StringRef)> Pred);
+
 private:
   std::unique_ptr<PassConceptT> Pass;
   bool EagerlyInvalidate;
diff --git a/llvm/include/llvm/IR/PassManagerInternal.h b/llvm/include/llvm/IR/PassManagerInternal.h
@@ -59,6 +59,13 @@ struct PassConcept {
   /// To opt-in, pass should implement `static bool isRequired()`. It's no-op
   /// to have `isRequired` always return false since that is the default.
   virtual bool isRequired() const = 0;
+
+  /// Polymorphic method to refurbish pass pipeline.
+  virtual void eraseIf(function_ref<bool(StringRef)> Pred) = 0;
+
+  /// There may be some empty PassManager after erasing,
+  /// use it to remove them.
+  virtual bool isEmpty() const = 0;
 };
 
 /// A template wrapper used to implement the polymorphic API.
@@ -114,6 +121,34 @@ struct PassModel : PassConcept<IRUnitT, AnalysisManagerT, ExtraArgTs...> {
 
   bool isRequired() const override { return passIsRequiredImpl<PassT>(); }
 
+  template <typename T>
+  using has_erase_if_t = decltype(std::declval<T &>().eraseIf(
+      std::declval<function_ref<bool(StringRef)>>()));
+
+  template <typename T>
+  std::enable_if_t<is_detected<has_erase_if_t, T>::value>
+  eraseIfImpl(function_ref<bool(StringRef)> Pred) {
+    Pass.eraseIf(Pred);
+  }
+
+  template <typename T>
+  std::enable_if_t<!is_detected<has_erase_if_t, T>::value>
+  eraseIfImpl(function_ref<bool(StringRef)>) {}
+
+  void eraseIf(function_ref<bool(StringRef)> Pred) override {
+    eraseIfImpl<PassT>(Pred);
+  }
+
+  template <typename T>
+  using has_is_empty_t = decltype(std::declval<T &>().isEmpty());
+
+  bool isEmpty() const override {
+    if constexpr (is_detected<has_is_empty_t, PassT>::value)
+      return Pass.isEmpty();
+    else
+      return false;
+  }
+
   PassT Pass;
 };
 
diff --git a/llvm/include/llvm/Transforms/Scalar/LoopPassManager.h b/llvm/include/llvm/Transforms/Scalar/LoopPassManager.h
@@ -134,6 +134,10 @@ class PassManager<Loop, LoopAnalysisManager, LoopStandardAnalysisResults &,
 
   static bool isRequired() { return true; }
 
+  /// Erase all passes that satisfy the predicate \p Pred.
+  /// For internal use only!
+  void eraseIf(function_ref<bool(StringRef)> Pred);
+
   size_t getNumLoopPasses() const { return LoopPasses.size(); }
   size_t getNumLoopNestPasses() const { return LoopNestPasses.size(); }
 
@@ -425,6 +429,10 @@ class FunctionToLoopPassAdaptor
 
   static bool isRequired() { return true; }
 
+  bool isEmpty() const { return Pass == nullptr; }
+
+  void eraseIf(function_ref<bool(StringRef)> Pred);
+
   bool isLoopNestMode() const { return LoopNestMode; }
 
 private:
diff --git a/llvm/lib/Analysis/CGSCCPassManager.cpp b/llvm/lib/Analysis/CGSCCPassManager.cpp
@@ -578,6 +578,17 @@ PreservedAnalyses CGSCCToFunctionPassAdaptor::run(LazyCallGraph::SCC &C,
   return PA;
 }
 
+void CGSCCToFunctionPassAdaptor::eraseIf(function_ref<bool(StringRef)> Pred) {
+  StringRef PassName = Pass->name();
+  if (PassName.contains("PassManager") || PassName.ends_with("PassAdaptor")) {
+    Pass->eraseIf(Pred);
+    if (Pass->isEmpty())
+      Pass.reset();
+  } else if (Pred(PassName)) {
+    Pass.reset();
+  }
+}
+
 bool CGSCCAnalysisManagerModuleProxy::Result::invalidate(
     Module &M, const PreservedAnalyses &PA,
     ModuleAnalysisManager::Invalidator &Inv) {
@@ -753,6 +764,18 @@ bool FunctionAnalysisManagerCGSCCProxy::Result::invalidate(
   return false;
 }
 
+void ModuleToPostOrderCGSCCPassAdaptor::eraseIf(
+    function_ref<bool(StringRef)> Pred) {
+  StringRef PassName = Pass->name();
+  if (PassName.contains("PassManager") || PassName.ends_with("PassAdaptor")) {
+    Pass->eraseIf(Pred);
+    if (Pass->isEmpty())
+      Pass.reset();
+  } else if (Pred(PassName)) {
+    Pass.reset();
+  }
+}
+
 } // end namespace llvm
 
 /// When a new SCC is created for the graph we first update the
diff --git a/llvm/lib/IR/PassManager.cpp b/llvm/lib/IR/PassManager.cpp
@@ -145,6 +145,17 @@ PreservedAnalyses ModuleToFunctionPassAdaptor::run(Module &M,
   return PA;
 }
 
+void ModuleToFunctionPassAdaptor::eraseIf(function_ref<bool(StringRef)> Pred) {
+  StringRef PassName = Pass->name();
+  if (PassName.contains("PassManager") || PassName.ends_with("PassAdaptor")) {
+    Pass->eraseIf(Pred);
+    if (Pass->isEmpty())
+      Pass.reset();
+  } else if (Pred(PassName)) {
+    Pass.reset();
+  }
+}
+
 template <>
 void llvm::printIRUnitNameForStackTrace<Module>(raw_ostream &OS,
                                                 const Module &IR) {
diff --git a/llvm/lib/Transforms/Scalar/LoopPassManager.cpp b/llvm/lib/Transforms/Scalar/LoopPassManager.cpp
@@ -63,6 +63,39 @@ void PassManager<Loop, LoopAnalysisManager, LoopStandardAnalysisResults &,
   }
 }
 
+void PassManager<Loop, LoopAnalysisManager, LoopStandardAnalysisResults &,
+                 LPMUpdater &>::eraseIf(function_ref<bool(StringRef)> Pred) {
+  std::vector<char> IsLoopNestPassVec(
+      static_cast<size_t>(IsLoopNestPass.size()));
+  for (unsigned Idx = 0, Sz = IsLoopNestPass.size(); Idx != Sz; ++Idx)
+    IsLoopNestPassVec[Idx] = IsLoopNestPass[Idx];
+
+  auto ILP = LoopPasses.begin();
+  auto ILNP = LoopNestPasses.begin();
+  for (auto I = IsLoopNestPassVec.begin(); I != IsLoopNestPassVec.end();) {
+    if (*I) {
+      if (Pred((*ILNP)->name())) {
+        I = IsLoopNestPassVec.erase(I);
+        ILNP = LoopNestPasses.erase(ILNP);
+        continue;
+      }
+      ++ILNP;
+    } else {
+      if (Pred((*ILP)->name())) {
+        I = IsLoopNestPassVec.erase(I);
+        ILP = LoopPasses.erase(ILP);
+        continue;
+      }
+      ++ILP;
+    }
+    ++I;
+  }
+
+  IsLoopNestPass.clear();
+  for (const auto I : IsLoopNestPassVec)
+    IsLoopNestPass.push_back(I);
+}
+
 // Run both loop passes and loop-nest passes on top-level loop \p L.
 PreservedAnalyses
 LoopPassManager::runWithLoopNestPasses(Loop &L, LoopAnalysisManager &AM,
@@ -363,6 +396,17 @@ PreservedAnalyses FunctionToLoopPassAdaptor::run(Function &F,
   return PA;
 }
 
+void FunctionToLoopPassAdaptor::eraseIf(function_ref<bool(StringRef)> Pred) {
+  StringRef PassName = Pass->name();
+  if (PassName.contains("PassManager") || PassName.ends_with("PassAdaptor")) {
+    Pass->eraseIf(Pred);
+    if (Pass->isEmpty())
+      Pass.reset();
+  } else if (Pred(PassName)) {
+    Pass.reset();
+  }
+}
+
 PrintLoopPass::PrintLoopPass() : OS(dbgs()) {}
 PrintLoopPass::PrintLoopPass(raw_ostream &OS, const std::string &Banner)
     : OS(OS), Banner(Banner) {}
