[RemoveDIs][DebugInfo] Update SROA to handle DPVAssigns (#78475)

SROA needs to update llvm.dbg.assign intrinsics when it migrates debug
info in response to alloca splitting; this patch updates the debug info
migration code to handle DPVAssigns as well, making use of generic code
to avoid duplication as much as possible.

llvm/include/llvm/IR/DebugInfo.h

@@ -236,6 +236,12 @@ inline SmallVector<DPValue *> getDPVAssignmentMarkers(const Instruction *Inst) {
   return {};
 }
 
+inline SmallVector<DPValue *> getDPVAssignmentMarkers(const Instruction *Inst) {
+  if (auto *ID = Inst->getMetadata(LLVMContext::MD_DIAssignID))
+    return cast<DIAssignID>(ID)->getAllDPValueUsers();
+  return {};
+}
+
 /// Delete the llvm.dbg.assign intrinsics linked to \p Inst.
 void deleteAssignmentMarkers(const Instruction *Inst);
 

llvm/lib/Transforms/Scalar/SROA.cpp

@@ -319,6 +319,29 @@ static DebugVariable getAggregateVariable(DbgVariableIntrinsic *DVI) {
   return DebugVariable(DVI->getVariable(), std::nullopt,
                        DVI->getDebugLoc().getInlinedAt());
 }
+static DebugVariable getAggregateVariable(DPValue *DPV) {
+  return DebugVariable(DPV->getVariable(), std::nullopt,
+                       DPV->getDebugLoc().getInlinedAt());
+}
+
+static DPValue *createLinkedAssign(DPValue *, DIBuilder &DIB,
+                                   Instruction *LinkedInstr, Value *NewValue,
+                                   DILocalVariable *Variable,
+                                   DIExpression *Expression, Value *Address,
+                                   DIExpression *AddressExpression,
+                                   const DILocation *DI) {
+  (void)DIB;
+  return DPValue::createLinkedDPVAssign(LinkedInstr, NewValue, Variable,
+                                        Expression, Address, AddressExpression,
+                                        DI);
+}
+static DbgAssignIntrinsic *createLinkedAssign(
+    DbgAssignIntrinsic *, DIBuilder &DIB, Instruction *LinkedInstr,
+    Value *NewValue, DILocalVariable *Variable, DIExpression *Expression,
+    Value *Address, DIExpression *AddressExpression, const DILocation *DI) {
+  return DIB.insertDbgAssign(LinkedInstr, NewValue, Variable, Expression,
+                             Address, AddressExpression, DI);
+}
 
 /// Find linked dbg.assign and generate a new one with the correct
 /// FragmentInfo. Link Inst to the new dbg.assign.  If Value is nullptr the
@@ -340,8 +363,9 @@ static void migrateDebugInfo(AllocaInst *OldAlloca, bool IsSplit,
                              Instruction *Inst, Value *Dest, Value *Value,
                              const DataLayout &DL) {
   auto MarkerRange = at::getAssignmentMarkers(OldInst);
+  auto DPVAssignMarkerRange = at::getDPVAssignmentMarkers(OldInst);
   // Nothing to do if OldInst has no linked dbg.assign intrinsics.
-  if (MarkerRange.empty())
+  if (MarkerRange.empty() && DPVAssignMarkerRange.empty())
     return;
 
   LLVM_DEBUG(dbgs() << "  migrateDebugInfo\n");
@@ -362,6 +386,9 @@ static void migrateDebugInfo(AllocaInst *OldAlloca, bool IsSplit,
   for (auto *DAI : at::getAssignmentMarkers(OldAlloca))
     BaseFragments[getAggregateVariable(DAI)] =
         DAI->getExpression()->getFragmentInfo();
+  for (auto *DPV : at::getDPVAssignmentMarkers(OldAlloca))
+    BaseFragments[getAggregateVariable(DPV)] =
+        DPV->getExpression()->getFragmentInfo();
 
   // The new inst needs a DIAssignID unique metadata tag (if OldInst has
   // one). It shouldn't already have one: assert this assumption.
@@ -371,7 +398,7 @@ static void migrateDebugInfo(AllocaInst *OldAlloca, bool IsSplit,
   DIBuilder DIB(*OldInst->getModule(), /*AllowUnresolved*/ false);
   assert(OldAlloca->isStaticAlloca());
 
-  for (DbgAssignIntrinsic *DbgAssign : MarkerRange) {
+  auto MigrateDbgAssign = [&](auto DbgAssign) {
     LLVM_DEBUG(dbgs() << "      existing dbg.assign is: " << *DbgAssign
                       << "\n");
     auto *Expr = DbgAssign->getExpression();
@@ -382,7 +409,7 @@ static void migrateDebugInfo(AllocaInst *OldAlloca, bool IsSplit,
       {
         auto R = BaseFragments.find(getAggregateVariable(DbgAssign));
         if (R == BaseFragments.end())
-          continue;
+          return;
         BaseFragment = R->second;
       }
       std::optional<DIExpression::FragmentInfo> CurrentFragment =
@@ -393,7 +420,7 @@ static void migrateDebugInfo(AllocaInst *OldAlloca, bool IsSplit,
           BaseFragment, CurrentFragment, NewFragment);
 
       if (Result == Skip)
-        continue;
+        return;
       if (Result == UseFrag && !(NewFragment == CurrentFragment)) {
         if (CurrentFragment) {
           // Rewrite NewFragment to be relative to the existing one (this is
@@ -425,9 +452,10 @@ static void migrateDebugInfo(AllocaInst *OldAlloca, bool IsSplit,
     }
 
     ::Value *NewValue = Value ? Value : DbgAssign->getValue();
-    auto *NewAssign = DIB.insertDbgAssign(
-        Inst, NewValue, DbgAssign->getVariable(), Expr, Dest,
-        DIExpression::get(Ctx, std::nullopt), DbgAssign->getDebugLoc());
+    auto *NewAssign = createLinkedAssign(
+        DbgAssign, DIB, Inst, NewValue, DbgAssign->getVariable(), Expr, Dest,
+        DIExpression::get(Expr->getContext(), std::nullopt),
+        DbgAssign->getDebugLoc());
 
     // If we've updated the value but the original dbg.assign has an arglist
     // then kill it now - we can't use the requested new value.
@@ -461,9 +489,11 @@ static void migrateDebugInfo(AllocaInst *OldAlloca, bool IsSplit,
     NewAssign->moveBefore(DbgAssign);
 
     NewAssign->setDebugLoc(DbgAssign->getDebugLoc());
-    LLVM_DEBUG(dbgs() << "Created new assign intrinsic: " << *NewAssign
-                      << "\n");
-  }
+    LLVM_DEBUG(dbgs() << "Created new assign: " << *NewAssign << "\n");
+  };
+
+  for_each(MarkerRange, MigrateDbgAssign);
+  for_each(DPVAssignMarkerRange, MigrateDbgAssign);
 }
 
 namespace {
@@ -3108,6 +3138,7 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
       // emit dbg.assign intrinsics for mem intrinsics storing through non-
       // constant geps, or storing a variable number of bytes.
       assert(at::getAssignmentMarkers(&II).empty() &&
+             at::getDPVAssignmentMarkers(&II).empty() &&
              "AT: Unexpected link to non-const GEP");
       deleteIfTriviallyDead(OldPtr);
       return false;
@@ -3254,11 +3285,13 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
       Value *AdjustedPtr = getNewAllocaSlicePtr(IRB, OldPtr->getType());
       if (IsDest) {
         // Update the address component of linked dbg.assigns.
-        for (auto *DAI : at::getAssignmentMarkers(&II)) {
-          if (llvm::is_contained(DAI->location_ops(), II.getDest()) ||
-              DAI->getAddress() == II.getDest())
-            DAI->replaceVariableLocationOp(II.getDest(), AdjustedPtr);
-        }
+        auto UpdateAssignAddress = [&](auto *DbgAssign) {
+          if (llvm::is_contained(DbgAssign->location_ops(), II.getDest()) ||
+              DbgAssign->getAddress() == II.getDest())
+            DbgAssign->replaceVariableLocationOp(II.getDest(), AdjustedPtr);
+        };
+        for_each(at::getAssignmentMarkers(&II), UpdateAssignAddress);
+        for_each(at::getDPVAssignmentMarkers(&II), UpdateAssignAddress);
         II.setDest(AdjustedPtr);
         II.setDestAlignment(SliceAlign);
       } else {
@@ -3842,6 +3875,7 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
                          DL);
       } else {
         assert(at::getAssignmentMarkers(Store).empty() &&
+               at::getDPVAssignmentMarkers(Store).empty() &&
                "AT: unexpected debug.assign linked to store through "
                "unbounded GEP");
       }
@@ -4861,10 +4895,22 @@ static void insertNewDbgInst(DIBuilder &DIB, DPValue *Orig, AllocaInst *NewAddr,
                              DIExpression *NewFragmentExpr,
                              Instruction *BeforeInst) {
   (void)DIB;
-  DPValue *New = new DPValue(ValueAsMetadata::get(NewAddr), Orig->getVariable(),
-                             NewFragmentExpr, Orig->getDebugLoc(),
-                             DPValue::LocationType::Declare);
-  BeforeInst->getParent()->insertDPValueBefore(New, BeforeInst->getIterator());
+  if (Orig->isDbgDeclare()) {
+    DPValue *DPV = DPValue::createDPVDeclare(
+        NewAddr, Orig->getVariable(), NewFragmentExpr, Orig->getDebugLoc());
+    BeforeInst->getParent()->insertDPValueBefore(DPV,
+                                                 BeforeInst->getIterator());
+    return;
+  }
+  if (!NewAddr->hasMetadata(LLVMContext::MD_DIAssignID)) {
+    NewAddr->setMetadata(LLVMContext::MD_DIAssignID,
+                         DIAssignID::getDistinct(NewAddr->getContext()));
+  }
+  auto *NewAssign = DPValue::createLinkedDPVAssign(
+      NewAddr, Orig->getValue(), Orig->getVariable(), NewFragmentExpr, NewAddr,
+      Orig->getAddressExpression(), Orig->getDebugLoc());
+  LLVM_DEBUG(dbgs() << "Created new DPVAssign: " << *NewAssign << "\n");
+  (void)NewAssign;
 }
 
 /// Walks the slices of an alloca and form partitions based on them,
@@ -5042,6 +5088,7 @@ bool SROA::splitAlloca(AllocaInst &AI, AllocaSlices &AS) {
   for_each(findDbgDeclares(&AI), MigrateOne);
   for_each(findDPVDeclares(&AI), MigrateOne);
   for_each(at::getAssignmentMarkers(&AI), MigrateOne);
+  for_each(at::getDPVAssignmentMarkers(&AI), MigrateOne);
 
   return Changed;
 }
@@ -5262,8 +5309,9 @@ std::pair<bool /*Changed*/, bool /*CFGChanged*/> SROA::runSROA(Function &F) {
          "Should not have modified the CFG when told to preserve it.");
 
   if (Changed && isAssignmentTrackingEnabled(*F.getParent())) {
-    for (auto &BB : F)
+    for (auto &BB : F) {
       RemoveRedundantDbgInstrs(&BB);
+    }
   }
 
   return {Changed, CFGChanged};

llvm/test/DebugInfo/Generic/assignment-tracking/instcombine/remove-redundant-dbg.ll

@@ -1,5 +1,7 @@
 ; RUN: opt -passes=sroa -S %s -o - \
 ; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
+; RUN: opt --try-experimental-debuginfo-iterators -passes=sroa -S %s -o - \
+; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
 
 ;; Check that sroa removes redundant debug intrinsics after it makes a
 ;; change. This has a significant positive impact on peak memory and compiler

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/after-inlining.ll

@@ -1,4 +1,5 @@
 ; RUN: opt %s -S -passes=sroa -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators %s -S -passes=sroa -o - | FileCheck %s
 
 ;; Check that SROA preserves the InlinedAt status of new dbg.assign intriniscs
 ;; it inserts.

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/alloca-single-slice.ll

@@ -1,5 +1,7 @@
 ; RUN: opt -passes=sroa,verify -S %s -o - \
 ; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
+; RUN: opt --try-experimental-debuginfo-iterators -passes=sroa,verify -S %s -o - \
+; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
 
 ; Check that single sliced allocas retain their assignment tracking debug info.
 

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/arglist.ll

@@ -1,4 +1,5 @@
 ; RUN: opt -passes=sroa -S %s -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators -passes=sroa -S %s -o - | FileCheck %s
 
 ;; Check that a dbg.assign for a promoted variable becomes a kill location if
 ;; it used an arglist.

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/complex.ll

@@ -1,5 +1,7 @@
 ; RUN: opt -passes=sroa -S -o - %s \
 ; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
+; RUN: opt --try-experimental-debuginfo-iterators -passes=sroa -S -o - %s \
+; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
 ;
 ;; Based on llvm/test/DebugInfo/ARM/sroa-complex.ll
 ;; generated from:

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/fail-fragment.ll

@@ -1,5 +1,7 @@
 ; RUN: opt -passes=sroa -S %s -o - \
 ; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
+; RUN: opt --try-experimental-debuginfo-iterators -passes=sroa -S %s -o - \
+; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
 
 ;; Check that a dbg.assign for a promoted variable becomes a kill location if
 ;; it used a fragment that can't be split (the first check directive below).

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/frag-2.ll

@@ -1,4 +1,5 @@
 ; RUN: opt -passes=sroa -S %s -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators -passes=sroa -S %s -o - | FileCheck %s
 
 ;; $ cat test.cpp
 ;; class a {

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/frag.ll

@@ -1,4 +1,5 @@
 ; RUN: opt %s -S -passes=sroa -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators %s -S -passes=sroa -o - | FileCheck %s
 
 ;; $ cat test.cpp
 ;; class c {

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/id.ll

@@ -1,4 +1,5 @@
 ; RUN: opt -passes=sroa -S %s -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators -passes=sroa -S %s -o - | FileCheck %s
 
 ;; Check that multiple dbg.assign intrinsics linked to a store that is getting
 ;; split (or at least that is touched by SROA, causing a replacement store to

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/memcpy.ll

@@ -1,5 +1,7 @@
 ; RUN: opt -passes=sroa,verify -S %s -o - \
 ; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
+; RUN: opt --try-experimental-debuginfo-iterators -passes=sroa,verify -S %s -o - \
+; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
 
 ;; Check that the new slices of an alloca and memcpy intructions get dbg.assign
 ;; intrinsics with the correct fragment info.

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/memmove-to-from-same-alloca.ll

@@ -1,5 +1,7 @@
 ; RUN: opt %s -passes=sroa -o - -S \
 ; RUN: | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators %s -passes=sroa -o - -S \
+; RUN: | FileCheck %s
 
 ;; Generated from this C++ source:
 ;; __attribute__((nodebug)) struct Blob {int P[6];} Glob;

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/remove-redundant-dbg.ll

@@ -1,5 +1,7 @@
 ; RUN: opt -passes=sroa -S %s -o - \
 ; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
+; RUN: opt --try-experimental-debuginfo-iterators -passes=sroa -S %s -o - \
+; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
 
 ;; Check that sroa removes redundant debug intrinsics after it makes a
 ;; change. This has a significant positive impact on peak memory and compiler

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/rewrite.ll

@@ -1,5 +1,7 @@
 ; RUN: opt -passes=sroa,verify -S %s -o - \
 ; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
+; RUN: opt --try-experimental-debuginfo-iterators -passes=sroa,verify -S %s -o - \
+; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
 
 ; Check that the new slices of an alloca and memset intructions get dbg.assign
 ; intrinsics with the correct fragment info. Ensure that only the

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/split-pre-fragmented-store-2.ll

@@ -1,5 +1,7 @@
 ; RUN: opt -S -passes=sroa -sroa-skip-mem2reg %s \
 ; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
+; RUN: opt --try-experimental-debuginfo-iterators -S -passes=sroa -sroa-skip-mem2reg %s \
+; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
 
 ;; NOTE: This is the same as split-pre-fragmented-store.ll except the base
 ;; alloca's dbg.assign has been altered to contain a fragment of the full

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/split-pre-fragmented-store.ll

@@ -1,5 +1,7 @@
 ; RUN: opt -S -passes=sroa -sroa-skip-mem2reg %s \
 ; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
+; RUN: opt --try-experimental-debuginfo-iterators -S -passes=sroa -sroa-skip-mem2reg %s \
+; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
 
 ;; IR hand-modified, originally generated from:
 ;; struct Pair { int a; int b; };

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/store.ll

@@ -1,5 +1,7 @@
 ; RUN: opt -passes=sroa,verify -S %s -o - \
 ; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
+; RUN: opt --try-experimental-debuginfo-iterators -passes=sroa,verify -S %s -o - \
+; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
 
 ; Check that the new slices of an alloca and memset intructions get dbg.assign
 ; intrinsics with the correct fragment info. Ensure that only the

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/unspecified-var-size.ll

@@ -1,4 +1,5 @@
 ; RUN: opt -S %s -passes=sroa -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators -S %s -passes=sroa -o - | FileCheck %s
 
 ;; $ cat test.cpp
 ;; #include <cstddef>

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/user-memcpy.ll

@@ -1,5 +1,7 @@
 ; RUN: opt -passes=sroa -S %s -o - \
 ; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
+; RUN: opt --try-experimental-debuginfo-iterators -passes=sroa -S %s -o - \
+; RUN: | FileCheck %s --implicit-check-not="call void @llvm.dbg"
 
 ;; Check that the fragments generated in SROA for a split alloca that has a
 ;; dbg.assign with non-zero-offset fragment are correct.

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/var-sized-fragment.ll

@@ -1,4 +1,5 @@
 ; RUN: opt -S %s -o - -passes=sroa | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators -S %s -o - -passes=sroa | FileCheck %s
 
 ;; SROA splits the alloca into two. Each slice already has a 32-bit variable
 ;; associated with it (the structured binding variables). Check SROA doesn't

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/vec-1.ll

@@ -1,4 +1,5 @@
 ; RUN: opt %s -S -passes=sroa -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators %s -S -passes=sroa -o - | FileCheck %s
 
 ;; Ensure that only the value-expression gets fragment info; that the
 ;; address-expression remains untouched.

llvm/test/DebugInfo/Generic/assignment-tracking/sroa/vec-2.ll

@@ -1,4 +1,5 @@
 ; RUN: opt %s -S -passes=sroa -o - | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators %s -S -passes=sroa -o - | FileCheck %s
 
 ;; $ cat test.cpp
 ;; class a {