[DebugInfo] Remap extracted DIAssignIDs in hotcoldsplit

[OCHyams]

Fixes issue #91814

When instructions are extracted into a new function the DIAssignID metadata
uses and attachments need to be remapped so that the stores and assignment
markers don't link to stores and assignment markers in the original function.

This matches inlining behaviour for DIAssignIDs.

[llvmbot]

@llvm/pr-subscribers-debuginfo

@llvm/pr-subscribers-llvm-transforms

Author: Orlando Cazalet-Hyams (OCHyams)

Changes

Fixes issue #91814

When instructions are extracted into a new function the DIAssignID metadata
uses and attachments need to be remapped so that the stores and assignment
markers don't link to stores and assignment markers in the original function.

This matches inlining behaviour for DIAssignIDs.

---

Full diff: https://github.com/llvm/llvm-project/pull/91940.diff

5 Files Affected:

• (modified) llvm/include/llvm/IR/DebugInfo.h (+4)
• (modified) llvm/lib/IR/DebugInfo.cpp (+24)
• (modified) llvm/lib/Transforms/Utils/CodeExtractor.cpp (+4-2)
• (modified) llvm/lib/Transforms/Utils/InlineFunction.cpp (+2-19)
• (added) llvm/test/DebugInfo/assignment-tracking/X86/hotcoldsplit.ll (+50)

diff --git a/llvm/include/llvm/IR/DebugInfo.h b/llvm/include/llvm/IR/DebugInfo.h
index 53cede5409e26..5b80218d6c5cc 100644
--- a/llvm/include/llvm/IR/DebugInfo.h
+++ b/llvm/include/llvm/IR/DebugInfo.h
@@ -268,6 +268,10 @@ bool calculateFragmentIntersect(
     uint64_t SliceSizeInBits, const DbgVariableRecord *DVRAssign,
     std::optional<DIExpression::FragmentInfo> &Result);
 
+/// Replace DIAssignID uses and attachments with IDs from \p Map.
+/// If an ID is unmapped a new ID is generated and added to \p Map.
+void remapAssignID(DenseMap<DIAssignID *, DIAssignID *> &Map, Instruction &I);
+
 /// Helper struct for trackAssignments, below. We don't use the similar
 /// DebugVariable class because trackAssignments doesn't (yet?) understand
 /// partial variables (fragment info) as input and want to make that clear and
diff --git a/llvm/lib/IR/DebugInfo.cpp b/llvm/lib/IR/DebugInfo.cpp
index 7976904b1fe9d..4c3f37ceaaa46 100644
--- a/llvm/lib/IR/DebugInfo.cpp
+++ b/llvm/lib/IR/DebugInfo.cpp
@@ -2130,6 +2130,30 @@ bool at::calculateFragmentIntersect(
                                         SliceSizeInBits, DVRAssign, Result);
 }
 
+/// Update inlined instructions' DIAssignID metadata. We need to do this
+/// otherwise a function inlined more than once into the same function
+/// will cause DIAssignID to be shared by many instructions.
+void at::remapAssignID(DenseMap<DIAssignID *, DIAssignID *> &Map,
+                       Instruction &I) {
+  auto GetNewID = [&Map](Metadata *Old) {
+    DIAssignID *OldID = cast<DIAssignID>(Old);
+    if (DIAssignID *NewID = Map.lookup(OldID))
+      return NewID;
+    DIAssignID *NewID = DIAssignID::getDistinct(OldID->getContext());
+    Map[OldID] = NewID;
+    return NewID;
+  };
+  // If we find a DIAssignID attachment or use, replace it with a new version.
+  for (DbgVariableRecord &DVR : filterDbgVars(I.getDbgRecordRange())) {
+    if (DVR.isDbgAssign())
+      DVR.setAssignId(GetNewID(DVR.getAssignID()));
+  }
+  if (auto *ID = I.getMetadata(LLVMContext::MD_DIAssignID))
+    I.setMetadata(LLVMContext::MD_DIAssignID, GetNewID(ID));
+  else if (auto *DAI = dyn_cast<DbgAssignIntrinsic>(&I))
+    DAI->setAssignId(GetNewID(DAI->getAssignID()));
+}
+
 /// Collect constant properies (base, size, offset) of \p StoreDest.
 /// Return std::nullopt if any properties are not constants or the
 /// offset from the base pointer is negative.
diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index 6988292ac7156..f2672b8e9118f 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -1678,8 +1678,9 @@ static void fixupDebugInfoPostExtraction(Function &OldFunc, Function &NewFunc,
     DVR->getMarker()->MarkedInstr->dropOneDbgRecord(DVR);
   DIB.finalizeSubprogram(NewSP);
 
-  // Fix up the scope information attached to the line locations in the new
-  // function.
+  // Fix up the scope information attached to the line locations and the
+  // debug assignment metadata in the new function.
+  DenseMap<DIAssignID *, DIAssignID *> AssignmentIDMap;
   for (Instruction &I : instructions(NewFunc)) {
     if (const DebugLoc &DL = I.getDebugLoc())
       I.setDebugLoc(
@@ -1695,6 +1696,7 @@ static void fixupDebugInfoPostExtraction(Function &OldFunc, Function &NewFunc,
       return MD;
     };
     updateLoopMetadataDebugLocations(I, updateLoopInfoLoc);
+    at::remapAssignID(AssignmentIDMap, I);
   }
   if (!TheCall.getDebugLoc())
     TheCall.setDebugLoc(DILocation::get(Ctx, 0, 0, OldSP));
diff --git a/llvm/lib/Transforms/Utils/InlineFunction.cpp b/llvm/lib/Transforms/Utils/InlineFunction.cpp
index 1aae561d8817b..c9c863f6abddd 100644
--- a/llvm/lib/Transforms/Utils/InlineFunction.cpp
+++ b/llvm/lib/Transforms/Utils/InlineFunction.cpp
@@ -1888,29 +1888,12 @@ static void trackInlinedStores(Function::iterator Start, Function::iterator End,
 /// otherwise a function inlined more than once into the same function
 /// will cause DIAssignID to be shared by many instructions.
 static void fixupAssignments(Function::iterator Start, Function::iterator End) {
-  // Map {Old, New} metadata. Not used directly - use GetNewID.
   DenseMap<DIAssignID *, DIAssignID *> Map;
-  auto GetNewID = [&Map](Metadata *Old) {
-    DIAssignID *OldID = cast<DIAssignID>(Old);
-    if (DIAssignID *NewID = Map.lookup(OldID))
-      return NewID;
-    DIAssignID *NewID = DIAssignID::getDistinct(OldID->getContext());
-    Map[OldID] = NewID;
-    return NewID;
-  };
   // Loop over all the inlined instructions. If we find a DIAssignID
   // attachment or use, replace it with a new version.
   for (auto BBI = Start; BBI != End; ++BBI) {
-    for (Instruction &I : *BBI) {
-      for (DbgVariableRecord &DVR : filterDbgVars(I.getDbgRecordRange())) {
-        if (DVR.isDbgAssign())
-          DVR.setAssignId(GetNewID(DVR.getAssignID()));
-      }
-      if (auto *ID = I.getMetadata(LLVMContext::MD_DIAssignID))
-        I.setMetadata(LLVMContext::MD_DIAssignID, GetNewID(ID));
-      else if (auto *DAI = dyn_cast<DbgAssignIntrinsic>(&I))
-        DAI->setAssignId(GetNewID(DAI->getAssignID()));
-    }
+    for (Instruction &I : *BBI)
+      at::remapAssignID(Map, I);
   }
 }
 #undef DEBUG_TYPE
diff --git a/llvm/test/DebugInfo/assignment-tracking/X86/hotcoldsplit.ll b/llvm/test/DebugInfo/assignment-tracking/X86/hotcoldsplit.ll
new file mode 100644
index 0000000000000..94bbfa91667fe
--- /dev/null
+++ b/llvm/test/DebugInfo/assignment-tracking/X86/hotcoldsplit.ll
@@ -0,0 +1,50 @@
+; RUN: opt %s -passes=hotcoldsplit -S | FileCheck %s
+
+;; Check the extracted DIAssignID gets remapped.
+
+; CHECK-LABEL: define void @_foo()
+; CHECK: common.ret:
+; CHECK-NEXT: #dbg_assign(i64 0, ![[#]], !DIExpression(DW_OP_LLVM_fragment, 0, 64), ![[ID1:[0-9]+]], ptr null, !DIExpression(), ![[#]])
+
+; CHECK-LABEL: define internal void @_foo.cold.1()
+; CHECK: store i64 0, ptr null, align 8, !DIAssignID ![[ID2:[0-9]+]]
+
+; CHECK-DAG: ![[ID1]] = distinct !DIAssignID()
+; CHECK-DAG: ![[ID2]] = distinct !DIAssignID()
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define void @_foo() !dbg !4 {
+entry:
+  br i1 false, label %if.then7, label %common.ret
+
+common.ret:                                       ; preds = %entry
+    #dbg_assign(i64 0, !7, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !12, ptr null, !DIExpression(), !13)
+  ret void
+
+if.then7:                                         ; preds = %entry
+  %call21 = load i1, ptr null, align 4294967296
+  store i64 0, ptr null, align 8, !DIAssignID !12
+  unreachable
+}
+
+declare void @llvm.dbg.assign(metadata, metadata, metadata, metadata, metadata, metadata)
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!3}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1, isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, enums: !2, splitDebugInlining: false, nameTableKind: None)
+!1 = !DIFile(filename: "file.cpp", directory: "foo")
+!2 = !{}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!4 = distinct !DISubprogram(name: "foo", linkageName: "_foo", scope: !5, file: !1, line: 425, type: !6, scopeLine: 425, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !2)
+!5 = !DINamespace(name: "llvm", scope: null)
+!6 = distinct !DISubroutineType(types: !2)
+!7 = !DILocalVariable(name: "Path", scope: !4, file: !1, line: 436, type: !8)
+!8 = !DIDerivedType(tag: DW_TAG_typedef, name: "string", scope: !9, file: !1, line: 79, baseType: !10)
+!9 = !DINamespace(name: "std", scope: null)
+!10 = distinct !DICompositeType(tag: DW_TAG_class_type, name: "basic_string<char, std::char_traits<char>, std::allocator<char> >", scope: !11, file: !1, line: 85, size: 256, flags: DIFlagTypePassByReference | DIFlagNonTrivial, elements: !2, templateParams: !2, identifier: "_ZTSNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE")
+!11 = !DINamespace(name: "__cxx11", scope: !9, exportSymbols: true)
+!12 = distinct !DIAssignID()
+!13 = !DILocation(line: 0, scope: !4)

[SLTozer]

LGTM. Just so I understand correctly, we already had the right behaviour for inlining but not for hot cold split, and this patch just moves that behaviour to a place where both inlining and hot cold splitting will use it?