Created owning and non-owning variants of MMIWP.

Author: matinraayai

llvm/include/llvm/CodeGen/CodeGenTargetMachineImpl.h

@@ -43,28 +43,24 @@ class CodeGenTargetMachineImpl : public TargetMachine {
   /// for generating a pipeline of CodeGen passes.
   virtual TargetPassConfig *createPassConfig(PassManagerBase &PM) override;
 
-  /// Creates a new \c MachineModuleInfo from this \c TargetMachine
-  std::unique_ptr<MachineModuleInfo> createMachineModuleInfo() const override;
-
   /// Add passes to the specified pass manager to get the specified file
   /// emitted. Typically, this will involve several steps of code generation.
   /// This method should return true if emission of this file type is not
   /// supported, or false on success.
-  /// \p MMI a \c MachineModuleInfo that holds the generated machine code
-  /// throughout the code generation process
+  /// \p MMI an optional, externally-managed \c MachineModuleInfo to
+  /// hold the generated machine code even after the \p PM is destroyed
   bool addPassesToEmitFile(PassManagerBase &PM, raw_pwrite_stream &Out,
                            raw_pwrite_stream *DwoOut, CodeGenFileType FileType,
-                           MachineModuleInfo *MMI,
-                           bool DisableVerify = true) override;
+                           bool DisableVerify = true,
+                           MachineModuleInfo *MMI = nullptr) override;
 
   /// Add passes to the specified pass manager to get machine code emitted with
   /// the MCJIT. This method returns true if machine code is not supported.
-  /// \p MMI a \c MachineModuleInfo that holds the generated machine code
-  /// throughout the code generation process
-  bool addPassesToEmitMC(PassManagerBase &PM,
-                         raw_pwrite_stream &Out,
-                         MachineModuleInfo *MMI,
-                         bool DisableVerify = true) override;
+  /// \p MMI an optional, externally-managed \c MachineModuleInfo to
+  /// hold the generated machine code even after the \p PM is destroyed
+  bool addPassesToEmitMC(PassManagerBase &PM, raw_pwrite_stream &Out,
+                         bool DisableVerify = true,
+                         MachineModuleInfo *MMI = nullptr) override;
 
   /// Adds an AsmPrinter pass to the pipeline that prints assembly or
   /// machine code from the MI representation.

llvm/include/llvm/CodeGen/MachineModuleInfo.h

@@ -163,19 +163,51 @@ class MachineModuleInfo {
   /// \}
 }; // End class MachineModuleInfo
 
+/// \brief Interface for pass that provide access to \c MachineModuleInfo
+/// being worked on
 class MachineModuleInfoWrapperPass : public ImmutablePass {
-  MachineModuleInfo &MMI;
 
 public:
   static char ID; // Pass identification, replacement for typeid
-  explicit MachineModuleInfoWrapperPass(MachineModuleInfo *MMI = nullptr);
+  MachineModuleInfoWrapperPass();
 
   // Initialization and Finalization
   bool doInitialization(Module &) override;
   bool doFinalization(Module &) override;
 
-  MachineModuleInfo &getMMI() { return MMI; }
-  const MachineModuleInfo &getMMI() const { return MMI; }
+  virtual MachineModuleInfo &getMMI() = 0;
+  virtual const MachineModuleInfo &getMMI() const = 0;
+};
+
+/// \brief a version of \c MachineModuleInfoWrapperPass that manages the
+/// lifetime of its \c MachineModuleInfo
+class OwningMachineModuleInfoWrapperPass : public MachineModuleInfoWrapperPass {
+  MachineModuleInfo MMI;
+
+public:
+  explicit OwningMachineModuleInfoWrapperPass(const TargetMachine &TM)
+      : MMI(&TM) {};
+
+  OwningMachineModuleInfoWrapperPass(const TargetMachine &TM,
+                                     MCContext &ExtContext)
+      : MMI(&TM, &ExtContext) {};
+
+  MachineModuleInfo &getMMI() override { return MMI; }
+  const MachineModuleInfo &getMMI() const override { return MMI; }
+};
+
+/// \brief a version of \c MachineModuleInfoWrapperPass that does not manage the
+/// lifetime of its \c MachineModuleInfo
+class NonOwningMachineModuleInfoWrapperPass
+    : public MachineModuleInfoWrapperPass {
+  MachineModuleInfo &MMI;
+
+public:
+  explicit NonOwningMachineModuleInfoWrapperPass(MachineModuleInfo &MMI)
+      : MMI(MMI) {};
+
+  MachineModuleInfo &getMMI() override { return MMI; }
+  const MachineModuleInfo &getMMI() const override { return MMI; }
 };
 
 /// An analysis that produces \c MachineModuleInfo for a module.

llvm/include/llvm/Target/TargetMachine.h

@@ -400,38 +400,30 @@ class TargetMachine {
   /// with the new pass manager. Only affects the "default" AAManager.
   virtual void registerDefaultAliasAnalyses(AAManager &) {}
 
-  /// For targets that utilize the target-independent code generator
-  /// (CodeGen), this method creates a new \c MachineModuleInfo from this
-  /// \c TargetMachine and returns it; For targets that don't use CodeGen,
-  /// returns nullptr
-  virtual std::unique_ptr<MachineModuleInfo> createMachineModuleInfo() const {
-    return nullptr;
-  }
-
   /// Add passes to the specified pass manager to get the specified file
   /// emitted. Typically, this will involve several steps of code generation.
   /// This method should return true if emission of this file type is not
   /// supported, or false on success.
   /// For targets that utilize the target-independent code generator
-  /// (CodeGen) to emit the file, \c MachineModuleInfo pointer will hold the
-  /// generated machine code; For targets that don't use CodeGen, it is set
-  /// to \c nullptr
+  /// (CodeGen), an externally-managed \c MachineModuleInfo can be provided
+  /// to ensure persistence of the generated machine code even after the pass
+  /// manager is destroyed
   virtual bool addPassesToEmitFile(PassManagerBase &, raw_pwrite_stream &,
                                    raw_pwrite_stream *, CodeGenFileType,
-                                   MachineModuleInfo *,
-                                   bool /*DisableVerify*/ = true) {
+                                   bool /*DisableVerify*/ = true,
+                                   MachineModuleInfo * = nullptr) {
     return true;
   }
 
   /// Add passes to the specified pass manager to get machine code emitted with
   /// the MCJIT. This method returns true if machine code is not supported.
   /// For targets that utilize the target-independent code generator
-  /// (CodeGen), \c MachineModuleInfo pointer will hold the generated
-  /// machine code; For targets that don't use CodeGen, it must be set
-  /// to \c nullptr
+  /// (CodeGen), an externally-managed \c MachineModuleInfo can be provided
+  /// to ensure persistence of the generated machine code even after the pass
+  /// manager is destroyed
   virtual bool addPassesToEmitMC(PassManagerBase &, raw_pwrite_stream &,
-                                 MachineModuleInfo *,
-                                 bool /*DisableVerify*/ = true) {
+                                 bool /*DisableVerify*/ = true,
+                                 MachineModuleInfo * = nullptr) {
     return true;
   }
 

llvm/lib/CodeGen/CodeGenTargetMachineImpl.cpp

@@ -208,20 +208,22 @@ CodeGenTargetMachineImpl::createMCStreamer(raw_pwrite_stream &Out,
   return std::move(AsmStreamer);
 }
 
-std::unique_ptr<MachineModuleInfo>
-CodeGenTargetMachineImpl::createMachineModuleInfo() const override {
-  return std::make_unique<MachineModuleInfo>(this);
-}
-
 bool CodeGenTargetMachineImpl::addPassesToEmitFile(
     PassManagerBase &PM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
-    CodeGenFileType FileType, MachineModuleInfo *MMI, bool DisableVerify) {
-  // Check if MMI is nullptr, or if MMI pointer is valid but the target
-  // machine of the MMI is not the same as this target machine
-  if (!MMI || &MMI->getTarget() != this)
-    return true;
+    CodeGenFileType FileType, bool DisableVerify, MachineModuleInfo *MMI) {
   // Add the wrapper pass for MMI
-  auto *MMIWP = new MachineModuleInfoWrapperPass(MMI);
+  MachineModuleInfoWrapperPass *MMIWP;
+  if (MMI) {
+    MMIWP = new NonOwningMachineModuleInfoWrapperPass(*MMI);
+  } else {
+    MMIWP = new OwningMachineModuleInfoWrapperPass(*this);
+    MMI = &MMIWP->getMMI();
+  }
+  // Check (for the case of externally-managed MMI) if the TM of MMI is
+  // the same as this target machine
+  if (&MMI->getTarget() != this)
+    return true;
+
   TargetPassConfig *PassConfig =
       addPassesToGenerateCode(*this, PM, DisableVerify, *MMIWP);
   if (!PassConfig)
@@ -240,21 +242,22 @@ bool CodeGenTargetMachineImpl::addPassesToEmitFile(
   return false;
 }
 
-/// addPassesToEmitMC - Add passes to the specified pass manager to get
-/// machine code emitted with the MCJIT. This method returns true if machine
-/// code is not supported. It fills the MCContext Ctx pointer which can be
-/// used to build custom MCStreamer.
-///
 bool CodeGenTargetMachineImpl::addPassesToEmitMC(PassManagerBase &PM,
                                                  raw_pwrite_stream &Out,
-                                                 MachineModuleInfo *MMI,
-                                                 bool DisableVerify) {
-  // Check if MMI is nullptr, or if MMI pointer is valid but the target
-  // machine of the MMI is not the same as this target machine
-  if (!MMI || &MMI->getTarget() != this)
+                                                 bool DisableVerify,
+                                                 MachineModuleInfo *MMI) {
+  // Add the wrapper pass for MMI
+  MachineModuleInfoWrapperPass *MMIWP;
+  if (MMI) {
+    MMIWP = new NonOwningMachineModuleInfoWrapperPass(*MMI);
+  } else {
+    MMIWP = new OwningMachineModuleInfoWrapperPass(*this);
+    MMI = &MMIWP->getMMI();
+  }
+  // Check (for the case of externally-managed MMI) if the TM of MMI is
+  // the same as this target machine
+  if (&MMI->getTarget() != this)
     return true;
-  // Add common CodeGen passes.
-  MachineModuleInfoWrapperPass *MMIWP = new MachineModuleInfoWrapperPass(MMI);
   TargetPassConfig *PassConfig =
       addPassesToGenerateCode(*this, PM, DisableVerify, *MMIWP);
   if (!PassConfig)

llvm/lib/CodeGen/MachineModuleInfo.cpp

@@ -139,12 +139,8 @@ FunctionPass *llvm::createFreeMachineFunctionPass() {
   return new FreeMachineFunction();
 }
 
-MachineModuleInfoWrapperPass::MachineModuleInfoWrapperPass(
-    MachineModuleInfo *MMI)
-    : ImmutablePass(ID), MMI([&] -> MachineModuleInfo & {
-        assert(MMI != nullptr, "MMI is nullptr");
-        return *MMI;
-      }()) {
+MachineModuleInfoWrapperPass::MachineModuleInfoWrapperPass()
+    : ImmutablePass(ID) {
   initializeMachineModuleInfoWrapperPassPass(*PassRegistry::getPassRegistry());
 }
 
@@ -179,6 +175,7 @@ static uint64_t getLocCookie(const SMDiagnostic &SMD, const SourceMgr &SrcMgr,
 }
 
 bool MachineModuleInfoWrapperPass::doInitialization(Module &M) {
+  MachineModuleInfo &MMI = getMMI();
   MMI.initialize();
   MMI.TheModule = &M;
   LLVMContext &Ctx = M.getContext();
@@ -196,7 +193,7 @@ bool MachineModuleInfoWrapperPass::doInitialization(Module &M) {
 }
 
 bool MachineModuleInfoWrapperPass::doFinalization(Module &M) {
-  MMI.finalize();
+  getMMI().finalize();
   return false;
 }
 

llvm/lib/ExecutionEngine/MCJIT/MCJIT.cpp

@@ -66,9 +66,8 @@ MCJIT::MCJIT(std::unique_ptr<Module> M, std::unique_ptr<TargetMachine> TM,
              std::shared_ptr<MCJITMemoryManager> MemMgr,
              std::shared_ptr<LegacyJITSymbolResolver> Resolver)
     : ExecutionEngine(TM->createDataLayout(), std::move(M)), TM(std::move(TM)),
-      Ctx(nullptr), MemMgr(std::move(MemMgr)),
-      Resolver(*this, std::move(Resolver)), Dyld(*this->MemMgr, this->Resolver),
-      ObjCache(nullptr) {
+      MemMgr(std::move(MemMgr)), Resolver(*this, std::move(Resolver)),
+      Dyld(*this->MemMgr, this->Resolver), ObjCache(nullptr) {
   // FIXME: We are managing our modules, so we do not want the base class
   // ExecutionEngine to manage them as well. To avoid double destruction
   // of the first (and only) module added in ExecutionEngine constructor
@@ -163,7 +162,7 @@ std::unique_ptr<MemoryBuffer> MCJIT::emitObject(Module *M) {
 
   // Turn the machine code intermediate representation into bytes in memory
   // that may be executed.
-  if (TM->addPassesToEmitMC(PM, Ctx, ObjStream, !getVerifyModules()))
+  if (TM->addPassesToEmitMC(PM, ObjStream, !getVerifyModules()))
     report_fatal_error("Target does not support MC emission!");
 
   // Initialize passes.

llvm/lib/ExecutionEngine/MCJIT/MCJIT.h

@@ -168,7 +168,6 @@ class MCJIT : public ExecutionEngine {
   };
 
   std::unique_ptr<TargetMachine> TM;
-  MCContext *Ctx;
   std::shared_ptr<MCJITMemoryManager> MemMgr;
   LinkingSymbolResolver Resolver;
   RuntimeDyld Dyld;

llvm/lib/ExecutionEngine/Orc/CompileUtils.cpp

@@ -46,8 +46,7 @@ Expected<SimpleCompiler::CompileResult> SimpleCompiler::operator()(Module &M) {
     raw_svector_ostream ObjStream(ObjBufferSV);
 
     legacy::PassManager PM;
-    MCContext *Ctx;
-    if (TM.addPassesToEmitMC(PM, Ctx, ObjStream))
+    if (TM.addPassesToEmitMC(PM, ObjStream))
       return make_error<StringError>("Target does not support MC emission",
                                      inconvertibleErrorCode());
     PM.run(M);

llvm/lib/LTO/ThinLTOCodeGenerator.cpp

@@ -331,9 +331,11 @@ static std::unique_ptr<MemoryBuffer> codegenModule(Module &TheModule,
   {
     raw_svector_ostream OS(OutputBuffer);
     legacy::PassManager PM;
+    std::unique_ptr<MachineModuleInfo> MMI = TM.createMachineModuleInfo();
 
-    // Setup the codegen now.
+    // Setup the codegen pipeline.
     if (TM.addPassesToEmitFile(PM, OS, nullptr, CodeGenFileType::ObjectFile,
+                               MMI.get(),
                                /* DisableVerify */ true))
       report_fatal_error("Failed to setup codegen");
 

llvm/lib/Target/DirectX/DirectXTargetMachine.cpp

@@ -133,8 +133,7 @@ void DirectXTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB) {
 
 bool DirectXTargetMachine::addPassesToEmitFile(
     PassManagerBase &PM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
-    CodeGenFileType FileType, bool DisableVerify,
-    MachineModuleInfoWrapperPass *MMIWP) {
+    CodeGenFileType FileType, bool DisableVerify, MachineModuleInfo *MMI) {
   TargetPassConfig *PassConfig = createPassConfig(PM);
   PassConfig->addCodeGenPrepare();
 
@@ -150,8 +149,17 @@ bool DirectXTargetMachine::addPassesToEmitFile(
       // globals don't pollute the DXIL.
       PM.add(createDXContainerGlobalsPass());
 
-      if (!MMIWP)
-        MMIWP = new MachineModuleInfoWrapperPass(this);
+      MachineModuleInfoWrapperPass *MMIWP;
+      if (MMI) {
+        MMIWP = new NonOwningMachineModuleInfoWrapperPass(*MMI);
+      } else {
+        MMIWP = new OwningMachineModuleInfoWrapperPass(*this);
+        MMI = &MMIWP->getMMI();
+      }
+      // Check (for the case of externally-managed MMI) if the TM of MMI is
+      // the same as this target machine
+      if (&MMI->getTarget() != this)
+        return true;
       PM.add(MMIWP);
       if (addAsmPrinter(PM, Out, DwoOut, FileType,
                         MMIWP->getMMI().getContext()))
@@ -166,9 +174,9 @@ bool DirectXTargetMachine::addPassesToEmitFile(
 }
 
 bool DirectXTargetMachine::addPassesToEmitMC(PassManagerBase &PM,
-                                             MCContext *&Ctx,
                                              raw_pwrite_stream &Out,
-                                             bool DisableVerify) {
+                                             bool DisableVerify,
+                                             MachineModuleInfo *MMI) {
   return true;
 }