[Assignment Tracking] Fix migrateDebuginfo in SROA

[Orlando]

Summary

Without this patch, migrateDebugInfo doesn't understand how to handle existing fragments that are smaller than the to-be-split store. This can occur if. e.g. a vector store (1 dbg.assign) is split (many dbg.assigns - 1 fragment for each scalar) and later those stores are re-vectorized (many dbg.assigns), and then SROA runs on that.

The approach taken in this patch is to drop intrinsics with fragments outside of the slice.

For example, starting with:

store <2 x float> %v, ptr %dest !DIAssignID !1
call void @llvm.dbg.assign(..., DIExpression(DW_OP_LLVM_fragment, 0, 32), !1, ...)
call void @llvm.dbg.assign(..., DIExpression(DW_OP_LLVM_fragment, 32, 32), !1, ...)

When visiting the slice of bits 0 to 31 we get:

store float v.extract.0, ptr %dest !DIAssignID !2
call void @llvm.dbg.assign(..., DIExpression(DW_OP_LLVM_fragment, 0, 32), !2, ...)

The other dbg.assign associated with the currently-split store is dropped for this split part. And visiting bits 32 to 63 we get the following:

store float v.extract.1, ptr %adjusted.dest !DIAssignID !3
call void @llvm.dbg.assign(..., DIExpression(DW_OP_LLVM_fragment, 32, 32), !3, ...)

I've added two tests that cover this case.

Implementing this meant re-writing the fragment-calculation part of migrateDebugInfo to work with the absolute offset of the new slice in terms of the base alloca (instead of the offset of the slice into the new alloca), the fragment (if any) of the variable associated with the base alloca, and the fragment associated with the split store. Because we need the offset into the base alloca for the variables being split, some careful wiring is required for memory intrinsics due to the fact that memory intrinsics can be split when either the source or dest allocas are split. In the case where the source alloca drives the splitting, we need to be careful to pass migrateDebugInfo the information in relation to the dest alloca.

[Orlando]

Additional note: I added -sroa-skip-mem2reg / SROASkipMem2Reg as a testing and debugging flag to improve testing of assignment tracking debug-info. It's useful to be able to verify that the SROA'd stores are linked to the resultant dbg.assigns (i.e. their DIAssignID metadata is correct). It's much easier to tell SROA "please don't mem2reg" than to come up with or reduce inputs that result in SROA scalaraizing something without those stores getting promoted by mem2reg.

[StephenTozer]

Looks broadly good to me, with some small suggestions.

llvm/lib/Transforms/Scalar/SROA.cpp

168–169

This looks like it would miss the case where Target.OffsetInBits + Target.SizeInBits > CurrentFragment->OffsetInBits + CurrentFragment->SizeInBits, i.e. the target fragment starts inside the current fragment but extends beyond it.

Optionally, I think the logic in this check might be easier to write and understand if we first set variables (Target|Current)Slice(Start|End), and then both this and the above check could be combined with TargetSliceStart >= CurrentSliceStart && TargetSliceEnd <= CurrentSliceEnd.

179

Slightly confusing that the name has changed to AbsoluteOffsetInBits but the description still says it is relative - is this an update error? If not, I think this might need a variable name that specifies what it is relative to.

246

With the guard against !NewFragment above, this could be simplified?

[Orlando]

Thanks @StephenTozer - I've addressed your suggestions. I've also removed AllocaSliceRewriter::RelativeOffset, which I had added in a previous patch.

llvm/lib/Transforms/Scalar/SROA.cpp

168–169

Good catch, I've added the helper methods startInBits and endInBits to DIExpression::FragmentInfo.

179

Ah you're right that is confusing - is that better? It is a little difficult to convey in a short-ish name exactly what it is + "InBits" (which I think is a useful suffix in the absence of bit/byte types).

To be clear, it's the slice's start bit in OldAlloca rather than its offset in the new alloca (if there is one).

246

Huh, you're right, nice one. Still have to use the !(x == y) form due to lack of operator!= overload though.

[StephenTozer]

llvm/include/llvm/IR/DebugInfoMetadata.h

2782

Nit so small you'd need an electron microscope (so feel free to ignore if you disagree), but imo if an illustrative example is being used to demonstrate that this returns the sum of the size and offset, they should both be non-zero values to make it clear that it's the sum and not just selecting the size.

llvm/lib/Transforms/Scalar/SROA.cpp

167–168

This check now subsumes the check above - since the fragment offset and size are always positive (we've got bigger problems if they aren't!), CurrentFragment->endInBits() <= Target.startInBits() will never be true if Target.startInBits() < CurrentFragment->startInBits() isn't (and same for the second condition).

179

YMMV but better to have a long-but-descriptive name than a short name that obscures or confuses the meaning of the variable - new one LGTM!

[Orlando]

+ Address comments

llvm/include/llvm/IR/DebugInfoMetadata.h

2782

This is reasonable and wise. Done.

llvm/lib/Transforms/Scalar/SROA.cpp

167–168

This is true - removed the first check and updated the comments.

[StephenTozer]

Thanks for the changes, LGTM.

[Orlando]

Thank you for the review!

llvm/include/llvm/IR/DebugInfoMetadata.h

public:
 
   /// Return whether there is exactly one operator and it is a DW_OP_deref;
   bool isDeref() const;
 
   /// Holds the characteristics of one fragment of a larger variable.
   struct FragmentInfo {
     uint64_t SizeInBits;
     uint64_t OffsetInBits;
+    /// Return the index of the first bit of the fragment.
+    uint64_t startInBits() const { return OffsetInBits; }
+    /// Return the index of the bit after the end of the fragment, e.g. for
+    /// fragment offset=16 and size=32 return their sum, 48.
+    uint64_t endInBits() const { return OffsetInBits + SizeInBits; }
   };
 
   /// Retrieve the details of this fragment expression.
   static std::optional<FragmentInfo> getFragmentInfo(expr_op_iterator Start,
                                                      expr_op_iterator End);
 
   /// Retrieve the details of this fragment expression.
   std::optional<FragmentInfo> getFragmentInfo() const {

llvm/lib/Transforms/Scalar/SROA.cpp

STATISTIC(
     "Number of stores rewritten into predicated loads to allow promotion");
 STATISTIC(NumDeleted, "Number of instructions deleted");
 STATISTIC(NumVectorized, "Number of vectorized aggregates");
 
 /// Hidden option to experiment with completely strict handling of inbounds
 /// GEPs.
 static cl::opt<bool> SROAStrictInbounds("sroa-strict-inbounds", cl::init(false),
                                         cl::Hidden);
+/// Disable running mem2reg during SROA in order to test or debug SROA.
+static cl::opt<bool> SROASkipMem2Reg("sroa-skip-mem2reg", cl::init(false),
+                                     cl::Hidden);
 namespace {
+
+/// Calculate the fragment of a variable to use when slicing a store
+/// based on the slice dimensions, existing fragment, and base storage
+/// fragment.
+/// Note that a returned value of std::nullopt indicates that there is
+/// no appropriate fragment available (rather than meaning use the whole
+/// variable, which is a common usage). Because the store is being sliced
+/// we always expect a fragment - there's never a case where the whole
+/// variable should be used.
+static std::optional<DIExpression::FragmentInfo>
+calculateFragment(uint64_t NewStorageSliceOffsetInBits,
+                  uint64_t NewStorageSliceSizeInBits,
+                  std::optional<DIExpression::FragmentInfo> StorageFragment,
+                  std::optional<DIExpression::FragmentInfo> CurrentFragment) {
+  DIExpression::FragmentInfo Target;
+  // If the base storage describes part of the variable apply the offset and
+  // the size constraint.
+  if (StorageFragment) {
+    Target.SizeInBits =
+        std::min(NewStorageSliceSizeInBits, StorageFragment->SizeInBits);
+    Target.OffsetInBits =
+        NewStorageSliceOffsetInBits + StorageFragment->OffsetInBits;
+  } else {
+    Target.SizeInBits = NewStorageSliceSizeInBits;
+    Target.OffsetInBits = NewStorageSliceOffsetInBits;
+  }
+
+  // No additional work to do if there isn't a fragment already, or there is
+  // but it already exactly describes the new assignment.
+  if (!CurrentFragment || *CurrentFragment == Target)
+    return Target;
+
+  // Reject the target fragment if it doesn't fit wholly within the current
+  // fragment. TODO: We could instead chop up the target to fit in the case of
+  // a partial overlap.
+  if (Target.startInBits() < CurrentFragment->startInBits() ||
+      Target.endInBits() > CurrentFragment->endInBits())
+    return std::nullopt;
+
+  // Target fits within the current fragment, return it.
+  return Target;
+}
+
+static DebugVariable getAggregateVariable(DbgVariableIntrinsic *DVI) {
+  return DebugVariable(DVI->getVariable(), std::nullopt,
+                       DVI->getDebugLoc().getInlinedAt());
+}
+
 /// 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
 /// value component is copied from the old dbg.assign to the new.
 /// \param OldAlloca             Alloca for the variable before splitting.
-/// \param RelativeOffsetInBits  Offset into \p OldAlloca relative to the
-///                              offset prior to splitting (change in offset).
+/// \param IsSplit               True if the store (not necessarily alloca)
+///                              is being split.
+/// \param OldAllocaOffsetInBits Offset of the slice taken from OldAlloca.
 /// \param SliceSizeInBits       New number of bits being written to.
 /// \param OldInst               Instruction that is being split.
 /// \param Inst                  New instruction performing this part of the
 ///                              split store.
 /// \param Dest                  Store destination.
 /// \param Value                 Stored value.
 /// \param DL                    Datalayout.
-static void migrateDebugInfo(AllocaInst *OldAlloca,
-                             uint64_t RelativeOffsetInBits,
+static void migrateDebugInfo(AllocaInst *OldAlloca, bool IsSplit,
+                             uint64_t OldAllocaOffsetInBits,
                              uint64_t SliceSizeInBits, Instruction *OldInst,
                              Instruction *Inst, Value *Dest, Value *Value,
                              const DataLayout &DL) {
   auto MarkerRange = at::getAssignmentMarkers(OldInst);
   // Nothing to do if OldInst has no linked dbg.assign intrinsics.
   if (MarkerRange.empty())
     return;
 
   LLVM_DEBUG(dbgs() << "  migrateDebugInfo\n");
   LLVM_DEBUG(dbgs() << "    OldAlloca: " << *OldAlloca << "\n");
-  LLVM_DEBUG(dbgs() << "    RelativeOffset: " << RelativeOffsetInBits << "\n");
+  LLVM_DEBUG(dbgs() << "    IsSplit: " << IsSplit << "\n");
+  LLVM_DEBUG(dbgs() << "    OldAllocaOffsetInBits: " << OldAllocaOffsetInBits
+                    << "\n");
   LLVM_DEBUG(dbgs() << "    SliceSizeInBits: " << SliceSizeInBits << "\n");
   LLVM_DEBUG(dbgs() << "    OldInst: " << *OldInst << "\n");
   LLVM_DEBUG(dbgs() << "    Inst: " << *Inst << "\n");
   LLVM_DEBUG(dbgs() << "    Dest: " << *Dest << "\n");
   if (Value)
     LLVM_DEBUG(dbgs() << "    Value: " << *Value << "\n");
 
+  /// Map of aggregate variables to their fragment associated with OldAlloca.
+  DenseMap<DebugVariable, std::optional<DIExpression::FragmentInfo>>
+      BaseFragments;
+  for (auto *DAI : at::getAssignmentMarkers(OldAlloca))
+    BaseFragments[getAggregateVariable(DAI)] =
+        DAI->getExpression()->getFragmentInfo();
+
   // The new inst needs a DIAssignID unique metadata tag (if OldInst has
   // one). It shouldn't already have one: assert this assumption.
   assert(!Inst->getMetadata(LLVMContext::MD_DIAssignID));
   DIAssignID *NewID = nullptr;
   auto &Ctx = Inst->getContext();
   DIBuilder DIB(*OldInst->getModule(), /*AllowUnresolved*/ false);
-  uint64_t AllocaSizeInBits = *OldAlloca->getAllocationSizeInBits(DL);
   assert(OldAlloca->isStaticAlloca());
 
   for (DbgAssignIntrinsic *DbgAssign : MarkerRange) {
     LLVM_DEBUG(dbgs() << "      existing dbg.assign is: " << *DbgAssign
                       << "\n");
     auto *Expr = DbgAssign->getExpression();
 
-    // Check if the dbg.assign already describes a fragment.
-    auto GetCurrentFragSize = [AllocaSizeInBits, DbgAssign,
-                               Expr]() -> uint64_t {
-      if (auto FI = Expr->getFragmentInfo())
-        return FI->SizeInBits;
-      if (auto VarSize = DbgAssign->getVariable()->getSizeInBits())
-        return *VarSize;
-      // The variable type has an unspecified size. This can happen in the
-      // case of DW_TAG_unspecified_type types, e.g.  std::nullptr_t. Because
-      // there is no fragment and we do not know the size of the variable type,
-      // we'll guess by looking at the alloca.
-      return AllocaSizeInBits;
-    };
-    uint64_t CurrentFragSize = GetCurrentFragSize();
-    bool MakeNewFragment = CurrentFragSize != SliceSizeInBits;
-    assert(MakeNewFragment || RelativeOffsetInBits == 0);
-
-    assert(SliceSizeInBits <= AllocaSizeInBits);
-    if (MakeNewFragment) {
-      assert(RelativeOffsetInBits + SliceSizeInBits <= CurrentFragSize);
+    if (IsSplit) {
+      std::optional<DIExpression::FragmentInfo> BaseFragment = std::nullopt;
+      {
+        auto R = BaseFragments.find(getAggregateVariable(DbgAssign));
+        if (R == BaseFragments.end())
+          continue;
+        BaseFragment = R->second;
+      }
+      std::optional<DIExpression::FragmentInfo> CurrentFragment =
+          Expr->getFragmentInfo();
+      std::optional<DIExpression::FragmentInfo> NewFragment =
+          calculateFragment(OldAllocaOffsetInBits, SliceSizeInBits,
+                            BaseFragment, CurrentFragment);
+      // Note that std::nullopt here means "skip this fragment" rather than
+      // "there is no fragment / use the whole variable".
+      if (!NewFragment)
+        continue;
+
+      if (!(NewFragment == CurrentFragment)) {
+        if (CurrentFragment) {
+          // Rewrite NewFragment to be relative to the existing one (this is
+          // what createFragmentExpression wants).  CalculateFragment has
+          // already resolved the size for us. FIXME: Should it return the
+          // relative fragment too?
+          NewFragment->OffsetInBits -= CurrentFragment->OffsetInBits;
+        }
+
         auto E = DIExpression::createFragmentExpression(
-          Expr, RelativeOffsetInBits, SliceSizeInBits);
+            Expr, NewFragment->OffsetInBits, NewFragment->SizeInBits);
         assert(E && "Failed to create fragment expr!");
         Expr = *E;
       }
+    }
 
     // If we haven't created a DIAssignID ID do that now and attach it to Inst.
     if (!NewID) {
       NewID = DIAssignID::getDistinct(Ctx);
       Inst->setMetadata(LLVMContext::MD_DIAssignID, NewID);
     }
 
     Value = Value ? Value : DbgAssign->getValue();
class llvm::sroa::AllocaSliceRewriter
   // the original alloca.
   uint64_t BeginOffset = 0;
   uint64_t EndOffset = 0;
 
   // The new offsets of the slice currently being rewritten relative to the
   // original alloca.
   uint64_t NewBeginOffset = 0, NewEndOffset = 0;
 
-  uint64_t RelativeOffset = 0;
   uint64_t SliceSize = 0;
   bool IsSplittable = false;
   bool IsSplit = false;
   Use *OldUse = nullptr;
   Instruction *OldPtr = nullptr;
 
   // Track post-rewrite users which are PHI nodes and Selects.
   SmallSetVector<PHINode *, 8> &PHIUsers;
bool visit(AllocaSlices::const_iterator I) {
     LLVM_DEBUG(dbgs() << "\n");
 
     // Compute the intersecting offset range.
     assert(BeginOffset < NewAllocaEndOffset);
     assert(EndOffset > NewAllocaBeginOffset);
     NewBeginOffset = std::max(BeginOffset, NewAllocaBeginOffset);
     NewEndOffset = std::min(EndOffset, NewAllocaEndOffset);
 
-    RelativeOffset = NewBeginOffset - BeginOffset;
     SliceSize = NewEndOffset - NewBeginOffset;
     LLVM_DEBUG(dbgs() << "   Begin:(" << BeginOffset << ", " << EndOffset
                       << ") NewBegin:(" << NewBeginOffset << ", "
                       << NewEndOffset << ") NewAllocaBegin:("
                       << NewAllocaBeginOffset << ", " << NewAllocaEndOffset
                       << ")\n");
-    assert(IsSplit || RelativeOffset == 0);
+    assert(IsSplit || NewBeginOffset == BeginOffset);
     OldUse = I->getUse();
     OldPtr = cast<Instruction>(OldUse->get());
 
     Instruction *OldUserI = cast<Instruction>(OldUse->getUser());
     IRB.SetInsertPoint(OldUserI);
     IRB.SetCurrentDebugLocation(OldUserI->getDebugLoc());
     IRB.getInserter().SetNamePrefix(
         Twine(NewAI.getName()) + "." + Twine(BeginOffset) + ".");
bool rewriteVectorizedStoreInst(Value *V, StoreInst &SI, Value *OldOp,
     StoreInst *Store = IRB.CreateAlignedStore(V, &NewAI, NewAI.getAlign());
     Store->copyMetadata(SI, {LLVMContext::MD_mem_parallel_loop_access,
                              LLVMContext::MD_access_group});
     if (AATags)
       Store->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
     Pass.DeadInsts.push_back(&SI);
 
     // NOTE: Careful to use OrigV rather than V.
-    migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &SI, Store,
-                     Store->getPointerOperand(), OrigV, DL);
+    migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8, &SI,
+                     Store, Store->getPointerOperand(), OrigV, DL);
     LLVM_DEBUG(dbgs() << "          to: " << *Store << "\n");
     return true;
   }
 
   bool rewriteIntegerStore(Value *V, StoreInst &SI, AAMDNodes AATags) {
     assert(IntTy && "We cannot extract an integer from the alloca");
     assert(!SI.isVolatile());
     if (DL.getTypeSizeInBits(V->getType()).getFixedValue() !=
         IntTy->getBitWidth()) {
       Value *Old = IRB.CreateAlignedLoad(NewAI.getAllocatedType(), &NewAI,
                                          NewAI.getAlign(), "oldload");
       Old = convertValue(DL, IRB, Old, IntTy);
       assert(BeginOffset >= NewAllocaBeginOffset && "Out of bounds offset");
       uint64_t Offset = BeginOffset - NewAllocaBeginOffset;
       V = insertInteger(DL, IRB, Old, SI.getValueOperand(), Offset, "insert");
     }
     V = convertValue(DL, IRB, V, NewAllocaTy);
     StoreInst *Store = IRB.CreateAlignedStore(V, &NewAI, NewAI.getAlign());
     Store->copyMetadata(SI, {LLVMContext::MD_mem_parallel_loop_access,
                              LLVMContext::MD_access_group});
     if (AATags)
       Store->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
 
-    migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &SI, Store,
-                     Store->getPointerOperand(), Store->getValueOperand(), DL);
+    migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8, &SI,
+                     Store, Store->getPointerOperand(),
+                     Store->getValueOperand(), DL);
 
     Pass.DeadInsts.push_back(&SI);
     LLVM_DEBUG(dbgs() << "          to: " << *Store << "\n");
     return true;
   }
 
   bool visitStoreInst(StoreInst &SI) {
     LLVM_DEBUG(dbgs() << "    original: " << SI << "\n");
NewSI->copyMetadata(SI, {LLVMContext::MD_mem_parallel_loop_access,
                              LLVMContext::MD_access_group});
     if (AATags)
       NewSI->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
     if (SI.isVolatile())
       NewSI->setAtomic(SI.getOrdering(), SI.getSyncScopeID());
     if (NewSI->isAtomic())
       NewSI->setAlignment(SI.getAlign());
 
-    migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &SI, NewSI,
-                     NewSI->getPointerOperand(), NewSI->getValueOperand(), DL);
+    migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8, &SI,
+                     NewSI, NewSI->getPointerOperand(),
+                     NewSI->getValueOperand(), DL);
 
     Pass.DeadInsts.push_back(&SI);
     deleteIfTriviallyDead(OldOp);
 
     LLVM_DEBUG(dbgs() << "          to: " << *NewSI << "\n");
     return NewSI->getPointerOperand() == &NewAI &&
            NewSI->getValueOperand()->getType() == NewAllocaTy &&
            !SI.isVolatile();
if (!CanContinue) {
       Type *SizeTy = II.getLength()->getType();
       Constant *Size = ConstantInt::get(SizeTy, NewEndOffset - NewBeginOffset);
       MemIntrinsic *New = cast<MemIntrinsic>(IRB.CreateMemSet(
           getNewAllocaSlicePtr(IRB, OldPtr->getType()), II.getValue(), Size,
           MaybeAlign(getSliceAlign()), II.isVolatile()));
       if (AATags)
         New->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
 
-      migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &II, New,
-                       New->getRawDest(), nullptr, DL);
+      migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8, &II,
+                       New, New->getRawDest(), nullptr, DL);
 
       LLVM_DEBUG(dbgs() << "          to: " << *New << "\n");
       return false;
     }
 
     // If we can represent this as a simple value, we have to build the actual
     // value to store, which requires expanding the byte present in memset to
     // a sensible representation for the alloca type. This is essentially
bool visitMemSetInst(MemSetInst &II) {
     Value *NewPtr = getPtrToNewAI(II.getDestAddressSpace(), II.isVolatile());
     StoreInst *New =
         IRB.CreateAlignedStore(V, NewPtr, NewAI.getAlign(), II.isVolatile());
     New->copyMetadata(II, {LLVMContext::MD_mem_parallel_loop_access,
                            LLVMContext::MD_access_group});
     if (AATags)
       New->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
 
-    migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &II, New,
-                     New->getPointerOperand(), V, DL);
+    migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8, &II,
+                     New, New->getPointerOperand(), V, DL);
 
     LLVM_DEBUG(dbgs() << "          to: " << *New << "\n");
     return !II.isVolatile();
   }
 
   bool visitMemTransferInst(MemTransferInst &II) {
     // Rewriting of memory transfer instructions can be a bit tricky. We break
     // them into two categories: split intrinsics and unsplit intrinsics.
if (EmitMemCpy) {
         SrcPtr = OurPtr;
         SrcAlign = SliceAlign;
       }
       CallInst *New = IRB.CreateMemCpy(DestPtr, DestAlign, SrcPtr, SrcAlign,
                                        Size, II.isVolatile());
       if (AATags)
         New->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
 
-      migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &II, New,
-                       DestPtr, nullptr, DL);
+      APInt Offset(DL.getIndexTypeSizeInBits(DestPtr->getType()), 0);
+      if (IsDest) {
+        migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8,
+                         &II, New, DestPtr, nullptr, DL);
+      } else if (AllocaInst *Base = dyn_cast<AllocaInst>(
+                     DestPtr->stripAndAccumulateConstantOffsets(
+                         DL, Offset, /*AllowNonInbounds*/ true))) {
+        migrateDebugInfo(Base, IsSplit, Offset.getZExtValue() * 8,
+                         SliceSize * 8, &II, New, DestPtr, nullptr, DL);
+      }
       LLVM_DEBUG(dbgs() << "          to: " << *New << "\n");
       return false;
     }
 
     bool IsWholeAlloca = NewBeginOffset == NewAllocaBeginOffset &&
                          NewEndOffset == NewAllocaEndOffset;
     uint64_t Size = NewEndOffset - NewBeginOffset;
     unsigned BeginIndex = VecTy ? getIndex(NewBeginOffset) : 0;
bool visitMemTransferInst(MemTransferInst &II) {
 
     StoreInst *Store = cast<StoreInst>(
         IRB.CreateAlignedStore(Src, DstPtr, DstAlign, II.isVolatile()));
     Store->copyMetadata(II, {LLVMContext::MD_mem_parallel_loop_access,
                              LLVMContext::MD_access_group});
     if (AATags)
       Store->setAAMetadata(AATags.shift(NewBeginOffset - BeginOffset));
 
-    migrateDebugInfo(&OldAI, RelativeOffset * 8, SliceSize * 8, &II, Store,
-                     DstPtr, Src, DL);
+    APInt Offset(DL.getIndexTypeSizeInBits(DstPtr->getType()), 0);
+    if (IsDest) {
+
+      migrateDebugInfo(&OldAI, IsSplit, NewBeginOffset * 8, SliceSize * 8, &II,
+                       Store, DstPtr, Src, DL);
+    } else if (AllocaInst *Base = dyn_cast<AllocaInst>(
+                   DstPtr->stripAndAccumulateConstantOffsets(
+                       DL, Offset, /*AllowNonInbounds*/ true))) {
+      migrateDebugInfo(Base, IsSplit, Offset.getZExtValue() * 8, SliceSize * 8,
+                       &II, Store, DstPtr, Src, DL);
+    }
+
     LLVM_DEBUG(dbgs() << "          to: " << *Store << "\n");
     return !II.isVolatile();
   }
 
   bool visitIntrinsicInst(IntrinsicInst &II) {
     assert((II.isLifetimeStartOrEnd() || II.isDroppable()) &&
            "Unexpected intrinsic!");
     LLVM_DEBUG(dbgs() << "    original: " << II << "\n");
void emitFunc(Type *Ty, Value *&Agg, Align Alignment, const Twine &Name) {
           IRB.CreateExtractValue(Agg, Indices, Name + ".extract");
       Value *InBoundsGEP =
           IRB.CreateInBoundsGEP(BaseTy, Ptr, GEPIndices, Name + ".gep");
       StoreInst *Store =
           IRB.CreateAlignedStore(ExtractValue, InBoundsGEP, Alignment);
 
       APInt Offset(
           DL.getIndexSizeInBits(Ptr->getType()->getPointerAddressSpace()), 0);
-      if (AATags &&
-          GEPOperator::accumulateConstantOffset(BaseTy, GEPIndices, DL, Offset))
+      GEPOperator::accumulateConstantOffset(BaseTy, GEPIndices, DL, Offset);
+      if (AATags)
         Store->setAAMetadata(AATags.shift(Offset.getZExtValue()));
 
       // migrateDebugInfo requires the base Alloca. Walk to it from this gep.
       // If we cannot (because there's an intervening non-const or unbounded
       // gep) then we wouldn't expect to see dbg.assign intrinsics linked to
       // this instruction.
-      APInt OffsetInBytes(DL.getTypeSizeInBits(Ptr->getType()), false);
-      Value *Base = InBoundsGEP->stripAndAccumulateInBoundsConstantOffsets(
-          DL, OffsetInBytes);
+      Value *Base = AggStore->getPointerOperand()->stripInBoundsOffsets();
       if (auto *OldAI = dyn_cast<AllocaInst>(Base)) {
         uint64_t SizeInBits =
             DL.getTypeSizeInBits(Store->getValueOperand()->getType());
-        migrateDebugInfo(OldAI, OffsetInBytes.getZExtValue() * 8, SizeInBits,
-                         AggStore, Store, Store->getPointerOperand(),
-                         Store->getValueOperand(), DL);
+        migrateDebugInfo(OldAI, /*IsSplit*/ true, Offset.getZExtValue() * 8,
+                         SizeInBits, AggStore, Store,
+                         Store->getPointerOperand(), Store->getValueOperand(),
+                         DL);
       } else {
         assert(at::getAssignmentMarkers(Store).empty() &&
                "AT: unexpected debug.assign linked to store through "
                "unbounded GEP");
       }
       LLVM_DEBUG(dbgs() << "          to: " << *Store << "\n");
     }
   };

 /// the PromotableAllocas list. If that list is empty, there is nothing to do.
 /// This function returns whether any promotion occurred.
 bool SROAPass::promoteAllocas(Function &F) {
   if (PromotableAllocas.empty())
     return false;
 
   NumPromoted += PromotableAllocas.size();
 
+  if (SROASkipMem2Reg) {
+    LLVM_DEBUG(dbgs() << "Not promoting allocas with mem2reg!\n");
+  } else {
     LLVM_DEBUG(dbgs() << "Promoting allocas with mem2reg...\n");
     PromoteMemToReg(PromotableAllocas, DTU->getDomTree(), AC);
+  }
+
   PromotableAllocas.clear();
   return true;
 }
 
 PreservedAnalyses SROAPass::runImpl(Function &F, DomTreeUpdater &RunDTU,
                                     AssumptionCache &RunAC) {
   LLVM_DEBUG(dbgs() << "SROA function: " << F.getName() << "\n");
   C = &F.getContext();

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

+; RUN: opt -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
+;; variable - the variable's size has been modified from 64 to 96 bits.
+;; This version of the test ensures that the behaviour being tested is still
+;; correct when the base alloca doesn't hold an entire variable.
+
+;; IR hand-modified, originally generated from:
+;; struct Pair { int a; int b; };
+;; Pair getVar();
+;; int fun() {
+;;   Pair var;
+;;   var = getVar();
+;;   return var.b;
+;; }
+;; Modification: split the dbg.assign linked the the memcpy(64 bits) into two,
+;; each describing a 32 bit fragment.
+;;
+;; Check that assignment tracking updates in SROA work when the store being
+;; split is described with one dbg.assign (covering different fragments). The
+;; store may have been already split and then merged again at some point.
+
+;; Alloca for var.a and associated dbg.assign:
+; CHECK: %var.sroa.0 = alloca i32, align 4, !DIAssignID ![[id_1:[0-9]+]]
+; CHECK-NEXT: call void @llvm.dbg.assign(metadata i1 undef, metadata ![[var:[0-9]+]], metadata !DIExpression(DW_OP_LLVM_fragment, 32, 32), metadata ![[id_1]], metadata ptr %var.sroa.0, metadata !DIExpression())
+
+;; Alloca for var.b and associated dbg.assign:
+; CHECK-NEXT: %var.sroa.1 = alloca i32, align 4, !DIAssignID ![[id_2:[0-9]+]]
+; CHECK-NEXT: call void @llvm.dbg.assign(metadata i1 undef, metadata ![[var]], metadata !DIExpression(DW_OP_LLVM_fragment, 64, 32), metadata ![[id_2]], metadata ptr %var.sroa.1, metadata !DIExpression())
+
+;; Store to var.b (split from store to var) and associated dbg.assigns. The
+;; dbg.assign for the fragment covering the (pre-split) assignment to var.a
+;; should not be linked to the store.
+; CHECK: store i32 %[[v:.*]], ptr %var.sroa.1,{{.*}}!DIAssignID ![[id_3:[0-9]+]]
+; CHECK-NEXT: call void @llvm.dbg.assign(metadata i32 %{{.*var\.sroa\.0.*}}, metadata ![[var]], metadata !DIExpression(DW_OP_LLVM_fragment, 32, 32), metadata ![[id_4:[0-9]+]], metadata ptr %var.sroa.0, metadata !DIExpression())
+; CHECK-NEXT: call void @llvm.dbg.assign(metadata i32 %[[v]], metadata ![[var]], metadata !DIExpression(DW_OP_LLVM_fragment, 64, 32), metadata ![[id_3]], metadata ptr %var.sroa.1, metadata !DIExpression())
+
+; CHECK-DAG: ![[id_1]] = distinct !DIAssignID()
+; CHECK-DAG: ![[id_2]] = distinct !DIAssignID()
+; CHECK-DAG: ![[id_3]] = distinct !DIAssignID()
+; CHECK-DAG: ![[id_4]] = distinct !DIAssignID()
+
+%struct.Tuple = type { i32, i32, i32 }
+
+define dso_local noundef i32 @_Z3funv() !dbg !9 {
+entry:
+  %var = alloca [2 x i32], align 4, !DIAssignID !19
+  call void @llvm.dbg.assign(metadata i1 undef, metadata !14, metadata !DIExpression(DW_OP_LLVM_fragment, 32, 64), metadata !19, metadata ptr %var, metadata !DIExpression()), !dbg !20
+  %ref.tmp = alloca %struct.Tuple, align 4
+  %call = call i64 @_Z6getVarv(), !dbg !22
+  store i64 %call, ptr %ref.tmp, align 4, !dbg !22
+  call void @llvm.memcpy.p0.p0.i64(ptr align 4 %var, ptr align 4 %ref.tmp, i64 8, i1 false), !dbg !23, !DIAssignID !29
+  call void @llvm.dbg.assign(metadata i1 undef, metadata !14, metadata !DIExpression(DW_OP_LLVM_fragment, 32, 32), metadata !29, metadata ptr %var, metadata !DIExpression()), !dbg !20
+  call void @llvm.dbg.assign(metadata i1 undef, metadata !14, metadata !DIExpression(DW_OP_LLVM_fragment, 64, 32), metadata !29, metadata ptr %var, metadata !DIExpression(DW_OP_plus, DW_OP_constu, 4)), !dbg !20
+  %b = getelementptr inbounds %struct.Tuple, ptr %var, i32 0, i32 1, !dbg !31
+  %0 = load i32, ptr %b, align 4, !dbg !31
+  ret i32 %0, !dbg !35
+}
+
+declare !dbg !36 i64 @_Z6getVarv()
+declare void @llvm.dbg.declare(metadata, metadata, metadata)
+declare void @llvm.memcpy.p0.p0.i64(ptr noalias nocapture writeonly, ptr noalias nocapture readonly, i64, i1 immarg)
+declare void @llvm.dbg.assign(metadata, metadata, metadata, metadata, metadata, metadata)
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!2, !3, !4, !5, !6, !7}
+!llvm.ident = !{!8}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1, producer: "clang version 16.0.0", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, splitDebugInlining: false, nameTableKind: None)
+!1 = !DIFile(filename: "test.cpp", directory: "/")
+!2 = !{i32 7, !"Dwarf Version", i32 5}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!4 = !{i32 1, !"wchar_size", i32 4}
+!5 = !{i32 8, !"PIC Level", i32 2}
+!6 = !{i32 7, !"PIE Level", i32 2}
+!7 = !{i32 7, !"uwtable", i32 2}
+!8 = !{!"clang version 16.0.0"}
+!9 = distinct !DISubprogram(name: "fun", linkageName: "_Z3funv", scope: !1, file: !1, line: 3, type: !10, scopeLine: 3, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !13)
+!10 = !DISubroutineType(types: !11)
+!11 = !{!12}
+!12 = !DIBasicType(name: "int", size: 32, encoding: DW_ATE_signed)
+!13 = !{!14}
+!14 = !DILocalVariable(name: "var", scope: !9, file: !1, line: 4, type: !15)
+!15 = distinct !DICompositeType(tag: DW_TAG_structure_type, name: "Tuple", file: !1, line: 1, size: 96, flags: DIFlagTypePassByValue, elements: !16, identifier: "_ZTS4Pair")
+!16 = !{!17, !18, !40}
+!17 = !DIDerivedType(tag: DW_TAG_member, name: "a", scope: !15, file: !1, line: 1, baseType: !12, size: 32)
+!18 = !DIDerivedType(tag: DW_TAG_member, name: "b", scope: !15, file: !1, line: 1, baseType: !12, size: 32, offset: 32)
+!19 = distinct !DIAssignID()
+!20 = !DILocation(line: 0, scope: !9)
+!21 = !DILocation(line: 4, column: 3, scope: !9)
+!22 = !DILocation(line: 5, column: 9, scope: !9)
+!23 = !DILocation(line: 5, column: 7, scope: !9)
+!29 = distinct !DIAssignID()
+!30 = !DILocation(line: 5, column: 3, scope: !9)
+!31 = !DILocation(line: 6, column: 14, scope: !9)
+!34 = !DILocation(line: 7, column: 1, scope: !9)
+!35 = !DILocation(line: 6, column: 3, scope: !9)
+!36 = !DISubprogram(name: "getVar", linkageName: "_Z6getVarv", scope: !1, file: !1, line: 2, type: !37, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized, retainedNodes: !39)
+!37 = !DISubroutineType(types: !38)
+!38 = !{!15}
+!39 = !{}
+!40 = !DIDerivedType(tag: DW_TAG_member, name: "c", scope: !15, file: !1, line: 1, baseType: !12, size: 32, offset: 64)

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

+; RUN: opt -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; };
+;; Pair getVar();
+;; int fun() {
+;;   Pair var;
+;;   var = getVar();
+;;   return var.b;
+;; }
+;; Modification: split the dbg.assign linked the the memcpy(64 bits) into two,
+;; each describing a 32 bit fragment.
+;;
+;; Check that assignment tracking updates in SROA work when the store being
+;; split is described with one dbg.assign (covering different fragments). The
+;; store may have been already split and then merged again at some point.
+
+;; Alloca for var.a and associated dbg.assign:
+; CHECK: %var.sroa.0 = alloca i32, align 4, !DIAssignID ![[id_1:[0-9]+]]
+; CHECK-NEXT: call void @llvm.dbg.assign(metadata i1 undef, metadata ![[var:[0-9]+]], metadata !DIExpression(DW_OP_LLVM_fragment, 0, 32), metadata ![[id_1]], metadata ptr %var.sroa.0, metadata !DIExpression())
+
+;; Alloca for var.b and associated dbg.assign:
+; CHECK-NEXT: %var.sroa.1 = alloca i32, align 4, !DIAssignID ![[id_2:[0-9]+]]
+; CHECK-NEXT: call void @llvm.dbg.assign(metadata i1 undef, metadata ![[var]], metadata !DIExpression(DW_OP_LLVM_fragment, 32, 32), metadata ![[id_2]], metadata ptr %var.sroa.1, metadata !DIExpression())
+
+;; Store to var.b (split from store to var) and associated dbg.assigns. The
+;; dbg.assign for the fragment covering the (pre-split) assignment to var.a
+;; should not be linked to the store.
+; CHECK: store i32 %[[v:.*]], ptr %var.sroa.1,{{.*}}!DIAssignID ![[id_3:[0-9]+]]
+; CHECK-NEXT: call void @llvm.dbg.assign(metadata i32 %{{.*var\.sroa\.0.*}}, metadata ![[var]], metadata !DIExpression(DW_OP_LLVM_fragment, 0, 32), metadata ![[id_4:[0-9]+]], metadata ptr %var.sroa.0, metadata !DIExpression())
+; CHECK-NEXT: call void @llvm.dbg.assign(metadata i32 %[[v]], metadata ![[var]], metadata !DIExpression(DW_OP_LLVM_fragment, 32, 32), metadata ![[id_3]], metadata ptr %var.sroa.1, metadata !DIExpression())
+
+; CHECK-DAG: ![[id_1]] = distinct !DIAssignID()
+; CHECK-DAG: ![[id_2]] = distinct !DIAssignID()
+; CHECK-DAG: ![[id_3]] = distinct !DIAssignID()
+; CHECK-DAG: ![[id_4]] = distinct !DIAssignID()
+
+%struct.Pair = type { i32, i32 }
+
+define dso_local noundef i32 @_Z3funv() !dbg !9 {
+entry:
+  %var = alloca %struct.Pair, align 4, !DIAssignID !19
+  call void @llvm.dbg.assign(metadata i1 undef, metadata !14, metadata !DIExpression(), metadata !19, metadata ptr %var, metadata !DIExpression()), !dbg !20
+  %ref.tmp = alloca %struct.Pair, align 4
+  %call = call i64 @_Z6getVarv(), !dbg !22
+  store i64 %call, ptr %ref.tmp, align 4, !dbg !22
+  call void @llvm.memcpy.p0.p0.i64(ptr align 4 %var, ptr align 4 %ref.tmp, i64 8, i1 false), !dbg !23, !DIAssignID !29
+  call void @llvm.dbg.assign(metadata i1 undef, metadata !14, metadata !DIExpression(DW_OP_LLVM_fragment, 0, 32), metadata !29, metadata ptr %var, metadata !DIExpression()), !dbg !20
+  call void @llvm.dbg.assign(metadata i1 undef, metadata !14, metadata !DIExpression(DW_OP_LLVM_fragment, 32, 32), metadata !29, metadata ptr %var, metadata !DIExpression(DW_OP_plus, DW_OP_constu, 4)), !dbg !20
+  %b = getelementptr inbounds %struct.Pair, ptr %var, i32 0, i32 1, !dbg !31
+  %0 = load i32, ptr %b, align 4, !dbg !31
+  ret i32 %0, !dbg !35
+}
+
+declare !dbg !36 i64 @_Z6getVarv()
+declare void @llvm.dbg.declare(metadata, metadata, metadata)
+declare void @llvm.memcpy.p0.p0.i64(ptr noalias nocapture writeonly, ptr noalias nocapture readonly, i64, i1 immarg)
+declare void @llvm.dbg.assign(metadata, metadata, metadata, metadata, metadata, metadata)
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!2, !3, !4, !5, !6, !7}
+!llvm.ident = !{!8}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1, producer: "clang version 16.0.0", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, splitDebugInlining: false, nameTableKind: None)
+!1 = !DIFile(filename: "test.cpp", directory: "/")
+!2 = !{i32 7, !"Dwarf Version", i32 5}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!4 = !{i32 1, !"wchar_size", i32 4}
+!5 = !{i32 8, !"PIC Level", i32 2}
+!6 = !{i32 7, !"PIE Level", i32 2}
+!7 = !{i32 7, !"uwtable", i32 2}
+!8 = !{!"clang version 16.0.0"}
+!9 = distinct !DISubprogram(name: "fun", linkageName: "_Z3funv", scope: !1, file: !1, line: 3, type: !10, scopeLine: 3, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !13)
+!10 = !DISubroutineType(types: !11)
+!11 = !{!12}
+!12 = !DIBasicType(name: "int", size: 32, encoding: DW_ATE_signed)
+!13 = !{!14}
+!14 = !DILocalVariable(name: "var", scope: !9, file: !1, line: 4, type: !15)
+!15 = distinct !DICompositeType(tag: DW_TAG_structure_type, name: "Pair", file: !1, line: 1, size: 64, flags: DIFlagTypePassByValue, elements: !16, identifier: "_ZTS4Pair")
+!16 = !{!17, !18}
+!17 = !DIDerivedType(tag: DW_TAG_member, name: "a", scope: !15, file: !1, line: 1, baseType: !12, size: 32)
+!18 = !DIDerivedType(tag: DW_TAG_member, name: "b", scope: !15, file: !1, line: 1, baseType: !12, size: 32, offset: 32)
+!19 = distinct !DIAssignID()
+!20 = !DILocation(line: 0, scope: !9)
+!21 = !DILocation(line: 4, column: 3, scope: !9)
+!22 = !DILocation(line: 5, column: 9, scope: !9)
+!23 = !DILocation(line: 5, column: 7, scope: !9)
+!29 = distinct !DIAssignID()
+!30 = !DILocation(line: 5, column: 3, scope: !9)
+!31 = !DILocation(line: 6, column: 14, scope: !9)
+!34 = !DILocation(line: 7, column: 1, scope: !9)
+!35 = !DILocation(line: 6, column: 3, scope: !9)
+!36 = !DISubprogram(name: "getVar", linkageName: "_Z6getVarv", scope: !1, file: !1, line: 2, type: !37, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized, retainedNodes: !39)
+!37 = !DISubroutineType(types: !38)
+!38 = !{!15}
+!39 = !{}