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

Author: paperchalice

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -16,11 +16,13 @@
 #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/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"
@@ -84,7 +86,9 @@ 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:
   /// Target machine description.
   TargetMachine &TM;
@@ -180,8 +184,6 @@ class AsmPrinter : public MachineFunctionPass {
   /// List of symbols to be inserted into PC sections.
   DenseMap<const MDNode *, SmallVector<const MCSymbol *>> PCSectionsSymbols;
 
-  static char ID;
-
 protected:
   MCSymbol *CurrentFnBegin = nullptr;
 
@@ -201,6 +203,10 @@ class AsmPrinter : public MachineFunctionPass {
 
   StackMaps SM;
 
+  MachineFunctionPass *MFPass = nullptr;
+  ModuleAnalysisManager *MAM = nullptr;
+  MachineFunctionAnalysisManager *MFAM = nullptr;
+
 private:
   /// If generated on the fly this own the instance.
   std::unique_ptr<MachineDominatorTree> OwnedMDT;
@@ -235,7 +241,7 @@ class AsmPrinter : public MachineFunctionPass {
   explicit AsmPrinter(TargetMachine &TM, std::unique_ptr<MCStreamer> Streamer);
 
 public:
-  ~AsmPrinter() override;
+  virtual ~AsmPrinter();
 
   DwarfDebug *getDwarfDebug() { return DD; }
   DwarfDebug *getDwarfDebug() const { return DD; }
@@ -375,24 +381,31 @@ class AsmPrinter : public MachineFunctionPass {
   // MachineFunctionPass Implementation.
   //===------------------------------------------------------------------===//
 
+  virtual StringRef getPassName() const;
+
   /// Record analysis usage.
-  void getAnalysisUsage(AnalysisUsage &AU) const override;
+  virtual void getAnalysisUsage(AnalysisUsage &AU) const;
 
   /// 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;
+  virtual bool doInitialization(Module &M);
 
   /// Shut down the asmprinter. If you override this in your pass, you must make
   /// sure to call it explicitly.
-  bool doFinalization(Module &M) override;
+  virtual bool doFinalization(Module &M);
 
   /// Emit the specified function out to the OutStreamer.
-  bool runOnMachineFunction(MachineFunction &MF) override {
+  virtual bool runOnMachineFunction(MachineFunction &MF) {
     SetupMachineFunction(MF);
     emitFunctionBody();
     return false;
   }
 
+  void setPass(MachineFunctionPass *P) { MFPass = P; }
+  void setMAM(ModuleAnalysisManager &AM) { MAM = &AM; }
+  void clearMAM() { MAM = nullptr; }
+  void setMFAM(MachineFunctionAnalysisManager &AM) { MFAM = &AM; }
+
   //===------------------------------------------------------------------===//
   // Coarse grained IR lowering routines.
   //===------------------------------------------------------------------===//
@@ -523,6 +536,7 @@ class AsmPrinter : public MachineFunctionPass {
 
   /// Emit the stack maps.
   void emitStackMaps();
+  void emitStackMaps(Module &M); // For new pass manager version.
 
   //===------------------------------------------------------------------===//
   // Overridable Hooks
@@ -922,6 +936,78 @@ class AsmPrinter : public MachineFunctionPass {
   }
 };
 
+class AsmPrinterInitializePass
+    : public PassInfoMixin<AsmPrinterInitializePass> {
+  IntrusiveRefCntPtr<AsmPrinter> Printer;
+
+public:
+  explicit AsmPrinterInitializePass(IntrusiveRefCntPtr<AsmPrinter> AP)
+      : Printer(std::move(AP)) {}
+
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+};
+
+class AsmPrinterPass : public PassInfoMixin<AsmPrinterPass> {
+  IntrusiveRefCntPtr<AsmPrinter> Printer;
+
+public:
+  explicit AsmPrinterPass(IntrusiveRefCntPtr<AsmPrinter> AP)
+      : Printer(std::move(AP)) {}
+
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+};
+
+class AsmPrinterFinalizePass : public PassInfoMixin<AsmPrinterFinalizePass> {
+  IntrusiveRefCntPtr<AsmPrinter> Printer;
+
+public:
+  explicit AsmPrinterFinalizePass(IntrusiveRefCntPtr<AsmPrinter> AP)
+      : Printer(std::move(AP)) {}
+
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &);
+};
+
+class AsmPrinterLegacy : public MachineFunctionPass {
+  std::unique_ptr<AsmPrinter> Printer;
+
+public:
+  static char ID;
+
+  explicit AsmPrinterLegacy(std::unique_ptr<AsmPrinter> AP);
+
+  AsmPrinter &getPrinter() { return *Printer; }
+
+  //===------------------------------------------------------------------===//
+  // MachineFunctionPass Implementation.
+  //===------------------------------------------------------------------===//
+
+  /// Record analysis usage.
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    MachineFunctionPass::getAnalysisUsage(AU);
+    Printer->getAnalysisUsage(AU);
+  }
+
+  /// 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 {
+    return Printer->doInitialization(M);
+  }
+
+  /// Shut down the asmprinter. If you override this in your pass, you must make
+  /// sure to call it explicitly.
+  bool doFinalization(Module &M) override { return Printer->doFinalization(M); }
+
+  /// Emit the specified function out to the OutStreamer.
+  bool runOnMachineFunction(MachineFunction &MF) override {
+    return Printer->runOnMachineFunction(MF);
+  }
+
+  StringRef getPassName() const override { return Printer->getPassName(); }
+};
+
+AsmPrinterLegacy *createAsmPrinterLegacy(std::unique_ptr<AsmPrinter> AP);
+
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_ASMPRINTER_H

llvm/include/llvm/MC/TargetRegistry.h

@@ -37,6 +37,7 @@
 namespace llvm {
 
 class AsmPrinter;
+class AsmPrinterLegacy;
 class MCAsmBackend;
 class MCAsmInfo;
 class MCAsmParser;
@@ -58,6 +59,8 @@ class MCTargetStreamer;
 class raw_ostream;
 class TargetMachine;
 class TargetOptions;
+
+AsmPrinterLegacy *createAsmPrinterLegacy(std::unique_ptr<AsmPrinter> AP);
 namespace mca {
 class CustomBehaviour;
 class InstrPostProcess;
@@ -166,8 +169,12 @@ class Target {
   // If it weren't for layering issues (this header is in llvm/Support, but
   // depends on MC?) this should take the Streamer by value rather than rvalue
   // reference.
-  using AsmPrinterCtorTy = AsmPrinter *(*)(
-      TargetMachine &TM, std::unique_ptr<MCStreamer> &&Streamer);
+  using AsmPrinterLegacyCtorTy =
+      AsmPrinterLegacy *(*)(TargetMachine &TM,
+                            std::unique_ptr<MCStreamer> &&Streamer);
+  using AsmPrinterCtorTy =
+      AsmPrinter *(*)(TargetMachine &TM,
+                      std::unique_ptr<MCStreamer> &&Streamer);
   using MCAsmBackendCtorTy = MCAsmBackend *(*)(const Target &T,
                                                const MCSubtargetInfo &STI,
                                                const MCRegisterInfo &MRI,
@@ -288,6 +295,7 @@ class Target {
 
   /// AsmPrinterCtorFn - Construction function for this target's AsmPrinter,
   /// if registered.
+  AsmPrinterLegacyCtorTy AsmPrinterLegacyCtorFn;
   AsmPrinterCtorTy AsmPrinterCtorFn;
 
   /// MCDisassemblerCtorFn - Construction function for this target's
@@ -487,6 +495,16 @@ class Target {
 
   /// createAsmPrinter - Create a target specific assembly printer pass.  This
   /// takes ownership of the MCStreamer object.
+  AsmPrinterLegacy *
+  createAsmPrinterLegacy(TargetMachine &TM,
+                         std::unique_ptr<MCStreamer> &&Streamer) const {
+    if (!AsmPrinterLegacyCtorFn)
+      return nullptr;
+    return AsmPrinterLegacyCtorFn(TM, std::move(Streamer));
+  }
+
+  /// Same as `createAsmPrinter`, but only create AsmPrinter, for new pass
+  /// manager.
   AsmPrinter *createAsmPrinter(TargetMachine &TM,
                                std::unique_ptr<MCStreamer> &&Streamer) const {
     if (!AsmPrinterCtorFn)
@@ -862,6 +880,9 @@ struct TargetRegistry {
   ///
   /// @param T - The target being registered.
   /// @param Fn - A function to construct an AsmPrinter for the target.
+  static void RegisterAsmPrinter(Target &T, Target::AsmPrinterLegacyCtorTy Fn) {
+    T.AsmPrinterLegacyCtorFn = Fn;
+  }
   static void RegisterAsmPrinter(Target &T, Target::AsmPrinterCtorTy Fn) {
     T.AsmPrinterCtorFn = Fn;
   }
@@ -1311,9 +1332,16 @@ template <class MCAsmParserImpl> struct RegisterMCAsmParser {
 template <class AsmPrinterImpl> struct RegisterAsmPrinter {
   RegisterAsmPrinter(Target &T) {
     TargetRegistry::RegisterAsmPrinter(T, &Allocator);
+    TargetRegistry::RegisterAsmPrinter(T, &AllocatorLegacy);
   }
 
 private:
+  static AsmPrinterLegacy *
+  AllocatorLegacy(TargetMachine &TM, std::unique_ptr<MCStreamer> &&Streamer) {
+    auto *P = new AsmPrinterImpl(TM, std::move(Streamer));
+    return createAsmPrinterLegacy(std::unique_ptr<AsmPrinter>(P));
+  }
+
   static AsmPrinter *Allocator(TargetMachine &TM,
                                std::unique_ptr<MCStreamer> &&Streamer) {
     return new AsmPrinterImpl(TM, std::move(Streamer));

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/CallBrPrepare.h"
 #include "llvm/CodeGen/CodeGenPrepare.h"
@@ -76,6 +77,7 @@
 #include "llvm/IRPrinter/IRPrintingPasses.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCTargetOptions.h"
+#include "llvm/Passes/PassBuilder.h"
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Error.h"
@@ -133,9 +135,8 @@ namespace llvm {
 template <typename DerivedT, typename TargetMachineT> class CodeGenPassBuilder {
 public:
   explicit CodeGenPassBuilder(TargetMachineT &TM,
-                              const CGPassBuilderOption &Opts,
-                              PassInstrumentationCallbacks *PIC)
-      : TM(TM), Opt(Opts), PIC(PIC) {
+                              const CGPassBuilderOption &Opts, PassBuilder &PB)
+      : TM(TM), Opt(Opts), PB(PB), PIC(PB.getPassInstrumentationCallbacks()) {
     // Target could set CGPassBuilderOption::MISchedPostRA to true to achieve
     //     substitutePass(&PostRASchedulerID, &PostMachineSchedulerID)
 
@@ -152,8 +153,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;
@@ -270,7 +271,9 @@ template <typename DerivedT, typename TargetMachineT> class CodeGenPassBuilder {
 
   TargetMachineT &TM;
   CGPassBuilderOption Opt;
+  PassBuilder &PB;
   PassInstrumentationCallbacks *PIC;
+  mutable IntrusiveRefCntPtr<AsmPrinter> PrinterImpl;
 
   template <typename TMC> TMC &getTM() const { return static_cast<TMC &>(TM); }
   CodeGenOptLevel getOptLevel() const { return TM.getOptLevel(); }
@@ -533,7 +536,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();
@@ -542,6 +545,14 @@ Error CodeGenPassBuilder<Derived, TargetMachineT>::buildPipeline(
   bool PrintAsm = TargetPassConfig::willCompleteCodeGenPipeline();
   bool PrintMIR = !PrintAsm && FileType != CodeGenFileType::Null;
 
+  if (PrintAsm) {
+    Expected<std::unique_ptr<MCStreamer>> MCStreamerOrErr =
+        TM.createMCStreamer(Out, DwoOut, FileType, Ctx);
+    if (auto Err = MCStreamerOrErr.takeError())
+      return Err;
+    PrinterImpl = PB.getAsmPrinter(std::move(*MCStreamerOrErr));
+  }
+
   {
     AddIRPass addIRPass(MPM, derived());
     addIRPass(RequireAnalysisPass<MachineModuleAnalysis, Module>());
@@ -550,29 +561,33 @@ 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));
+  }
 
   return verifyStartStop(*StartStopInfo);
 }
@@ -680,6 +695,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());