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

[Orlando]

Summary

Add Instruction::mergeDIAssignID(ArrayRef<const Instruction* > SourceInstructions) which merges the DIAssignID metadata attachments on the 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.

[jmorse]

LGTM

llvm/include/llvm/IR/Instruction.h

515

Worth adding the rider that if this instruction has no DIAssignID, then it is given one, no? (I think that's what the 3rd unit test is doing?)

llvm/lib/IR/DebugInfo.cpp

848

"Uninserted instruction merged"?

866

(Style guide says ++It)

[Orlando]

+ Address review comments

[jmorse]

LGTM

llvm/include/llvm/IR/Instruction.h

public:
   /// Drop the instruction's debug location. This does not guarantee removal
   /// of the !dbg source location attachment, as it must set a line 0 location
   /// with scope information attached on call instructions. To guarantee
   /// removal of the !dbg attachment, use the \ref setDebugLoc() API.
   /// Note: it is undefined behavior to call this on an instruction not
   /// 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;
   MDNode *getMetadataImpl(StringRef Kind) const;
   void
   getAllMetadataImpl(SmallVectorImpl<std::pair<unsigned, MDNode *>> &) const;
 
   /// Update the LLVMContext ID-to-Instruction(s) mapping. If \p ID is nullptr

llvm/lib/IR/DebugInfo.cpp

unsigned llvm::getDebugMetadataVersionFromModule(const Module &M) {
   return 0;
 }
 
 void Instruction::applyMergedLocation(const DILocation *LocA,
                                       const DILocation *LocB) {
   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() {
   const DebugLoc &DL = getDebugLoc();
   if (!DL)
     return;
 
   // If this isn't a call, drop the location to allow a location from a

llvm/lib/Transforms/Utils/Local.cpp

for (const auto &MD : Metadata) {
     MDNode *KMD = MD.second;
 
     switch (Kind) {
       default:
         K->setMetadata(Kind, nullptr); // Remove unknown metadata
         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;
       case LLVMContext::MD_alias_scope:
         K->setMetadata(Kind, MDNode::getMostGenericAliasScope(JMD, KMD));
         break;
       case LLVMContext::MD_noalias:
       case LLVMContext::MD_mem_parallel_loop_access:

llvm/unittests/IR/DebugInfoTest.cpp

TEST(AssignmentTrackingTest, Utils) {
   EXPECT_EQ(Fun1.begin()->size(), 2u);
   // fun2's alloca should have the same DIAssignID and remain linked to its
   // llvm.dbg.assign.
   EXPECT_EQ(Fun2ID, cast_or_null<DIAssignID>(
                         Fun2Alloca.getMetadata(LLVMContext::MD_DIAssignID)));
   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