[llvm] f2a322b - [SROA] prevent crash on large memset length (PR50910)

[Sanjay Patel via llvm-commits]

Author: Sanjay Patel
Date: 2021-07-31T14:07:30-04:00
New Revision: f2a322bfcfbc62b5523f32c4eded6faf2cad2e24

URL: https://github.com/llvm/llvm-project/commit/f2a322bfcfbc62b5523f32c4eded6faf2cad2e24
DIFF: https://github.com/llvm/llvm-project/commit/f2a322bfcfbc62b5523f32c4eded6faf2cad2e24.diff

LOG: [SROA] prevent crash on large memset length (PR50910)

I don't know much about this pass, but we need a stronger
check on the memset length arg to avoid an assert. The
current code was added with D59000.
The test is reduced from:
https://llvm.org/PR50910

Differential Revision: https://reviews.llvm.org/D106462

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/SROA.cpp
    llvm/test/Transforms/SROA/slice-width.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 5ec01454e5b2f..fe160d5415bd2 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -2811,10 +2811,11 @@ class llvm::sroa::AllocaSliceRewriter
       if (BeginOffset > NewAllocaBeginOffset ||
           EndOffset < NewAllocaEndOffset)
         return false;
+      // Length must be in range for FixedVectorType.
       auto *C = cast<ConstantInt>(II.getLength());
-      if (C->getBitWidth() > 64)
+      const uint64_t Len = C->getLimitedValue();
+      if (Len > std::numeric_limits<unsigned>::max())
         return false;
-      const auto Len = C->getZExtValue();
       auto *Int8Ty = IntegerType::getInt8Ty(NewAI.getContext());
       auto *SrcTy = FixedVectorType::get(Int8Ty, Len);
       return canConvertValue(DL, SrcTy, AllocaTy) &&

diff  --git a/llvm/test/Transforms/SROA/slice-width.ll b/llvm/test/Transforms/SROA/slice-width.ll
index a801de68217ff..b15e66b462c0f 100644
--- a/llvm/test/Transforms/SROA/slice-width.ll
+++ b/llvm/test/Transforms/SROA/slice-width.ll
@@ -145,3 +145,16 @@ define void @PR50888() {
   call void @llvm.memset.p0i8.i64(i8* align 16 %array, i8 0, i64 ptrtoint (void ()* @PR50888 to i64), i1 false)
   ret void
 }
+
+; Don't crash on out-of-bounds length.
+
+define void @PR50910() {
+; CHECK-LABEL: @PR50910(
+; CHECK-NEXT:    [[T1:%.*]] = alloca i8, i64 1, align 8
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 8 [[T1]], i8 0, i64 1, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %t1 = alloca i8, i64 1, align 8
+  call void @llvm.memset.p0i8.i64(i8* align 8 %t1, i8 0, i64 4294967296, i1 false)
+  ret void
+}