[NewPM/Codegen] Move MachineModuleInfo ownership outside of analysis (#80937)

With the legacy pass manager, MachineModuleInfoWrapperPass owned the
MachineModuleInfo used in the codegen pipeline. It can do this since
it's an ImmutablePass that doesn't get invalidated.

However, with the new pass manager, it is legal for the
ModuleAnalysisManager to clear all of its analyses, regardless of if the
analysis does not want to be invalidated. So we must move ownership of
the MachineModuleInfo outside of the analysis (this is similar to
PassInstrumentation). For now, make the PassBuilder user register a
MachineModuleAnalysis that returns a reference to a MachineModuleInfo
that the user owns. Perhaps we can find a better place to own the
MachineModuleInfo to make using the codegen pass manager less cumbersome
in the future.

Author: aeubanks

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()