[CodeGen][NewPM] Port AsmPrinter to new pass manager

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -16,13 +16,15 @@
 #define LLVM_CODEGEN_ASMPRINTER_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/StaticDataProfileInfo.h"
 #include "llvm/BinaryFormat/Dwarf.h"
 #include "llvm/CodeGen/DwarfStringPoolEntry.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/CodeGen/StackMaps.h"
 #include "llvm/DebugInfo/CodeView/CodeView.h"
 #include "llvm/IR/InlineAsm.h"
@@ -86,7 +88,10 @@ class RemarkStreamer;
 }
 
 /// This class is intended to be used as a driving class for all asm writers.
-class AsmPrinter : public MachineFunctionPass {
+/// Use lightweight RefCountedBase here because AsmPrinter is shared only in
+/// pass manager.
+class AsmPrinter : public RefCountedBase<AsmPrinter>,
+                   public MachineFunctionPass {
 public:
   /// Target machine description.
   TargetMachine &TM;
@@ -209,6 +214,11 @@ class AsmPrinter : public MachineFunctionPass {
 
   StackMaps SM;
 
+  /// If one of these pass managers is not null, then it is in new pass manager.
+  ModuleAnalysisManager *MAM = nullptr;
+  MachineFunctionAnalysisManager *MFAM = nullptr;
+  bool inNewPassManager() const { return MAM || MFAM; }
+
 private:
   /// If generated on the fly this own the instance.
   std::unique_ptr<MachineDominatorTree> OwnedMDT;
@@ -244,7 +254,7 @@ class AsmPrinter : public MachineFunctionPass {
              char &ID = AsmPrinter::ID);
 
 public:
-  ~AsmPrinter() override;
+  virtual ~AsmPrinter();
 
   DwarfDebug *getDwarfDebug() { return DD; }
   DwarfDebug *getDwarfDebug() const { return DD; }
@@ -388,23 +398,43 @@ class AsmPrinter : public MachineFunctionPass {
   // MachineFunctionPass Implementation.
   //===------------------------------------------------------------------===//
 
+  virtual StringRef getPassName() const override;
+
   /// Record analysis usage.
-  void getAnalysisUsage(AnalysisUsage &AU) const override;
+  virtual void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   /// Set up the AsmPrinter when we are working on a new module. If your pass
   /// overrides this, it must make sure to explicitly call this implementation.
-  bool doInitialization(Module &M) override;
+  /// TODO: Keep only the new pass manager doInitialization.
+  virtual bool doInitialization(Module &M) override;
+  virtual void doInitialization(Module &M, ModuleAnalysisManager &MAM) {
+    this->MAM = &MAM;
+    doInitialization(M);
+    this->MAM = nullptr;
+  }
 
   /// Shut down the asmprinter. If you override this in your pass, you must make
   /// sure to call it explicitly.
-  bool doFinalization(Module &M) override;
+  /// TODO: Keep only the new pass manager doFinalization.
+  virtual bool doFinalization(Module &M) override;
+  virtual void doFinalization(Module &M, ModuleAnalysisManager &MAM) {
+    this->MAM = &MAM;
+    doFinalization(M);
+    this->MAM = nullptr;
+  }
 
   /// Emit the specified function out to the OutStreamer.
-  bool runOnMachineFunction(MachineFunction &MF) override {
+  /// TODO: Keep only the new pass manager run.
+  virtual bool runOnMachineFunction(MachineFunction &MF) override {
     SetupMachineFunction(MF);
     emitFunctionBody();
     return false;
   }
+  virtual void run(MachineFunction &MF, MachineFunctionAnalysisManager &MFAM) {
+    this->MFAM = &MFAM;
+    SetupMachineFunction(MF);
+    emitFunctionBody();
+  }
 
   //===------------------------------------------------------------------===//
   // Coarse grained IR lowering routines.
@@ -539,6 +569,7 @@ class AsmPrinter : public MachineFunctionPass {
 
   /// Emit the stack maps.
   void emitStackMaps();
+  void emitStackMaps(Module &M); // For new pass manager version.
 
   //===------------------------------------------------------------------===//
   // Overridable Hooks
@@ -948,6 +979,43 @@ class AsmPrinter : public MachineFunctionPass {
   }
 };
 
+class AsmPrinterInitializePass
+    : public PassInfoMixin<AsmPrinterInitializePass> {
+  IntrusiveRefCntPtr<AsmPrinter> Printer;
+
+public:
+  AsmPrinterInitializePass(IntrusiveRefCntPtr<AsmPrinter> Printer)
+      : Printer(Printer) {}
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
+    Printer->doInitialization(M, MAM);
+    return PreservedAnalyses::all();
+  }
+};
+
+class AsmPrinterPass : public PassInfoMixin<AsmPrinterPass> {
+  IntrusiveRefCntPtr<AsmPrinter> Printer;
+
+public:
+  AsmPrinterPass(IntrusiveRefCntPtr<AsmPrinter> Printer) : Printer(Printer) {}
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM) {
+    Printer->run(MF, MFAM);
+    return PreservedAnalyses::all();
+  }
+};
+
+class AsmPrinterFinalizePass : public PassInfoMixin<AsmPrinterFinalizePass> {
+  IntrusiveRefCntPtr<AsmPrinter> Printer;
+
+public:
+  AsmPrinterFinalizePass(IntrusiveRefCntPtr<AsmPrinter> Printer)
+      : Printer(Printer) {}
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
+    Printer->doFinalization(M, MAM);
+    return PreservedAnalyses::all();
+  }
+};
+
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_ASMPRINTER_H

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -22,6 +22,7 @@
 #include "llvm/Analysis/ScopedNoAliasAA.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/TypeBasedAliasAnalysis.h"
+#include "llvm/CodeGen/AsmPrinter.h"
 #include "llvm/CodeGen/AssignmentTrackingAnalysis.h"
 #include "llvm/CodeGen/BranchFoldingPass.h"
 #include "llvm/CodeGen/CallBrPrepare.h"
@@ -103,6 +104,8 @@
 #include "llvm/IRPrinter/IRPrintingPasses.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCTargetOptions.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Passes/PassBuilder.h"
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Error.h"
@@ -179,8 +182,8 @@ template <typename DerivedT, typename TargetMachineT> class CodeGenPassBuilder {
   }
 
   Error buildPipeline(ModulePassManager &MPM, raw_pwrite_stream &Out,
-                      raw_pwrite_stream *DwoOut,
-                      CodeGenFileType FileType) const;
+                      raw_pwrite_stream *DwoOut, CodeGenFileType FileType,
+                      MCContext &Ctx) const;
 
   PassInstrumentationCallbacks *getPassInstrumentationCallbacks() const {
     return PIC;
@@ -298,6 +301,7 @@ template <typename DerivedT, typename TargetMachineT> class CodeGenPassBuilder {
   TargetMachineT &TM;
   CGPassBuilderOption Opt;
   PassInstrumentationCallbacks *PIC;
+  mutable IntrusiveRefCntPtr<AsmPrinter> PrinterImpl;
 
   template <typename TMC> TMC &getTM() const { return static_cast<TMC &>(TM); }
   CodeGenOptLevel getOptLevel() const { return TM.getOptLevel(); }
@@ -560,7 +564,7 @@ template <typename DerivedT, typename TargetMachineT> class CodeGenPassBuilder {
 template <typename Derived, typename TargetMachineT>
 Error CodeGenPassBuilder<Derived, TargetMachineT>::buildPipeline(
     ModulePassManager &MPM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
-    CodeGenFileType FileType) const {
+    CodeGenFileType FileType, MCContext &Ctx) const {
   auto StartStopInfo = TargetPassConfig::getStartStopInfo(*PIC);
   if (!StartStopInfo)
     return StartStopInfo.takeError();
@@ -569,6 +573,17 @@ Error CodeGenPassBuilder<Derived, TargetMachineT>::buildPipeline(
   bool PrintAsm = TargetPassConfig::willCompleteCodeGenPipeline();
   bool PrintMIR = !PrintAsm && FileType != CodeGenFileType::Null;
 
+  if (PrintAsm) {
+    AddIRPass addIRPass(MPM, derived());
+    Expected<std::unique_ptr<MCStreamer>> MCStreamerOrErr =
+        TM.createMCStreamer(Out, DwoOut, FileType, Ctx);
+    if (auto Err = MCStreamerOrErr.takeError())
+      return Err;
+    PrinterImpl =
+        TM.getTarget().createAsmPrinter(TM, std::move(*MCStreamerOrErr));
+    addIRPass(AsmPrinterInitializePass(PrinterImpl));
+  }
+
   {
     AddIRPass addIRPass(MPM, derived());
     addIRPass(RequireAnalysisPass<MachineModuleAnalysis, Module>());
@@ -577,29 +592,35 @@ Error CodeGenPassBuilder<Derived, TargetMachineT>::buildPipeline(
     addISelPasses(addIRPass);
   }
 
-  AddMachinePass addPass(MPM, derived());
+  // Ensure we destruct `addPass`.
+  {
+    AddMachinePass addPass(MPM, derived());
 
-  if (PrintMIR)
-    addPass(PrintMIRPreparePass(Out), /*Force=*/true);
+    if (PrintMIR)
+      addPass(PrintMIRPreparePass(Out), /*Force=*/true);
 
-  if (auto Err = addCoreISelPasses(addPass))
-    return std::move(Err);
+    if (auto Err = addCoreISelPasses(addPass))
+      return std::move(Err);
 
-  if (auto Err = derived().addMachinePasses(addPass))
-    return std::move(Err);
+    if (auto Err = derived().addMachinePasses(addPass))
+      return std::move(Err);
 
-  if (!Opt.DisableVerify)
-    addPass(MachineVerifierPass());
+    if (!Opt.DisableVerify)
+      addPass(MachineVerifierPass());
 
-  if (PrintAsm) {
-    derived().addAsmPrinter(
-        addPass, [this, &Out, DwoOut, FileType](MCContext &Ctx) {
-          return this->TM.createMCStreamer(Out, DwoOut, FileType, Ctx);
-        });
+    if (PrintAsm)
+      addPass(AsmPrinterPass(PrinterImpl));
+
+    if (PrintMIR)
+      addPass(PrintMIRPass(Out), /*Force=*/true);
   }
 
-  if (PrintMIR)
-    addPass(PrintMIRPass(Out), /*Force=*/true);
+  {
+    AddIRPass addIRPass(MPM, derived());
+    addIRPass(AsmPrinterFinalizePass(PrinterImpl));
+  }
+
+  PrinterImpl.reset();
 
   return verifyStartStop(*StartStopInfo);
 }
@@ -707,6 +728,10 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addIRPasses(
     addPass(ExpandMemCmpPass(&TM));
   }
 
+  // This should be the last IR module pass.
+  // if (TargetPassConfig::willCompleteCodeGenPipeline())
+  //   addPass(AsmPrinterInitializePass(PrinterImpl));
+
   // Run GC lowering passes for builtin collectors
   // TODO: add a pass insertion point here
   addPass(GCLoweringPass());

llvm/include/llvm/Target/TargetMachine.h

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