From f381cd069965dabfeb277f30a4e532d7fd498f6e Mon Sep 17 00:00:00 2001
From: Alexey Bataev <a.bataev@outlook.com>
Date: Tue, 3 Sep 2024 04:52:47 -0700
Subject: [PATCH] [SLP]Fix PR107036: Check if the type of the user is sizable
 before requesting its size.

Only some instructions should be considered as potentially reducing the
size of the operands types, not all instructions should be considered.

Fixes https://github.com/llvm/llvm-project/issues/107036
---
 .../Transforms/Vectorize/SLPVectorizer.cpp    |  9 +++---
 .../X86/minbw-user-non-sizable.ll             | 31 +++++++++++++++++++
 2 files changed, 36 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/Transforms/SLPVectorizer/X86/minbw-user-non-sizable.ll

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 58137cd6f543f2..f58803fc56a2d7 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -16055,15 +16055,16 @@ void BoUpSLP::computeMinimumValueSizes() {
                     const TreeEntry *UserTE = E.UserTreeIndices.back().UserTE;
                     if (TE == UserTE || !TE)
                       return false;
+                    if (!isa<CastInst, BinaryOperator, FreezeInst, PHINode,
+                             SelectInst>(U) ||
+                        !isa<CastInst, BinaryOperator, FreezeInst, PHINode,
+                             SelectInst>(UserTE->getMainOp()))
+                      return true;
                     unsigned UserTESz = DL->getTypeSizeInBits(
                         UserTE->Scalars.front()->getType());
                     auto It = MinBWs.find(TE);
                     if (It != MinBWs.end() && It->second.first > UserTESz)
                       return true;
-                    // The size of icmp is always 1 and should not be
-                    // considered.
-                    if (TE->getOpcode() == Instruction::ICmp)
-                      return true;
                     return DL->getTypeSizeInBits(U->getType()) > UserTESz;
                   }));
         })) {
diff --git a/llvm/test/Transforms/SLPVectorizer/X86/minbw-user-non-sizable.ll b/llvm/test/Transforms/SLPVectorizer/X86/minbw-user-non-sizable.ll
new file mode 100644
index 00000000000000..7e7d4352e27733
--- /dev/null
+++ b/llvm/test/Transforms/SLPVectorizer/X86/minbw-user-non-sizable.ll
@@ -0,0 +1,31 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S --passes=slp-vectorizer -mtriple=x86_64-unknown-linux-gnu < %s -slp-threshold=-100 | FileCheck %s
+
+define void @test(ptr %i) {
+; CHECK-LABEL: define void @test(
+; CHECK-SAME: ptr [[I:%.*]]) {
+; CHECK-NEXT:  [[BB:.*]]:
+; CHECK-NEXT:    br label %[[BB2:.*]]
+; CHECK:       [[BB2]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = phi <2 x i32> [ [[TMP3:%.*]], %[[BB2]] ], [ zeroinitializer, %[[BB]] ]
+; CHECK-NEXT:    store <2 x i32> [[TMP0]], ptr [[I]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <2 x i32> [[TMP0]], <2 x i32> <i32 0, i32 poison>, <2 x i32> <i32 2, i32 1>
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc <2 x i32> [[TMP1]] to <2 x i1>
+; CHECK-NEXT:    [[TMP3]] = select <2 x i1> [[TMP2]], <2 x i32> zeroinitializer, <2 x i32> zeroinitializer
+; CHECK-NEXT:    br label %[[BB2]]
+;
+bb:
+  %i1 = getelementptr i8, ptr %i, i64 4
+  br label %bb2
+
+bb2:
+  %i3 = phi i32 [ %i6, %bb2 ], [ 0, %bb ]
+  %i4 = phi i32 [ %i8, %bb2 ], [ 0, %bb ]
+  store i32 %i3, ptr %i
+  store i32 %i4, ptr %i1
+  %i5 = trunc i32 0 to i1
+  %i6 = select i1 %i5, i32 0, i32 0
+  %i7 = trunc i32 %i4 to i1
+  %i8 = select i1 %i7, i32 0, i32 0
+  br label %bb2
+}