[DebugInfo] Handle additional types of stores in assignment tracking

llvm/lib/CodeGen/AssignmentTrackingAnalysis.cpp

@@ -23,6 +23,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/PrintPasses.h"
@@ -1103,6 +1104,8 @@ class AssignmentTrackingLowering {
   using UntaggedStoreAssignmentMap =
       DenseMap<const Instruction *,
                SmallVector<std::pair<VariableID, at::AssignmentInfo>>>;
+  using UnknownStoreAssignmentMap =
+      DenseMap<const Instruction *, SmallVector<VariableID>>;
 
 private:
   /// The highest numbered VariableID for partially promoted variables plus 1,
@@ -1113,6 +1116,9 @@ class AssignmentTrackingLowering {
   /// Map untagged stores to the variable fragments they assign to. Used by
   /// processUntaggedInstruction.
   UntaggedStoreAssignmentMap UntaggedStoreVars;
+  /// Map untagged unknown stores (e.g. strided/masked store intrinsics)
+  /// to the variables they may assign to. Used by processUntaggedInstruction.
+  UnknownStoreAssignmentMap UnknownStoreVars;
 
   // Machinery to defer inserting dbg.values.
   using InstInsertMap = MapVector<VarLocInsertPt, SmallVector<VarLocInfo>>;
@@ -1355,6 +1361,8 @@ class AssignmentTrackingLowering {
   /// Update \p LiveSet after encountering an instruciton without a DIAssignID
   /// attachment, \p I.
   void processUntaggedInstruction(Instruction &I, BlockInfo *LiveSet);
+  void processUnknownStoreToVariable(Instruction &I, VariableID &Var,
+                                     BlockInfo *LiveSet);
   void processDbgAssign(AssignRecord Assign, BlockInfo *LiveSet);
   void processDbgVariableRecord(DbgVariableRecord &DVR, BlockInfo *LiveSet);
   void processDbgValue(
@@ -1604,6 +1612,45 @@ void AssignmentTrackingLowering::processNonDbgInstruction(
     processUntaggedInstruction(I, LiveSet);
 }
 
+void AssignmentTrackingLowering::processUnknownStoreToVariable(
+    Instruction &I, VariableID &Var, BlockInfo *LiveSet) {
+  // We may have assigned to some unknown fragment of the variable, so
+  // treat the memory assignment as unknown for now.
+  addMemDef(LiveSet, Var, Assignment::makeNoneOrPhi());
+  // If we weren't already using a memory location, we don't need to do
+  // anything more.
+  if (getLocKind(LiveSet, Var) != LocKind::Mem)
+    return;
+  // If there is a live debug value for this variable, fall back to using
+  // that.
+  Assignment DbgAV = LiveSet->getAssignment(BlockInfo::Debug, Var);
+  if (DbgAV.Status != Assignment::NoneOrPhi && DbgAV.Source) {
+    LLVM_DEBUG(dbgs() << "Switching to fallback debug value: ";
+               DbgAV.dump(dbgs()); dbgs() << "\n");
+    setLocKind(LiveSet, Var, LocKind::Val);
+    emitDbgValue(LocKind::Val, DbgAV.Source, &I);
+    return;
+  }
+  // Otherwise, find a suitable insert point, before the next instruction or
+  // DbgRecord after I.
+  auto InsertBefore = getNextNode(&I);
+  assert(InsertBefore && "Shouldn't be inserting after a terminator");
+
+  // Get DILocation for this assignment.
+  DebugVariable V = FnVarLocs->getVariable(Var);
+  DILocation *InlinedAt = const_cast<DILocation *>(V.getInlinedAt());
+  const DILocation *DILoc = DILocation::get(
+      Fn.getContext(), 0, 0, V.getVariable()->getScope(), InlinedAt);
+
+  VarLocInfo VarLoc;
+  VarLoc.VariableID = static_cast<VariableID>(Var);
+  VarLoc.Expr = DIExpression::get(I.getContext(), {});
+  VarLoc.Values = RawLocationWrapper(
+      ValueAsMetadata::get(PoisonValue::get(Type::getInt1Ty(I.getContext()))));
+  VarLoc.DL = DILoc;
+  InsertBeforeMap[InsertBefore].push_back(VarLoc);
+}
+
 void AssignmentTrackingLowering::processUntaggedInstruction(
     Instruction &I, AssignmentTrackingLowering::BlockInfo *LiveSet) {
   // Interpret stack stores that are not tagged as an assignment in memory for
@@ -1619,8 +1666,21 @@ void AssignmentTrackingLowering::processUntaggedInstruction(
   // "early", for example.
   assert(!I.hasMetadata(LLVMContext::MD_DIAssignID));
   auto It = UntaggedStoreVars.find(&I);
-  if (It == UntaggedStoreVars.end())
+  if (It == UntaggedStoreVars.end()) {
+    // It is possible that we have an untagged unknown store, i.e. one that
+    // cannot be represented as a simple (base, offset, size) - in this case we
+    // should undef the memory location of the variable, as if we had a tagged
+    // store that did not match the current assignment.
+    // FIXME: It should be possible to support these stores, but it would
+    // require more extensive changes to our representation of assignments.
+    if (auto UnhandledStoreIt = UnknownStoreVars.find(&I);
+        UnhandledStoreIt != UnknownStoreVars.end()) {
+      LLVM_DEBUG(dbgs() << "Processing untagged unknown store " << I << "\n");
+      for (auto &Var : UnhandledStoreIt->second)
+        processUnknownStoreToVariable(I, Var, LiveSet);
+    }
     return; // No variables associated with the store destination.
+  }
 
   LLVM_DEBUG(dbgs() << "processUntaggedInstruction on UNTAGGED INST " << I
                     << "\n");
@@ -2123,6 +2183,26 @@ getUntaggedStoreAssignmentInfo(const Instruction &I, const DataLayout &Layout) {
   return std::nullopt;
 }
 
+AllocaInst *getUnknownStore(const Instruction &I, const DataLayout &Layout) {
+  auto *II = dyn_cast<IntrinsicInst>(&I);
+  if (!II)
+    return nullptr;
+  Intrinsic::ID ID = II->getIntrinsicID();
+  if (ID != Intrinsic::experimental_vp_strided_store &&
+      ID != Intrinsic::masked_store && ID != Intrinsic::vp_scatter &&
+      ID != Intrinsic::masked_scatter && ID != Intrinsic::vp_store &&
+      ID != Intrinsic::masked_compressstore)
+    return nullptr;
+  Value *MemOp = II->getArgOperand(1);
+  // We don't actually use the constant offset for now, but we may in future,
+  // and the non-accumulating versions do not support a vector of pointers.
+  APInt Offset(Layout.getIndexTypeSizeInBits(MemOp->getType()), 0);
+  Value *Base = MemOp->stripAndAccumulateConstantOffsets(Layout, Offset, true);
+  // For Base pointers that are not an alloca instruction we don't need to do
+  // anything, and simply return nullptr.
+  return dyn_cast<AllocaInst>(Base);
+}
+
 DbgDeclareInst *DynCastToDbgDeclare(DbgVariableIntrinsic *DVI) {
   return dyn_cast<DbgDeclareInst>(DVI);
 }
@@ -2145,14 +2225,16 @@ DbgVariableRecord *DynCastToDbgDeclare(DbgVariableRecord *DVR) {
 /// subsequent variables are either stack homed or fully promoted.
 ///
 /// Finally, populate UntaggedStoreVars with a mapping of untagged stores to
-/// the stored-to variable fragments.
+/// the stored-to variable fragments, and UnknownStoreVars with a mapping
+/// of untagged unknown stores to the stored-to variable aggregates.
 ///
 /// These tasks are bundled together to reduce the number of times we need
 /// to iterate over the function as they can be achieved together in one pass.
 static AssignmentTrackingLowering::OverlapMap buildOverlapMapAndRecordDeclares(
     Function &Fn, FunctionVarLocsBuilder *FnVarLocs,
     const DenseSet<DebugAggregate> &VarsWithStackSlot,
     AssignmentTrackingLowering::UntaggedStoreAssignmentMap &UntaggedStoreVars,
+    AssignmentTrackingLowering::UnknownStoreAssignmentMap &UnknownStoreVars,
     unsigned &TrackedVariablesVectorSize) {
   DenseSet<DebugVariable> Seen;
   // Map of Variable: [Fragments].
@@ -2161,7 +2243,8 @@ static AssignmentTrackingLowering::OverlapMap buildOverlapMapAndRecordDeclares(
   // - dbg.declare    -> add single location variable record
   // - dbg.*          -> Add fragments to FragmentMap
   // - untagged store -> Add fragments to FragmentMap and update
-  //                     UntaggedStoreVars.
+  //                     UntaggedStoreVars, or add to UnknownStoreVars if
+  //                     we can't determine the fragment overlap.
   // We need to add fragments for untagged stores too so that we can correctly
   // clobber overlapped fragment locations later.
   SmallVector<DbgDeclareInst *> InstDeclares;
@@ -2224,6 +2307,25 @@ static AssignmentTrackingLowering::OverlapMap buildOverlapMapAndRecordDeclares(
           HandleDbgAssignForStore(DAI);
         for (DbgVariableRecord *DVR : at::getDVRAssignmentMarkers(Info->Base))
           HandleDbgAssignForStore(DVR);
+      } else if (auto *AI = getUnknownStore(I, Fn.getDataLayout())) {
+        // Find markers linked to this alloca.
+        auto HandleDbgAssignForUnknownStore = [&](auto *Assign) {
+          // Because we can't currently represent the fragment info for this
+          // store, we treat it as an unusable store to the whole variable.
+          DebugVariable DV =
+              DebugVariable(Assign->getVariable(), std::nullopt,
+                            Assign->getDebugLoc().getInlinedAt());
+          DebugAggregate DA = {DV.getVariable(), DV.getInlinedAt()};
+          if (!VarsWithStackSlot.contains(DA))
+            return;
+
+          // Cache this info for later.
+          UnknownStoreVars[&I].push_back(FnVarLocs->insertVariable(DV));
+        };
+        for (DbgAssignIntrinsic *DAI : at::getAssignmentMarkers(AI))
+          HandleDbgAssignForUnknownStore(DAI);
+        for (DbgVariableRecord *DVR : at::getDVRAssignmentMarkers(AI))
+          HandleDbgAssignForUnknownStore(DVR);
       }
     }
   }
@@ -2298,7 +2400,7 @@ bool AssignmentTrackingLowering::run(FunctionVarLocsBuilder *FnVarLocsBuilder) {
   // neither does LiveDebugVariables) because that is difficult to do and
   // appears to be rare occurance.
   VarContains = buildOverlapMapAndRecordDeclares(
-      Fn, FnVarLocs, *VarsWithStackSlot, UntaggedStoreVars,
+      Fn, FnVarLocs, *VarsWithStackSlot, UntaggedStoreVars, UnknownStoreVars,
       TrackedVariablesVectorSize);
 
   // Prepare for traversal.

llvm/lib/IR/Verifier.cpp

@@ -4956,8 +4956,12 @@ void Verifier::visitProfMetadata(Instruction &I, MDNode *MD) {
 
 void Verifier::visitDIAssignIDMetadata(Instruction &I, MDNode *MD) {
   assert(I.hasMetadata(LLVMContext::MD_DIAssignID));
+  // DIAssignID metadata must be attached to either an alloca or some form of
+  // store/memory-writing instruction.
+  // FIXME: We allow all intrinsic insts here to avoid trying to enumerate all
+  // possible store intrinsics.
   bool ExpectedInstTy =
-      isa<AllocaInst>(I) || isa<StoreInst>(I) || isa<MemIntrinsic>(I);
+      isa<AllocaInst>(I) || isa<StoreInst>(I) || isa<IntrinsicInst>(I);
   CheckDI(ExpectedInstTy, "!DIAssignID attached to unexpected instruction kind",
           I, MD);
   // Iterate over the MetadataAsValue uses of the DIAssignID - these should

llvm/test/CodeGen/RISCV/di-assignment-tracking-vector.ll

@@ -0,0 +1,65 @@
+; RUN: llc -mtriple=riscv64 < %s | FileCheck %s --implicit-check-not=DEBUG_VALUE
+
+;; Verify that tagged and untagged non-contiguous stores are handled correctly
+;; by assignment tracking.
+;; * The store to "i" is untagged, and results in the memory location being
+;;   dropped in favour of the debug value 1010 after the store.
+;; * The store to "j" is tagged with a corresponding dbg_assign, which allows
+;;   us to keep using the memory location.
+
+; CHECK-LABEL: foo:
+; CHECK-NEXT:  .Lfunc_begin0:
+; CHECK:       # %bb.0
+; CHECK:         addi    a1, sp, 48
+; CHECK-NEXT:    #DEBUG_VALUE: foo:i <- [DW_OP_deref] $x12
+; CHECK-NEXT:    #DEBUG_VALUE: foo:j <- [DW_OP_deref] $x12
+; CHECK:         vsse32.v
+; CHECK-NEXT:    #DEBUG_VALUE: foo:i <- 1010
+; CHECK-NEXT:    vsse32.v
+
+
+target datalayout = "e-m:e-p:64:64-i64:64-i128:128-n32:64-S128"
+target triple = "riscv64-unknown-linux-gnu"
+
+define void @foo() #0 !dbg !5 {
+entry:
+  %i = alloca i64, align 8, !DIAssignID !6
+  %j = alloca i64, align 8, !DIAssignID !12
+  %sar_height.i = getelementptr i8, ptr %i, i64 24
+  store ptr %sar_height.i, ptr null, align 8
+  %vui.i = getelementptr i8, ptr %i, i64 44
+  %0 = load i32, ptr %vui.i, align 4
+  %sar_width.i = getelementptr i8, ptr %i, i64 20
+  %i_sar_width.i = getelementptr i8, ptr %i, i64 48
+  %j_sar_width.j = getelementptr i8, ptr %j, i64 48
+    #dbg_assign(i32 1010, !7, !DIExpression(), !6, ptr %i_sar_width.i, !DIExpression(), !9)
+    #dbg_assign(i32 2121, !17, !DIExpression(), !12, ptr %i_sar_width.i, !DIExpression(), !9)
+  %1 = load <2 x i32>, ptr %sar_width.i, align 4
+  call void @llvm.experimental.vp.strided.store.v2i32.p0.i64(<2 x i32> %1, ptr align 4 %i_sar_width.i, i64 -4, <2 x i1> splat (i1 true), i32 2)
+  call void @llvm.experimental.vp.strided.store.v2i32.p0.i64(<2 x i32> %1, ptr align 4 %j_sar_width.j, i64 -4, <2 x i1> splat (i1 true), i32 2), !DIAssignID !13
+    #dbg_assign(i32 1010, !7, !DIExpression(), !14, ptr %i_sar_width.i, !DIExpression(), !9)
+    #dbg_assign(i32 2121, !17, !DIExpression(), !13, ptr %i_sar_width.i, !DIExpression(), !9)
+  ret void
+}
+
+attributes #0 = { "target-features"="+v" }
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!3, !4}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C11, file: !1, isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, producer: "clang version 21.0.0git")
+!1 = !DIFile(filename: "test.c", directory: "/")
+!2 = !{}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!4 = !{i32 7, !"debug-info-assignment-tracking", i1 true}
+!5 = distinct !DISubprogram(name: "foo", linkageName: "foo", scope: !1, file: !1, line: 1, scopeLine: 1, type: !10, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !11)
+!6 = distinct !DIAssignID()
+!7 = !DILocalVariable(name: "i", scope: !5, file: !1, line: 7, type: !8)
+!8 = !DIBasicType(name: "int32_t", size: 32, encoding: DW_ATE_signed)
+!9 = !DILocation(line: 5, scope: !5)
+!10 = !DISubroutineType(types: !2)
+!11 = !{!7, !17}
+!12 = distinct !DIAssignID()
+!13 = distinct !DIAssignID()
+!14 = distinct !DIAssignID()
+!17 = !DILocalVariable(name: "j", scope: !5, file: !1, line: 7, type: !8)

llvm/test/Verifier/diassignid-vector-stores.ll

@@ -0,0 +1,39 @@
+; RUN: llvm-as -disable-output <%s 2>&1 | FileCheck %s --implicit-check-not="attached to unexpected instruction kind"
+;; Check that we allow intrinsics to have !DIAssignID attachments, but we do not
+;; allow non-intrinsic calls to have them.
+;; FIXME: Ideally we would also not allow non-store intrinsics, e.g. the
+;; llvm.vp.load intrinsic in this test.
+
+; CHECK: !DIAssignID attached to unexpected instruction kind
+; CHECK-NEXT: call void @g()
+
+declare void @g()
+
+define void @f() !dbg !5 {
+  call void @llvm.vp.store.v2i8.p0(<2 x i8> poison, ptr poison, <2 x i1> poison, i32 poison), !DIAssignID !6
+  call void @llvm.vp.scatter.v2i8.v2p0(<2 x i8> poison, <2 x ptr> poison, <2 x i1> poison, i32 poison), !DIAssignID !7
+  call void @llvm.experimental.vp.strided.store.v2i8.i64(<2 x i8> poison, ptr poison, i64 poison, <2 x i1> poison, i32 poison), !DIAssignID !8
+  call void @llvm.masked.store.v2i8.p0(<2 x i8> poison, ptr poison, i32 1, <2 x i1> poison), !DIAssignID !9
+  call void @llvm.masked.scatter.v2i8.v2p0(<2 x i8> poison, <2 x ptr> poison, i32 1, <2 x i1> poison), !DIAssignID !10
+  %r = call <2 x i8> @llvm.vp.load.v2i8.p0(ptr poison, <2 x i1> poison, i32 poison), !DIAssignID !11
+  call void @g(), !DIAssignID !12
+  ret void
+}
+
+!llvm.module.flags = !{!0}
+!llvm.dbg.cu = !{!1}
+
+!0 = !{i32 2, !"Debug Info Version", i32 3}
+!1 = distinct !DICompileUnit(language: DW_LANG_Swift, producer: "clang",
+                             file: !2, emissionKind: 2)
+!2 = !DIFile(filename: "path/to/file", directory: "/path/to/dir")
+!3 = !{null}
+!4 = !DISubroutineType(types: !3)
+!5 = distinct !DISubprogram(name: "f", scope: !2, file: !2, line: 1, type: !4, scopeLine: 2, unit: !1)
+!6 = distinct !DIAssignID()
+!7 = distinct !DIAssignID()
+!8 = distinct !DIAssignID()
+!9 = distinct !DIAssignID()
+!10 = distinct !DIAssignID()
+!11 = distinct !DIAssignID()
+!12 = distinct !DIAssignID()