[CodeGen] Port PrintMIR to new pass manager (#79440)

The legacy version print machine functions to a string stream, then
output the module and string in `doFinalization`. This patch break
`MIRPrintingPass` into two parts `PrintMIRPreparePass` and
`PrintMIRPass`. `PrintMIRPreparePass` output the original IR in yaml
string, `PrintMIRPass` just print the machine function, so we can avoid
the `doFinalization`.

Author: paperchalice

llvm/include/llvm/CodeGen/MIRPrinter.h

@@ -14,14 +14,33 @@
 #ifndef LLVM_CODEGEN_MIRPRINTER_H
 #define LLVM_CODEGEN_MIRPRINTER_H
 
+#include "llvm/CodeGen/MachinePassManager.h"
+#include "llvm/Support/raw_ostream.h"
+
 namespace llvm {
 
 class MachineBasicBlock;
 class MachineFunction;
 class Module;
-class raw_ostream;
 template <typename T> class SmallVectorImpl;
 
+class PrintMIRPreparePass : public MachinePassInfoMixin<PrintMIRPreparePass> {
+  raw_ostream &OS;
+
+public:
+  PrintMIRPreparePass(raw_ostream &OS = errs()) : OS(OS) {}
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MFAM);
+};
+
+class PrintMIRPass : public MachinePassInfoMixin<PrintMIRPass> {
+  raw_ostream &OS;
+
+public:
+  PrintMIRPass(raw_ostream &OS = errs()) : OS(OS) {}
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+};
+
 /// Print LLVM IR using the MIR serialization format to the given output stream.
 void printMIR(raw_ostream &OS, const Module &M);
 

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -126,8 +126,8 @@ MACHINE_FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis,
 #define MACHINE_FUNCTION_PASS(NAME, PASS_NAME, CONSTRUCTOR)
 #endif
 // MACHINE_FUNCTION_PASS("free-machine-function", FreeMachineFunctionPass, ())
-// MACHINE_FUNCTION_PASS("mir-printer", PrintMIRPass, ())
 MACHINE_FUNCTION_PASS("no-op-machine-function", NoOpMachineFunctionPass, ())
+MACHINE_FUNCTION_PASS("print", PrintMIRPass, ())
 #undef MACHINE_FUNCTION_PASS
 
 // After a pass is converted to new pass manager, its entry should be moved from
@@ -205,7 +205,6 @@ DUMMY_MACHINE_FUNCTION_PASS("machineinstr-printer", MachineFunctionPrinterPass,
                             (OS, Banner))
 DUMMY_MACHINE_FUNCTION_PASS("machinelicm", MachineLICMPass, ())
 DUMMY_MACHINE_FUNCTION_PASS("machineverifier", MachineVerifierPass, (Banner))
-DUMMY_MACHINE_FUNCTION_PASS("mir-printer", PrintMIRPass, ())
 DUMMY_MACHINE_FUNCTION_PASS("mirfs-discriminators", MIRAddFSDiscriminatorsPass,
                             (P))
 DUMMY_MACHINE_FUNCTION_PASS("opt-phis", OptimizePHIsPass, ())

llvm/lib/CodeGen/MIRPrintingPass.cpp

@@ -15,11 +15,20 @@
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/InitializePasses.h"
-#include "llvm/Support/Debug.h"
-#include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
 
+PreservedAnalyses PrintMIRPreparePass::run(Module &M, ModuleAnalysisManager &) {
+  printMIR(OS, M);
+  return PreservedAnalyses::all();
+}
+
+PreservedAnalyses PrintMIRPass::run(MachineFunction &MF,
+                                    MachineFunctionAnalysisManager &) {
+  printMIR(OS, MF);
+  return PreservedAnalyses::all();
+}
+
 namespace {
 
 /// This pass prints out the LLVM IR to an output stream using the MIR

llvm/lib/Passes/PassBuilder.cpp

@@ -89,6 +89,7 @@
 #include "llvm/CodeGen/InterleavedLoadCombine.h"
 #include "llvm/CodeGen/JMCInstrumenter.h"
 #include "llvm/CodeGen/LowerEmuTLS.h"
+#include "llvm/CodeGen/MIRPrinter.h"
 #include "llvm/CodeGen/SafeStack.h"
 #include "llvm/CodeGen/SelectOptimize.h"
 #include "llvm/CodeGen/ShadowStackGCLowering.h"

llvm/tools/llc/NewPMDriver.cpp

@@ -168,6 +168,9 @@ int llvm::compileModuleWithNewPM(
 
   MachineFunctionAnalysisManager MFAM(FAM, MAM);
 
+  ModulePassManager MPM;
+  MachineFunctionPassManager MFPM;
+
   if (!PassPipeline.empty()) {
     // Construct a custom pass pipeline that starts after instruction
     // selection.
@@ -177,21 +180,18 @@ int llvm::compileModuleWithNewPM(
       return 1;
     }
 
-    MachineFunctionPassManager MFPM;
     ExitOnErr(PB.parsePassPipeline(MFPM, PassPipeline));
+    MPM.addPass(PrintMIRPreparePass(*OS));
     MFPM.addPass(PrintMIRPass(*OS));
     MFPM.addPass(FreeMachineFunctionPass());
 
     auto &MMI = MFAM.getResult<MachineModuleAnalysis>(*M);
     if (MIR->parseMachineFunctions(*M, MMI))
       return 1;
 
-    RunPasses(BOS.get(), Out.get(), M.get(), Context, Buffer, nullptr, nullptr,
-              MFPM, MFAM);
+    RunPasses(BOS.get(), Out.get(), M.get(), Context, Buffer, &MPM, &MAM, MFPM,
+              MFAM);
   } else {
-    ModulePassManager MPM;
-    MachineFunctionPassManager MFPM;
-
     ExitOnErr(LLVMTM.buildCodeGenPipeline(MPM, MFPM, MFAM, *OS,
                                           DwoOut ? &DwoOut->os() : nullptr,
                                           FileType, Opt, &PIC));