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

[spatel]

Summary

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

[spatel]

Rebased - I found another bug ( https://llvm.org/PR50888 ) near here and patched it with f14495dc75d7 .

[spatel]

Ping.

[lebedev.ri]

I suppose we could instead clamp to the maximal remaining alloca size,
but it's UB so it doesn't seem worth spending too much time on it.
Seems fine to me.

llvm/lib/Transforms/Scalar/SROA.cpp

bool visitMemSetInst(MemSetInst &II) {
     Type *ScalarTy = AllocaTy->getScalarType();
 
     const bool CanContinue = [&]() {
       if (VecTy || IntTy)
         return true;
       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) &&
              DL.isLegalInteger(DL.getTypeSizeInBits(ScalarTy).getFixedSize());
     }();
 
     // If this doesn't map cleanly onto the alloca type, and that type isn't
     // a single value type, just emit a memset.

llvm/test/Transforms/SROA/slice-width.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -sroa -S | FileCheck %s
 target datalayout = "e-p:64:64:64-p1:16:16:16-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f32:32:32-f64:64:64-f80:128-v64:64:64-v128:128:128-a0:0:64-n8:16:32:64"
 
 declare void @llvm.memcpy.p0i8.p0i8.i32(i8* nocapture, i8* nocapture, i32, i1) nounwind
 declare void @llvm.memset.p0i8.i32(i8* nocapture, i8, i32, i1) nounwind
+declare void @llvm.memset.p0i8.i64(i8* nocapture, i8, i64, i1) nounwind
 
 ; This tests that allocas are not split into slices that are not byte width multiple
 define void @no_split_on_non_byte_width(i32) {
 ; CHECK-LABEL: @no_split_on_non_byte_width(
 ; CHECK-NEXT:    [[ARG_SROA_0:%.*]] = alloca i8, align 8
 ; CHECK-NEXT:    [[ARG_SROA_0_0_EXTRACT_TRUNC:%.*]] = trunc i32 [[TMP0:%.*]] to i8
 ; CHECK-NEXT:    store i8 [[ARG_SROA_0_0_EXTRACT_TRUNC]], i8* [[ARG_SROA_0]], align 8
 ; CHECK-NEXT:    [[ARG_SROA_3_0_EXTRACT_SHIFT:%.*]] = lshr i32 [[TMP0]], 8
entry:
   %tmp2 = alloca %S.vec3float, align 4
   %2 = bitcast %S.vec3float* %tmp2 to i8*
   %3 = bitcast %S.vec3float* %tmp1 to i8*
   call void @llvm.memcpy.p0i8.p0i8.i32(i8* align 4 %2, i8* align 4 %3, i32 12, i1 false)
 
   %result = call i32 @memcpy_vec3float_helper(%S.vec3float* %tmp2)
   ret i32 %result
 }
+
+; 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
+}