[Assignment Tracking][SROA] Fix fragment when slice size equals variable size

[Orlando]

Summary

Correctly handle the case of splitting an alloca which backs contiguous distinct variables, where a slice's size equals the size of a backed variable.

We need to ensure that we don't generate fragments expressions with fragments of the same size as the variable as this is a verifier error.

Prior to this patch a fragment expression would be created in this situation. e.g. splitting an alloca i64 with two adjacent 32-bit variables into two 32-bit allocas, the new dbg.assign expressions would contain (DW_OP_LLVM_fragment, 0, 32) and (DW_OP_LLVM_fragment, 32, 32) even though those fragments cover each variable entirely.

[jmorse]

I've got some quibbles about the code layout and wording, but otherwise this LGTM.

llvm/lib/Transforms/Scalar/SROA.cpp

153–160

I feel this could be more terse, something like "If this slice extracts the entirety of an independent variable from a larger alloca, do not produce a fragment expression, as the variable is not fragmented". IMO it's better to be terse than to give too much detail, YMMV.

256–261

IMO: using if UseNoFrag ...else if Skip... would allow UseNoFrag and Skip to be handled up front, leaving UseFrag as the only onwards path which avoids indenting the rest of the block below. We can rely on optimisations to make this just as fast as a switch, IMO.

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

8

22

Does the upper 32 bits not get a dbg.value(poison) at least? Best to test for that not having an expression if it does get generated, as that's the purpose of this test.

[Orlando]

+ Address some comments and leave some inline replies

llvm/lib/Transforms/Scalar/SROA.cpp

153–160

SGTM

256–261

Both UseNoFrag and UseFrag need to reach code outside the outer if (f (IsSplit)), but we only want the fragment-making code to occur when we UseFrag. I've restructured it to avoid a new additional layer of indentation - are you happy with this?

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

22

Unfortunately the upper 32 bits (which are tracked with a dbg.declare) get incorrect debug-info. See https://github.com/llvm/llvm-project/issues/61981. This issue is independent of assignment tracking.

[Orlando]

Fix typo and undo formatting change that was previously required due to additional indentation.

[jmorse]

llvm/lib/Transforms/Scalar/SROA.cpp

256–261

SGTM, it was only the extra indentation that bothered me.

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

22

Unfortunately the upper 32 bits (which are tracked with a dbg.declare) get incorrect debug-info. See https://github.com/llvm/llvm-project/issues/61981. This issue is independent of assignment tracking.

Sure, but the fragment information for that is orthogonal surely, we can test for the wrong location but right fragment information, with a "FIXME" saying the location is wrong and a link to the bug? That'll be updated when the bug is fixed, and means we still achieve coverage of the upper bits of this alloca that's being split up either way.

[Orlando]

+ Check debug intrinsic for upper bits (variable b) in the test doesn't get a fragment expression

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

22

SGTM, how does this look?

[Orlando]

+ Update comment in the test.

[jmorse]

LGTM

llvm/lib/Transforms/Scalar/SROA.cpp

 /// 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,
+/// Results:
+/// UseFrag - Use Target as the new fragment.
+/// UseNoFrag - The new slice already covers the whole variable.
+/// Skip - The new alloca slice doesn't include this variable.
+enum FragCalcResult { UseFrag, UseNoFrag, Skip };
+static FragCalcResult
+calculateFragment(DILocalVariable *Variable,
+                  uint64_t NewStorageSliceOffsetInBits,
                   uint64_t NewStorageSliceSizeInBits,
                   std::optional<DIExpression::FragmentInfo> StorageFragment,
-                  std::optional<DIExpression::FragmentInfo> CurrentFragment) {
-  DIExpression::FragmentInfo Target;
+                  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;
   }
 
+  // If this slice extracts the entirety of an independent variable from a
+  // larger alloca, do not produce a fragment expression, as the variable is
+  // not fragmented.
+  if (!CurrentFragment) {
+    if (auto Size = Variable->getSizeInBits()) {
+      // Treat the current fragment as covering the whole variable.
+      CurrentFragment =  DIExpression::FragmentInfo(*Size, 0);
+      if (Target == CurrentFragment)
+        return UseNoFrag;
+    }
+  }
+
   // 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;
+    return UseFrag;
 
   // 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;
+    return Skip;
 
   // Target fits within the current fragment, return it.
-  return Target;
+  return UseFrag;
 }
 
 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
if (IsSplit) {
       {
         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;
+      DIExpression::FragmentInfo NewFragment;
+      FragCalcResult Result = calculateFragment(
+          DbgAssign->getVariable(), OldAllocaOffsetInBits, SliceSizeInBits,
+          BaseFragment, CurrentFragment, NewFragment);
 
-      if (!(NewFragment == CurrentFragment)) {
+      if (Result == Skip)
+        continue;
+      if (Result == UseFrag && !(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;
+          NewFragment.OffsetInBits -= CurrentFragment->OffsetInBits;
         }
         // Add the new fragment info to the existing expression if possible.
         if (auto E = DIExpression::createFragmentExpression(
-                Expr, NewFragment->OffsetInBits, NewFragment->SizeInBits)) {
+                Expr, NewFragment.OffsetInBits, NewFragment.SizeInBits)) {
           Expr = *E;
         } else {
           // Otherwise, add the new fragment info to an empty expression and
           // discard the value component of this dbg.assign as the value cannot
           // be computed with the new fragment.
           Expr = *DIExpression::createFragmentExpression(
               DIExpression::get(Expr->getContext(), std::nullopt),
-              NewFragment->OffsetInBits, NewFragment->SizeInBits);
+              NewFragment.OffsetInBits, NewFragment.SizeInBits);
           SetKillLocation = true;
         }
       }
     }
 
     // If we haven't created a DIAssignID ID do that now and attach it to Inst.
     if (!NewID) {
       NewID = DIAssignID::getDistinct(Ctx);

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

+; RUN: opt -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
+;; mistakenly create 32 bit fragments for them.
+
+;; NOTE: The upper 32 bits are currently
+;; not tracked with assignment tracking (due to the dbg.declare containing an
+;; expression).
+;; The dbg intrinsics for the unnamed aggregate variable have been commented
+;; out to reduce test clutter.
+
+;; From C++ source:
+;; class two {public:int a; int b;}
+;; two get();
+;; int fun() {
+;;   auto [a, b] = get();
+;;   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.0.0.extract.trunc, metadata ![[A:[0-9]+]], metadata !DIExpression())
+; CHECK: ![[A]] = !DILocalVariable(name: "a",
+; CHECK: ![[B]] = !DILocalVariable(name: "b",
+
+%class.two = type { i32, i32 }
+
+define dso_local noundef i32 @_Z3funv() #0 !dbg !10 {
+entry:
+  %0 = alloca %class.two, align 4, !DIAssignID !23
+  ;;call void @llvm.dbg.assign(metadata i1 undef, metadata !17, metadata !DIExpression(), metadata !23, metadata ptr %0, metadata !DIExpression()), !dbg !24
+  call void @llvm.dbg.assign(metadata i1 undef, metadata !15, metadata !DIExpression(), metadata !23, metadata ptr %0, metadata !DIExpression()), !dbg !24
+  call void @llvm.dbg.declare(metadata ptr %0, metadata !16, metadata !DIExpression(DW_OP_plus_uconst, 4)), !dbg !26
+  %call = call i64 @_Z3getv()
+  store i64 %call, ptr %0, align 4, !DIAssignID !28
+  ;;call void @llvm.dbg.assign(metadata i64 %call, metadata !17, metadata !DIExpression(), metadata !28, metadata ptr %0, metadata !DIExpression()), !dbg !24
+  call void @llvm.dbg.assign(metadata i64 %call, metadata !15, metadata !DIExpression(), metadata !28, metadata ptr %0, metadata !DIExpression()), !dbg !24
+  %a = getelementptr inbounds %class.two, ptr %0, i32 0, i32 0
+  %1 = load i32, ptr %a, align 4
+  ret i32 %1
+}
+
+declare void @llvm.dbg.declare(metadata, metadata, metadata) #2
+declare !dbg !38 i64 @_Z3getv() #3
+declare void @llvm.dbg.assign(metadata, metadata, metadata, metadata, metadata, metadata) #2
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!2, !3, !4, !5, !6, !7, !8}
+!llvm.ident = !{!9}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1, producer: "clang version 17.0.0", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, splitDebugInlining: false, nameTableKind: None)
+!1 = !DIFile(filename: "test.cpp", directory: "h")
+!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 = !{i32 7, !"debug-info-assignment-tracking", i1 true}
+!9 = !{!"clang version 17.0.0"}
+!10 = distinct !DISubprogram(name: "fun", linkageName: "_Z3funv", scope: !1, file: !1, line: 3, type: !11, scopeLine: 3, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !14)
+!11 = !DISubroutineType(types: !12)
+!12 = !{!13}
+!13 = !DIBasicType(name: "int", size: 32, encoding: DW_ATE_signed)
+!14 = !{!15, !16, !17}
+!15 = !DILocalVariable(name: "a", scope: !10, file: !1, line: 4, type: !13)
+!16 = !DILocalVariable(name: "b", scope: !10, file: !1, line: 4, type: !13)
+!17 = !DILocalVariable(scope: !10, file: !1, line: 4, type: !18)
+!18 = distinct !DICompositeType(tag: DW_TAG_class_type, name: "two", file: !19, line: 1, size: 64, flags: DIFlagTypePassByValue, elements: !20, identifier: "_ZTS3two")
+!19 = !DIFile(filename: "./include.h", directory: "/")
+!20 = !{!21, !22}
+!21 = !DIDerivedType(tag: DW_TAG_member, name: "a", scope: !18, file: !19, line: 1, baseType: !13, size: 32, flags: DIFlagPublic)
+!22 = !DIDerivedType(tag: DW_TAG_member, name: "b", scope: !18, file: !19, line: 1, baseType: !13, size: 32, offset: 32, flags: DIFlagPublic)
+!23 = distinct !DIAssignID()
+!24 = !DILocation(line: 0, scope: !10)
+!26 = !DILocation(line: 4, column: 12, scope: !10)
+!28 = distinct !DIAssignID()
+!38 = !DISubprogram(name: "get", linkageName: "_Z3getv", scope: !1, file: !1, line: 2, type: !39, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized, retainedNodes: !41)
+!39 = !DISubroutineType(types: !40)
+!40 = !{!18}
+!41 = !{}