[DebugInfo] Handle additional types of stores in assignment tracking (#129070)

Fixes #126417.

Currently, assignment tracking recognizes allocas, stores, and mem
intrinsics as valid instructions to tag with DIAssignID, with allocas
representing the allocation for a variable and the others representing
instructions that may assign to the variable. There are other intrinsics
that can perform these assignments however, and if we transform a store
instruction into one of these intrinsics and correctly transfer the
DIAssignID over, this results in a verifier error. The
AssignmentTrackingAnalysis pass also does not know how to handle these
intrinsics if they are untagged, as it does not know how to extract
assignment information (base address, offset, size) from them.

This patch adds _some_ support for some intrinsics that may perform
assignments: masked store/scatter, and vp store/strided store/scatter.
This patch does not add support for extracting assignment information
from these, as they may store with either non-constant size or to
non-contiguous blocks of memory; instead it adds support for recognizing
untagged stores with "unknown" assignment info, for which we assume that
the memory location of the associated variable should not be used, as we
can't determine which fragments of it should or should not be used.

In principle, it should be possible to handle the more complex cases
mentioned above, but it would require more substantial changes to
AssignmentTrackingAnalysis, and it is mostly only needed as a fallback
if the DIAssignID is not preserved on these alternative stores.

Author: SLTozer

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"
@@ -1102,6 +1103,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,
@@ -1112,6 +1115,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>>;
@@ -1354,6 +1360,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(
@@ -1603,6 +1611,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
@@ -1618,8 +1665,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");
@@ -2122,6 +2182,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);
 }
@@ -2144,14 +2224,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].
@@ -2160,7 +2242,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;
@@ -2223,6 +2306,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);
       }
     }
   }
@@ -2297,7 +2399,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

@@ -4966,8 +4966,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()