[NewPM][CodeGen] Refactoring CodeGenPassBuilder

- Remove redundant `std::move`.
- Use a unified function to add passes.
- Support pass substitution.

Author: paperchalice

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -468,7 +468,7 @@ class LLVMTargetMachine : public TargetMachine {
   virtual Error buildCodeGenPipeline(ModulePassManager &, raw_pwrite_stream &,
                                      raw_pwrite_stream *, CodeGenFileType,
                                      const CGPassBuilderOption &,
-                                     PassInstrumentationCallbacks *) {
+                                     PassBuilder &) {
     return make_error<StringError>("buildCodeGenPipeline is not overridden",
                                    inconvertibleErrorCode());
   }

llvm/include/llvm/Target/TargetMachine.h

@@ -24,23 +24,22 @@ class X86CodeGenPassBuilder
 public:
   explicit X86CodeGenPassBuilder(X86TargetMachine &TM,
                                  const CGPassBuilderOption &Opts,
-                                 PassInstrumentationCallbacks *PIC)
-      : CodeGenPassBuilder(TM, Opts, PIC) {}
-  void addPreISel(AddIRPass &addPass) const;
-  void addAsmPrinter(AddMachinePass &, CreateMCStreamer) const;
-  Error addInstSelector(AddMachinePass &) const;
+                                 PassBuilder &PB)
+      : CodeGenPassBuilder(TM, Opts, PB) {}
+  void addPreISel();
+  void addAsmPrinter(CreateMCStreamer);
+  Error addInstSelector();
 };
 
-void X86CodeGenPassBuilder::addPreISel(AddIRPass &addPass) const {
+void X86CodeGenPassBuilder::addPreISel() {
   // TODO: Add passes pre instruction selection.
 }
 
-void X86CodeGenPassBuilder::addAsmPrinter(AddMachinePass &addPass,
-                                          CreateMCStreamer) const {
+void X86CodeGenPassBuilder::addAsmPrinter(CreateMCStreamer) {
   // TODO: Add AsmPrinter.
 }
 
-Error X86CodeGenPassBuilder::addInstSelector(AddMachinePass &) const {
+Error X86CodeGenPassBuilder::addInstSelector() {
   // TODO: Add instruction selector.
   return Error::success();
 }
@@ -49,8 +48,7 @@ Error X86CodeGenPassBuilder::addInstSelector(AddMachinePass &) const {
 
 Error X86TargetMachine::buildCodeGenPipeline(
     ModulePassManager &MPM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
-    CodeGenFileType FileType, const CGPassBuilderOption &Opt,
-    PassInstrumentationCallbacks *PIC) {
-  auto CGPB = X86CodeGenPassBuilder(*this, Opt, PIC);
+    CodeGenFileType FileType, const CGPassBuilderOption &Opt, PassBuilder &PB) {
+  auto CGPB = X86CodeGenPassBuilder(*this, Opt, PB);
   return CGPB.buildPipeline(MPM, Out, DwoOut, FileType);
 }

llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp

@@ -61,7 +61,7 @@ class X86TargetMachine final : public LLVMTargetMachine {
   Error buildCodeGenPipeline(ModulePassManager &, raw_pwrite_stream &,
                              raw_pwrite_stream *, CodeGenFileType,
                              const CGPassBuilderOption &,
-                             PassInstrumentationCallbacks *) override;
+                             PassBuilder &) override;
 
   bool isJIT() const { return IsJIT; }
 

llvm/lib/Target/X86/X86TargetMachine.h

@@ -1,5 +1,5 @@
 ; RUN: llc -mtriple=x86_64-pc-linux-gnu -enable-new-pm -print-pipeline-passes -start-before=mergeicmps -stop-after=gc-lowering -filetype=null %s | FileCheck --match-full-lines %s --check-prefix=NULL
 ; RUN: llc -mtriple=x86_64-pc-linux-gnu -enable-new-pm -print-pipeline-passes -start-before=mergeicmps -stop-after=gc-lowering -o /dev/null %s | FileCheck --match-full-lines %s --check-prefix=OBJ
 
-; NULL: function(mergeicmps,expand-memcmp,gc-lowering)
+; NULL: function(mergeicmps,expand-memcmp,gc-lowering,invalidate<machine-function-info>)
 ; OBJ: function(mergeicmps,expand-memcmp,gc-lowering),PrintMIRPreparePass,function(machine-function(print),invalidate<machine-function-info>)

llvm/test/tools/llc/new-pm/start-stop.ll

@@ -159,7 +159,7 @@ int llvm::compileModuleWithNewPM(
       return 1;
   } else {
     ExitOnErr(LLVMTM.buildCodeGenPipeline(
-        MPM, *OS, DwoOut ? &DwoOut->os() : nullptr, FileType, Opt, &PIC));
+        MPM, *OS, DwoOut ? &DwoOut->os() : nullptr, FileType, Opt, PB));
   }
 
   if (PrintPipelinePasses) {

llvm/tools/llc/NewPMDriver.cpp

@@ -25,6 +25,7 @@ add_llvm_unittest(CodeGenTests
   AMDGPUMetadataTest.cpp
   AsmPrinterDwarfTest.cpp
   CCStateTest.cpp
+  CodeGenPassBuilderTest.cpp
   DIEHashTest.cpp
   DIETest.cpp
   DwarfStringPoolEntryRefTest.cpp

llvm/unittests/CodeGen/CMakeLists.txt

@@ -0,0 +1,126 @@
+//===- llvm/unittest/CodeGen/CodeGenPassBuilderTest.cpp -------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Passes/CodeGenPassBuilder.h"
+#include "llvm/CodeGen/MachinePassManager.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetLowering.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Target/TargetMachine.h"
+
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+class TestTargetMachine : public LLVMTargetMachine {
+public:
+  TestTargetMachine()
+      : LLVMTargetMachine(Target(), "", Triple(""), "", "", TargetOptions(),
+                          Reloc::Static, CodeModel::Small,
+                          CodeGenOptLevel::Default) {}
+};
+
+TestTargetMachine &createTargetMachine() {
+  static TestTargetMachine TM;
+  return TM;
+}
+
+struct DisabledMachineFunctionPass
+    : public PassInfoMixin<DisabledMachineFunctionPass> {
+  PreservedAnalyses run(MachineFunction &, MachineFunctionAnalysisManager &) {
+    return PreservedAnalyses::all();
+  }
+};
+
+struct ReplacedMachineFunctionPass
+    : public PassInfoMixin<ReplacedMachineFunctionPass> {
+  PreservedAnalyses run(MachineFunction &, MachineFunctionAnalysisManager &) {
+    return PreservedAnalyses::all();
+  }
+};
+
+class TestCodeGenPassBuilder
+    : public CodeGenPassBuilder<TestCodeGenPassBuilder, TestTargetMachine> {
+
+public:
+  explicit TestCodeGenPassBuilder()
+      : CodeGenPassBuilder(createTargetMachine(), CGPassBuilderOption(),
+                           nullptr) {
+    // Declare disabled passes in constructor.
+    disablePass<NoOpModulePass>(3); // Disable the third NoOpModulePass.
+    disablePass<DisabledMachineFunctionPass>();
+  }
+
+  // Override substitutePass is also OK.
+  // template <typename PassT> auto substitutePass() {
+  //   if constexpr (std::is_same_v<PassT, ReplacedMachineFunctionPass>)
+  //     return NoOpMachineFunctionPass();
+  //   else
+  //     return;
+  // }
+
+  void buildTestPipeline(ModulePassManager &MPM) {
+    addPass(NoOpModulePass());
+    addPass<NoOpModulePass, NoOpFunctionPass, NoOpModulePass,
+            DisabledMachineFunctionPass, NoOpFunctionPass,
+            NoOpMachineFunctionPass, ReplacedMachineFunctionPass>();
+    mergePassManager();
+    MPM.addPass(std::move(getMPM()));
+    getMPM() = ModulePassManager();
+  }
+};
+
+class CodeGenPassBuilderTest : public testing::Test {
+public:
+  CodeGenPassBuilderTest() {
+    PIC.addClassToPassName(NoOpModulePass::name(), "no-op-module");
+    PIC.addClassToPassName(NoOpFunctionPass::name(), "no-op-function");
+    PIC.addClassToPassName(NoOpMachineFunctionPass::name(),
+                           "no-op-machine-function");
+    PIC.addClassToPassName(DisabledMachineFunctionPass::name(), "disabled");
+    PIC.addClassToPassName(ReplacedMachineFunctionPass::name(), "replaced");
+  }
+
+  void buildPipeline(ModulePassManager &MPM) {
+    TestCodeGenPassBuilder PB;
+    PB.buildTestPipeline(MPM);
+  }
+
+  std::string getPipelineText(ModulePassManager &MPM) {
+    std::string PipelineText;
+    raw_string_ostream OS(PipelineText);
+    MPM.printPipeline(
+        OS, [&](StringRef S) { return PIC.getPassNameForClassName(S); });
+    return PipelineText;
+  }
+
+  PassInstrumentationCallbacks PIC;
+};
+
+} // namespace
+
+using PassBuilderBase =
+    CodeGenPassBuilder<TestCodeGenPassBuilder, TestTargetMachine>;
+
+// Add a specialization to substitute a pass.
+template <>
+template <>
+auto PassBuilderBase::substitutePass<ReplacedMachineFunctionPass>() {
+  return NoOpMachineFunctionPass();
+}
+
+TEST_F(CodeGenPassBuilderTest, Basic) {
+  ModulePassManager MPM;
+  buildPipeline(MPM);
+  EXPECT_EQ(getPipelineText(MPM),
+            "no-op-module,no-op-module,function(no-op-function,no-op-function,"
+            "machine-function(no-op-machine-function,no-op-machine-function))");
+}