Support eraseIf in pass manager

paperchalice · paperchalice · commit 7ca2eed1117d · 2024-11-19T21:57:20.000+08:00
diff --git a/llvm/include/llvm/Analysis/CGSCCPassManager.h b/llvm/include/llvm/Analysis/CGSCCPassManager.h
@@ -311,7 +311,10 @@ struct CGSCCUpdateResult {
 /// pass over the module to enable a \c FunctionAnalysisManager to be used
 /// within this run safely.
 class ModuleToPostOrderCGSCCPassAdaptor
-    : public PassInfoMixin<ModuleToPostOrderCGSCCPassAdaptor> {
+    : public PassInfoMixin<ModuleToPostOrderCGSCCPassAdaptor>,
+      public AdaptorMixin<ModuleToPostOrderCGSCCPassAdaptor> {
+  friend AdaptorMixin<ModuleToPostOrderCGSCCPassAdaptor>;
+
 public:
   using PassConceptT =
       detail::PassConcept<LazyCallGraph::SCC, CGSCCAnalysisManager,
@@ -441,7 +444,10 @@ LazyCallGraph::SCC &updateCGAndAnalysisManagerForCGSCCPass(
 /// pass over the SCC to enable a \c FunctionAnalysisManager to be used
 /// within this run safely.
 class CGSCCToFunctionPassAdaptor
-    : public PassInfoMixin<CGSCCToFunctionPassAdaptor> {
+    : public PassInfoMixin<CGSCCToFunctionPassAdaptor>,
+      public AdaptorMixin<CGSCCToFunctionPassAdaptor> {
+  friend AdaptorMixin<CGSCCToFunctionPassAdaptor>;
+
 public:
   using PassConceptT = detail::PassConcept<Function, FunctionAnalysisManager>;
 
diff --git a/llvm/include/llvm/CodeGen/MachinePassManager.h b/llvm/include/llvm/CodeGen/MachinePassManager.h
@@ -191,7 +191,10 @@ class FunctionAnalysisManagerMachineFunctionProxy
 };
 
 class FunctionToMachineFunctionPassAdaptor
-    : public PassInfoMixin<FunctionToMachineFunctionPassAdaptor> {
+    : public PassInfoMixin<FunctionToMachineFunctionPassAdaptor>,
+      public AdaptorMixin<FunctionToMachineFunctionPassAdaptor> {
+  friend AdaptorMixin<FunctionToMachineFunctionPassAdaptor>;
+
 public:
   using PassConceptT =
       detail::PassConcept<MachineFunction, MachineFunctionAnalysisManager>;
diff --git a/llvm/include/llvm/IR/PassManager.h b/llvm/include/llvm/IR/PassManager.h
@@ -218,6 +218,22 @@ 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();) {
+      auto &P = *I;
+      P->eraseIf(Pred);
+      bool IsSpecial = P->name().ends_with("PassAdaptor") ||
+                       P->name().contains("PassManager");
+      bool PredResult = Pred(P->name());
+      if ((!IsSpecial && PredResult) || (IsSpecial && P->isEmpty()))
+        I = Passes.erase(I);
+      else
+        ++I;
+    }
+  }
+
 protected:
   using PassConceptT =
       detail::PassConcept<IRUnitT, AnalysisManagerT, ExtraArgTs...>;
@@ -797,6 +813,28 @@ extern template class OuterAnalysisManagerProxy<ModuleAnalysisManager,
 using ModuleAnalysisManagerFunctionProxy =
     OuterAnalysisManagerProxy<ModuleAnalysisManager, Function>;
 
+/// Simple mix-in for pass adaptor. If adaptor contains only a single pass
+/// instance in it, then it can inherit this mix-in to get default `isEmpty()`
+/// and `eraseIf` implementation. This mix-in must have access to the `Pass`
+/// member in adaptor.
+template <typename DerivedT> struct AdaptorMixin {
+  bool isEmpty() const { return derived().Pass == nullptr; }
+
+  void eraseIf(function_ref<bool(StringRef)> Pred) {
+    StringRef PassName = derived().Pass->name();
+    if (PassName.contains("PassManager") || PassName.ends_with("PassAdaptor")) {
+      derived().Pass->eraseIf(Pred);
+      if (derived().Pass->isEmpty())
+        derived().Pass.reset();
+    } else if (Pred(PassName)) {
+      derived().Pass.reset();
+    }
+  }
+
+private:
+  DerivedT &derived() { return *static_cast<Derived *>(this); }
+};
+
 /// Trivial adaptor that maps from a module to its functions.
 ///
 /// Designed to allow composition of a FunctionPass(Manager) and
@@ -821,7 +859,10 @@ using ModuleAnalysisManagerFunctionProxy =
 /// analyses are not invalidated while the function passes are running, so they
 /// may be stale.  Function analyses will not be stale.
 class ModuleToFunctionPassAdaptor
-    : public PassInfoMixin<ModuleToFunctionPassAdaptor> {
+    : public PassInfoMixin<ModuleToFunctionPassAdaptor>,
+      public AdaptorMixin<ModuleToFunctionPassAdaptor> {
+  friend AdaptorMixin<ModuleToFunctionPassAdaptor>;
+
 public:
   using PassConceptT = detail::PassConcept<Function, FunctionAnalysisManager>;
 
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,33 @@ 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();
+    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(); }
 
@@ -399,7 +403,10 @@ std::optional<PreservedAnalyses> LoopPassManager::runSinglePass(
 /// \fn createLoopFunctionToLoopPassAdaptor to see when loop mode and loop-nest
 /// mode are used.
 class FunctionToLoopPassAdaptor
-    : public PassInfoMixin<FunctionToLoopPassAdaptor> {
+    : public PassInfoMixin<FunctionToLoopPassAdaptor>,
+      public AdaptorMixin<FunctionToLoopPassAdaptor> {
+  friend AdaptorMixin<FunctionToLoopPassAdaptor>;
+
 public:
   using PassConceptT =
       detail::PassConcept<Loop, LoopAnalysisManager,
diff --git a/llvm/lib/Transforms/Scalar/LoopPassManager.cpp b/llvm/lib/Transforms/Scalar/LoopPassManager.cpp
@@ -62,6 +62,45 @@ void PassManager<Loop, LoopAnalysisManager, LoopStandardAnalysisResults &,
   }
 }
 
+void PassManager<Loop, LoopAnalysisManager, LoopStandardAnalysisResults &,
+                 LPMUpdater &>::eraseIf(function_ref<bool(StringRef)> Pred) {
+  assert(LoopPasses.size() + LoopNestPasses.size() == IsLoopNestPass.size() &&
+         "Wrong precondition!");
+
+  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);
+
+  assert(LoopPasses.size() + LoopNestPasses.size() == IsLoopNestPass.size() &&
+         "Wrong postcondition!");
+}
+
 // Run both loop passes and loop-nest passes on top-level loop \p L.
 PreservedAnalyses
 LoopPassManager::runWithLoopNestPasses(Loop &L, LoopAnalysisManager &AM,
