[Assignment Tracking] Fix fragments for assignments to variables smaller than the alloca

[Orlando]

Summary

Prior to this patch the trackAssignments function would attribute all stores to an alloca to all variables linked to the alloca. This is wrong in the case where the alloca contains variables which are smaller than the alloca, and caused erroneous fragment information to be generated.

Now stores outside the variable bounds are discarded, and we check whether a fragment is needed based on whether the store covers the entire variable as opposed to whether it covers the entire alloca (except for variables of unknown size).

Note that trackAssignments doesn't yet understand whole variables sitting at anything other than offset 0 in an alloca - those variables are still tracked using dbg.declares.

Fixes https://lab.llvm.org/buildbot/#/builders/70/builds/36007

[jmorse]

What would be the implications of a store that partially overlaps an AT tracked variable -- as a result of this patch, we would lose some of that information (to avoid a crash)? It feels like an extremely unlikely edge case, so doesn't worry me too much, but is the ultimate fix to truncate the assignment here, or just not produce that kind of debug-info in the first place? And if the latter, should we assert for that, or just conclude "SROA really hates what you're doing in this situation".

Otherwise LGTM

llvm/lib/IR/DebugInfo.cpp

1849

all variables that _reach_ here? Otherwise it's unclear what context "here" is in.

1850

[Orlando]

What would be the implications of a store that partially overlaps an AT tracked variable -- as a result of this patch, we would lose some of that information (to avoid a crash)?

Yeah that's right.

It feels like an extremely unlikely edge case, so doesn't worry me too much, but is the ultimate fix to truncate the assignment here, or just not produce that kind of debug-info in the first place? And if the latter, should we assert for that, or just conclude "SROA really hates what you're doing in this situation".

I think truncating would make sense too - I've added a note in the comments but I think you're right that the situation should be extremely rare.

Thanks for the review.

llvm/lib/IR/DebugInfo.cpp

 }
 
 std::optional<AssignmentInfo> at::getAssignmentInfo(const DataLayout &DL,
                                                     const AllocaInst *AI) {
   uint64_t SizeInBits = DL.getTypeSizeInBits(AI->getAllocatedType());
   return getAssignmentInfoImpl(DL, AI, SizeInBits);
 }
 
+/// Returns nullptr if the assignment shouldn't be attributed to this variable.
 static CallInst *emitDbgAssign(AssignmentInfo Info, Value *Val, Value *Dest,
                                Instruction &StoreLikeInst,
                                const VarRecord &VarRec, DIBuilder &DIB) {
   auto *ID = StoreLikeInst.getMetadata(LLVMContext::MD_DIAssignID);
   assert(ID && "Store instruction must have DIAssignID metadata");
   (void)ID;
 
+  bool StoreToWholeVariable = Info.StoreToWholeAlloca;
+  if (auto Size = VarRec.Var->getSizeInBits()) {
+    // Discard stores to bits outside this variable. NOTE: trackAssignments
+    // doesn't understand base expressions yet, so all variables here are
+    // guarenteed to start at offset 0 in the alloca.
+    if (Info.OffsetInBits + Info.SizeInBits > *Size)
+      return nullptr;
+    // FIXME: As noted above - only variables at offset 0 are handled
+    // currently.
+    StoreToWholeVariable = Info.OffsetInBits == /*VarOffsetInAlloca*/ 0 &&
+                           Info.SizeInBits == *Size;
+  }
+
   DIExpression *Expr =
       DIExpression::get(StoreLikeInst.getContext(), std::nullopt);
-  if (!Info.StoreToWholeAlloca) {
+  if (!StoreToWholeVariable) {
     auto R = DIExpression::createFragmentExpression(Expr, Info.OffsetInBits,
                                                     Info.SizeInBits);
     assert(R.has_value() && "failed to create fragment expression");
     Expr = *R;
   }
   DIExpression *AddrExpr =
       DIExpression::get(StoreLikeInst.getContext(), std::nullopt);
   return DIB.insertDbgAssign(&StoreLikeInst, Val, VarRec.Var, Expr, Dest,
for (Instruction &I : *BBI) {
         ID = DIAssignID::getDistinct(Ctx);
         I.setMetadata(LLVMContext::MD_DIAssignID, ID);
       }
 
       for (const VarRecord &R : LocalIt->second) {
         auto *Assign =
             emitDbgAssign(*Info, ValueComponent, DestComponent, I, R, DIB);
         (void)Assign;
-        LLVM_DEBUG(errs() << " > INSERT: " << *Assign << "\n");
+        LLVM_DEBUG(if (Assign) errs() << " > INSERT: " << *Assign << "\n");
       }
     }
   }
 }
 
 bool AssignmentTrackingPass::runOnFunction(Function &F) {
   // No value in assignment tracking without optimisations.
   if (F.hasFnAttribute(Attribute::OptimizeNone))

llvm/test/DebugInfo/Generic/assignment-tracking/inline/shared-alloca.ll

+; RUN: opt -S %s -passes=inline -o - \
+; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
+
+;; The dbg.assign linked to the large alloca describes a variable sitting at
+;; offset 0, size 64. Check:
+;; A) an inlined store to the alloca outside of the bits 0-64 is not attributed
+;;    to the variable.
+;; B) a fragment expression is not created if the inlined store is the size of
+;;    the variable.
+;; C) a fragment expression is created if the inlined store is not the size of
+;;    the variable.
+
+;; Test A:
+; CHECK: %0 = alloca %"struct.llvm::detail::DenseMapPair", i32 0, align 8, !DIAssignID ![[ID1:[0-9]+]]
+; CHECK: call void @llvm.dbg.assign(metadata i1 undef, metadata ![[#]], metadata !DIExpression(), metadata ![[ID1]], metadata ptr %0, metadata !DIExpression())
+
+;; Test B:
+;; CHECK: store i64 1, ptr %0, align 4, !DIAssignID ![[ID2:[0-9]+]]
+;; CHECK: call void @llvm.dbg.assign(metadata i64 1, metadata ![[#]], metadata !DIExpression(), metadata ![[ID2]], metadata ptr %0, metadata !DIExpression())
+
+;; Test C:
+;; CHECK: store i32 2, ptr %0, align 4, !DIAssignID ![[ID3:[0-9]+]]
+;; CHECK: call void @llvm.dbg.assign(metadata i32 2, metadata ![[#]], metadata !DIExpression(DW_OP_LLVM_fragment, 0, 32), metadata ![[ID3]], metadata ptr %0, metadata !DIExpression())
+
+%"struct.llvm::detail::DenseMapPair" = type { %"struct.std::pair" }
+%"struct.std::pair" = type { ptr, %"class.llvm::SmallVector" }
+%"class.llvm::SmallVector" = type { %"class.llvm::SmallVectorImpl", %"struct.llvm::SmallVectorStorage" }
+%"class.llvm::SmallVectorImpl" = type { %"class.llvm::SmallVectorTemplateBase" }
+%"class.llvm::SmallVectorTemplateBase" = type { %"class.llvm::SmallVectorTemplateCommon" }
+%"class.llvm::SmallVectorTemplateCommon" = type { %"class.llvm::SmallVectorBase" }
+%"class.llvm::SmallVectorBase" = type { ptr, i64, i64 }
+%"struct.llvm::SmallVectorStorage" = type { [40 x i8] }
+
+define void @_Z6verifyv() {
+entry:
+  %0 = alloca %"struct.llvm::detail::DenseMapPair", i32 0, align 8, !DIAssignID !5
+  call void @llvm.dbg.assign(metadata i1 undef, metadata !6, metadata !DIExpression(), metadata !5, metadata ptr %0, metadata !DIExpression()), !dbg !14
+  call void @_ZN4llvm6detail12DenseMapPairIP4SCEVNS_11SmallVectorI6FoldIDLj40EEEEC2ERKS7_(ptr %0)
+  ret void
+}
+
+define linkonce_odr void @_ZN4llvm6detail12DenseMapPairIP4SCEVNS_11SmallVectorI6FoldIDLj40EEEEC2ERKS7_(ptr %this) {
+entry:
+  %second.i = getelementptr %"struct.std::pair", ptr %this, i64 0, i32 1
+  store i32 0, ptr %second.i, align 4
+  %test = getelementptr %"struct.std::pair", ptr %this, i64 0, i32 0
+  store i64 1, ptr %test
+  store i32 2, ptr %test
+  ret void
+}
+
+declare void @llvm.dbg.assign(metadata, metadata, metadata, metadata, metadata, metadata)
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!3, !4}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1, producer: "clang version 17.0.0", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, enums: !2, retainedTypes: !2, globals: !2, imports: !2, splitDebugInlining: false, nameTableKind: None)
+!1 = !DIFile(filename: "reduce.cpp", directory: "")
+!2 = !{}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!4 = !{i32 7, !"debug-info-assignment-tracking", i1 true}
+!5 = distinct !DIAssignID()
+!6 = !DILocalVariable(name: "a", scope: !7, file: !1, line: 10, type: !12)
+!7 = distinct !DILexicalBlock(scope: !8, file: !1, line: 10, column: 3)
+!8 = distinct !DILexicalBlock(scope: !9, file: !1, line: 10, column: 3)
+!9 = distinct !DISubprogram(name: "verify", linkageName: "_Z6verifyv", scope: !1, file: !1, line: 9, type: !10, scopeLine: 9, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !2)
+!10 = !DISubroutineType(types: !11)
+!11 = !{null}
+!12 = !DIDerivedType(tag: DW_TAG_pointer_type, baseType: !13, size: 64)
+!13 = distinct !DICompositeType(tag: DW_TAG_class_type, name: "SCEV", file: !1, line: 3, flags: DIFlagFwdDecl | DIFlagNonTrivial, identifier: "_ZTS4SCEV")
+!14 = !DILocation(line: 0, scope: !7)
+