[MIRPrinter] Add machine metadata support.

- Distinct metadata needs generating in the codegen to attach correct
  AAInfo on the loads/stores after lowering, merging, and other relevant
  transformations.
- This patch adds 'MachhineModuleSlotTracker' to help assign slot
  numbers to these newly generated unnamed metadata nodes.
- To help 'MachhineModuleSlotTracker' track machine metadata, the
  original 'SlotTracker' is rebased from 'AbstractSlotTrackerStorage',
  which provides basic interfaces to create/retrive metadata slots. In
  addition, once LLVM IR is processsed, additional hooks are also
  introduced to help collect machine metadata and assign them slot
  numbers.
- Finally, if there is any such machine metadata, 'MIRPrinter' outputs
  an additional 'machineMetadataNodes' field containing all the
  definition of those nodes.

Reviewed By: arsenm

Differential Revision: https://reviews.llvm.org/D103205

include/llvm/CodeGen/MIRYamlMapping.h

@@ -705,6 +705,7 @@ struct MachineFunction {
   std::vector<CallSiteInfo> CallSitesInfo;
   std::vector<DebugValueSubstitution> DebugValueSubstitutions;
   MachineJumpTable JumpTableInfo;
+  std::vector<StringValue> MachineMetadataNodes;
   BlockStringValue Body;
 };
 
@@ -739,6 +740,9 @@ template <> struct MappingTraits<MachineFunction> {
     YamlIO.mapOptional("machineFunctionInfo", MF.MachineFuncInfo);
     if (!YamlIO.outputting() || !MF.JumpTableInfo.Entries.empty())
       YamlIO.mapOptional("jumpTable", MF.JumpTableInfo, MachineJumpTable());
+    if (!YamlIO.outputting() || !MF.MachineMetadataNodes.empty())
+      YamlIO.mapOptional("machineMetadataNodes", MF.MachineMetadataNodes,
+                         std::vector<StringValue>());
     YamlIO.mapOptional("body", MF.Body, BlockStringValue());
   }
 };

include/llvm/CodeGen/MachineModuleSlotTracker.h

@@ -0,0 +1,45 @@
+//===-- llvm/CodeGen/MachineModuleInfo.h ------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_MACHINEMODULESLOTTRACKER_H
+#define LLVM_CODEGEN_MACHINEMODULESLOTTRACKER_H
+
+#include "llvm/IR/ModuleSlotTracker.h"
+
+namespace llvm {
+
+class AbstractSlotTrackerStorage;
+class Function;
+class MachineModuleInfo;
+class MachineFunction;
+class Module;
+
+class MachineModuleSlotTracker : public ModuleSlotTracker {
+  const Function &TheFunction;
+  const MachineModuleInfo &TheMMI;
+  unsigned MDNStartSlot, MDNEndSlot;
+
+  void processMachineFunctionMetadata(AbstractSlotTrackerStorage *AST,
+                                      const MachineFunction &MF);
+  void processMachineModule(AbstractSlotTrackerStorage *AST, const Module *M,
+                            bool ShouldInitializeAllMetadata);
+  void processMachineFunction(AbstractSlotTrackerStorage *AST,
+                              const Function *F,
+                              bool ShouldInitializeAllMetadata);
+
+public:
+  MachineModuleSlotTracker(const MachineFunction *MF,
+                           bool ShouldInitializeAllMetadata = true);
+  ~MachineModuleSlotTracker();
+
+  void collectMachineMDNodes(MachineMDNodeListType &L) const;
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_MACHINEMODULESLOTTRACKER_H

include/llvm/IR/ModuleSlotTracker.h

@@ -9,7 +9,10 @@
 #ifndef LLVM_IR_MODULESLOTTRACKER_H
 #define LLVM_IR_MODULESLOTTRACKER_H
 
+#include <functional>
 #include <memory>
+#include <utility>
+#include <vector>
 
 namespace llvm {
 
@@ -17,6 +20,18 @@ class Module;
 class Function;
 class SlotTracker;
 class Value;
+class MDNode;
+
+/// Abstract interface of slot tracker storage.
+class AbstractSlotTrackerStorage {
+public:
+  virtual ~AbstractSlotTrackerStorage();
+
+  virtual unsigned getNextMetadataSlot() = 0;
+
+  virtual void createMetadataSlot(const MDNode *) = 0;
+  virtual int getMetadataSlot(const MDNode *) = 0;
+};
 
 /// Manage lifetime of a slot tracker for printing IR.
 ///
@@ -36,6 +51,11 @@ class ModuleSlotTracker {
   const Function *F = nullptr;
   SlotTracker *Machine = nullptr;
 
+  std::function<void(AbstractSlotTrackerStorage *, const Module *, bool)>
+      ProcessModuleHookFn;
+  std::function<void(AbstractSlotTrackerStorage *, const Function *, bool)>
+      ProcessFunctionHookFn;
+
 public:
   /// Wrap a preinitialized SlotTracker.
   ModuleSlotTracker(SlotTracker &Machine, const Module *M,
@@ -52,7 +72,7 @@ public:
                              bool ShouldInitializeAllMetadata = true);
 
   /// Destructor to clean up storage.
-  ~ModuleSlotTracker();
+  virtual ~ModuleSlotTracker();
 
   /// Lazily creates a slot tracker.
   SlotTracker *getMachine();
@@ -72,6 +92,16 @@ public:
   /// this method.
   /// Return -1 if the value is not in the function's SlotTracker.
   int getLocalSlot(const Value *V);
+
+  void setProcessHook(
+      std::function<void(AbstractSlotTrackerStorage *, const Module *, bool)>);
+  void setProcessHook(std::function<void(AbstractSlotTrackerStorage *,
+                                         const Function *, bool)>);
+
+  using MachineMDNodeListType =
+      std::vector<std::pair<unsigned, const MDNode *>>;
+
+  void collectMDNodes(MachineMDNodeListType &L, unsigned LB, unsigned UB) const;
 };
 
 } // end namespace llvm

lib/CodeGen/CMakeLists.txt

@@ -91,6 +91,7 @@ add_llvm_component_library(LLVMCodeGen
   MachineLoopUtils.cpp
   MachineModuleInfo.cpp
   MachineModuleInfoImpls.cpp
+  MachineModuleSlotTracker.cpp
   MachineOperand.cpp
   MachineOptimizationRemarkEmitter.cpp
   MachineOutliner.cpp

lib/CodeGen/MIRPrinter.cpp

@@ -29,13 +29,14 @@
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
+#include "llvm/CodeGen/MachineModuleSlotTracker.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/PseudoSourceValue.h"
+#include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
-#include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DebugInfo.h"
@@ -135,6 +136,9 @@ public:
   void convertCallSiteObjects(yaml::MachineFunction &YMF,
                               const MachineFunction &MF,
                               ModuleSlotTracker &MST);
+  void convertMachineMetadataNodes(yaml::MachineFunction &YMF,
+                                   const MachineFunction &MF,
+                                   MachineModuleSlotTracker &MST);
 
 private:
   void initRegisterMaskIds(const MachineFunction &MF);
@@ -215,7 +219,7 @@ void MIRPrinter::print(const MachineFunction &MF) {
       MachineFunctionProperties::Property::FailedISel);
 
   convert(YamlMF, MF.getRegInfo(), MF.getSubtarget().getRegisterInfo());
-  ModuleSlotTracker MST(MF.getFunction().getParent());
+  MachineModuleSlotTracker MST(&MF);
   MST.incorporateFunction(MF.getFunction());
   convert(MST, YamlMF.FrameInfo, MF.getFrameInfo());
   convertStackObjects(YamlMF, MF, MST);
@@ -243,6 +247,10 @@ void MIRPrinter::print(const MachineFunction &MF) {
     IsNewlineNeeded = true;
   }
   StrOS.flush();
+  // Convert machine metadata collected during the print of the machine
+  // function.
+  convertMachineMetadataNodes(YamlMF, MF, MST);
+
   yaml::Output Out(OS);
   if (!SimplifyMIR)
       Out.setWriteDefaultValues(true);
@@ -525,6 +533,19 @@ void MIRPrinter::convertCallSiteObjects(yaml::MachineFunction &YMF,
              });
 }
 
+void MIRPrinter::convertMachineMetadataNodes(yaml::MachineFunction &YMF,
+                                             const MachineFunction &MF,
+                                             MachineModuleSlotTracker &MST) {
+  MachineModuleSlotTracker::MachineMDNodeListType MDList;
+  MST.collectMachineMDNodes(MDList);
+  for (auto &MD : MDList) {
+    std::string NS;
+    raw_string_ostream StrOS(NS);
+    MD.second->print(StrOS, MST, MF.getFunction().getParent());
+    YMF.MachineMetadataNodes.push_back(StrOS.str());
+  }
+}
+
 void MIRPrinter::convert(yaml::MachineFunction &MF,
                          const MachineConstantPool &ConstantPool) {
   unsigned ID = 0;

lib/CodeGen/MachineModuleSlotTracker.cpp

@@ -0,0 +1,81 @@
+//===-- llvm/CodeGen/MachineModuleInfo.cpp ----------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/MachineModuleSlotTracker.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+
+using namespace llvm;
+
+void MachineModuleSlotTracker::processMachineFunctionMetadata(
+    AbstractSlotTrackerStorage *AST, const MachineFunction &MF) {
+  // Create metadata created within the backend.
+  for (const MachineBasicBlock &MBB : MF)
+    for (const MachineInstr &MI : MBB.instrs())
+      for (const MachineMemOperand *MMO : MI.memoperands()) {
+        AAMDNodes AAInfo = MMO->getAAInfo();
+        if (AAInfo.TBAA)
+          AST->createMetadataSlot(AAInfo.TBAA);
+        if (AAInfo.TBAAStruct)
+          AST->createMetadataSlot(AAInfo.TBAAStruct);
+        if (AAInfo.Scope)
+          AST->createMetadataSlot(AAInfo.Scope);
+        if (AAInfo.NoAlias)
+          AST->createMetadataSlot(AAInfo.NoAlias);
+      }
+}
+
+void MachineModuleSlotTracker::processMachineModule(
+    AbstractSlotTrackerStorage *AST, const Module *M,
+    bool ShouldInitializeAllMetadata) {
+  if (ShouldInitializeAllMetadata) {
+    for (const Function &F : *M) {
+      if (&F != &TheFunction)
+        continue;
+      MDNStartSlot = AST->getNextMetadataSlot();
+      if (auto *MF = TheMMI.getMachineFunction(F))
+        processMachineFunctionMetadata(AST, *MF);
+      MDNEndSlot = AST->getNextMetadataSlot();
+      break;
+    }
+  }
+}
+
+void MachineModuleSlotTracker::processMachineFunction(
+    AbstractSlotTrackerStorage *AST, const Function *F,
+    bool ShouldInitializeAllMetadata) {
+  if (!ShouldInitializeAllMetadata && F == &TheFunction) {
+    MDNStartSlot = AST->getNextMetadataSlot();
+    if (auto *MF = TheMMI.getMachineFunction(*F))
+      processMachineFunctionMetadata(AST, *MF);
+    MDNEndSlot = AST->getNextMetadataSlot();
+  }
+}
+
+void MachineModuleSlotTracker::collectMachineMDNodes(
+    MachineMDNodeListType &L) const {
+  collectMDNodes(L, MDNStartSlot, MDNEndSlot);
+}
+
+MachineModuleSlotTracker::MachineModuleSlotTracker(
+    const MachineFunction *MF, bool ShouldInitializeAllMetadata)
+    : ModuleSlotTracker(MF->getFunction().getParent(),
+                        ShouldInitializeAllMetadata),
+      TheFunction(MF->getFunction()), TheMMI(MF->getMMI()), MDNStartSlot(0),
+      MDNEndSlot(0) {
+  setProcessHook([this](AbstractSlotTrackerStorage *AST, const Module *M,
+                        bool ShouldInitializeAllMetadata) {
+    this->processMachineModule(AST, M, ShouldInitializeAllMetadata);
+  });
+  setProcessHook([this](AbstractSlotTrackerStorage *AST, const Function *F,
+                        bool ShouldInitializeAllMetadata) {
+    this->processMachineFunction(AST, F, ShouldInitializeAllMetadata);
+  });
+}
+
+MachineModuleSlotTracker::~MachineModuleSlotTracker() = default;

lib/IR/AsmWriter.cpp

@@ -713,6 +713,8 @@ void TypePrinting::printStructBody(StructType *STy, raw_ostream &OS) {
     OS << '>';
 }
 
+AbstractSlotTrackerStorage::~AbstractSlotTrackerStorage() {}
+
 namespace llvm {
 
 //===----------------------------------------------------------------------===//
@@ -720,7 +722,7 @@ namespace llvm {
 //===----------------------------------------------------------------------===//
 /// This class provides computation of slot numbers for LLVM Assembly writing.
 ///
-class SlotTracker {
+class SlotTracker : public AbstractSlotTrackerStorage {
 public:
   /// ValueMap - A mapping of Values to slot numbers.
   using ValueMap = DenseMap<const Value *, unsigned>;
@@ -734,6 +736,11 @@ private:
   bool FunctionProcessed = false;
   bool ShouldInitializeAllMetadata;
 
+  std::function<void(AbstractSlotTrackerStorage *, const Module *, bool)>
+      ProcessModuleHookFn;
+  std::function<void(AbstractSlotTrackerStorage *, const Function *, bool)>
+      ProcessFunctionHookFn;
+
   /// The summary index for which we are holding slot numbers.
   const ModuleSummaryIndex *TheIndex = nullptr;
 
@@ -788,11 +795,22 @@ public:
   SlotTracker(const SlotTracker &) = delete;
   SlotTracker &operator=(const SlotTracker &) = delete;
 
+  ~SlotTracker() = default;
+
+  void setProcessHook(
+      std::function<void(AbstractSlotTrackerStorage *, const Module *, bool)>);
+  void setProcessHook(std::function<void(AbstractSlotTrackerStorage *,
+                                         const Function *, bool)>);
+
+  unsigned getNextMetadataSlot() override { return mdnNext; }
+
+  void createMetadataSlot(const MDNode *N) override;
+
   /// Return the slot number of the specified value in it's type
   /// plane.  If something is not in the SlotTracker, return -1.
   int getLocalSlot(const Value *V);
   int getGlobalSlot(const GlobalValue *V);
-  int getMetadataSlot(const MDNode *N);
+  int getMetadataSlot(const MDNode *N) override;
   int getAttributeGroupSlot(AttributeSet AS);
   int getModulePathSlot(StringRef Path);
   int getGUIDSlot(GlobalValue::GUID GUID);
@@ -893,6 +911,10 @@ SlotTracker *ModuleSlotTracker::getMachine() {
   MachineStorage =
       std::make_unique<SlotTracker>(M, ShouldInitializeAllMetadata);
   Machine = MachineStorage.get();
+  if (ProcessModuleHookFn)
+    Machine->setProcessHook(ProcessModuleHookFn);
+  if (ProcessFunctionHookFn)
+    Machine->setProcessHook(ProcessFunctionHookFn);
   return Machine;
 }
 
@@ -915,6 +937,18 @@ int ModuleSlotTracker::getLocalSlot(const Value *V) {
   return Machine->getLocalSlot(V);
 }
 
+void ModuleSlotTracker::setProcessHook(
+    std::function<void(AbstractSlotTrackerStorage *, const Module *, bool)>
+        Fn) {
+  ProcessModuleHookFn = Fn;
+}
+
+void ModuleSlotTracker::setProcessHook(
+    std::function<void(AbstractSlotTrackerStorage *, const Function *, bool)>
+        Fn) {
+  ProcessFunctionHookFn = Fn;
+}
+
 static SlotTracker *createSlotTracker(const Value *V) {
   if (const Argument *FA = dyn_cast<Argument>(V))
     return new SlotTracker(FA->getParent());
@@ -1025,6 +1059,9 @@ void SlotTracker::processModule() {
       CreateAttributeSetSlot(FnAttrs);
   }
 
+  if (ProcessModuleHookFn)
+    ProcessModuleHookFn(this, TheModule, ShouldInitializeAllMetadata);
+
   ST_DEBUG("end processModule!\n");
 }
 
@@ -1065,6 +1102,9 @@ void SlotTracker::processFunction() {
     }
   }
 
+  if (ProcessFunctionHookFn)
+    ProcessFunctionHookFn(this, TheFunction, ShouldInitializeAllMetadata);
+
   FunctionProcessed = true;
 
   ST_DEBUG("end processFunction!\n");
@@ -1155,6 +1195,21 @@ int SlotTracker::getGlobalSlot(const GlobalValue *V) {
   return MI == mMap.end() ? -1 : (int)MI->second;
 }
 
+void SlotTracker::setProcessHook(
+    std::function<void(AbstractSlotTrackerStorage *, const Module *, bool)>
+        Fn) {
+  ProcessModuleHookFn = Fn;
+}
+
+void SlotTracker::setProcessHook(
+    std::function<void(AbstractSlotTrackerStorage *, const Function *, bool)>
+        Fn) {
+  ProcessFunctionHookFn = Fn;
+}
+
+/// getMetadataSlot - Get the slot number of a MDNode.
+void SlotTracker::createMetadataSlot(const MDNode *N) { CreateMetadataSlot(N); }
+
 /// getMetadataSlot - Get the slot number of a MDNode.
 int SlotTracker::getMetadataSlot(const MDNode *N) {
   // Check for uninitialized state and do lazy initialization.
@@ -4767,6 +4822,17 @@ void ModuleSummaryIndex::print(raw_ostream &ROS, bool IsForDebug) const {
   W.printModuleSummaryIndex();
 }
 
+void ModuleSlotTracker::collectMDNodes(MachineMDNodeListType &L, unsigned LB,
+                                       unsigned UB) const {
+  SlotTracker *ST = MachineStorage.get();
+  if (!ST)
+    return;
+
+  for (auto &I : llvm::make_range(ST->mdn_begin(), ST->mdn_end()))
+    if (I.second >= LB && I.second < UB)
+      L.push_back(std::make_pair(I.second, I.first));
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 // Value::dump - allow easy printing of Values from the debugger.
 LLVM_DUMP_METHOD

unittests/CMakeLists.txt

@@ -33,6 +33,7 @@ add_subdirectory(LineEditor)
 add_subdirectory(Linker)
 add_subdirectory(MC)
 add_subdirectory(MI)
+add_subdirectory(MIR)
 add_subdirectory(Object)
 add_subdirectory(ObjectYAML)
 add_subdirectory(Option)

unittests/CodeGen/CMakeLists.txt

@@ -5,6 +5,7 @@ set(LLVM_LINK_COMPONENTS
   AsmPrinter
   CodeGen
   Core
+  FileCheck
   MC
   MIRParser
   Passes

unittests/MIR/CMakeLists.txt

@@ -0,0 +1,16 @@
+set(LLVM_LINK_COMPONENTS
+  ${LLVM_TARGETS_TO_BUILD}
+  CodeGen
+  Core
+  FileCheck
+  MC
+  MIRParser
+  Support
+  Target
+  )
+
+add_llvm_unittest(MIRTests
+  MachineMetadata.cpp
+  )
+
+target_link_libraries(MIRTests PRIVATE LLVMTestingSupport)

unittests/MIR/MachineMetadata.cpp

@@ -0,0 +1,481 @@
+//===- MachineInstrBundleIteratorTest.cpp ---------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/MIRParser/MIRParser.h"
+#include "llvm/CodeGen/MIRPrinter.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineMemOperand.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/CodeGen/MachineModuleSlotTracker.h"
+#include "llvm/CodeGen/MachineOperand.h"
+#include "llvm/CodeGen/TargetFrameLowering.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetLowering.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/FileCheck/FileCheck.h"
+#include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/ModuleSlotTracker.h"
+#include "llvm/MC/MCAsmInfo.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/TargetRegistry.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Target/TargetMachine.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+class MachineMetadataTest : public testing::Test {
+public:
+  MachineMetadataTest() {}
+
+protected:
+  LLVMContext Context;
+  std::unique_ptr<Module> M;
+  std::unique_ptr<MIRParser> MIR;
+
+  static void SetUpTestCase() {
+    InitializeAllTargetInfos();
+    InitializeAllTargets();
+    InitializeAllTargetMCs();
+  }
+
+  void SetUp() override { M = std::make_unique<Module>("Dummy", Context); }
+
+  void addHooks(ModuleSlotTracker &MST, const MachineOperand &MO) {
+    // Setup hooks to assign slot numbers for the specified machine metadata.
+    MST.setProcessHook([&MO](AbstractSlotTrackerStorage *AST, const Module *M,
+                             bool ShouldInitializeAllMetadata) {
+      if (ShouldInitializeAllMetadata) {
+        if (MO.isMetadata())
+          AST->createMetadataSlot(MO.getMetadata());
+      }
+    });
+    MST.setProcessHook([&MO](AbstractSlotTrackerStorage *AST, const Function *F,
+                             bool ShouldInitializeAllMetadata) {
+      if (!ShouldInitializeAllMetadata) {
+        if (MO.isMetadata())
+          AST->createMetadataSlot(MO.getMetadata());
+      }
+    });
+  }
+
+  std::unique_ptr<LLVMTargetMachine>
+  createTargetMachine(std::string TT, StringRef CPU, StringRef FS) {
+    std::string Error;
+    const Target *T = TargetRegistry::lookupTarget(TT, Error);
+    if (!T)
+      return nullptr;
+    TargetOptions Options;
+    return std::unique_ptr<LLVMTargetMachine>(static_cast<LLVMTargetMachine *>(
+        T->createTargetMachine(TT, CPU, FS, Options, None, None)));
+  }
+
+  std::unique_ptr<Module> parseMIR(const TargetMachine &TM, StringRef MIRCode,
+                                   const char *FnName, MachineModuleInfo &MMI) {
+    SMDiagnostic Diagnostic;
+    std::unique_ptr<MemoryBuffer> MBuffer = MemoryBuffer::getMemBuffer(MIRCode);
+    MIR = createMIRParser(std::move(MBuffer), Context);
+    if (!MIR)
+      return nullptr;
+
+    std::unique_ptr<Module> Mod = MIR->parseIRModule();
+    if (!Mod)
+      return nullptr;
+
+    Mod->setDataLayout(TM.createDataLayout());
+
+    if (MIR->parseMachineFunctions(*Mod, MMI)) {
+      M.reset();
+      return nullptr;
+    }
+
+    return Mod;
+  }
+};
+
+// Helper to dump the printer output into a string.
+static std::string print(std::function<void(raw_ostream &OS)> PrintFn) {
+  std::string Str;
+  raw_string_ostream OS(Str);
+  PrintFn(OS);
+  OS.flush();
+  return Str;
+}
+
+TEST_F(MachineMetadataTest, TrivialHook) {
+  // Verify that post-process hook is invoked to assign slot numbers for
+  // machine metadata.
+  ASSERT_TRUE(M);
+
+  // Create a MachineOperand with a metadata and print it.
+  Metadata *MDS = MDString::get(Context, "foo");
+  MDNode *Node = MDNode::get(Context, MDS);
+  MachineOperand MO = MachineOperand::CreateMetadata(Node);
+
+  // Checking some preconditions on the newly created
+  // MachineOperand.
+  ASSERT_TRUE(MO.isMetadata());
+  ASSERT_EQ(MO.getMetadata(), Node);
+
+  ModuleSlotTracker MST(M.get());
+  addHooks(MST, MO);
+
+  // Print a MachineOperand containing a metadata node.
+  EXPECT_EQ("!0", print([&](raw_ostream &OS) {
+              MO.print(OS, MST, LLT{}, /*OpIdx*/ ~0U, /*PrintDef=*/false,
+                       /*IsStandalone=*/false,
+                       /*ShouldPrintRegisterTies=*/false, /*TiedOperandIdx=*/0,
+                       /*TRI=*/nullptr,
+                       /*IntrinsicInfo=*/nullptr);
+            }));
+  // Print the definition of that metadata node.
+  EXPECT_EQ("!0 = !{!\"foo\"}",
+            print([&](raw_ostream &OS) { Node->print(OS, MST); }));
+}
+
+TEST_F(MachineMetadataTest, BasicHook) {
+  // Verify that post-process hook is invoked to assign slot numbers for
+  // machine metadata. When both LLVM IR and machine IR contain metadata,
+  // ensure that machine metadata is always assigned after LLVM IR.
+  ASSERT_TRUE(M);
+
+  // Create a MachineOperand with a metadata and print it.
+  Metadata *MachineMDS = MDString::get(Context, "foo");
+  MDNode *MachineNode = MDNode::get(Context, MachineMDS);
+  MachineOperand MO = MachineOperand::CreateMetadata(MachineNode);
+
+  // Checking some preconditions on the newly created
+  // MachineOperand.
+  ASSERT_TRUE(MO.isMetadata());
+  ASSERT_EQ(MO.getMetadata(), MachineNode);
+
+  // Create metadata in LLVM IR.
+  NamedMDNode *MD = M->getOrInsertNamedMetadata("namedmd");
+  Metadata *MDS = MDString::get(Context, "bar");
+  MDNode *Node = MDNode::get(Context, MDS);
+  MD->addOperand(Node);
+
+  ModuleSlotTracker MST(M.get());
+  addHooks(MST, MO);
+
+  // Print a MachineOperand containing a metadata node.
+  EXPECT_EQ("!1", print([&](raw_ostream &OS) {
+              MO.print(OS, MST, LLT{}, /*OpIdx*/ ~0U, /*PrintDef=*/false,
+                       /*IsStandalone=*/false,
+                       /*ShouldPrintRegisterTies=*/false, /*TiedOperandIdx=*/0,
+                       /*TRI=*/nullptr,
+                       /*IntrinsicInfo=*/nullptr);
+            }));
+  // Print the definition of these unnamed metadata nodes.
+  EXPECT_EQ("!0 = !{!\"bar\"}",
+            print([&](raw_ostream &OS) { Node->print(OS, MST); }));
+  EXPECT_EQ("!1 = !{!\"foo\"}",
+            print([&](raw_ostream &OS) { MachineNode->print(OS, MST); }));
+}
+
+static bool checkOutput(std::string CheckString, std::string Output) {
+  auto CheckBuffer = MemoryBuffer::getMemBuffer(CheckString, "");
+  auto OutputBuffer = MemoryBuffer::getMemBuffer(Output, "Output", false);
+
+  SmallString<4096> CheckFileBuffer;
+  FileCheckRequest Req;
+  FileCheck FC(Req);
+  StringRef CheckFileText =
+      FC.CanonicalizeFile(*CheckBuffer.get(), CheckFileBuffer);
+
+  SourceMgr SM;
+  SM.AddNewSourceBuffer(MemoryBuffer::getMemBuffer(CheckFileText, "CheckFile"),
+                        SMLoc());
+  Regex PrefixRE = FC.buildCheckPrefixRegex();
+  if (FC.readCheckFile(SM, CheckFileText, PrefixRE))
+    return false;
+
+  auto OutBuffer = OutputBuffer->getBuffer();
+  SM.AddNewSourceBuffer(std::move(OutputBuffer), SMLoc());
+  return FC.checkInput(SM, OutBuffer);
+}
+
+TEST_F(MachineMetadataTest, MMSlotTrackerAArch64) {
+  auto TM = createTargetMachine(Triple::normalize("aarch64--"), "", "");
+  if (!TM)
+    GTEST_SKIP();
+
+  StringRef MIRString = R"MIR(
+--- |
+  define i32 @test0(i32* %p) {
+    %r = load i32, i32* %p, align 4
+    ret i32 %r
+  }
+...
+---
+name:            test0
+liveins:
+  - { reg: '$x0', virtual-reg: '%0' }
+body:             |
+  bb.0 (%ir-block.0):
+    liveins: $x0
+
+  %0:gpr64common = COPY $x0
+  %1:gpr32 = LDRWui %0, 0 :: (load 4 from %ir.p)
+...
+)MIR";
+
+  MachineModuleInfo MMI(TM.get());
+  M = parseMIR(*TM, MIRString, "test0", MMI);
+  ASSERT_TRUE(M);
+
+  auto *MF = MMI.getMachineFunction(*M->getFunction("test0"));
+  auto *MBB = MF->getBlockNumbered(0);
+
+  auto &MI = MBB->back();
+  ASSERT_TRUE(MI.hasOneMemOperand());
+
+  // Create and attached scoped AA metadata on that instruction with one MMO.
+  MDBuilder MDB(Context);
+  MDNode *Domain = MDB.createAnonymousAliasScopeDomain("domain");
+  MDNode *Scope0 = MDB.createAnonymousAliasScope(Domain, "scope0");
+  MDNode *Scope1 = MDB.createAnonymousAliasScope(Domain, "scope1");
+  MDNode *Set0 = MDNode::get(Context, {Scope0});
+  MDNode *Set1 = MDNode::get(Context, {Scope1});
+
+  AAMDNodes AAInfo;
+  AAInfo.TBAA = AAInfo.TBAAStruct = nullptr;
+  AAInfo.Scope = Set0;
+  AAInfo.NoAlias = Set1;
+
+  auto *OldMMO = MI.memoperands().front();
+  auto *NewMMO = MF->getMachineMemOperand(OldMMO, AAInfo);
+  MI.setMemRefs(*MF, NewMMO);
+
+  MachineModuleSlotTracker MST(MF);
+  // Print that MI with new machine metadata, which slot numbers should be
+  // assigned.
+  EXPECT_EQ("%1:gpr32 = LDRWui %0, 0 :: (load 4 from %ir.p, "
+            "!alias.scope !0, !noalias !3)",
+            print([&](raw_ostream &OS) {
+              MI.print(OS, MST, /*IsStandalone=*/false, /*SkipOpers=*/false,
+                       /*SkipDebugLoc=*/false, /*AddNewLine=*/false);
+            }));
+
+  std::vector<const MDNode *> Generated{Domain, Scope0, Scope1, Set0, Set1};
+  // Examine machine metadata collected. They should match ones
+  // afore-generated.
+  std::vector<const MDNode *> Collected;
+  MachineModuleSlotTracker::MachineMDNodeListType MDList;
+  MST.collectMachineMDNodes(MDList);
+  for (auto &MD : MDList)
+    Collected.push_back(MD.second);
+
+  std::sort(Generated.begin(), Generated.end());
+  std::sort(Collected.begin(), Collected.end());
+  EXPECT_EQ(Collected, Generated);
+
+  // FileCheck the output from MIR printer.
+  std::string Output = print([&](raw_ostream &OS) { printMIR(OS, *MF); });
+  std::string CheckString = R"(
+CHECK: machineMetadataNodes:
+CHECK-DAG: ![[MMDOMAIN:[0-9]+]] = distinct !{!{{[0-9]+}}, !"domain"}
+CHECK-DAG: ![[MMSCOPE0:[0-9]+]] = distinct !{!{{[0-9]+}}, ![[MMDOMAIN]], !"scope0"}
+CHECK-DAG: ![[MMSCOPE1:[0-9]+]] = distinct !{!{{[0-9]+}}, ![[MMDOMAIN]], !"scope1"}
+CHECK-DAG: ![[MMSET0:[0-9]+]] = !{![[MMSCOPE0]]}
+CHECK-DAG: ![[MMSET1:[0-9]+]] = !{![[MMSCOPE1]]}
+CHECK: body:
+CHECK: %1:gpr32 = LDRWui %0, 0 :: (load 4 from %ir.p, !alias.scope ![[MMSET0]], !noalias ![[MMSET1]])
+)";
+  EXPECT_TRUE(checkOutput(CheckString, Output));
+}
+
+TEST_F(MachineMetadataTest, MMSlotTrackerX64) {
+  auto TM = createTargetMachine(Triple::normalize("x86_64--"), "", "");
+  if (!TM)
+    GTEST_SKIP();
+
+  StringRef MIRString = R"MIR(
+--- |
+  define i32 @test0(i32* %p) {
+    %r = load i32, i32* %p, align 4
+    ret i32 %r
+  }
+...
+---
+name:            test0
+liveins:
+  - { reg: '$rdi', virtual-reg: '%0' }
+body:             |
+  bb.0 (%ir-block.0):
+    liveins: $rdi
+
+  %0:gr64 = COPY $rdi
+  %1:gr32 = MOV32rm %0, 1, $noreg, 0, $noreg :: (load 4 from %ir.p)
+...
+)MIR";
+
+  MachineModuleInfo MMI(TM.get());
+  M = parseMIR(*TM, MIRString, "test0", MMI);
+  ASSERT_TRUE(M);
+
+  auto *MF = MMI.getMachineFunction(*M->getFunction("test0"));
+  auto *MBB = MF->getBlockNumbered(0);
+
+  auto &MI = MBB->back();
+  ASSERT_FALSE(MI.memoperands_empty());
+  ASSERT_TRUE(MI.hasOneMemOperand());
+
+  // Create and attached scoped AA metadata on that instruction with one MMO.
+  MDBuilder MDB(Context);
+  MDNode *Domain = MDB.createAnonymousAliasScopeDomain("domain");
+  MDNode *Scope0 = MDB.createAnonymousAliasScope(Domain, "scope0");
+  MDNode *Scope1 = MDB.createAnonymousAliasScope(Domain, "scope1");
+  MDNode *Set0 = MDNode::get(Context, {Scope0});
+  MDNode *Set1 = MDNode::get(Context, {Scope1});
+
+  AAMDNodes AAInfo;
+  AAInfo.TBAA = AAInfo.TBAAStruct = nullptr;
+  AAInfo.Scope = Set0;
+  AAInfo.NoAlias = Set1;
+
+  auto *OldMMO = MI.memoperands().front();
+  auto *NewMMO = MF->getMachineMemOperand(OldMMO, AAInfo);
+  MI.setMemRefs(*MF, NewMMO);
+
+  MachineModuleSlotTracker MST(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 4 from %ir.p, "
+            "!alias.scope !0, !noalias !3)",
+            print([&](raw_ostream &OS) {
+              MI.print(OS, MST, /*IsStandalone=*/false, /*SkipOpers=*/false,
+                       /*SkipDebugLoc=*/false, /*AddNewLine=*/false);
+            }));
+
+  std::vector<const MDNode *> Generated{Domain, Scope0, Scope1, Set0, Set1};
+  // Examine machine metadata collected. They should match ones
+  // afore-generated.
+  std::vector<const MDNode *> Collected;
+  MachineModuleSlotTracker::MachineMDNodeListType MDList;
+  MST.collectMachineMDNodes(MDList);
+  for (auto &MD : MDList)
+    Collected.push_back(MD.second);
+
+  std::sort(Generated.begin(), Generated.end());
+  std::sort(Collected.begin(), Collected.end());
+  EXPECT_EQ(Collected, Generated);
+
+  // FileCheck the output from MIR printer.
+  std::string Output = print([&](raw_ostream &OS) { printMIR(OS, *MF); });
+  std::string CheckString = R"(
+CHECK: machineMetadataNodes:
+CHECK-DAG: ![[MMDOMAIN:[0-9]+]] = distinct !{!{{[0-9]+}}, !"domain"}
+CHECK-DAG: ![[MMSCOPE0:[0-9]+]] = distinct !{!{{[0-9]+}}, ![[MMDOMAIN]], !"scope0"}
+CHECK-DAG: ![[MMSCOPE1:[0-9]+]] = distinct !{!{{[0-9]+}}, ![[MMDOMAIN]], !"scope1"}
+CHECK-DAG: ![[MMSET0:[0-9]+]] = !{![[MMSCOPE0]]}
+CHECK-DAG: ![[MMSET1:[0-9]+]] = !{![[MMSCOPE1]]}
+CHECK: body:
+CHECK: %1:gr32 = MOV32rm %0, 1, $noreg, 0, $noreg :: (load 4 from %ir.p, !alias.scope ![[MMSET0]], !noalias ![[MMSET1]])
+)";
+  EXPECT_TRUE(checkOutput(CheckString, Output));
+}
+
+TEST_F(MachineMetadataTest, MMSlotTrackerAMDGPU) {
+  auto TM = createTargetMachine(Triple::normalize("amdgcn-amd-amdhsa"),
+                                "gfx1010", "");
+  if (!TM)
+    GTEST_SKIP();
+
+  StringRef MIRString = R"MIR(
+--- |
+  define i32 @test0(i32* %p) {
+    %r = load i32, i32* %p, align 4
+    ret i32 %r
+  }
+...
+---
+name:            test0
+liveins:
+  - { reg: '$vgpr0', virtual-reg: '%0' }
+  - { reg: '$vgpr1', virtual-reg: '%1' }
+  - { reg: '$sgpr30_sgpr31', virtual-reg: '%2' }
+body:             |
+  bb.0 (%ir-block.0):
+    liveins: $vgpr0, $vgpr1, $sgpr30_sgpr31
+
+    %2:sreg_64 = COPY $sgpr30_sgpr31
+    %1:vgpr_32 = COPY $vgpr1
+    %0:vgpr_32 = COPY $vgpr0
+    %8:vreg_64 = REG_SEQUENCE %0, %subreg.sub0, %1, %subreg.sub1
+    %6:vreg_64 = COPY %8
+    %5:vgpr_32 = FLAT_LOAD_DWORD killed %6, 0, 0, implicit $exec, implicit $flat_scr :: (load 4 from %ir.p)
+...
+)MIR";
+
+  MachineModuleInfo MMI(TM.get());
+  M = parseMIR(*TM, MIRString, "test0", MMI);
+  ASSERT_TRUE(M);
+
+  auto *MF = MMI.getMachineFunction(*M->getFunction("test0"));
+  auto *MBB = MF->getBlockNumbered(0);
+
+  auto &MI = MBB->back();
+  ASSERT_FALSE(MI.memoperands_empty());
+  ASSERT_TRUE(MI.hasOneMemOperand());
+
+  // Create and attached scoped AA metadata on that instruction with one MMO.
+  MDBuilder MDB(Context);
+  MDNode *Domain = MDB.createAnonymousAliasScopeDomain("domain");
+  MDNode *Scope0 = MDB.createAnonymousAliasScope(Domain, "scope0");
+  MDNode *Scope1 = MDB.createAnonymousAliasScope(Domain, "scope1");
+  MDNode *Set0 = MDNode::get(Context, {Scope0});
+  MDNode *Set1 = MDNode::get(Context, {Scope1});
+
+  AAMDNodes AAInfo;
+  AAInfo.TBAA = AAInfo.TBAAStruct = nullptr;
+  AAInfo.Scope = Set0;
+  AAInfo.NoAlias = Set1;
+
+  auto *OldMMO = MI.memoperands().front();
+  auto *NewMMO = MF->getMachineMemOperand(OldMMO, AAInfo);
+  MI.setMemRefs(*MF, NewMMO);
+
+  MachineModuleSlotTracker MST(MF);
+  // Print that MI with new machine metadata, which slot numbers should be
+  // assigned.
+  EXPECT_EQ(
+      "%5:vgpr_32 = FLAT_LOAD_DWORD killed %4, 0, 0, implicit $exec, implicit "
+      "$flat_scr :: (load 4 from %ir.p, !alias.scope !0, !noalias !3)",
+      print([&](raw_ostream &OS) {
+        MI.print(OS, MST, /*IsStandalone=*/false, /*SkipOpers=*/false,
+                 /*SkipDebugLoc=*/false, /*AddNewLine=*/false);
+      }));
+
+  std::vector<const MDNode *> Generated{Domain, Scope0, Scope1, Set0, Set1};
+  // Examine machine metadata collected. They should match ones
+  // afore-generated.
+  std::vector<const MDNode *> Collected;
+  MachineModuleSlotTracker::MachineMDNodeListType MDList;
+  MST.collectMachineMDNodes(MDList);
+  for (auto &MD : MDList)
+    Collected.push_back(MD.second);
+
+  std::sort(Generated.begin(), Generated.end());
+  std::sort(Collected.begin(), Collected.end());
+  EXPECT_EQ(Collected, Generated);
+
+  // FileCheck the output from MIR printer.
+  std::string Output = print([&](raw_ostream &OS) { printMIR(OS, *MF); });
+  std::string CheckString = R"(
+CHECK: machineMetadataNodes:
+CHECK-DAG: ![[MMDOMAIN:[0-9]+]] = distinct !{!{{[0-9]+}}, !"domain"}
+CHECK-DAG: ![[MMSCOPE0:[0-9]+]] = distinct !{!{{[0-9]+}}, ![[MMDOMAIN]], !"scope0"}
+CHECK-DAG: ![[MMSCOPE1:[0-9]+]] = distinct !{!{{[0-9]+}}, ![[MMDOMAIN]], !"scope1"}
+CHECK-DAG: ![[MMSET0:[0-9]+]] = !{![[MMSCOPE0]]}
+CHECK-DAG: ![[MMSET1:[0-9]+]] = !{![[MMSCOPE1]]}
+CHECK: body:
+CHECK: %5:vgpr_32 = FLAT_LOAD_DWORD killed %4, 0, 0, implicit $exec, implicit $flat_scr :: (load 4 from %ir.p, !alias.scope ![[MMSET0]], !noalias ![[MMSET1]])
+)";
+  EXPECT_TRUE(checkOutput(CheckString, Output));
+}