[llvm] 6aaf1e7 - [LoopIdiom] Fix store size SCEV type.

[Clement Courbet via llvm-commits]

Author: Clement Courbet
Date: 2021-10-11T09:39:06+02:00
New Revision: 6aaf1e7ea931a1865cc97ebf8fcb084772f8142e

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

LOG: [LoopIdiom] Fix store size SCEV type.

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.

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

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

Modified: 
    llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp b/llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp
index 021c37748936e..4b7d63c9efd06 100644
--- a/llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp
@@ -786,7 +786,8 @@ bool LoopIdiomRecognize::processLoopStores(SmallVectorImpl<StoreInst *> &SL,
 
     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,

diff  --git a/llvm/test/Transforms/LoopIdiom/memset-pr52104.ll b/llvm/test/Transforms/LoopIdiom/memset-pr52104.ll
new file mode 100644
index 0000000000000..025c392a5734b
--- /dev/null
+++ b/llvm/test/Transforms/LoopIdiom/memset-pr52104.ll
@@ -0,0 +1,43 @@
+; 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
+}