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

Author: vikramRH

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"
@@ -87,7 +89,8 @@ class RemarkStreamer;
 }
 
 /// This class is intended to be used as a driving class for all asm writers.
-class LLVM_ABI AsmPrinter : public MachineFunctionPass {
+class LLVM_ABI AsmPrinter : public RefCountedBase<AsmPrinter>,
+                            public MachineFunctionPass {
 public:
   /// Target machine description.
   TargetMachine &TM;
@@ -247,6 +250,11 @@ class LLVM_ABI 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;
@@ -282,7 +290,7 @@ class LLVM_ABI AsmPrinter : public MachineFunctionPass {
              char &ID = AsmPrinter::ID);
 
 public:
-  ~AsmPrinter() override;
+  virtual ~AsmPrinter();
 
   DwarfDebug *getDwarfDebug() { return DD; }
   DwarfDebug *getDwarfDebug() const { return DD; }
@@ -437,23 +445,43 @@ class LLVM_ABI 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.
@@ -590,6 +618,7 @@ class LLVM_ABI AsmPrinter : public MachineFunctionPass {
 
   /// Emit the stack maps.
   void emitStackMaps();
+  void emitStackMaps(Module &M); // For new pass manager version.
 
   //===------------------------------------------------------------------===//
   // Overridable Hooks
@@ -999,6 +1028,43 @@ class LLVM_ABI 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

@@ -23,6 +23,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"
@@ -107,6 +108,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"
@@ -190,8 +193,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;
@@ -356,6 +359,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(); }
@@ -621,7 +625,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();
@@ -630,6 +634,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>(),
@@ -641,29 +656,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);
 }
@@ -776,6 +797,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

@@ -487,7 +487,7 @@ class LLVM_ABI 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());