[llvm] [SLP] Fix crash of shuffle poison (PR #106857)

[via llvm-commits]

https://github.com/tcwzxx updated https://github.com/llvm/llvm-project/pull/106857

>From 1eaf29ef5a5ad318eb62387fd6ae784579114f44 Mon Sep 17 00:00:00 2001
From: tcwzxx <tcwzxx at gmail.com>
Date: Sat, 31 Aug 2024 23:18:29 +0800
Subject: [PATCH] Fix crash of shuffle poison

---
 .../Transforms/Vectorize/SLPVectorizer.cpp    |  6 +++-
 .../crash_extractelement_poison.ll            | 31 +++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/SLPVectorizer/crash_extractelement_poison.ll

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 3d41c978281351..4e786f0c77d17a 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -11204,6 +11204,10 @@ BoUpSLP::tryToGatherSingleRegisterExtractElements(
       UndefVectorExtracts.push_back(I);
       continue;
     }
+    if (Idx >= VecTy->getNumElements()) {
+      UndefVectorExtracts.push_back(I);
+      continue;
+    }
     SmallBitVector ExtractMask(VecTy->getNumElements(), true);
     ExtractMask.reset(*Idx);
     if (isUndefVector(EI->getVectorOperand(), ExtractMask).all()) {
@@ -11251,7 +11255,7 @@ BoUpSLP::tryToGatherSingleRegisterExtractElements(
   // shuffle of a single/two vectors the scalars are extracted from.
   std::optional<TTI::ShuffleKind> Res =
       isFixedVectorShuffle(GatheredExtracts, Mask);
-  if (!Res) {
+  if (!Res || all_of(Mask, [](int Idx) { return Idx == PoisonMaskElem; })) {
     // TODO: try to check other subsets if possible.
     // Restore the original VL if attempt was not successful.
     copy(SavedVL, VL.begin());
diff --git a/llvm/test/Transforms/SLPVectorizer/crash_extractelement_poison.ll b/llvm/test/Transforms/SLPVectorizer/crash_extractelement_poison.ll
new file mode 100644
index 00000000000000..32ea3984cc4cb8
--- /dev/null
+++ b/llvm/test/Transforms/SLPVectorizer/crash_extractelement_poison.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 < %s | FileCheck %s
+
+define void @test(i8 %0, i8 %1) {
+; CHECK-LABEL: define void @test(
+; CHECK-SAME: i8 [[TMP0:%.*]], i8 [[TMP1:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[L:%.*]] = load <4 x i8>, ptr getelementptr (i8, ptr null, i32 8), align 1
+; CHECK-NEXT:    [[LI15:%.*]] = extractelement <4 x i8> [[L]], i64 15
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <4 x i8> poison, i8 [[TMP0]], i32 0
+; CHECK-NEXT:    [[TMP3:%.*]] = insertelement <4 x i8> [[TMP2]], i8 [[TMP1]], i32 1
+; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <4 x i8> [[TMP3]], i8 [[LI15]], i32 3
+; CHECK-NEXT:    [[TMP5:%.*]] = shufflevector <4 x i8> [[TMP4]], <4 x i8> poison, <4 x i32> <i32 0, i32 1, i32 0, i32 3>
+; CHECK-NEXT:    [[TMP6:%.*]] = icmp ne <4 x i8> [[TMP5]], zeroinitializer
+; CHECK-NEXT:    ret void
+;
+entry:
+  %l = load <4 x i8>, ptr getelementptr (i8, ptr null, i32 8), align 1
+  %li15 = extractelement <4 x i8> %l, i64 15
+  %2 = icmp ne i8 %0, 0
+  %3 = icmp ne i8 %1, 0
+  %4 = icmp ne i8 %0, 0
+  %.i15 = icmp ne i8 %li15, 0
+
+  %i0244 = insertelement <4 x i1> zeroinitializer, i1 %2, i64 0
+  %i1245 = insertelement <4 x i1> %i0244, i1 %3, i64 1
+  %i2246 = insertelement <4 x i1> %i1245, i1 %4, i64 2
+  %14 = insertelement <4 x i1> %i2246, i1 %.i15, i64 3
+  ret void
+}