[MIRPrinter] Add machine metadata support.

[hliao]

Summary

• 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.

[hliao]

Revise that machine tracker a bit to ensure machine metadata is alway per-function state.

• The printer of the machine function is done on a per-function wise. That makes the scanning of machine data through module unnecessary and problematic as MF may be removed deliberately under certain cases. I found that issue when validating the parser support.

[nikic]

Unfortunately I don't know anything about MIR -- adding @arsenm who I think should be familiar with this, or at least know someone who is :)

[arsenm]

Thanks for doing this, this has been painful forever.

Can you add some plain mir tests in test/CodeGen/MIR?

llvm/include/llvm/CodeGen/MIRYamlMapping.h

707

Is this just for the metadata nodes attached to the function itself?

llvm/unittests/CodeGen/MachineMetadata.cpp

123

EXPECT_EQ

155

EXPECT_EQ

327–328

Only need the hasOneMemOperand

[hliao]

Separate that test into MIR.

[hliao]

the patch is revised following comments.

llvm/include/llvm/CodeGen/MIRYamlMapping.h

707

yes, as 'printMIR' inherently works per function wise, the metadata here refers to ones attached to the specified MF. So far, metadata generated in the backend are all unnamed ones and only referenced within that function only.
Furthermore, our MIR/MF in some tools may be removed in a function-pass ('createFreeMachineFunctionPass'). That makes the whole machine module inconsistent between prints of different MFs.

[hliao]

Fix warngins from clang-tidy.

[arsenm]

llvm/include/llvm/CodeGen/MIRYamlMapping.h

707

So the intent is you do not need the IR section to use the metadata? I had assumed this would directly read from the IR. I'm all for breaking the IR dependency but wondering what happens if the IR section metadata is inconsistent with the MIR section metadata

[hliao]

llvm/include/llvm/CodeGen/MIRYamlMapping.h

707

those metadata is generated directly within the backend instead of reading from IR. For instance, in D102255, new scoped AA metadata is generated when lowering memcpy into machine instructions. In D102821, when two stores are merged, potentially we need a new scoped AA metadata as the new scope should be the union of that 2 stores and the new scoped no-alias should be the intersect. In short, with the newly introduced scoped AA metadata, the backend may need to generate new metadata, which could not be referenced directly from IR.

[hliao]

Kindly PING

[hliao]

Kindly PING again

[arsenm]

Needs standalone MIR tests (i.e. round trip tests not bundled in a unit test)

Also I would like to see some tests where there are inconsistencies between the metadata present in the IR section, and the explicit metadata in the MIR

[hliao]

In D103205#2811695, @arsenm wrote:

Needs standalone MIR tests (i.e. round trip tests not bundled in a unit test)

Also I would like to see some tests where there are inconsistencies between the metadata present in the IR section, and the explicit metadata in the MIR

round-trip needs parser support (under review @ https://reviews.llvm.org/D103282), where cases with mixed metadata are also included. Shall we merge them into one?

[hliao]

PING for the further comment on the round-trip test, which requires MIR parser support reviewed @ D103282

llvm/include/llvm/CodeGen/MIRYamlMapping.h

struct MachineFunction {
   MachineFrameInfo FrameInfo;
   std::vector<FixedMachineStackObject> FixedStackObjects;
   std::vector<MachineStackObject> StackObjects;
   std::vector<MachineConstantPoolValue> Constants; /// Constant pool.
   std::unique_ptr<MachineFunctionInfo> MachineFuncInfo;
   std::vector<CallSiteInfo> CallSitesInfo;
   std::vector<DebugValueSubstitution> DebugValueSubstitutions;
   MachineJumpTable JumpTableInfo;
+  std::vector<StringValue> MachineMetadataNodes;
   BlockStringValue Body;
 };
 
 template <> struct MappingTraits<MachineFunction> {
   static void mapping(IO &YamlIO, MachineFunction &MF) {
     YamlIO.mapRequired("name", MF.Name);
     YamlIO.mapOptional("alignment", MF.Alignment, None);
     YamlIO.mapOptional("exposesReturnsTwice", MF.ExposesReturnsTwice, false);
YamlIO.mapOptional("callSites", MF.CallSitesInfo,
                        std::vector<CallSiteInfo>());
     YamlIO.mapOptional("debugValueSubstitutions", MF.DebugValueSubstitutions,
                        std::vector<DebugValueSubstitution>());
     YamlIO.mapOptional("constants", MF.Constants,
                        std::vector<MachineConstantPoolValue>());
     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());
   }
 };
 
 } // end namespace yaml
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_MIRYAMLMAPPING_H

llvm/include/llvm/CodeGen/MachineModuleSlotTracker.h

+//===-- 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

llvm/include/llvm/IR/ModuleSlotTracker.h

 //===-- llvm/IR/ModuleSlotTracker.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_IR_MODULESLOTTRACKER_H
 #define LLVM_IR_MODULESLOTTRACKER_H
 
+#include <functional>
 #include <memory>
+#include <utility>
+#include <vector>
 
 namespace llvm {
 
 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.
 ///
 /// Wrapper around the \a SlotTracker used internally by \a AsmWriter.  This
 /// class allows callers to share the cost of incorporating the metadata in a
 /// module or a function.
 ///
 /// If the IR changes from underneath \a ModuleSlotTracker, strings like
 /// "<badref>" will be printed, or, worse, the wrong slots entirely.
 class ModuleSlotTracker {
   /// Storage for a slot tracker.
   std::unique_ptr<SlotTracker> MachineStorage;
   bool ShouldCreateStorage = false;
   bool ShouldInitializeAllMetadata = false;
 
   const Module *M = nullptr;
   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,
                     const Function *F = nullptr);
 
   /// Construct a slot tracker from a module.
   ///
   /// If \a M is \c nullptr, uses a null slot tracker.  Otherwise, initializes
   /// a slot tracker, and initializes all metadata slots.  \c
   /// ShouldInitializeAllMetadata defaults to true because this is expected to
   /// be shared between multiple callers, and otherwise MDNode references will
   /// not match up.
   explicit ModuleSlotTracker(const Module *M,
                              bool ShouldInitializeAllMetadata = true);
 
   /// Destructor to clean up storage.
-  ~ModuleSlotTracker();
+  virtual ~ModuleSlotTracker();
 
   /// Lazily creates a slot tracker.
   SlotTracker *getMachine();
 
   const Module *getModule() const { return M; }
   const Function *getCurrentFunction() const { return F; }
 
   /// Incorporate the given function.
   ///
   /// Purge the currently incorporated function and incorporate \c F.  If \c F
   /// is currently incorporated, this is a no-op.
   void incorporateFunction(const Function &F);
 
   /// Return the slot number of the specified local value.
   ///
   /// A function that defines this value should be incorporated prior to calling
   /// 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
 
 #endif

llvm/lib/CodeGen/CMakeLists.txt

add_llvm_component_library(LLVMCodeGen
   MachineFunctionSplitter.cpp
   MachineInstrBundle.cpp
   MachineInstr.cpp
   MachineLICM.cpp
   MachineLoopInfo.cpp
   MachineLoopUtils.cpp
   MachineModuleInfo.cpp
   MachineModuleInfoImpls.cpp
+  MachineModuleSlotTracker.cpp
   MachineOperand.cpp
   MachineOptimizationRemarkEmitter.cpp
   MachineOutliner.cpp
   MachinePassManager.cpp
   MachinePipeliner.cpp
   MachinePostDominators.cpp
   MachineRegionInfo.cpp
   MachineRegisterInfo.cpp

llvm/lib/CodeGen/MIRPrinter.cpp

 #include "llvm/CodeGen/MIRYamlMapping.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineConstantPool.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #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/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"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/IRPrintingPasses.h"
 #include "llvm/IR/InstrTypes.h"
void convert(yaml::MachineFunction &MF,
                const MachineConstantPool &ConstantPool);
   void convert(ModuleSlotTracker &MST, yaml::MachineJumpTable &YamlJTI,
                const MachineJumpTableInfo &JTI);
   void convertStackObjects(yaml::MachineFunction &YMF,
                            const MachineFunction &MF, ModuleSlotTracker &MST);
   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);
 };
 
 /// This class prints out the machine instructions using the MIR serialization
 /// format.
 class MIPrinter {
void MIRPrinter::print(const MachineFunction &MF) {
   YamlMF.RegBankSelected = MF.getProperties().hasProperty(
       MachineFunctionProperties::Property::RegBankSelected);
   YamlMF.Selected = MF.getProperties().hasProperty(
       MachineFunctionProperties::Property::Selected);
   YamlMF.FailedISel = MF.getProperties().hasProperty(
       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);
   convertCallSiteObjects(YamlMF, MF, MST);
   for (auto &Sub : MF.DebugValueSubstitutions)
     YamlMF.DebugValueSubstitutions.push_back({Sub.first.first, Sub.first.second,
                                               Sub.second.first,
                                               Sub.second.second});
void MIRPrinter::print(const MachineFunction &MF) {
   for (const auto &MBB : MF) {
     if (IsNewlineNeeded)
       StrOS << "\n";
     MIPrinter(StrOS, MST, RegisterMaskIds, StackObjectOperandMapping)
         .print(MBB);
     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);
   Out << YamlMF;
 }
 
 static void printCustomRegMask(const uint32_t *RegMask, raw_ostream &OS,
                                const TargetRegisterInfo *TRI) {
void MIRPrinter::convertCallSiteObjects(yaml::MachineFunction &YMF,
   llvm::sort(YMF.CallSitesInfo.begin(), YMF.CallSitesInfo.end(),
              [](yaml::CallSiteInfo A, yaml::CallSiteInfo B) {
                if (A.CallLocation.BlockNum == B.CallLocation.BlockNum)
                  return A.CallLocation.Offset < B.CallLocation.Offset;
                return A.CallLocation.BlockNum < B.CallLocation.BlockNum;
              });
 }
 
+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;
   for (const MachineConstantPoolEntry &Constant : ConstantPool.getConstants()) {
     std::string Str;
     raw_string_ostream StrOS(Str);
     if (Constant.isMachineConstantPoolEntry()) {
       Constant.Val.MachineCPVal->print(StrOS);

llvm/lib/CodeGen/MachineModuleSlotTracker.cpp

+//===-- 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;

llvm/lib/IR/AsmWriter.cpp

if (STy->getNumElements() == 0) {
     }
 
     OS << " }";
   }
   if (STy->isPacked())
     OS << '>';
 }
 
+AbstractSlotTrackerStorage::~AbstractSlotTrackerStorage() {}
+
 namespace llvm {
 
 //===----------------------------------------------------------------------===//
 // SlotTracker Class: Enumerate slot numbers for unnamed values
 //===----------------------------------------------------------------------===//
 /// 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>;
 
 private:
   /// TheModule - The module for which we are holding slot numbers.
   const Module* TheModule;
 
   /// TheFunction - The function for which we are holding slot numbers.
   const Function* TheFunction = nullptr;
   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;
 
   /// mMap - The slot map for the module level data.
   ValueMap mMap;
   unsigned mNext = 0;
 
   /// fMap - The slot map for the function level data.
explicit SlotTracker(const Function *F,
                        bool ShouldInitializeAllMetadata = false);
 
   /// Construct from a module summary index.
   explicit SlotTracker(const ModuleSummaryIndex *Index);
 
   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);
   int getTypeIdSlot(StringRef Id);
 
   /// If you'd like to deal with a function instead of just a module, use
   /// this method to get its data into the SlotTracker.
   void incorporateFunction(const Function *F) {

 SlotTracker *ModuleSlotTracker::getMachine() {
   if (!ShouldCreateStorage)
     return Machine;
 
   ShouldCreateStorage = false;
   MachineStorage =
       std::make_unique<SlotTracker>(M, ShouldInitializeAllMetadata);
   Machine = MachineStorage.get();
+  if (processModuleHookFn)
+    Machine->setProcessHook(processModuleHookFn);
+  if (processFunctionHookFn)
+    Machine->setProcessHook(processFunctionHookFn);
   return Machine;
 }
 
 void ModuleSlotTracker::incorporateFunction(const Function &F) {
   // Using getMachine() may lazily create the slot tracker.
   if (!getMachine())
     return;
 
   // Nothing to do if this is the right function already.
   if (this->F == &F)
     return;
   if (this->F)
     Machine->purgeFunction();
   Machine->incorporateFunction(&F);
   this->F = &F;
 }
 
 int ModuleSlotTracker::getLocalSlot(const Value *V) {
   assert(F && "No function incorporated");
   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());
 
   if (const Instruction *I = dyn_cast<Instruction>(V))
     if (I->getParent())
       return new SlotTracker(I->getParent()->getParent());
 
for (const Function &F : *TheModule) {
 
     // Add all the function attributes to the table.
     // FIXME: Add attributes of other objects?
     AttributeSet FnAttrs = F.getAttributes().getFnAttributes();
     if (FnAttrs.hasAttributes())
       CreateAttributeSetSlot(FnAttrs);
   }
 
+  if (processModuleHookFn)
+    processModuleHookFn(this, TheModule, ShouldInitializeAllMetadata);
+
   ST_DEBUG("end processModule!\n");
 }
 
 // Process the arguments, basic blocks, and instructions  of a function.
 void SlotTracker::processFunction() {
   ST_DEBUG("begin processFunction!\n");
   fNext = 0;
 
for (auto &I : BB) {
         // Add all the call attributes to the table.
         AttributeSet Attrs = Call->getAttributes().getFnAttributes();
         if (Attrs.hasAttributes())
           CreateAttributeSetSlot(Attrs);
       }
     }
   }
 
+  if (processFunctionHookFn)
+    processFunctionHookFn(this, TheFunction, ShouldInitializeAllMetadata);
+
   FunctionProcessed = true;
 
   ST_DEBUG("end processFunction!\n");
 }
 
 // Iterate through all the GUID in the index and create slots for them.
 int SlotTracker::processIndex() {
   ST_DEBUG("begin processIndex!\n");
int SlotTracker::getGlobalSlot(const GlobalValue *V) {
   // Check for uninitialized state and do lazy initialization.
   initializeIfNeeded();
 
   // Find the value in the module map
   ValueMap::iterator MI = mMap.find(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.
   initializeIfNeeded();
 
   // Find the MDNode in the module map
   mdn_iterator MI = mdnMap.find(N);
   return MI == mdnMap.end() ? -1 : (int)MI->second;

 
 void ModuleSummaryIndex::print(raw_ostream &ROS, bool IsForDebug) const {
   SlotTracker SlotTable(this);
   formatted_raw_ostream OS(ROS);
   AssemblyWriter W(OS, SlotTable, this, IsForDebug);
   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
 void Value::dump() const { print(dbgs(), /*IsForDebug=*/true); dbgs() << '\n'; }
 
 // Type::dump - allow easy printing of Types from the debugger.
 LLVM_DUMP_METHOD
 void Type::dump() const { print(dbgs(), /*IsForDebug=*/true); dbgs() << '\n'; }

llvm/unittests/CodeGen/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   ${LLVM_TARGETS_TO_BUILD}
   Analysis
   AsmParser
   AsmPrinter
   CodeGen
   Core
+  FileCheck
   MC
   MIRParser
   Passes
   SelectionDAG
   Support
   Target
   )
 
add_llvm_unittest(CodeGenTests
   MachineInstrTest.cpp
   MachineOperandTest.cpp
   PassManagerTest.cpp
   ScalableVectorMVTsTest.cpp
   SelectionDAGAddressAnalysisTest.cpp
   TypeTraitsTest.cpp
   TargetOptionsTest.cpp
   TestAsmPrinter.cpp
+  MachineMetadata.cpp
   )
 
 add_subdirectory(GlobalISel)
 
 target_link_libraries(CodeGenTests PRIVATE LLVMTestingSupport)

llvm/unittests/CodeGen/MachineMetadata.cpp

+//===- 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_TRUE(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_TRUE(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_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: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));
+}