[SROA] Fix incorrect offsets for structured binding variables

llvm/include/llvm/IR/DebugInfo.h

@@ -226,7 +226,7 @@ void RAUW(DIAssignID *Old, DIAssignID *New);
 /// Remove all Assignment Tracking related intrinsics and metadata from \p F.
 void deleteAll(Function *F);
 
-/// Calculate the fragment of the variable in \p DAI covered
+/// Calculate the fragment of the variable in \p DVI covered
 /// from (Dest + SliceOffsetInBits) to
 ///   to (Dest + SliceOffsetInBits + SliceSizeInBits)
 ///
@@ -235,10 +235,15 @@ void deleteAll(Function *F);
 /// variable size) in DAI.
 ///
 /// Result contains a zero-sized fragment if there's no intersect.
+/// \p DVI may be either a DbgDeclareInst or a DbgAssignIntrinsic.
+///
+/// \p NewExprOffsetInBits is set to the difference between the first bit of
+/// memory the fragment describes and the first bit of the slice.
 bool calculateFragmentIntersect(
     const DataLayout &DL, const Value *Dest, uint64_t SliceOffsetInBits,
-    uint64_t SliceSizeInBits, const DbgAssignIntrinsic *DAI,
-    std::optional<DIExpression::FragmentInfo> &Result);
+    uint64_t SliceSizeInBits, const DbgVariableIntrinsic *DVI,
+    std::optional<DIExpression::FragmentInfo> &Result,
+    uint64_t &NewExprOffsetInBits);
 
 /// Helper struct for trackAssignments, below. We don't use the similar
 /// DebugVariable class because trackAssignments doesn't (yet?) understand

llvm/lib/CodeGen/AssignmentTrackingAnalysis.cpp

@@ -1981,12 +1981,15 @@ static AssignmentTrackingLowering::OverlapMap buildOverlapMapAndRecordDeclares(
         for (DbgAssignIntrinsic *DAI : at::getAssignmentMarkers(Info->Base)) {
           std::optional<DIExpression::FragmentInfo> FragInfo;
 
+          uint64_t NewExprOffsetInBits;
           // Skip this assignment if the affected bits are outside of the
           // variable fragment.
-          if (!at::calculateFragmentIntersect(
-                  I.getModule()->getDataLayout(), Info->Base,
-                  Info->OffsetInBits, Info->SizeInBits, DAI, FragInfo) ||
-              (FragInfo && FragInfo->SizeInBits == 0))
+          if (!at::calculateFragmentIntersect(I.getModule()->getDataLayout(),
+                                              Info->Base, Info->OffsetInBits,
+                                              Info->SizeInBits, DAI, FragInfo,
+                                              NewExprOffsetInBits) ||
+              (FragInfo && FragInfo->SizeInBits == 0) ||
+              NewExprOffsetInBits > 0)
             continue;
 
           // FragInfo from calculateFragmentIntersect is nullopt if the

llvm/lib/IR/DebugInfo.cpp

@@ -1782,8 +1782,34 @@ void at::deleteAll(Function *F) {
 
 bool at::calculateFragmentIntersect(
     const DataLayout &DL, const Value *Dest, uint64_t SliceOffsetInBits,
-    uint64_t SliceSizeInBits, const DbgAssignIntrinsic *DAI,
-    std::optional<DIExpression::FragmentInfo> &Result) {
+    uint64_t SliceSizeInBits, const DbgVariableIntrinsic *DVI,
+    std::optional<DIExpression::FragmentInfo> &Result,
+    uint64_t &NewExprOffsetInBits) {
+
+  // Only dbg.assign and dbg.declares are allowed because this function
+  // deals with memory locations. This isn't comprehensive because dbg.values
+  // are able to describe memory locations; support for dbg.values can be added
+  // if/when needed.
+  assert(isa<DbgAssignIntrinsic>(DVI) || isa<DbgDeclareInst>(DVI));
+
+  // There isn't a shared interface to get the "address" parts out of a
+  // dbg.declare and dbg.assign, so provide some wrappers now.
+  auto GetAddress = [DVI]() -> const Value * {
+    if (const auto *DAI = dyn_cast<DbgAssignIntrinsic>(DVI))
+      return DAI->getAddress();
+    return cast<DbgDeclareInst>(DVI)->getAddress();
+  };
+  auto IsKillLocation = [DVI]() -> bool {
+    if (const auto *DAI = dyn_cast<DbgAssignIntrinsic>(DVI))
+      return DAI->isKillAddress();
+    return cast<DbgDeclareInst>(DVI)->isKillLocation();
+  };
+  auto GetExpression = [DVI]() -> const DIExpression * {
+    if (const auto *DAI = dyn_cast<DbgAssignIntrinsic>(DVI))
+      return DAI->getAddressExpression();
+    return cast<DbgDeclareInst>(DVI)->getExpression();
+  };
+
   // There are multiple offsets at play in this function, so let's break it
   // down. Starting with how variables may be stored in allocas:
   //
@@ -1874,23 +1900,23 @@ bool at::calculateFragmentIntersect(
   //      SliceOfVariable ∩ DAI_fragment = (offset: 144, size: 16)
   // SliceOfVariable tells us the bits of the variable described by DAI that are
   // affected by the DSE.
-  if (DAI->isKillAddress())
+  if (IsKillLocation())
     return false;
 
-  DIExpression::FragmentInfo VarFrag = DAI->getFragmentOrEntireVariable();
+  DIExpression::FragmentInfo VarFrag = DVI->getFragmentOrEntireVariable();
   if (VarFrag.SizeInBits == 0)
     return false; // Variable size is unknown.
 
   // Calculate the difference between Dest and the dbg.assign address +
   // address-modifying expression.
   int64_t PointerOffsetInBits;
   {
-    auto DestOffsetInBytes = DAI->getAddress()->getPointerOffsetFrom(Dest, DL);
+    auto DestOffsetInBytes = GetAddress()->getPointerOffsetFrom(Dest, DL);
     if (!DestOffsetInBytes)
       return false; // Can't calculate difference in addresses.
 
     int64_t ExprOffsetInBytes;
-    if (!DAI->getAddressExpression()->extractIfOffset(ExprOffsetInBytes))
+    if (!GetExpression()->extractIfOffset(ExprOffsetInBytes))
       return false;
 
     int64_t PointerOffsetInBytes = *DestOffsetInBytes + ExprOffsetInBytes;
@@ -1899,11 +1925,19 @@ bool at::calculateFragmentIntersect(
 
   // Adjust the slice offset so that we go from describing the a slice
   // of memory to a slice of the variable.
-  int64_t NewOffsetInBits =
+  int64_t AdjustedSliceOffsetInBits =
       SliceOffsetInBits + VarFrag.OffsetInBits - PointerOffsetInBits;
-  if (NewOffsetInBits < 0)
-    return false; // Fragment offsets can only be positive.
-  DIExpression::FragmentInfo SliceOfVariable(SliceSizeInBits, NewOffsetInBits);
+  NewExprOffsetInBits = std::max(0l, -AdjustedSliceOffsetInBits);
+  AdjustedSliceOffsetInBits = std::max(0l, AdjustedSliceOffsetInBits);
+
+  // Check if the variable fragment sits outside this memory slice.
+  if (NewExprOffsetInBits >= SliceSizeInBits) {
+    Result = {0, 0};
+    return true;
+  }
+
+  DIExpression::FragmentInfo SliceOfVariable(
+      SliceSizeInBits - NewExprOffsetInBits, AdjustedSliceOffsetInBits);
   // Intersect the variable slice with DAI's fragment to trim it down to size.
   DIExpression::FragmentInfo TrimmedSliceOfVariable =
       DIExpression::FragmentInfo::intersect(SliceOfVariable, VarFrag);

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

@@ -530,11 +530,12 @@ static void shortenAssignment(Instruction *Inst, Value *OriginalDest,
   SmallVector<DbgAssignIntrinsic *> Linked(LinkedRange.begin(),
                                            LinkedRange.end());
   for (auto *DAI : Linked) {
+    uint64_t NewExprOffsetInBits;
     std::optional<DIExpression::FragmentInfo> NewFragment;
     if (!at::calculateFragmentIntersect(DL, OriginalDest, DeadSliceOffsetInBits,
-                                        DeadSliceSizeInBits, DAI,
-                                        NewFragment) ||
-        !NewFragment) {
+                                        DeadSliceSizeInBits, DAI, NewFragment,
+                                        NewExprOffsetInBits) ||
+        !NewFragment || NewExprOffsetInBits > 0) {
       // We couldn't calculate the intersecting fragment for some reason. Be
       // cautious and unlink the whole assignment from the store.
       DAI->setKillAddress();

llvm/lib/Transforms/Scalar/SROA.cpp

@@ -4813,54 +4813,43 @@ bool SROAPass::splitAlloca(AllocaInst &AI, AllocaSlices &AS) {
   for (auto *DbgAssign : at::getAssignmentMarkers(&AI))
     DbgVariables.push_back(DbgAssign);
   for (DbgVariableIntrinsic *DbgVariable : DbgVariables) {
-    auto *Expr = DbgVariable->getExpression();
     DIBuilder DIB(*AI.getModule(), /*AllowUnresolved*/ false);
-    uint64_t AllocaSize =
-        DL.getTypeSizeInBits(AI.getAllocatedType()).getFixedValue();
-    for (auto Fragment : Fragments) {
-      // Create a fragment expression describing the new partition or reuse AI's
-      // expression if there is only one partition.
-      auto *FragmentExpr = Expr;
-      if (Fragment.Size < AllocaSize || Expr->isFragment()) {
-        // If this alloca is already a scalar replacement of a larger aggregate,
-        // Fragment.Offset describes the offset inside the scalar.
-        auto ExprFragment = Expr->getFragmentInfo();
-        uint64_t Offset = ExprFragment ? ExprFragment->OffsetInBits : 0;
-        uint64_t Start = Offset + Fragment.Offset;
-        uint64_t Size = Fragment.Size;
-        if (ExprFragment) {
-          uint64_t AbsEnd =
-              ExprFragment->OffsetInBits + ExprFragment->SizeInBits;
-          if (Start >= AbsEnd) {
-            // No need to describe a SROAed padding.
-            continue;
-          }
-          Size = std::min(Size, AbsEnd - Start);
-        }
-        // The new, smaller fragment is stenciled out from the old fragment.
-        if (auto OrigFragment = FragmentExpr->getFragmentInfo()) {
-          assert(Start >= OrigFragment->OffsetInBits &&
-                 "new fragment is outside of original fragment");
-          Start -= OrigFragment->OffsetInBits;
-        }
 
-        // The alloca may be larger than the variable.
-        auto VarSize = DbgVariable->getVariable()->getSizeInBits();
-        if (VarSize) {
-          if (Size > *VarSize)
-            Size = *VarSize;
-          if (Size == 0 || Start + Size > *VarSize)
-            continue;
-        }
-
-        // Avoid creating a fragment expression that covers the entire variable.
-        if (!VarSize || *VarSize != Size) {
-          if (auto E =
-                  DIExpression::createFragmentExpression(Expr, Start, Size))
-            FragmentExpr = *E;
-          else
-            continue;
-        }
+    for (auto Fragment : Fragments) {
+      uint64_t OffestFromNewAllocaInBits;
+      std::optional<DIExpression::FragmentInfo> NewDbgFragment;
+
+      // Drop debug info for this variable fragment if we can't compute an
+      // intersect between it and the alloca slice.
+      if (!at::calculateFragmentIntersect(
+              DL, &AI, Fragment.Offset, Fragment.Size, DbgVariable,
+              NewDbgFragment, OffestFromNewAllocaInBits))
+        continue; // Do not migrate this fragment to this slice.
+
+      // Zero sized fragment indicates there's no intersect between the variable
+      // fragment and the alloca slice. Skip this slice for this variable
+      // fragment.
+      if (NewDbgFragment && !NewDbgFragment->SizeInBits)
+        continue; // Do not migrate this fragment to this slice.
+
+      // No fragment indicates DbgVariable's variable or fragment exactly
+      // overlaps the slice; copy its fragment (or nullopt if there isn't one).
+      if (!NewDbgFragment)
+        NewDbgFragment = DbgVariable->getFragment();
+
+      // calculateFragmentIntersect fails if DbgVariable's expression is not a
+      // trivial offset expression, meaning it must contains only an offset and
+      // fragment. Start from scratch; add the fragment and then the offset.
+      // If calculateFragmentIntersect gets smarter this needs updating
+      // (sroa-alloca-offset.ll should start failing) - we'd need to copy the
+      // other parts of the original expression.
+      DIExpression *NewExpr = DIExpression::get(AI.getContext(), {});
+      if (NewDbgFragment)
+        NewExpr = *DIExpression::createFragmentExpression(
+            NewExpr, NewDbgFragment->OffsetInBits, NewDbgFragment->SizeInBits);
+      if (OffestFromNewAllocaInBits > 0) {
+        int64_t OffsetInBytes = (OffestFromNewAllocaInBits + 7) / 8;
+        NewExpr = DIExpression::prepend(NewExpr, /*flags=*/0, OffsetInBytes);
       }
 
       // Remove any existing intrinsics on the new alloca describing
@@ -4884,14 +4873,14 @@ bool SROAPass::splitAlloca(AllocaInst &AI, AllocaSlices &AS) {
         }
         auto *NewAssign = DIB.insertDbgAssign(
             Fragment.Alloca, DbgAssign->getValue(), DbgAssign->getVariable(),
-            FragmentExpr, Fragment.Alloca, DbgAssign->getAddressExpression(),
+            NewExpr, Fragment.Alloca, DbgAssign->getAddressExpression(),
             DbgAssign->getDebugLoc());
         NewAssign->setDebugLoc(DbgAssign->getDebugLoc());
         LLVM_DEBUG(dbgs() << "Created new assign intrinsic: " << *NewAssign
                           << "\n");
       } else {
-        DIB.insertDeclare(Fragment.Alloca, DbgVariable->getVariable(),
-                          FragmentExpr, DbgVariable->getDebugLoc(), &AI);
+        DIB.insertDeclare(Fragment.Alloca, DbgVariable->getVariable(), NewExpr,
+                          DbgVariable->getDebugLoc(), &AI);
       }
     }
   }

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

@@ -18,10 +18,7 @@
 ;;   return a;
 ;; }
 
-;; FIXME: Variable 'b' gets an incorrect location (value and expression) - see
-;; llvm.org/PR61981. This check just ensures that no fragment info is added to
-;; the dbg.value.
-; CHECK: dbg.value(metadata i32 %.sroa.0.0.extract.trunc, metadata ![[B:[0-9]+]], metadata !DIExpression(DW_OP_plus_uconst, 4))
+; CHECK: dbg.value(metadata i32 %.sroa.2.0.extract.trunc, metadata ![[B:[0-9]+]], metadata !DIExpression())
 ; CHECK: dbg.value(metadata i32 %.sroa.0.0.extract.trunc, metadata ![[A:[0-9]+]], metadata !DIExpression())
 ; CHECK: ![[A]] = !DILocalVariable(name: "a",
 ; CHECK: ![[B]] = !DILocalVariable(name: "b",