[LoopIdiom] Fix store size SCEV type.

[courbet]

Summary

We were using the type of the loop back edge count to represent the
store size. This failed for small loop counts (e.g. in the added test,
the loop count was an i2).

Use the index type instead.

Fixes PR52104.

[nikic]

LGTM

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

for (SetVector<StoreInst *>::iterator it = Heads.begin(), e = Heads.end();
 
     // Check to see if the stride matches the size of the stores.  If so, then
     // we know that every byte is touched in the loop.
     if (StoreSize != Stride && StoreSize != -Stride)
       continue;
 
     bool IsNegStride = StoreSize == -Stride;
 
-    const SCEV *StoreSizeSCEV = SE->getConstant(BECount->getType(), StoreSize);
+    Type *IntIdxTy = DL->getIndexType(StorePtr->getType());
+    const SCEV *StoreSizeSCEV = SE->getConstant(IntIdxTy, StoreSize);
     if (processLoopStridedStore(StorePtr, StoreSizeSCEV,
                                 MaybeAlign(HeadStore->getAlignment()),
                                 StoredVal, HeadStore, AdjacentStores, StoreEv,
                                 BECount, IsNegStride)) {
       TransformedStores.insert(AdjacentStores.begin(), AdjacentStores.end());
       Changed = true;
     }
   }

llvm/test/Transforms/LoopIdiom/memset-pr52104.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -basic-aa -loop-idiom < %s -S | FileCheck %s
+
+define void @f(i32* nocapture nonnull align 4 dereferenceable(20) %0, i32 %1) local_unnamed_addr #0 align 32 {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:    [[TMP3:%.*]] = trunc i32 [[TMP1:%.*]] to i2
+; CHECK-NEXT:    [[TMP4:%.*]] = zext i2 [[TMP3]] to i64
+; CHECK-NEXT:    [[SCEVGEP:%.*]] = getelementptr i32, i32* [[TMP0:%.*]], i64 [[TMP4]]
+; CHECK-NEXT:    [[SCEVGEP1:%.*]] = bitcast i32* [[SCEVGEP]] to i8*
+; CHECK-NEXT:    [[TMP5:%.*]] = sub i2 -1, [[TMP3]]
+; CHECK-NEXT:    [[TMP6:%.*]] = zext i2 [[TMP5]] to i64
+; CHECK-NEXT:    [[TMP7:%.*]] = shl nuw nsw i64 [[TMP6]], 2
+; CHECK-NEXT:    [[TMP8:%.*]] = add nuw nsw i64 [[TMP7]], 4
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 4 [[SCEVGEP1]], i8 0, i64 [[TMP8]], i1 false)
+; CHECK-NEXT:    br label [[TMP9:%.*]]
+; CHECK:       9:
+; CHECK-NEXT:    [[TMP10:%.*]] = phi i32 [ [[TMP14:%.*]], [[TMP9]] ], [ [[TMP1]], [[TMP2:%.*]] ]
+; CHECK-NEXT:    [[TMP11:%.*]] = and i32 [[TMP10]], 3
+; CHECK-NEXT:    [[TMP12:%.*]] = zext i32 [[TMP11]] to i64
+; CHECK-NEXT:    [[TMP13:%.*]] = getelementptr inbounds i32, i32* [[TMP0]], i64 [[TMP12]]
+; CHECK-NEXT:    [[TMP14]] = add nsw i32 [[TMP10]], 1
+; CHECK-NEXT:    [[TMP15:%.*]] = and i32 [[TMP14]], 3
+; CHECK-NEXT:    [[TMP16:%.*]] = icmp eq i32 [[TMP15]], 0
+; CHECK-NEXT:    br i1 [[TMP16]], label [[TMP17:%.*]], label [[TMP9]]
+; CHECK:       17:
+; CHECK-NEXT:    ret void
+;
+  br label %3
+
+3:                                                ; preds = %3, %1
+  %4 = phi i32 [ %8, %3 ], [ %1, %2 ]
+  %5 = and i32 %4, 3
+  %6 = zext i32 %5 to i64
+  %7 = getelementptr inbounds i32, i32* %0, i64 %6
+  store i32 0, i32* %7, align 4
+  %8 = add nsw i32 %4, 1
+  %9 = and i32 %8, 3
+  %10 = icmp eq i32 %9, 0
+  br i1 %10, label %11, label %3
+
+11:                                               ; preds = %4
+  ret void
+}