Reapply [Assignment Tracking][5/*] Add core infrastructure for instruction reference

Previously reverted in 41f5a00. Fold in
D133576 previously reverted in d29d5ff.

---

The Assignment Tracking debug-info feature is outlined in this RFC:

https://discourse.llvm.org/t/
rfc-assignment-tracking-a-better-way-of-specifying-variable-locations-in-ir

Overview
It's possible to find intrinsics linked to an instruction by looking at the
MetadataAsValue uses of the attached DIAssignID. That covers instruction ->
intrinsic(s) lookup. Add a global DIAssignID -> instruction(s) map which gives
us the ability to perform intrinsic -> instruction(s) lookup. Add plumbing to
keep the map up to date through optimisations and add utility functions
including two that perform those lookups. Finally, add a unittest.

Details
In llvm/lib/IR/LLVMContextImpl.h add AssignmentIDToInstrs which maps DIAssignID
* attachments to Instruction *s. Because the DIAssignID * is the key we can't
use a TrackingMDNodeRef for it, and therefore cannot easily update the mapping
when a temporary DIAssignID is replaced.

Temporary DIAssignID's are only used in IR parsing to deal with metadata
forward references. Update llvm/lib/AsmParser/LLParser.cpp to avoid using
temporary DIAssignID's for attachments.

In llvm/lib/IR/Metadata.cpp add Instruction::updateDIAssignIDMapping which is
called to remove or add an entry (or both) to AssignmentIDToInstrs. Call this
from Instruction::setMetadata and add a call to setMetadata in Intruction's
dtor that explicitly unsets the DIAssignID so that the mappging gets updated.

In llvm/lib/IR/DebugInfo.cpp and DebugInfo.h add utility functions:

    getAssignmentInsts(const DbgAssignIntrinsic *DAI)
    getAssignmentMarkers(const Instruction *Inst)
    RAUW(DIAssignID *Old, DIAssignID *New)
    deleteAll(Function *F)
    deleteAssignmentMarkers(const Instruction *Inst)

These core utils are tested in llvm/unittests/IR/DebugInfoTest.cpp.

Reviewed By: jmorse

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

Author: OCHyams

llvm/include/llvm/AsmParser/LLParser.h

@@ -108,6 +108,12 @@ namespace llvm {
 
     SmallVector<Instruction*, 64> InstsWithTBAATag;
 
+    /// DIAssignID metadata does not support temporary RAUW so we cannot use
+    /// the normal metadata forward reference resolution method. Instead,
+    /// non-temporary DIAssignID are attached to instructions (recorded here)
+    /// then replaced later.
+    DenseMap<MDNode *, SmallVector<Instruction *, 2>> TempDIAssignIDAttachments;
+
     // Type resolution handling data structures.  The location is set when we
     // have processed a use of the type but not a definition yet.
     StringMap<std::pair<Type*, LocTy> > NamedTypes;

llvm/include/llvm/IR/DebugInfo.h

@@ -21,7 +21,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/ADT/iterator_range.h"
-#include "llvm/IR/DebugInfoMetadata.h"
+#include "llvm/IR/IntrinsicInst.h"
 
 namespace llvm {
 
@@ -159,6 +159,70 @@ class DebugInfoFinder {
   SmallPtrSet<const MDNode *, 32> NodesSeen;
 };
 
+/// Assignment Tracking (at).
+namespace at {
+//
+// Utilities for enumerating storing instructions from an assignment ID.
+//
+/// A range of instructions.
+using AssignmentInstRange =
+    iterator_range<SmallVectorImpl<Instruction *>::iterator>;
+/// Return a range of instructions (typically just one) that have \p ID
+/// as an attachment.
+/// Iterators invalidated by adding or removing DIAssignID metadata to/from any
+/// instruction (including by deleting or cloning instructions).
+AssignmentInstRange getAssignmentInsts(DIAssignID *ID);
+/// Return a range of instructions (typically just one) that perform the
+/// assignment that \p DAI encodes.
+/// Iterators invalidated by adding or removing DIAssignID metadata to/from any
+/// instruction (including by deleting or cloning instructions).
+inline AssignmentInstRange getAssignmentInsts(const DbgAssignIntrinsic *DAI) {
+  return getAssignmentInsts(DAI->getAssignID());
+}
+
+//
+// Utilities for enumerating llvm.dbg.assign intrinsic from an assignment ID.
+//
+/// High level: this is an iterator for llvm.dbg.assign intrinsics.
+/// Implementation details: this is a wrapper around Value's User iterator that
+/// dereferences to a DbgAssignIntrinsic ptr rather than a User ptr.
+class DbgAssignIt
+    : public iterator_adaptor_base<DbgAssignIt, Value::user_iterator,
+                                   typename std::iterator_traits<
+                                       Value::user_iterator>::iterator_category,
+                                   DbgAssignIntrinsic *, std::ptrdiff_t,
+                                   DbgAssignIntrinsic **,
+                                   DbgAssignIntrinsic *&> {
+public:
+  DbgAssignIt(Value::user_iterator It) : iterator_adaptor_base(It) {}
+  DbgAssignIntrinsic *operator*() const { return cast<DbgAssignIntrinsic>(*I); }
+};
+/// A range of llvm.dbg.assign intrinsics.
+using AssignmentMarkerRange = iterator_range<DbgAssignIt>;
+/// Return a range of dbg.assign intrinsics which use \ID as an operand.
+/// Iterators invalidated by deleting an intrinsic contained in this range.
+AssignmentMarkerRange getAssignmentMarkers(DIAssignID *ID);
+/// Return a range of dbg.assign intrinsics for which \p Inst performs the
+/// assignment they encode.
+/// Iterators invalidated by deleting an intrinsic contained in this range.
+inline AssignmentMarkerRange getAssignmentMarkers(const Instruction *Inst) {
+  if (auto *ID = Inst->getMetadata(LLVMContext::MD_DIAssignID))
+    return getAssignmentMarkers(cast<DIAssignID>(ID));
+  else
+    return make_range(Value::user_iterator(), Value::user_iterator());
+}
+
+/// Delete the llvm.dbg.assign intrinsics linked to \p Inst.
+void deleteAssignmentMarkers(const Instruction *Inst);
+
+/// Replace all uses (and attachments) of \p Old with \p New.
+void RAUW(DIAssignID *Old, DIAssignID *New);
+
+/// Remove all Assignment Tracking related intrinsics and metadata from \p F.
+void deleteAll(Function *F);
+
+} // end namespace at
+
 /// Return true if assignment tracking is enabled.
 bool getEnableAssignmentTracking();
 } // end namespace llvm

llvm/include/llvm/IR/Instruction.h

@@ -515,6 +515,10 @@ class Instruction : public User,
   void
   getAllMetadataImpl(SmallVectorImpl<std::pair<unsigned, MDNode *>> &) const;
 
+  /// Update the LLVMContext ID-to-Instruction(s) mapping. If \p ID is nullptr
+  /// then clear the mapping for this instruction.
+  void updateDIAssignIDMapping(DIAssignID *ID);
+
 public:
   //===--------------------------------------------------------------------===//
   // Predicates and helper methods.

llvm/include/llvm/module.modulemap

@@ -286,6 +286,7 @@ module LLVM_intrinsic_gen {
   module IR_PatternMatch { header "IR/PatternMatch.h" export * }
   module IR_SafepointIRVerifier { header "IR/SafepointIRVerifier.h" export * }
   module IR_Statepoint { header "IR/Statepoint.h" export * }
+  module IR_DebugInfo { header "IR/DebugInfo.h" export * }
 
   export *
 }

llvm/lib/AsmParser/LLParser.cpp

@@ -853,7 +853,18 @@ bool LLParser::parseStandaloneMetadata() {
   // See if this was forward referenced, if so, handle it.
   auto FI = ForwardRefMDNodes.find(MetadataID);
   if (FI != ForwardRefMDNodes.end()) {
-    FI->second.first->replaceAllUsesWith(Init);
+    auto *ToReplace = FI->second.first.get();
+    // DIAssignID has its own special forward-reference "replacement" for
+    // attachments (the temporary attachments are never actually attached).
+    if (isa<DIAssignID>(Init)) {
+      for (auto *Inst : TempDIAssignIDAttachments[ToReplace]) {
+        assert(!Inst->getMetadata(LLVMContext::MD_DIAssignID) &&
+               "Inst unexpectedly already has DIAssignID attachment");
+        Inst->setMetadata(LLVMContext::MD_DIAssignID, Init);
+      }
+    }
+
+    ToReplace->replaceAllUsesWith(Init);
     ForwardRefMDNodes.erase(FI);
 
     assert(NumberedMetadata[MetadataID] == Init && "Tracking VH didn't work");
@@ -2082,7 +2093,11 @@ bool LLParser::parseInstructionMetadata(Instruction &Inst) {
     if (parseMetadataAttachment(MDK, N))
       return true;
 
-    Inst.setMetadata(MDK, N);
+    if (MDK == LLVMContext::MD_DIAssignID)
+      TempDIAssignIDAttachments[N].push_back(&Inst);
+    else
+      Inst.setMetadata(MDK, N);
+
     if (MDK == LLVMContext::MD_tbaa)
       InstsWithTBAATag.push_back(&Inst);
 

llvm/lib/IR/DebugInfo.cpp

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm-c/DebugInfo.h"
+#include "LLVMContextImpl.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
@@ -37,6 +38,7 @@
 #include <utility>
 
 using namespace llvm;
+using namespace llvm::at;
 using namespace llvm::dwarf;
 
 static cl::opt<bool>
@@ -1632,3 +1634,70 @@ LLVMMetadataKind LLVMGetMetadataKind(LLVMMetadataRef Metadata) {
     return (LLVMMetadataKind)LLVMGenericDINodeMetadataKind;
   }
 }
+
+AssignmentInstRange at::getAssignmentInsts(DIAssignID *ID) {
+  assert(ID && "Expected non-null ID");
+  LLVMContext &Ctx = ID->getContext();
+  auto &Map = Ctx.pImpl->AssignmentIDToInstrs;
+
+  auto MapIt = Map.find(ID);
+  if (MapIt == Map.end())
+    return make_range(nullptr, nullptr);
+
+  return make_range(MapIt->second.begin(), MapIt->second.end());
+}
+
+AssignmentMarkerRange at::getAssignmentMarkers(DIAssignID *ID) {
+  assert(ID && "Expected non-null ID");
+  LLVMContext &Ctx = ID->getContext();
+
+  auto *IDAsValue = MetadataAsValue::getIfExists(Ctx, ID);
+
+  // The ID is only used wrapped in MetadataAsValue(ID), so lets check that
+  // one of those already exists first.
+  if (!IDAsValue)
+    return make_range(Value::user_iterator(), Value::user_iterator());
+
+  return make_range(IDAsValue->user_begin(), IDAsValue->user_end());
+}
+
+void at::deleteAssignmentMarkers(const Instruction *Inst) {
+  auto Range = getAssignmentMarkers(Inst);
+  if (Range.empty())
+    return;
+  SmallVector<DbgAssignIntrinsic *> ToDelete(Range.begin(), Range.end());
+  for (auto *DAI : ToDelete)
+    DAI->eraseFromParent();
+}
+
+void at::RAUW(DIAssignID *Old, DIAssignID *New) {
+  // Replace MetadataAsValue uses.
+  if (auto *OldIDAsValue =
+          MetadataAsValue::getIfExists(Old->getContext(), Old)) {
+    auto *NewIDAsValue = MetadataAsValue::get(Old->getContext(), New);
+    OldIDAsValue->replaceAllUsesWith(NewIDAsValue);
+  }
+
+  // Replace attachments.
+  AssignmentInstRange InstRange = getAssignmentInsts(Old);
+  // Use intermediate storage for the instruction ptrs because the
+  // getAssignmentInsts range iterators will be invalidated by adding and
+  // removing DIAssignID attachments.
+  SmallVector<Instruction *> InstVec(InstRange.begin(), InstRange.end());
+  for (auto *I : InstVec)
+    I->setMetadata(LLVMContext::MD_DIAssignID, New);
+}
+
+void at::deleteAll(Function *F) {
+  SmallVector<DbgAssignIntrinsic *, 12> ToDelete;
+  for (BasicBlock &BB : *F) {
+    for (Instruction &I : BB) {
+      if (auto *DAI = dyn_cast<DbgAssignIntrinsic>(&I))
+        ToDelete.push_back(DAI);
+      else
+        I.setMetadata(LLVMContext::MD_DIAssignID, nullptr);
+    }
+  }
+  for (auto *DAI : ToDelete)
+    DAI->eraseFromParent();
+}

llvm/lib/IR/Instruction.cpp

@@ -55,6 +55,10 @@ Instruction::~Instruction() {
   //   instructions in a BasicBlock are deleted).
   if (isUsedByMetadata())
     ValueAsMetadata::handleRAUW(this, UndefValue::get(getType()));
+
+  // Explicitly remove DIAssignID metadata to clear up ID -> Instruction(s)
+  // mapping in LLVMContext.
+  setMetadata(LLVMContext::MD_DIAssignID, nullptr);
 }
 
 

llvm/lib/IR/LLVMContextImpl.h

@@ -1499,6 +1499,11 @@ class LLVMContextImpl {
   /// Collection of metadata used in this context.
   DenseMap<const Value *, MDAttachments> ValueMetadata;
 
+  /// Map DIAssignID -> Instructions with that attachment.
+  /// Managed by Instruction via Instruction::updateDIAssignIDMapping.
+  /// Query using the at:: functions defined in DebugInfo.h.
+  DenseMap<DIAssignID *, SmallVector<Instruction *, 1>> AssignmentIDToInstrs;
+
   /// Collection of per-GlobalObject sections used in this context.
   DenseMap<const GlobalObject *, StringRef> GlobalObjectSections;
 

llvm/lib/IR/Metadata.cpp

@@ -1425,6 +1425,37 @@ void Instruction::dropUnknownNonDebugMetadata(ArrayRef<unsigned> KnownIDs) {
   }
 }
 
+void Instruction::updateDIAssignIDMapping(DIAssignID *ID) {
+  auto &IDToInstrs = getContext().pImpl->AssignmentIDToInstrs;
+  if (const DIAssignID *CurrentID =
+          cast_or_null<DIAssignID>(getMetadata(LLVMContext::MD_DIAssignID))) {
+    // Nothing to do if the ID isn't changing.
+    if (ID == CurrentID)
+      return;
+
+    // Unmap this instruction from its current ID.
+    auto InstrsIt = IDToInstrs.find(CurrentID);
+    assert(InstrsIt != IDToInstrs.end() &&
+           "Expect existing attachment to be mapped");
+
+    auto &InstVec = InstrsIt->second;
+    auto *InstIt = std::find(InstVec.begin(), InstVec.end(), this);
+    assert(InstIt != InstVec.end() &&
+           "Expect instruction to be mapped to attachment");
+    // The vector contains a ptr to this. If this is the only element in the
+    // vector, remove the ID:vector entry, otherwise just remove the
+    // instruction from the vector.
+    if (InstVec.size() == 1)
+      IDToInstrs.erase(InstrsIt);
+    else
+      InstVec.erase(InstIt);
+  }
+
+  // Map this instruction to the new ID.
+  if (ID)
+    IDToInstrs[ID].push_back(this);
+}
+
 void Instruction::setMetadata(unsigned KindID, MDNode *Node) {
   if (!Node && !hasMetadata())
     return;
@@ -1435,6 +1466,16 @@ void Instruction::setMetadata(unsigned KindID, MDNode *Node) {
     return;
   }
 
+  // Update DIAssignID to Instruction(s) mapping.
+  if (KindID == LLVMContext::MD_DIAssignID) {
+    // The DIAssignID tracking infrastructure doesn't support RAUWing temporary
+    // nodes with DIAssignIDs. The cast_or_null below would also catch this, but
+    // having a dedicated assert helps make this obvious.
+    assert((!Node || !Node->isTemporary()) &&
+           "Temporary DIAssignIDs are invalid");
+    updateDIAssignIDMapping(cast_or_null<DIAssignID>(Node));
+  }
+
   Value::setMetadata(KindID, Node);
 }
 

llvm/lib/IR/Verifier.cpp

@@ -68,6 +68,7 @@
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/DerivedTypes.h"
@@ -4557,6 +4558,10 @@ void Verifier::visitDIAssignIDMetadata(Instruction &I, MDNode *MD) {
       CheckDI(isa<DbgAssignIntrinsic>(User),
               "!DIAssignID should only be used by llvm.dbg.assign intrinsics",
               MD, User);
+      // All of the dbg.assign intrinsics should be in the same function as I.
+      if (auto *DAI = dyn_cast<DbgAssignIntrinsic>(User))
+        CheckDI(DAI->getFunction() == I.getFunction(),
+                "dbg.assign not in same function as inst", DAI, &I);
     }
   }
 }
@@ -6017,6 +6022,10 @@ void Verifier::visitDbgIntrinsic(StringRef Kind, DbgVariableIntrinsic &DII) {
     CheckDI(isa<DIExpression>(DAI->getRawAddressExpression()),
             "invalid llvm.dbg.assign intrinsic address expression", &DII,
             DAI->getRawAddressExpression());
+    // All of the linked instructions should be in the same function as DII.
+    for (Instruction *I : at::getAssignmentInsts(DAI))
+      CheckDI(DAI->getFunction() == I->getFunction(),
+              "inst not in same function as dbg.assign", I, DAI);
   }
 
   // Ignore broken !dbg attachments; they're checked elsewhere.