From eb07600f8eff793664a4d6b181e6f31de1f5f973 Mon Sep 17 00:00:00 2001
From: paperchalice <liujunchang97@outlook.com>
Date: Mon, 8 Apr 2024 20:52:37 +0800
Subject: [NewPM] support `disablePass`, `insertPass` in `CodeGenPassBuilder`
 (#87147)

Currently only machine function passes support `insertPass`, but it
seems to be enough, all targets tune their pipelines when adding machine
function passes.
---
 llvm/include/llvm/Passes/CodeGenPassBuilder.h | 37 +++++++++++++++++----------
 1 file changed, 24 insertions(+), 13 deletions(-)

diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index 00eb9b0..ab231dbb 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -176,9 +176,6 @@ protected:
       // Add Function Pass
       if constexpr (is_detected<is_function_pass_t, PassT>::value) {
         FPM.addPass(std::forward<PassT>(Pass));
-
-        for (auto &C : PB.AfterCallbacks)
-          C(Name);
       } else {
         // Add Module Pass
         if (!FPM.isEmpty()) {
@@ -187,9 +184,6 @@ protected:
         }
 
         MPM.addPass(std::forward<PassT>(Pass));
-
-        for (auto &C : PB.AfterCallbacks)
-          C(Name);
       }
     }
 
@@ -222,9 +216,6 @@ protected:
       // Add Function Pass
       if constexpr (is_detected<is_machine_function_pass_t, PassT>::value) {
         MFPM.addPass(std::forward<PassT>(Pass));
-
-        for (auto &C : PB.AfterCallbacks)
-          C(Name);
       } else {
         // Add Module Pass
         if (!MFPM.isEmpty()) {
@@ -234,10 +225,10 @@ protected:
         }
 
         MPM.addPass(std::forward<PassT>(Pass));
-
-        for (auto &C : PB.AfterCallbacks)
-          C(Name);
       }
+
+      for (auto &C : PB.AfterCallbacks)
+        C(Name, MFPM);
     }
 
   private:
@@ -461,6 +452,24 @@ protected:
   Error addRegAssignmentFast(AddMachinePass &) const;
   Error addRegAssignmentOptimized(AddMachinePass &) const;
 
+  /// Allow the target to disable a specific pass by default.
+  /// Backend can declare unwanted passes in constructor.
+  template <typename... PassTs> void disablePass() {
+    BeforeCallbacks.emplace_back(
+        [](StringRef Name) { return ((Name != PassTs::name()) && ...); });
+  }
+
+  /// Insert InsertedPass pass after TargetPass pass.
+  /// Only machine function passes are supported.
+  template <typename TargetPassT, typename InsertedPassT>
+  void insertPass(InsertedPassT &&Pass) {
+    AfterCallbacks.emplace_back(
+        [&](StringRef Name, MachineFunctionPassManager &MFPM) mutable {
+          if (Name == TargetPassT::name())
+            MFPM.addPass(std::forward<InsertedPassT>(Pass));
+        });
+  }
+
 private:
   DerivedT &derived() { return static_cast<DerivedT &>(*this); }
   const DerivedT &derived() const {
@@ -480,7 +489,9 @@ private:
 
   mutable SmallVector<llvm::unique_function<bool(StringRef)>, 4>
       BeforeCallbacks;
-  mutable SmallVector<llvm::unique_function<void(StringRef)>, 4> AfterCallbacks;
+  mutable SmallVector<
+      llvm::unique_function<void(StringRef, MachineFunctionPassManager &)>, 4>
+      AfterCallbacks;
 
   /// Helper variable for `-start-before/-start-after/-stop-before/-stop-after`
   mutable bool Started = true;
-- 
cgit v1.1