[CodeGen][NPM] Avoid MachineModuleInfo in MachineModuleSlotTracker (#171257)

This work is a continuation of
[PR#140530](https://github.com/llvm/llvm-project/pull/140530),
co-authored by @optimisan.

The PR refactors `MachineModuleSlotTracker` to support both the Legacy
Pass Manager and New Pass Manager by removing its direct dependency on
`MachineModuleInfo`.

`MachineModuleSlotTracker` requires `MachineModuleInfo` to obtain
`MachineFunction` instances when printing MIR. But
`MachineFunctionAnalysis` provides `MachineFunction` for NPM.

This patch refactors `MachineModuleSlotTracker` to use a function
callback (`MFGetterFnT`) instead of directly depending on
`MachineModuleInfo`.

---------

Co-authored-by: vikhegde <vikram.hegde@amd.com>

llvm/include/llvm/CodeGen/MIRPrinter.h

@@ -48,11 +48,14 @@ public:
 /// Print LLVM IR using the MIR serialization format to the given output stream.
 LLVM_ABI void printMIR(raw_ostream &OS, const Module &M);
 
-/// Print a machine function using the MIR serialization format to the given
-/// output stream.
+/// Print MIR using Legacy Pass Manager (uses MachineModuleInfo).
 LLVM_ABI void printMIR(raw_ostream &OS, const MachineModuleInfo &MMI,
                        const MachineFunction &MF);
 
+/// Print MIR using New Pass Manager (uses FunctionAnalysisManager).
+LLVM_ABI void printMIR(raw_ostream &OS, FunctionAnalysisManager &FAM,
+                       const MachineFunction &MF);
+
 /// Determine a possible list of successors of a basic block based on the
 /// basic block machine operand being used inside the block. This should give
 /// you the correct list of successor blocks in most cases except for things

llvm/include/llvm/CodeGen/MachineModuleSlotTracker.h

@@ -9,6 +9,7 @@
 #ifndef LLVM_CODEGEN_MACHINEMODULESLOTTRACKER_H
 #define LLVM_CODEGEN_MACHINEMODULESLOTTRACKER_H
 
+#include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/IR/ModuleSlotTracker.h"
 #include "llvm/Support/Compiler.h"
 
@@ -20,9 +21,11 @@ class MachineModuleInfo;
 class MachineFunction;
 class Module;
 
+using MFGetterFnT = function_ref<MachineFunction *(const Function &)>;
+
 class LLVM_ABI MachineModuleSlotTracker : public ModuleSlotTracker {
   const Function &TheFunction;
-  const MachineModuleInfo &TheMMI;
+  const MachineFunction *TheMF;
   unsigned MDNStartSlot = 0, MDNEndSlot = 0;
 
   void processMachineFunctionMetadata(AbstractSlotTrackerStorage *AST,
@@ -34,8 +37,7 @@ class LLVM_ABI MachineModuleSlotTracker : public ModuleSlotTracker {
                               bool ShouldInitializeAllMetadata);
 
 public:
-  MachineModuleSlotTracker(const MachineModuleInfo &MMI,
-                           const MachineFunction *MF,
+  MachineModuleSlotTracker(MFGetterFnT Fn, const MachineFunction *MF,
                            bool ShouldInitializeAllMetadata = true);
   ~MachineModuleSlotTracker() override;
 

llvm/lib/CodeGen/MIRPrinter.cpp

@@ -25,9 +25,11 @@
 #include "llvm/CodeGen/MachineConstantPool.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionAnalysis.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineModuleSlotTracker.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
@@ -102,8 +104,8 @@ struct MFPrintState {
   /// Synchronization scope names registered with LLVMContext.
   SmallVector<StringRef, 8> SSNs;
 
-  MFPrintState(const MachineModuleInfo &MMI, const MachineFunction &MF)
-      : MST(MMI, &MF) {}
+  MFPrintState(MFGetterFnT Fn, const MachineFunction &MF)
+      : MST(std::move(Fn), &MF) {}
 };
 
 } // end anonymous namespace
@@ -169,9 +171,10 @@ static void convertCalledGlobals(yaml::MachineFunction &YMF,
                                  const MachineFunction &MF,
                                  MachineModuleSlotTracker &MST);
 
-static void printMF(raw_ostream &OS, const MachineModuleInfo &MMI,
+static void printMF(raw_ostream &OS, MFGetterFnT Fn,
                     const MachineFunction &MF) {
-  MFPrintState State(MMI, MF);
+  MFPrintState State(std::move(Fn), MF);
+
   State.RegisterMaskIds = initRegisterMaskIds(MF);
 
   yaml::MachineFunction YamlMF;
@@ -1020,5 +1023,17 @@ void llvm::printMIR(raw_ostream &OS, const Module &M) {
 
 void llvm::printMIR(raw_ostream &OS, const MachineModuleInfo &MMI,
                     const MachineFunction &MF) {
-  printMF(OS, MMI, MF);
+  printMF(OS, [&](const Function &F) { return MMI.getMachineFunction(F); }, MF);
+}
+
+void llvm::printMIR(raw_ostream &OS, FunctionAnalysisManager &FAM,
+                    const MachineFunction &MF) {
+  printMF(
+      OS,
+      [&](const Function &F) {
+        return &FAM.getResult<MachineFunctionAnalysis>(
+                       const_cast<Function &>(F))
+                    .getMF();
+      },
+      MF);
 }

llvm/lib/CodeGen/MIRPrintingPass.cpp

@@ -27,12 +27,10 @@ PreservedAnalyses PrintMIRPreparePass::run(Module &M, ModuleAnalysisManager &) {
 
 PreservedAnalyses PrintMIRPass::run(MachineFunction &MF,
                                     MachineFunctionAnalysisManager &MFAM) {
-  auto &MAMP = MFAM.getResult<ModuleAnalysisManagerMachineFunctionProxy>(MF);
-  Module *M = MF.getFunction().getParent();
-  const MachineModuleInfo &MMI =
-      MAMP.getCachedResult<MachineModuleAnalysis>(*M)->getMMI();
+  auto &FAM = MFAM.getResult<FunctionAnalysisManagerMachineFunctionProxy>(MF)
+                  .getManager();
 
-  printMIR(OS, MMI, MF);
+  printMIR(OS, FAM, MF);
   return PreservedAnalyses::all();
 }
 
@@ -59,10 +57,10 @@ struct MIRPrintingPass : public MachineFunctionPass {
     std::string Str;
     raw_string_ostream StrOS(Str);
 
-    const MachineModuleInfo &MMI =
-        getAnalysis<MachineModuleInfoWrapperPass>().getMMI();
+    MachineModuleInfo *MMI =
+        &getAnalysis<MachineModuleInfoWrapperPass>().getMMI();
 
-    printMIR(StrOS, MMI, MF);
+    printMIR(StrOS, *MMI, MF);
     MachineFunctions.append(Str);
     return false;
   }

llvm/lib/CodeGen/MachineModuleSlotTracker.cpp

@@ -39,8 +39,8 @@ void MachineModuleSlotTracker::processMachineModule(
       if (&F != &TheFunction)
         continue;
       MDNStartSlot = AST->getNextMetadataSlot();
-      if (auto *MF = TheMMI.getMachineFunction(F))
-        processMachineFunctionMetadata(AST, *MF);
+      if (TheMF)
+        processMachineFunctionMetadata(AST, *TheMF);
       MDNEndSlot = AST->getNextMetadataSlot();
       break;
     }
@@ -52,8 +52,8 @@ void MachineModuleSlotTracker::processMachineFunction(
     bool ShouldInitializeAllMetadata) {
   if (!ShouldInitializeAllMetadata && F == &TheFunction) {
     MDNStartSlot = AST->getNextMetadataSlot();
-    if (auto *MF = TheMMI.getMachineFunction(*F))
-      processMachineFunctionMetadata(AST, *MF);
+    if (TheMF)
+      processMachineFunctionMetadata(AST, *TheMF);
     MDNEndSlot = AST->getNextMetadataSlot();
   }
 }
@@ -64,11 +64,10 @@ void MachineModuleSlotTracker::collectMachineMDNodes(
 }
 
 MachineModuleSlotTracker::MachineModuleSlotTracker(
-    const MachineModuleInfo &MMI, const MachineFunction *MF,
-    bool ShouldInitializeAllMetadata)
+    MFGetterFnT Fn, const MachineFunction *MF, bool ShouldInitializeAllMetadata)
     : ModuleSlotTracker(MF->getFunction().getParent(),
                         ShouldInitializeAllMetadata),
-      TheFunction(MF->getFunction()), TheMMI(MMI) {
+      TheFunction(MF->getFunction()), TheMF(Fn(MF->getFunction())) {
   setProcessHook([this](AbstractSlotTrackerStorage *AST, const Module *M,
                         bool ShouldInitializeAllMetadata) {
     this->processMachineModule(AST, M, ShouldInitializeAllMetadata);

llvm/unittests/MIR/MachineMetadata.cpp

@@ -250,7 +250,8 @@ body:             |
   auto *NewMMO = MF->getMachineMemOperand(OldMMO, AAInfo);
   MI.setMemRefs(*MF, NewMMO);
 
-  MachineModuleSlotTracker MST(MMI, MF);
+  MachineModuleSlotTracker MST(
+      [&](const Function &F) { return MMI.getMachineFunction(F); }, MF);
   // Print that MI with new machine metadata, which slot numbers should be
   // assigned.
   EXPECT_EQ("%1:gpr32 = LDRWui %0, 0 :: (load (s32) from %ir.p, "
@@ -400,7 +401,8 @@ body:             |
   auto *NewMMO = MF->getMachineMemOperand(OldMMO, AAInfo);
   MI.setMemRefs(*MF, NewMMO);
 
-  MachineModuleSlotTracker MST(MMI, MF);
+  MachineModuleSlotTracker MST(
+      [&](const Function &F) { return MMI.getMachineFunction(F); }, MF);
   // Print that MI with new machine metadata, which slot numbers should be
   // assigned.
   EXPECT_EQ("%1:gr32 = MOV32rm %0, 1, $noreg, 0, $noreg :: (load (s32) from %ir.p, "
@@ -498,7 +500,8 @@ body:             |
   auto *NewMMO = MF->getMachineMemOperand(OldMMO, AAInfo);
   MI.setMemRefs(*MF, NewMMO);
 
-  MachineModuleSlotTracker MST(MMI, MF);
+  MachineModuleSlotTracker MST(
+      [&](const Function &F) { return MMI.getMachineFunction(F); }, MF);
   // Print that MI with new machine metadata, which slot numbers should be
   // assigned.
   EXPECT_EQ(