[Assignment Tracking][8/*] Add DIAssignID merging utilities

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

Add method:

  Instruction::mergeDIAssignID(
      ArrayRef<const Instruction* > SourceInstructions)

which merges the DIAssignID metadata attachments on `SourceInstructions` and
`this` and replaces uses of the original IDs with the new shared one.

This is used when stores are merged, for example sinking stores out of a
if-diamond CFG or vectorizing contiguous stores.

Reviewed By: jmorse

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

diff --git a/llvm/include/llvm/IR/Instruction.h b/llvm/include/llvm/IR/Instruction.h
index 131a741..9fa8a0a 100644 (file)
--- a/llvm/include/llvm/IR/Instruction.h
+++ b/llvm/include/llvm/IR/Instruction.h
@@ -508,6 +508,17 @@ public:
   /// currently inserted into a function.
   void dropLocation();
 
+  /// Merge the DIAssignID metadata from this instruction and those attached to
+  /// instructions in \p SourceInstructions. This process performs a RAUW on
+  /// the MetadataAsValue uses of the merged DIAssignID nodes. Not every
+  /// instruction in \p SourceInstructions needs to have DIAssignID
+  /// metadata. If none of them do then nothing happens. If this instruction
+  /// does not have a DIAssignID attachment but at least one in \p
+  /// SourceInstructions does then the merged one will be attached to
+  /// it. However, instructions without attachments in \p SourceInstructions
+  /// are not modified.
+  void mergeDIAssignID(ArrayRef<const Instruction *> SourceInstructions);
+
 private:
   // These are all implemented in Metadata.cpp.
   MDNode *getMetadataImpl(unsigned KindID) const;

diff --git a/llvm/lib/IR/DebugInfo.cpp b/llvm/lib/IR/DebugInfo.cpp
index 74fab1f..b30b51e 100644 (file)
--- a/llvm/lib/IR/DebugInfo.cpp
+++ b/llvm/lib/IR/DebugInfo.cpp
@@ -840,6 +840,36 @@ void Instruction::applyMergedLocation(const DILocation *LocA,
   setDebugLoc(DILocation::getMergedLocation(LocA, LocB));
 }
 
+void Instruction::mergeDIAssignID(
+    ArrayRef<const Instruction *> SourceInstructions) {
+  // Replace all uses (and attachments) of all the DIAssignIDs
+  // on SourceInstructions with a single merged value.
+  Function *Fn = getFunction();
+  assert(Fn && "Uninserted instruction merged");
+  // Collect up the DIAssignID tags.
+  SmallVector<DIAssignID *, 4> IDs;
+  for (const Instruction *I : SourceInstructions) {
+    if (auto *MD = I->getMetadata(LLVMContext::MD_DIAssignID))
+      IDs.push_back(cast<DIAssignID>(MD));
+    assert(Fn == I->getFunction() &&
+           "Merging with instruction from another function not allowed");
+  }
+
+  // Add this instruction's DIAssignID too, if it has one.
+  if (auto *MD = getMetadata(LLVMContext::MD_DIAssignID))
+    IDs.push_back(cast<DIAssignID>(MD));
+
+  if (IDs.empty())
+    return; // No DIAssignID tags to process.
+
+  DIAssignID *MergeID = IDs[0];
+  for (auto It = std::next(IDs.begin()), End = IDs.end(); It != End; ++It) {
+    if (*It != MergeID)
+      at::RAUW(*It, MergeID);
+  }
+  setMetadata(LLVMContext::MD_DIAssignID, MergeID);
+}
+
 void Instruction::updateLocationAfterHoist() { dropLocation(); }
 
 void Instruction::dropLocation() {

diff --git a/llvm/lib/Transforms/Utils/Local.cpp b/llvm/lib/Transforms/Utils/Local.cpp
index e31e691..f3c2371 100644 (file)
--- a/llvm/lib/Transforms/Utils/Local.cpp
+++ b/llvm/lib/Transforms/Utils/Local.cpp
@@ -2583,6 +2583,9 @@ void llvm::combineMetadata(Instruction *K, const Instruction *J,
         break;
       case LLVMContext::MD_dbg:
         llvm_unreachable("getAllMetadataOtherThanDebugLoc returned a MD_dbg");
+      case LLVMContext::MD_DIAssignID:
+        K->mergeDIAssignID(J);
+        break;
       case LLVMContext::MD_tbaa:
         K->setMetadata(Kind, MDNode::getMostGenericTBAA(JMD, KMD));
         break;

diff --git a/llvm/unittests/IR/DebugInfoTest.cpp b/llvm/unittests/IR/DebugInfoTest.cpp
index e58b4f5..970de06 100644 (file)
--- a/llvm/unittests/IR/DebugInfoTest.cpp
+++ b/llvm/unittests/IR/DebugInfoTest.cpp
@@ -514,4 +514,157 @@ TEST(AssignmentTrackingTest, Utils) {
   EXPECT_FALSE(at::getAssignmentMarkers(&Fun2Alloca).empty());
 }
 
+TEST(AssignmentTrackingTest, InstrMethods) {
+  // Test the assignment tracking Instruction methods.
+  // This includes:
+  //     Instruction::mergeDIAssignID
+
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"(
+    define dso_local void @fun() #0 !dbg !8 {
+    entry:
+      %Local = alloca [2 x i32], align 4, !DIAssignID !12
+      call void @llvm.dbg.assign(metadata i1 undef, metadata !13, metadata !DIExpression(), metadata !12, metadata [2 x i32]* %Local, metadata !DIExpression()), !dbg !18
+      %arrayidx = getelementptr inbounds [2 x i32], [2 x i32]* %Local, i64 0, i64 0, !dbg !19
+      store i32 5, i32* %arrayidx, align 4, !dbg !20, !DIAssignID !21
+      call void @llvm.dbg.assign(metadata i32 5, metadata !13, metadata !DIExpression(DW_OP_LLVM_fragment, 0, 32), metadata !21, metadata i32* %arrayidx, metadata !DIExpression()), !dbg !18
+      %arrayidx1 = getelementptr inbounds [2 x i32], [2 x i32]* %Local, i64 0, i64 1, !dbg !22
+      store i32 6, i32* %arrayidx1, align 4, !dbg !23, !DIAssignID !24
+      call void @llvm.dbg.assign(metadata i32 6, metadata !13, metadata !DIExpression(DW_OP_LLVM_fragment, 32, 32), metadata !24, metadata i32* %arrayidx1, metadata !DIExpression()), !dbg !18
+      ret void, !dbg !25
+    }
+
+    declare void @llvm.dbg.assign(metadata, metadata, metadata, metadata, metadata, metadata) #1
+
+    !llvm.dbg.cu = !{!0}
+    !llvm.module.flags = !{!2, !3, !4, !5, !6}
+    !llvm.ident = !{!7}
+
+    !0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1, producer: "clang version 14.0.0", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, splitDebugInlining: false, nameTableKind: None)
+    !1 = !DIFile(filename: "test.cpp", directory: "/")
+    !2 = !{i32 7, !"Dwarf Version", i32 5}
+    !3 = !{i32 2, !"Debug Info Version", i32 3}
+    !4 = !{i32 1, !"wchar_size", i32 4}
+    !5 = !{i32 7, !"uwtable", i32 1}
+    !6 = !{i32 7, !"frame-pointer", i32 2}
+    !7 = !{!"clang version 14.0.0"}
+    !8 = distinct !DISubprogram(name: "fun", linkageName: "fun", scope: !1, file: !1, line: 1, type: !9, scopeLine: 1, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !11)
+    !9 = !DISubroutineType(types: !10)
+    !10 = !{null}
+    !11 = !{}
+    !12 = distinct !DIAssignID()
+    !13 = !DILocalVariable(name: "Local", scope: !8, file: !1, line: 2, type: !14)
+    !14 = !DICompositeType(tag: DW_TAG_array_type, baseType: !15, size: 64, elements: !16)
+    !15 = !DIBasicType(name: "int", size: 32, encoding: DW_ATE_signed)
+    !16 = !{!17}
+    !17 = !DISubrange(count: 2)
+    !18 = !DILocation(line: 0, scope: !8)
+    !19 = !DILocation(line: 3, column: 3, scope: !8)
+    !20 = !DILocation(line: 3, column: 12, scope: !8)
+    !21 = distinct !DIAssignID()
+    !22 = !DILocation(line: 4, column: 3, scope: !8)
+    !23 = !DILocation(line: 4, column: 12, scope: !8)
+    !24 = distinct !DIAssignID()
+    !25 = !DILocation(line: 5, column: 1, scope: !8)
+  )");
+
+  // Check the test IR isn't malformed.
+  ASSERT_TRUE(M);
+  Function &Fun = *M->getFunction("fun");
+  SmallVector<Instruction *> Stores;
+  for (auto &BB : Fun) {
+    for (auto &I : BB) {
+      if (isa<StoreInst>(&I))
+        Stores.push_back(&I);
+    }
+  }
+
+  // The test requires (at least) 2 stores.
+  ASSERT_TRUE(Stores.size() == 2);
+  // Use SetVectors to check that the attachments and markers are unique
+  // (another test requirement).
+  SetVector<Metadata *> OrigIDs;
+  SetVector<DbgAssignIntrinsic *> Markers;
+  for (const Instruction *SI : Stores) {
+    Metadata *ID = SI->getMetadata(LLVMContext::MD_DIAssignID);
+    ASSERT_TRUE(OrigIDs.insert(ID));
+    ASSERT_TRUE(ID != nullptr);
+    auto Range = at::getAssignmentMarkers(SI);
+    ASSERT_TRUE(std::distance(Range.begin(), Range.end()) == 1);
+    ASSERT_TRUE(Markers.insert(*Range.begin()));
+  }
+
+  // Test 1 - mergeDIAssignID.
+  //
+  // Input            store0->mergeDIAssignID(store1)
+  // -----            -------------------------
+  // store0 !x        store0 !x
+  // dbg.assign0 !x   dbg.assign !x
+  // store1 !y        store1 !x
+  // dbg.assign1 !y   dbg.assign1 !x
+  {
+    Stores[0]->mergeDIAssignID(Stores[1]);
+    // Check that the stores share the same ID.
+    Metadata *NewID0 = Stores[0]->getMetadata(LLVMContext::MD_DIAssignID);
+    Metadata *NewID1 = Stores[1]->getMetadata(LLVMContext::MD_DIAssignID);
+    EXPECT_NE(NewID0, nullptr);
+    EXPECT_EQ(NewID0, NewID1);
+    EXPECT_EQ(Markers[0]->getAssignID(), NewID0);
+    EXPECT_EQ(Markers[1]->getAssignID(), NewID0);
+  }
+
+  // Test 2 - mergeDIAssignID.
+  //
+  // Input            store0->mergeDIAssignID(store1)
+  // -----            -------------------------
+  // store0 !x        store0 !x
+  // dbg.assign0 !x   dbg.assign !x
+  // store1           store1
+  {
+    Stores[1]->setMetadata(LLVMContext::MD_DIAssignID, nullptr);
+    Stores[0]->mergeDIAssignID(Stores[1]);
+    // Check that store1 doesn't get a new ID.
+    Metadata *NewID0 = Stores[0]->getMetadata(LLVMContext::MD_DIAssignID);
+    Metadata *NewID1 = Stores[1]->getMetadata(LLVMContext::MD_DIAssignID);
+    EXPECT_NE(NewID0, nullptr);
+    EXPECT_EQ(NewID1, nullptr);
+    EXPECT_EQ(Markers[0]->getAssignID(), NewID0);
+  }
+
+  // Test 3 - mergeDIAssignID.
+  //
+  // Input            store1->mergeDIAssignID(store0)
+  // -----            -------------------------
+  // store0 !x        store0 !x
+  // dbg.assign0 !x   dbg.assign !x
+  // store1           store1 !x
+  {
+    Stores[1]->setMetadata(LLVMContext::MD_DIAssignID, nullptr);
+    Stores[1]->mergeDIAssignID(Stores[0]);
+    // Check that the stores share the same ID (note store1 starts with none).
+    Metadata *NewID0 = Stores[0]->getMetadata(LLVMContext::MD_DIAssignID);
+    Metadata *NewID1 = Stores[1]->getMetadata(LLVMContext::MD_DIAssignID);
+    EXPECT_NE(NewID0, nullptr);
+    EXPECT_EQ(NewID0, NewID1);
+    EXPECT_EQ(Markers[0]->getAssignID(), NewID0);
+  }
+
+  // Test 4 - mergeDIAssignID.
+  //
+  // Input            store1->mergeDIAssignID(store0)
+  // -----            -------------------------
+  // store0 !x        store0 !x
+  // dbg.assign0 !x   dbg.assign !x
+  // store1 !x        store1 !x
+  {
+    Stores[0]->mergeDIAssignID(Stores[1]);
+    // Check that the stores share the same ID.
+    Metadata *NewID0 = Stores[0]->getMetadata(LLVMContext::MD_DIAssignID);
+    Metadata *NewID1 = Stores[1]->getMetadata(LLVMContext::MD_DIAssignID);
+    EXPECT_NE(NewID0, nullptr);
+    EXPECT_EQ(NewID0, NewID1);
+    EXPECT_EQ(Markers[0]->getAssignID(), NewID0);
+  }
+}
+
 } // end namespace