[llvm] Fix handling of very large allocation in llvm.objectsize expansion

serge-sans-paille · serge-sans-paille · commit a7e4478078b7 · 2024-11-12T21:22:24.000+01:00
The internal structure used to carry intermediate computations hold
signed values. If an object size happens to overflow signed values, we
can get invalid result, so make sure this situation never happens.

This is not very limitative as static allocation of such large values
should scarcely happen.
diff --git a/llvm/include/llvm/Analysis/MemoryBuiltins.h b/llvm/include/llvm/Analysis/MemoryBuiltins.h
@@ -223,6 +223,10 @@ struct SizeOffsetAPInt : public SizeOffsetType<APInt, SizeOffsetAPInt> {
 
 /// OffsetSpan - Used internally by \p ObjectSizeOffsetVisitor. Represents a
 /// point in memory as a pair of allocated bytes before and after it.
+///
+/// \c Before and \c After fields are signed values. It makes it possible to
+/// represent out-of-bound access, e.g. as a result of a GEP, at the expense of
+/// not being able to represent very large allocation.
 struct OffsetSpan {
   APInt Before; /// Number of allocated bytes before this point.
   APInt After;  /// Number of allocated bytes after this point.
@@ -255,7 +259,7 @@ class ObjectSizeOffsetVisitor
   SmallDenseMap<Instruction *, OffsetSpan, 8> SeenInsts;
   unsigned InstructionsVisited;
 
-  APInt align(APInt Size, MaybeAlign Align);
+  APInt alignAndCap(APInt Size, MaybeAlign Align);
 
   static OffsetSpan unknown() { return OffsetSpan(); }
 
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -668,10 +668,14 @@ STATISTIC(ObjectVisitorArgument,
 STATISTIC(ObjectVisitorLoad,
           "Number of load instructions with unsolved size and offset");
 
-APInt ObjectSizeOffsetVisitor::align(APInt Size, MaybeAlign Alignment) {
+/// Align \p Size according to \p Alignment. If \p Size is greater than
+/// getSignedMaxValue(), set it as unknown as we can only represent signed value
+/// in OffsetSpan.
+APInt ObjectSizeOffsetVisitor::alignAndCap(APInt Size, MaybeAlign Alignment) {
   if (Options.RoundToAlign && Alignment)
-    return APInt(IntTyBits, alignTo(Size.getZExtValue(), *Alignment));
-  return Size;
+    Size = APInt(IntTyBits, alignTo(Size.getZExtValue(), *Alignment));
+
+  return Size.isNegative() ? APInt() : Size;
 }
 
 ObjectSizeOffsetVisitor::ObjectSizeOffsetVisitor(const DataLayout &DL,
@@ -733,8 +737,19 @@ OffsetSpan ObjectSizeOffsetVisitor::computeImpl(Value *V) {
       ORT.After = APInt();
   }
   // If the computed bound is "unknown" we cannot add the stripped offset.
-  return {(ORT.knownBefore() ? ORT.Before + Offset : ORT.Before),
-          (ORT.knownAfter() ? ORT.After - Offset : ORT.After)};
+  if (ORT.knownBefore()) {
+    bool Overflow;
+    ORT.Before = ORT.Before.sadd_ov(Offset, Overflow);
+    if (Overflow)
+      ORT.Before = APInt();
+  }
+  if (ORT.knownAfter()) {
+    bool Overflow;
+    ORT.After = ORT.After.ssub_ov(Offset, Overflow);
+    if (Overflow)
+      ORT.After = APInt();
+  }
+  return ORT;
 }
 
 OffsetSpan ObjectSizeOffsetVisitor::computeValue(Value *V) {
@@ -780,8 +795,9 @@ OffsetSpan ObjectSizeOffsetVisitor::visitAllocaInst(AllocaInst &I) {
   if (!isUIntN(IntTyBits, ElemSize.getKnownMinValue()))
     return ObjectSizeOffsetVisitor::unknown();
   APInt Size(IntTyBits, ElemSize.getKnownMinValue());
+
   if (!I.isArrayAllocation())
-    return OffsetSpan(Zero, align(Size, I.getAlign()));
+    return OffsetSpan(Zero, alignAndCap(Size, I.getAlign()));
 
   Value *ArraySize = I.getArraySize();
   if (const ConstantInt *C = dyn_cast<ConstantInt>(ArraySize)) {
@@ -791,8 +807,9 @@ OffsetSpan ObjectSizeOffsetVisitor::visitAllocaInst(AllocaInst &I) {
 
     bool Overflow;
     Size = Size.umul_ov(NumElems, Overflow);
+
     return Overflow ? ObjectSizeOffsetVisitor::unknown()
-                    : OffsetSpan(Zero, align(Size, I.getAlign()));
+                    : OffsetSpan(Zero, alignAndCap(Size, I.getAlign()));
   }
   return ObjectSizeOffsetVisitor::unknown();
 }
@@ -806,12 +823,16 @@ OffsetSpan ObjectSizeOffsetVisitor::visitArgument(Argument &A) {
   }
 
   APInt Size(IntTyBits, DL.getTypeAllocSize(MemoryTy));
-  return OffsetSpan(Zero, align(Size, A.getParamAlign()));
+  return OffsetSpan(Zero, alignAndCap(Size, A.getParamAlign()));
 }
 
 OffsetSpan ObjectSizeOffsetVisitor::visitCallBase(CallBase &CB) {
-  if (std::optional<APInt> Size = getAllocSize(&CB, TLI))
+  if (std::optional<APInt> Size = getAllocSize(&CB, TLI)) {
+    // Very large unsigned value cannot be represented as OffsetSpan.
+    if (Size->isNegative())
+      return ObjectSizeOffsetVisitor::unknown();
     return OffsetSpan(Zero, *Size);
+  }
   return ObjectSizeOffsetVisitor::unknown();
 }
 
@@ -852,7 +873,7 @@ OffsetSpan ObjectSizeOffsetVisitor::visitGlobalVariable(GlobalVariable &GV) {
     return ObjectSizeOffsetVisitor::unknown();
 
   APInt Size(IntTyBits, DL.getTypeAllocSize(GV.getValueType()));
-  return OffsetSpan(Zero, align(Size, GV.getAlign()));
+  return OffsetSpan(Zero, alignAndCap(Size, GV.getAlign()));
 }
 
 OffsetSpan ObjectSizeOffsetVisitor::visitIntToPtrInst(IntToPtrInst &) {
@@ -944,7 +965,11 @@ OffsetSpan ObjectSizeOffsetVisitor::findLoadOffsetRange(
       if (!C)
         return Unknown();
 
-      return Known({APInt(C->getValue().getBitWidth(), 0), C->getValue()});
+      APInt CSize = C->getValue();
+      if (CSize.isNegative())
+        return Unknown();
+
+      return Known({APInt(CSize.getBitWidth(), 0), CSize});
     }
 
     return Unknown();
diff --git a/llvm/test/Transforms/LowerConstantIntrinsics/builtin-object-size-phi.ll b/llvm/test/Transforms/LowerConstantIntrinsics/builtin-object-size-phi.ll
@@ -118,6 +118,20 @@ if.end:
   ret i64 %size
 }
 
+define dso_local i64 @pick_max_large(i1 %c) local_unnamed_addr {
+; CHECK-LABEL: @pick_max_large(
+; CHECK-NEXT:    [[BUFFER:%.*]] = alloca i8, i64 -7, align 1
+; CHECK-NEXT:    [[S:%.*]] = select i1 [[C:%.*]], ptr null, ptr [[BUFFER]]
+; CHECK-NEXT:    ret i64 -1
+;
+  %buffer = alloca i8, i64 -7 ; Actually a very large positive integer
+  %s = select i1 %c, ptr null, ptr %buffer
+  %objsize = tail call i64 @llvm.objectsize.i64.p0(ptr %s, i1 false, i1 false, i1 false)
+  ret i64 %objsize
+
+}
+
+
 define i64 @pick_negative_offset(i32 %n) {
 ; CHECK-LABEL: @pick_negative_offset(
 ; CHECK-NEXT:  entry:
diff --git a/llvm/test/Transforms/LowerConstantIntrinsics/objectsize_basic.ll b/llvm/test/Transforms/LowerConstantIntrinsics/objectsize_basic.ll
@@ -195,6 +195,62 @@ define i64 @out_of_bound_gep() {
   ret i64 %objsize
 }
 
+define i64 @wrapping_gep() {
+; CHECK-LABEL: @wrapping_gep(
+; CHECK-NEXT:    [[OBJ:%.*]] = alloca i8, i64 4, align 1
+; CHECK-NEXT:    [[SLIDE:%.*]] = getelementptr i8, ptr [[OBJ]], i64 9223372036854775807
+; CHECK-NEXT:    [[SLIDE_BIS:%.*]] = getelementptr i8, ptr [[SLIDE]], i64 9223372036854775807
+; CHECK-NEXT:    ret i64 0
+;
+  %obj = alloca i8, i64 4
+  %slide = getelementptr i8, ptr %obj, i64 9223372036854775807
+  %slide.bis = getelementptr i8, ptr %slide, i64 9223372036854775807
+  %objsize = call i64 @llvm.objectsize.i64(ptr %slide.bis, i1 false, i1 false, i1 false)
+  ret i64 %objsize
+}
+
+define i64 @wrapping_gep_neg() {
+; CHECK-LABEL: @wrapping_gep_neg(
+; CHECK-NEXT:    [[OBJ:%.*]] = alloca i8, i64 9223372036854775807, align 1
+; CHECK-NEXT:    [[SLIDE:%.*]] = getelementptr i8, ptr [[OBJ]], i64 9223372036854775807
+; CHECK-NEXT:    [[SLIDE_BIS:%.*]] = getelementptr i8, ptr [[SLIDE]], i64 3
+; CHECK-NEXT:    [[SLIDE_TER:%.*]] = getelementptr i8, ptr [[SLIDE_BIS]], i64 -4
+; CHECK-NEXT:    ret i64 1
+;
+  %obj = alloca i8, i64 9223372036854775807
+  %slide = getelementptr i8, ptr %obj, i64 9223372036854775807
+  %slide.bis = getelementptr i8, ptr %slide, i64 3
+  %slide.ter = getelementptr i8, ptr %slide.bis, i64 -4
+  %objsize = call i64 @llvm.objectsize.i64(ptr %slide.ter, i1 false, i1 false, i1 false)
+  ret i64 %objsize
+}
+
+; We don't analyze allocations larger than platform's ptrdiff_t
+define i64 @large_alloca() {
+; CHECK-LABEL: @large_alloca(
+; CHECK-NEXT:    [[OBJ:%.*]] = alloca i8, i64 -9223372036854775808, align 1
+; CHECK-NEXT:    [[SLIDE:%.*]] = getelementptr i8, ptr [[OBJ]], i64 9223372036854775807
+; CHECK-NEXT:    ret i64 -1
+;
+  %obj = alloca i8, i64 9223372036854775808
+  %slide = getelementptr i8, ptr %obj, i64 9223372036854775807
+  %objsize = call i64 @llvm.objectsize.i64(ptr %slide, i1 false, i1 false, i1 false)
+  ret i64 %objsize
+}
+
+; We don't analyze allocations larger than platform's ptrdiff_t
+define i64 @large_malloc() {
+; CHECK-LABEL: @large_malloc(
+; CHECK-NEXT:    [[OBJ:%.*]] = call ptr @malloc(i64 -9223372036854775808)
+; CHECK-NEXT:    [[SLIDE:%.*]] = getelementptr i8, ptr [[OBJ]], i64 9223372036854775807
+; CHECK-NEXT:    ret i64 -1
+;
+  %obj = call ptr @malloc(i64 9223372036854775808)
+  %slide = getelementptr i8, ptr %obj, i64 9223372036854775807
+  %objsize = call i64 @llvm.objectsize.i64(ptr %slide, i1 false, i1 false, i1 false)
+  ret i64 %objsize
+}
+
 define i64 @out_of_bound_negative_gep() {
 ; CHECK-LABEL: @out_of_bound_negative_gep(
 ; CHECK-NEXT:    [[OBJ:%.*]] = alloca i8, i32 4, align 1
