[CodeGen] Allow CodeGenPassBuilder to add module pass after function pass (#77084)

paperchalice · web-flow · commit ae1c1ed6af8d · 2024-01-12T08:37:12.000+08:00
In fact, there are several backends, e.g. AArch64, AMDGPU etc. add
module pass after function pass, this patch removes this constraint.
This patch also adds a simple unit test for `CodeGenPassBuilder`.
diff --git a/llvm/include/llvm/CodeGen/CodeGenPassBuilder.h b/llvm/include/llvm/CodeGen/CodeGenPassBuilder.h
@@ -175,45 +175,33 @@ template <typename DerivedT> class CodeGenPassBuilder {
   // Function object to maintain state while adding codegen IR passes.
   class AddIRPass {
   public:
-    AddIRPass(ModulePassManager &MPM, bool DebugPM, bool Check = true)
-        : MPM(MPM) {
-      if (Check)
-        AddingFunctionPasses = false;
-    }
+    AddIRPass(ModulePassManager &MPM) : MPM(MPM) {}
     ~AddIRPass() {
-      MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
-    }
-
-    // Add Function Pass
-    template <typename PassT>
-    std::enable_if_t<is_detected<is_function_pass_t, PassT>::value>
-    operator()(PassT &&Pass) {
-      if (AddingFunctionPasses && !*AddingFunctionPasses)
-        AddingFunctionPasses = true;
-      FPM.addPass(std::forward<PassT>(Pass));
+      if (!FPM.isEmpty())
+        MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
     }
 
-    // Add Module Pass
-    template <typename PassT>
-    std::enable_if_t<is_detected<is_module_pass_t, PassT>::value &&
-                     !is_detected<is_function_pass_t, PassT>::value>
-    operator()(PassT &&Pass) {
-      assert((!AddingFunctionPasses || !*AddingFunctionPasses) &&
-             "could not add module pass after adding function pass");
-      MPM.addPass(std::forward<PassT>(Pass));
+    template <typename PassT> void operator()(PassT &&Pass) {
+      static_assert((is_detected<is_function_pass_t, PassT>::value ||
+                     is_detected<is_module_pass_t, PassT>::value) &&
+                    "Only module pass and function pass are supported.");
+
+      // Add Function Pass
+      if constexpr (is_detected<is_function_pass_t, PassT>::value) {
+        FPM.addPass(std::forward<PassT>(Pass));
+      } else {
+        // Add Module Pass
+        if (!FPM.isEmpty()) {
+          MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+          FPM = FunctionPassManager();
+        }
+        MPM.addPass(std::forward<PassT>(Pass));
+      }
     }
 
   private:
     ModulePassManager &MPM;
     FunctionPassManager FPM;
-    // The codegen IR pipeline are mostly function passes with the exceptions of
-    // a few loop and module passes. `AddingFunctionPasses` make sures that
-    // we could only add module passes at the beginning of the pipeline. Once
-    // we begin adding function passes, we could no longer add module passes.
-    // This special-casing introduces less adaptor passes. If we have the need
-    // of adding module passes after function passes, we could change the
-    // implementation to accommodate that.
-    std::optional<bool> AddingFunctionPasses;
   };
 
   // Function object to maintain state while adding codegen machine passes.
@@ -488,7 +476,7 @@ Error CodeGenPassBuilder<Derived>::buildPipeline(
     ModulePassManager &MPM, MachineFunctionPassManager &MFPM,
     raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
     CodeGenFileType FileType) const {
-  AddIRPass addIRPass(MPM, Opt.DebugPM);
+  AddIRPass addIRPass(MPM);
   // `ProfileSummaryInfo` is always valid.
   addIRPass(RequireAnalysisPass<ProfileSummaryAnalysis, Module>());
   addIRPass(RequireAnalysisPass<CollectorMetadataAnalysis, Module>());
@@ -627,8 +615,8 @@ void CodeGenPassBuilder<Derived>::addIRPasses(AddIRPass &addPass) const {
 
   // Run loop strength reduction before anything else.
   if (getOptLevel() != CodeGenOptLevel::None && !Opt.DisableLSR) {
-    addPass(createFunctionToLoopPassAdaptor(
-        LoopStrengthReducePass(), /*UseMemorySSA*/ true, Opt.DebugPM));
+    addPass(createFunctionToLoopPassAdaptor(LoopStrengthReducePass(),
+                                            /*UseMemorySSA=*/true));
     // FIXME: use -stop-after so we could remove PrintLSR
     if (Opt.PrintLSR)
       addPass(PrintFunctionPass(dbgs(), "\n\n*** Code after LSR ***\n"));
@@ -647,9 +635,6 @@ void CodeGenPassBuilder<Derived>::addIRPasses(AddIRPass &addPass) const {
   // Run GC lowering passes for builtin collectors
   // TODO: add a pass insertion point here
   addPass(GCLoweringPass());
-  // FIXME: `ShadowStackGCLoweringPass` now is a
-  // module pass, so it will trigger assertion.
-  // See comment of `AddingFunctionPasses`
   addPass(ShadowStackGCLoweringPass());
   addPass(LowerConstantIntrinsicsPass());
 
diff --git a/llvm/include/llvm/Passes/PassBuilder.h b/llvm/include/llvm/Passes/PassBuilder.h
@@ -743,6 +743,90 @@ bool parseAnalysisUtilityPasses(
 
   return false;
 }
+
+// These are special since they are only for testing purposes.
+
+/// No-op module pass which does nothing.
+struct NoOpModulePass : PassInfoMixin<NoOpModulePass> {
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &) {
+    return PreservedAnalyses::all();
+  }
+};
+
+/// No-op module analysis.
+class NoOpModuleAnalysis : public AnalysisInfoMixin<NoOpModuleAnalysis> {
+  friend AnalysisInfoMixin<NoOpModuleAnalysis>;
+  static AnalysisKey Key;
+
+public:
+  struct Result {};
+  Result run(Module &, ModuleAnalysisManager &) { return Result(); }
+};
+
+/// No-op CGSCC pass which does nothing.
+struct NoOpCGSCCPass : PassInfoMixin<NoOpCGSCCPass> {
+  PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &,
+                        LazyCallGraph &, CGSCCUpdateResult &UR) {
+    return PreservedAnalyses::all();
+  }
+};
+
+/// No-op CGSCC analysis.
+class NoOpCGSCCAnalysis : public AnalysisInfoMixin<NoOpCGSCCAnalysis> {
+  friend AnalysisInfoMixin<NoOpCGSCCAnalysis>;
+  static AnalysisKey Key;
+
+public:
+  struct Result {};
+  Result run(LazyCallGraph::SCC &, CGSCCAnalysisManager &, LazyCallGraph &G) {
+    return Result();
+  }
+};
+
+/// No-op function pass which does nothing.
+struct NoOpFunctionPass : PassInfoMixin<NoOpFunctionPass> {
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &) {
+    return PreservedAnalyses::all();
+  }
+};
+
+/// No-op function analysis.
+class NoOpFunctionAnalysis : public AnalysisInfoMixin<NoOpFunctionAnalysis> {
+  friend AnalysisInfoMixin<NoOpFunctionAnalysis>;
+  static AnalysisKey Key;
+
+public:
+  struct Result {};
+  Result run(Function &, FunctionAnalysisManager &) { return Result(); }
+};
+
+/// No-op loop nest pass which does nothing.
+struct NoOpLoopNestPass : PassInfoMixin<NoOpLoopNestPass> {
+  PreservedAnalyses run(LoopNest &L, LoopAnalysisManager &,
+                        LoopStandardAnalysisResults &, LPMUpdater &) {
+    return PreservedAnalyses::all();
+  }
+};
+
+/// No-op loop pass which does nothing.
+struct NoOpLoopPass : PassInfoMixin<NoOpLoopPass> {
+  PreservedAnalyses run(Loop &L, LoopAnalysisManager &,
+                        LoopStandardAnalysisResults &, LPMUpdater &) {
+    return PreservedAnalyses::all();
+  }
+};
+
+/// No-op loop analysis.
+class NoOpLoopAnalysis : public AnalysisInfoMixin<NoOpLoopAnalysis> {
+  friend AnalysisInfoMixin<NoOpLoopAnalysis>;
+  static AnalysisKey Key;
+
+public:
+  struct Result {};
+  Result run(Loop &, LoopAnalysisManager &, LoopStandardAnalysisResults &) {
+    return Result();
+  }
+};
 }
 
 #endif
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
@@ -300,109 +300,13 @@ cl::opt<bool> PrintPipelinePasses(
              "(best-effort only)."));
 } // namespace llvm
 
-namespace {
-
-// The following passes/analyses have custom names, otherwise their name will
-// include `(anonymous namespace)`. These are special since they are only for
-// testing purposes and don't live in a header file.
-
-/// No-op module pass which does nothing.
-struct NoOpModulePass : PassInfoMixin<NoOpModulePass> {
-  PreservedAnalyses run(Module &M, ModuleAnalysisManager &) {
-    return PreservedAnalyses::all();
-  }
-
-  static StringRef name() { return "NoOpModulePass"; }
-};
-
-/// No-op module analysis.
-class NoOpModuleAnalysis : public AnalysisInfoMixin<NoOpModuleAnalysis> {
-  friend AnalysisInfoMixin<NoOpModuleAnalysis>;
-  static AnalysisKey Key;
-
-public:
-  struct Result {};
-  Result run(Module &, ModuleAnalysisManager &) { return Result(); }
-  static StringRef name() { return "NoOpModuleAnalysis"; }
-};
-
-/// No-op CGSCC pass which does nothing.
-struct NoOpCGSCCPass : PassInfoMixin<NoOpCGSCCPass> {
-  PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &,
-                        LazyCallGraph &, CGSCCUpdateResult &UR) {
-    return PreservedAnalyses::all();
-  }
-  static StringRef name() { return "NoOpCGSCCPass"; }
-};
-
-/// No-op CGSCC analysis.
-class NoOpCGSCCAnalysis : public AnalysisInfoMixin<NoOpCGSCCAnalysis> {
-  friend AnalysisInfoMixin<NoOpCGSCCAnalysis>;
-  static AnalysisKey Key;
-
-public:
-  struct Result {};
-  Result run(LazyCallGraph::SCC &, CGSCCAnalysisManager &, LazyCallGraph &G) {
-    return Result();
-  }
-  static StringRef name() { return "NoOpCGSCCAnalysis"; }
-};
-
-/// No-op function pass which does nothing.
-struct NoOpFunctionPass : PassInfoMixin<NoOpFunctionPass> {
-  PreservedAnalyses run(Function &F, FunctionAnalysisManager &) {
-    return PreservedAnalyses::all();
-  }
-  static StringRef name() { return "NoOpFunctionPass"; }
-};
-
-/// No-op function analysis.
-class NoOpFunctionAnalysis : public AnalysisInfoMixin<NoOpFunctionAnalysis> {
-  friend AnalysisInfoMixin<NoOpFunctionAnalysis>;
-  static AnalysisKey Key;
-
-public:
-  struct Result {};
-  Result run(Function &, FunctionAnalysisManager &) { return Result(); }
-  static StringRef name() { return "NoOpFunctionAnalysis"; }
-};
-
-/// No-op loop nest pass which does nothing.
-struct NoOpLoopNestPass : PassInfoMixin<NoOpLoopNestPass> {
-  PreservedAnalyses run(LoopNest &L, LoopAnalysisManager &,
-                        LoopStandardAnalysisResults &, LPMUpdater &) {
-    return PreservedAnalyses::all();
-  }
-  static StringRef name() { return "NoOpLoopNestPass"; }
-};
-
-/// No-op loop pass which does nothing.
-struct NoOpLoopPass : PassInfoMixin<NoOpLoopPass> {
-  PreservedAnalyses run(Loop &L, LoopAnalysisManager &,
-                        LoopStandardAnalysisResults &, LPMUpdater &) {
-    return PreservedAnalyses::all();
-  }
-  static StringRef name() { return "NoOpLoopPass"; }
-};
-
-/// No-op loop analysis.
-class NoOpLoopAnalysis : public AnalysisInfoMixin<NoOpLoopAnalysis> {
-  friend AnalysisInfoMixin<NoOpLoopAnalysis>;
-  static AnalysisKey Key;
-
-public:
-  struct Result {};
-  Result run(Loop &, LoopAnalysisManager &, LoopStandardAnalysisResults &) {
-    return Result();
-  }
-  static StringRef name() { return "NoOpLoopAnalysis"; }
-};
-
 AnalysisKey NoOpModuleAnalysis::Key;
 AnalysisKey NoOpCGSCCAnalysis::Key;
 AnalysisKey NoOpFunctionAnalysis::Key;
 AnalysisKey NoOpLoopAnalysis::Key;
 
+namespace {
+
 /// Whether or not we should populate a PassInstrumentationCallbacks's class to
 /// pass name map.
 ///
diff --git a/llvm/unittests/CodeGen/CMakeLists.txt b/llvm/unittests/CodeGen/CMakeLists.txt
@@ -7,13 +7,16 @@ set(LLVM_LINK_COMPONENTS
   CodeGenTypes
   Core
   FileCheck
+  IRPrinter
   MC
   MIRParser
   Passes
+  ScalarOpts
   SelectionDAG
   Support
   Target
   TargetParser
+  TransformUtils
   )
 
 add_llvm_unittest(CodeGenTests
@@ -22,6 +25,7 @@ add_llvm_unittest(CodeGenTests
   AMDGPUMetadataTest.cpp
   AsmPrinterDwarfTest.cpp
   CCStateTest.cpp
+  CodeGenPassBuilderTest.cpp
   DIEHashTest.cpp
   DIETest.cpp
   DwarfStringPoolEntryRefTest.cpp
diff --git a/llvm/unittests/CodeGen/CodeGenPassBuilderTest.cpp b/llvm/unittests/CodeGen/CodeGenPassBuilderTest.cpp
