[llvm] [NewPM] support `disablePass`, `insertPass` in `CodeGenPassBuilder` (PR #87147)

[via llvm-commits]

https://github.com/paperchalice updated https://github.com/llvm/llvm-project/pull/87147

>From f14724946269902cd3f2a1a3e8ae33de44a0c153 Mon Sep 17 00:00:00 2001
From: PaperChalice <liujunchang97 at outlook.com>
Date: Sat, 30 Mar 2024 17:09:32 +0800
Subject: [PATCH] [NewPM] support `disablePass`, `insertPass` in
 `CodeGenPassBuilder` 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 00eb9b096a9356..ab231dbbb6a68c 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -176,9 +176,6 @@ template <typename DerivedT> class CodeGenPassBuilder {
       // 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 @@ template <typename DerivedT> class CodeGenPassBuilder {
         }
 
         MPM.addPass(std::forward<PassT>(Pass));
-
-        for (auto &C : PB.AfterCallbacks)
-          C(Name);
       }
     }
 
@@ -222,9 +216,6 @@ template <typename DerivedT> class CodeGenPassBuilder {
       // 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 @@ template <typename DerivedT> class CodeGenPassBuilder {
         }
 
         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 @@ template <typename DerivedT> class CodeGenPassBuilder {
   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 @@ template <typename DerivedT> class CodeGenPassBuilder {
 
   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;