[Assignment Tracking] Check getTypeSizeInBits result for scalable vector types

[Orlando]

Summary

Without this patch, in getAssignmentInfo the result of getTypeSizeInBits is cast to uint64_t, which a) is an operation that will eventually be unsupported by the API according to the comments, and b) causes an assertion failure if the type is a scalable vector. Don't cast the TypeSize to uint64_t and check isScalable before getting the fixed size.

This can result in incorrect variable locations, see llvm.org/PR62346 (but is better than crashing).

[DavidSpickett]

I'm not qualified to review this, others are.

[paulwalker-arm]

llvm/lib/IR/DebugInfo.cpp

1950

Is the use of "Scalar" here accurate? given it's only scalable vectors that are being rejected.

Only a suggestion but what about changing getAssignmentInfoImpl to take a TypeSize instead of an uint64_t and then reject scalable sizes within that? This way if/when scalable support is required the structure will already be in place.

[Orlando]

+ Address review comment

llvm/lib/IR/DebugInfo.cpp

1950

Thanks, I agree that way is better - fixed.

[Orlando]

Thanks!

llvm/lib/IR/DebugInfo.cpp

else
     Result = TrimmedSliceOfVariable;
   return true;
 }
 
 /// Collect constant properies (base, size, offset) of \p StoreDest.
 /// Return std::nullopt if any properties are not constants.
 static std::optional<AssignmentInfo>
 getAssignmentInfoImpl(const DataLayout &DL, const Value *StoreDest,
-                      uint64_t SizeInBits) {
+                      TypeSize SizeInBits) {
+  if (SizeInBits.isScalable())
+    return std::nullopt;
   APInt GEPOffset(DL.getIndexTypeSizeInBits(StoreDest->getType()), 0);
   const Value *Base = StoreDest->stripAndAccumulateConstantOffsets(
       DL, GEPOffset, /*AllowNonInbounds*/ true);
   uint64_t OffsetInBytes = GEPOffset.getLimitedValue();
   // Check for overflow.
   if (OffsetInBytes == UINT64_MAX)
     return std::nullopt;
   if (const auto *Alloca = dyn_cast<AllocaInst>(Base))
     return AssignmentInfo(DL, Alloca, OffsetInBytes * 8, SizeInBits);
   return std::nullopt;
 }
 
 std::optional<AssignmentInfo> at::getAssignmentInfo(const DataLayout &DL,
                                                     const MemIntrinsic *I) {
   const Value *StoreDest = I->getRawDest();
   // Assume 8 bit bytes.
   auto *ConstLengthInBytes = dyn_cast<ConstantInt>(I->getLength());
   if (!ConstLengthInBytes)
     // We can't use a non-const size, bail.
     return std::nullopt;
   uint64_t SizeInBits = 8 * ConstLengthInBytes->getZExtValue();
-  return getAssignmentInfoImpl(DL, StoreDest, SizeInBits);
+  return getAssignmentInfoImpl(DL, StoreDest, TypeSize::getFixed(SizeInBits));
 }
 
 std::optional<AssignmentInfo> at::getAssignmentInfo(const DataLayout &DL,
                                                     const StoreInst *SI) {
-  const Value *StoreDest = SI->getPointerOperand();
-  uint64_t SizeInBits = DL.getTypeSizeInBits(SI->getValueOperand()->getType());
-  return getAssignmentInfoImpl(DL, StoreDest, SizeInBits);
+  TypeSize SizeInBits = DL.getTypeSizeInBits(SI->getValueOperand()->getType());
+  return getAssignmentInfoImpl(DL, SI->getPointerOperand(), SizeInBits);
 }
 
 std::optional<AssignmentInfo> at::getAssignmentInfo(const DataLayout &DL,
                                                     const AllocaInst *AI) {
-  uint64_t SizeInBits = DL.getTypeSizeInBits(AI->getAllocatedType());
+  TypeSize SizeInBits = DL.getTypeSizeInBits(AI->getAllocatedType());
   return getAssignmentInfoImpl(DL, AI, SizeInBits);
 }
 
 /// Returns nullptr if the assignment shouldn't be attributed to this variable.
 static CallInst *emitDbgAssign(AssignmentInfo Info, Value *Val, Value *Dest,
                                Instruction &StoreLikeInst,
                                const VarRecord &VarRec, DIBuilder &DIB) {
   auto *ID = StoreLikeInst.getMetadata(LLVMContext::MD_DIAssignID);

llvm/test/DebugInfo/assignment-tracking/AArch64/lit.local.cfg

+if not 'AArch64' in config.root.targets:
+    config.unsupported = True

llvm/test/DebugInfo/assignment-tracking/AArch64/scalable-vectors.ll

+; RUN: llc %s -stop-after=finalize-isel -o - | FileCheck %s
+
+;; Hand written. Check AssignmentTrackingAnalysis doesn't try to get the size
+;; of scalable vectors (which causes an assertion failure).
+
+; CHECK: DBG_VALUE %stack.0, $noreg, ![[#]], !DIExpression(DW_OP_deref)
+; CHECK: DBG_VALUE i256 0, $noreg, ![[#]], !DIExpression()
+;; FIXME: We should reinstate (or poison) the variable at the store. See
+;; llvm.org/PR62346.
+
+target triple = "aarch64"
+
+define dso_local void @square(ptr %0) local_unnamed_addr #0 !dbg !9 {
+entry:
+   %1 = alloca <32 x i8>, !DIAssignID !18
+   call void @llvm.dbg.assign(metadata i1 poison, metadata !19, metadata !DIExpression(), metadata !18, metadata ptr %1, metadata !DIExpression()), !dbg !21
+   %2 = load <vscale x 8 x i8>, ptr %0, align 1
+   call void @llvm.dbg.assign(metadata i256 0, metadata !19, metadata !DIExpression(), metadata !22, metadata ptr %1, metadata !DIExpression()), !dbg !21
+   store <vscale x 8 x i8> %2, ptr %1, align 1
+   ret void
+}
+
+declare void @llvm.dbg.assign(metadata, metadata, metadata, metadata, metadata, metadata)
+
+attributes #0 = { vscale_range(1,16) "target-cpu"="generic" "target-features"="+sve" }
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!2, !3, !7}
+!llvm.ident = !{!8}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C11, file: !1, producer: "clang version 17.0.0", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, splitDebugInlining: false, nameTableKind: None)
+!1 = !DIFile(filename: "/app/example.c", directory: "/app")
+!2 = !{i32 7, !"Dwarf Version", i32 4}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!7 = !{i32 7, !"debug-info-assignment-tracking", i1 true}
+!8 = !{!"clang version 17.0.0"}
+!9 = distinct !DISubprogram(name: "square", scope: !10, file: !10, line: 1, type: !11, scopeLine: 1, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !17)
+!10 = !DIFile(filename: "example.c", directory: "/app")
+!11 = !DISubroutineType(types: !12)
+!12 = !{null, !13, !14, !16}
+!13 = !DIDerivedType(tag: DW_TAG_restrict_type, baseType: !14)
+!14 = !DIDerivedType(tag: DW_TAG_pointer_type, baseType: !15, size: 64)
+!15 = !DIBasicType(name: "unsigned char", size: 8, encoding: DW_ATE_unsigned_char)
+!16 = !DIBasicType(name: "int", size: 32, encoding: DW_ATE_signed)
+!17 = !{}
+!18 =  distinct !DIAssignID()
+!19 = !DILocalVariable(name: "i", scope: !9, file: !10, line: 2, type: !20)
+!20 = !DIBasicType(name: "cats", size: 256, encoding: DW_ATE_signed)
+!21 = !DILocation(line: 0, scope: !9)
+!22 =  distinct !DIAssignID()
+