Created owning and non-owning variants of MMIWP.

matinraayai · matinraayai · commit 8715ed6fbcfb · 2025-04-05T23:52:01.000-04:00
diff --git a/llvm/include/llvm/CodeGen/CodeGenTargetMachineImpl.h b/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.
diff --git a/llvm/include/llvm/CodeGen/MachineModuleInfo.h b/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.
diff --git a/llvm/include/llvm/Target/TargetMachine.h b/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;
   }
 
diff --git a/llvm/lib/CodeGen/CodeGenTargetMachineImpl.cpp b/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)
diff --git a/llvm/lib/CodeGen/MachineModuleInfo.cpp b/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;
 }
 
diff --git a/llvm/lib/ExecutionEngine/MCJIT/MCJIT.cpp b/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.
diff --git a/llvm/lib/ExecutionEngine/MCJIT/MCJIT.h b/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;
diff --git a/llvm/lib/ExecutionEngine/Orc/CompileUtils.cpp b/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);
diff --git a/llvm/lib/LTO/ThinLTOCodeGenerator.cpp b/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");
 
diff --git a/llvm/lib/Target/DirectX/DirectXTargetMachine.cpp b/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;
 }
 
diff --git a/llvm/lib/Target/DirectX/DirectXTargetMachine.h b/llvm/lib/Target/DirectX/DirectXTargetMachine.h
@@ -33,10 +33,10 @@ class DirectXTargetMachine : public CodeGenTargetMachineImpl {
   bool addPassesToEmitFile(PassManagerBase &PM, raw_pwrite_stream &Out,
                            raw_pwrite_stream *DwoOut, CodeGenFileType FileType,
                            bool DisableVerify,
-                           MachineModuleInfoWrapperPass *MMIWP) override;
+                           MachineModuleInfo *MMIWP) override;
 
-  bool addPassesToEmitMC(PassManagerBase &PM, MCContext *&Ctx,
-                         raw_pwrite_stream &Out, bool DisableVerify) override;
+  bool addPassesToEmitMC(PassManagerBase &PM, raw_pwrite_stream &Out,
+                         bool DisableVerify, MachineModuleInfo *MMIWP) override;
 
   const DirectXSubtarget *getSubtargetImpl(const Function &) const override;
 
diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.h b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.h
@@ -52,8 +52,8 @@ class NVPTXTargetMachine : public CodeGenTargetMachineImpl {
   TargetPassConfig *createPassConfig(PassManagerBase &PM) override;
 
   // Emission of machine code through MCJIT is not supported.
-  bool addPassesToEmitMC(PassManagerBase &, MCContext *&, raw_pwrite_stream &,
-                         bool = true) override {
+  bool addPassesToEmitMC(PassManagerBase &, raw_pwrite_stream &,
+                         bool = true, MachineModuleInfo *) override {
     return true;
   }
   TargetLoweringObjectFile *getObjFileLowering() const override {
diff --git a/llvm/lib/Target/SPIRV/SPIRVAPI.cpp b/llvm/lib/Target/SPIRV/SPIRVAPI.cpp
@@ -121,15 +121,14 @@ SPIRVTranslate(Module *M, std::string &SpirvObj, std::string &ErrMsg,
   TargetLibraryInfoImpl TLII(M->getTargetTriple());
   legacy::PassManager PM;
   PM.add(new TargetLibraryInfoWrapperPass(TLII));
-  std::unique_ptr<MachineModuleInfoWrapperPass> MMIWP(
-      new MachineModuleInfoWrapperPass(Target.get()));
-  const_cast<TargetLoweringObjectFile *>(Target->getObjFileLowering())
-      ->Initialize(MMIWP.get()->getMMI().getContext(), *Target);
+
+  MachineModuleInfo MMI(Target.get());
+  Target->getObjFileLowering()->Initialize(MMI.getContext(), *Target);
 
   SmallString<4096> OutBuffer;
   raw_svector_ostream OutStream(OutBuffer);
   if (Target->addPassesToEmitFile(PM, OutStream, nullptr,
-                                  CodeGenFileType::ObjectFile)) {
+                                  CodeGenFileType::ObjectFile, true, &MMI)) {
     ErrMsg = "Target machine cannot emit a file of this type";
     return false;
   }
diff --git a/offload/plugins-nextgen/common/src/JIT.cpp b/offload/plugins-nextgen/common/src/JIT.cpp
@@ -175,11 +175,12 @@ void JITEngine::codegen(TargetMachine *TM, TargetLibraryInfoImpl *TLII,
                         Module &M, raw_pwrite_stream &OS) {
   legacy::PassManager PM;
   PM.add(new TargetLibraryInfoWrapperPass(*TLII));
-  MachineModuleInfoWrapperPass *MMIWP = new MachineModuleInfoWrapperPass(TM);
+  MachineModuleInfo MMI(TM);
   TM->addPassesToEmitFile(PM, OS, nullptr,
                           TT.isNVPTX() ? CodeGenFileType::AssemblyFile
                                        : CodeGenFileType::ObjectFile,
-                          /*DisableVerify=*/false, MMIWP);
+                          /*DisableVerify=*/false,
+                          &MMI);
 
   PM.run(M);
 }

Original file line number	Diff line number	Diff line change
`@@ -139,12 +139,8 @@ FunctionPass *llvm::createFreeMachineFunctionPass() {`
`139`	`139`	`return new FreeMachineFunction();`
`140`	`140`	`}`
`141`	`141`
`142`		`-MachineModuleInfoWrapperPass::MachineModuleInfoWrapperPass(`
`143`		`- MachineModuleInfo *MMI)`
`144`		`- : ImmutablePass(ID), MMI([&] -> MachineModuleInfo & {`
`145`		`- assert(MMI != nullptr, "MMI is nullptr");`
`146`		`- return *MMI;`
`147`		`- }()) {`
	`142`	`+MachineModuleInfoWrapperPass::MachineModuleInfoWrapperPass()`
	`143`	`+ : ImmutablePass(ID) {`
`148`	`144`	`initializeMachineModuleInfoWrapperPassPass(*PassRegistry::getPassRegistry());`
`149`	`145`	`}`
`150`	`146`
`@@ -179,6 +175,7 @@ static uint64_t getLocCookie(const SMDiagnostic &SMD, const SourceMgr &SrcMgr,`
`179`	`175`	`}`
`180`	`176`
`181`	`177`	`bool MachineModuleInfoWrapperPass::doInitialization(Module &M) {`
	`178`	`+ MachineModuleInfo &MMI = getMMI();`
`182`	`179`	`MMI.initialize();`
`183`	`180`	`MMI.TheModule = &M;`
`184`	`181`	`LLVMContext &Ctx = M.getContext();`
`@@ -196,7 +193,7 @@ bool MachineModuleInfoWrapperPass::doInitialization(Module &M) {`
`196`	`193`	`}`
`197`	`194`
`198`	`195`	`bool MachineModuleInfoWrapperPass::doFinalization(Module &M) {`
`199`		`- MMI.finalize();`
	`196`	`+ getMMI().finalize();`
`200`	`197`	`return false;`
`201`	`198`	`}`
`202`	`199`