[NewPM/Codegen] Move MachineModuleInfo ownership outside of analysis

llvm/include/llvm/CodeGen/MachineModuleInfo.h

@@ -218,21 +218,31 @@ class MachineModuleInfoWrapperPass : public ImmutablePass {
   const MachineModuleInfo &getMMI() const { return MMI; }
 };
 
-/// An analysis that produces \c MachineInfo for a module.
+/// An analysis that produces \c MachineModuleInfo for a module.
+/// This does not produce its own MachineModuleInfo because we need a consistent
+/// MachineModuleInfo to keep ownership of MachineFunctions regardless of
+/// analysis invalidation/clearing. So something outside the analysis
+/// infrastructure must own the MachineModuleInfo.
 class MachineModuleAnalysis : public AnalysisInfoMixin<MachineModuleAnalysis> {
   friend AnalysisInfoMixin<MachineModuleAnalysis>;
   static AnalysisKey Key;
 
-  const LLVMTargetMachine *TM;
+  MachineModuleInfo &MMI;
 
 public:
-  /// Provide the result type for this analysis pass.
-  using Result = MachineModuleInfo;
+  class Result {
+    MachineModuleInfo &MMI;
+    Result(MachineModuleInfo &MMI) : MMI(MMI) {}
+    friend class MachineModuleAnalysis;
 
-  MachineModuleAnalysis(const LLVMTargetMachine *TM) : TM(TM) {}
+  public:
+    MachineModuleInfo &getMMI() { return MMI; }
+  };
+
+  MachineModuleAnalysis(MachineModuleInfo &MMI) : MMI(MMI) {}
 
   /// Run the analysis pass and produce machine module information.
-  MachineModuleInfo run(Module &M, ModuleAnalysisManager &);
+  Result run(Module &M, ModuleAnalysisManager &);
 };
 
 } // end namespace llvm

llvm/lib/CodeGen/MachineModuleInfo.cpp

@@ -237,11 +237,10 @@ bool MachineModuleInfoWrapperPass::doFinalization(Module &M) {
 
 AnalysisKey MachineModuleAnalysis::Key;
 
-MachineModuleInfo MachineModuleAnalysis::run(Module &M,
-                                             ModuleAnalysisManager &) {
-  MachineModuleInfo MMI(TM);
+MachineModuleAnalysis::Result
+MachineModuleAnalysis::run(Module &M, ModuleAnalysisManager &) {
   MMI.TheModule = &M;
-  MMI.DbgInfoAvailable = !DisableDebugInfoPrinting &&
-                         !M.debug_compile_units().empty();
-  return MMI;
+  MMI.DbgInfoAvailable =
+      !DisableDebugInfoPrinting && !M.debug_compile_units().empty();
+  return Result(MMI);
 }

llvm/lib/CodeGen/MachinePassManager.cpp

@@ -29,7 +29,7 @@ Error MachineFunctionPassManager::run(Module &M,
   // because we don't run any module pass in codegen pipeline. This is very
   // important because the codegen state is stored in MMI which is the analysis
   // result of MachineModuleAnalysis. MMI should not be recomputed.
-  auto &MMI = MFAM.getResult<MachineModuleAnalysis>(M);
+  auto &MMI = MFAM.getResult<MachineModuleAnalysis>(M).getMMI();
 
   (void)RequireCodeGenSCCOrder;
   assert(!RequireCodeGenSCCOrder && "not implemented");

llvm/tools/llc/NewPMDriver.cpp

@@ -147,6 +147,8 @@ int llvm::compileModuleWithNewPM(
   Opt.DebugPM = DebugPM;
   Opt.RegAlloc = RegAlloc;
 
+  MachineModuleInfo MMI(&LLVMTM);
+
   PassInstrumentationCallbacks PIC;
   StandardInstrumentations SI(Context, Opt.DebugPM);
   SI.registerCallbacks(PIC);
@@ -164,7 +166,7 @@ int llvm::compileModuleWithNewPM(
   PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
 
   FAM.registerPass([&] { return TargetLibraryAnalysis(TLII); });
-  MAM.registerPass([&] { return MachineModuleAnalysis(&LLVMTM); });
+  MAM.registerPass([&] { return MachineModuleAnalysis(MMI); });
 
   MachineFunctionAnalysisManager MFAM(FAM, MAM);
 
@@ -185,7 +187,7 @@ int llvm::compileModuleWithNewPM(
     MFPM.addPass(PrintMIRPass(*OS));
     MFPM.addPass(FreeMachineFunctionPass());
 
-    auto &MMI = MFAM.getResult<MachineModuleAnalysis>(*M);
+    auto &MMI = MFAM.getResult<MachineModuleAnalysis>(*M).getMMI();
     if (MIR->parseMachineFunctions(*M, MMI))
       return 1;
 

llvm/unittests/CodeGen/PassManagerTest.cpp

@@ -118,7 +118,8 @@ struct TestMachineFunctionPass : public PassInfoMixin<TestMachineFunctionPass> {
 
     // Query module analysis result.
     MachineModuleInfo &MMI =
-        MFAM.getResult<MachineModuleAnalysis>(*MF.getFunction().getParent());
+        MFAM.getResult<MachineModuleAnalysis>(*MF.getFunction().getParent())
+            .getMMI();
     // 1 + 1 + 1 = 3
     Count += (MMI.getModule() == MF.getFunction().getParent());
 
@@ -144,7 +145,7 @@ struct TestMachineModulePass : public PassInfoMixin<TestMachineModulePass> {
       : Count(Count), MachineModulePassCount(MachineModulePassCount) {}
 
   Error run(Module &M, MachineFunctionAnalysisManager &MFAM) {
-    MachineModuleInfo &MMI = MFAM.getResult<MachineModuleAnalysis>(M);
+    MachineModuleInfo &MMI = MFAM.getResult<MachineModuleAnalysis>(M).getMMI();
     // + 1
     Count += (MMI.getModule() == &M);
     MachineModulePassCount.push_back(Count);
@@ -209,6 +210,7 @@ TEST_F(PassManagerTest, Basic) {
   LLVMTargetMachine *LLVMTM = static_cast<LLVMTargetMachine *>(TM.get());
   M->setDataLayout(TM->createDataLayout());
 
+  MachineModuleInfo MMI(LLVMTM);
   LoopAnalysisManager LAM;
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;
@@ -220,7 +222,7 @@ TEST_F(PassManagerTest, Basic) {
 
   FAM.registerPass([&] { return TestFunctionAnalysis(); });
   FAM.registerPass([&] { return PassInstrumentationAnalysis(); });
-  MAM.registerPass([&] { return MachineModuleAnalysis(LLVMTM); });
+  MAM.registerPass([&] { return MachineModuleAnalysis(MMI); });
   MAM.registerPass([&] { return PassInstrumentationAnalysis(); });
 
   MachineFunctionAnalysisManager MFAM;

llvm/unittests/MIR/PassBuilderCallbacksTest.cpp

@@ -298,6 +298,7 @@ class MachineFunctionCallbacksTest : public testing::Test {
   }
 
   std::unique_ptr<LLVMTargetMachine> TM;
+  std::unique_ptr<MachineModuleInfo> MMI;
 
   LLVMContext Context;
   std::unique_ptr<Module> M;
@@ -355,9 +356,9 @@ class MachineFunctionCallbacksTest : public testing::Test {
             TripleName, "", "", TargetOptions(), std::nullopt)));
     if (!TM)
       GTEST_SKIP();
-    MachineModuleInfo MMI(TM.get());
-    M = parseMIR(*TM, MIRString, MMI);
-    AM.registerPass([&] { return MachineModuleAnalysis(TM.get()); });
+    MMI = std::make_unique<MachineModuleInfo>(TM.get());
+    M = parseMIR(*TM, MIRString, *MMI);
+    AM.registerPass([&] { return MachineModuleAnalysis(*MMI); });
   }
 
   MachineFunctionCallbacksTest()