[llvm] 24a043a - [SLP] Fix crash of shuffle poison (#106857)

[via llvm-commits]

Author: tcwzxx
Date: 2024-09-01T20:24:09+08:00
New Revision: 24a043a6ff052ad7c123e67d020d688eea1a7efa

URL: https://github.com/llvm/llvm-project/commit/24a043a6ff052ad7c123e67d020d688eea1a7efa
DIFF: https://github.com/llvm/llvm-project/commit/24a043a6ff052ad7c123e67d020d688eea1a7efa.diff

LOG: [SLP] Fix crash of shuffle poison (#106857)

When the shuffle masks are `PoisonMaskElem`, there is not need to check
the cost of `SK_ExtractSubvector`. It is free. Otherwise, it will cause
the compiler to crash.

Assertion `(Idx + EltsPerVector) <= alignTo(NumElts, EltsPerVector) &&
"SK_ExtractSubvector index out of range"' failed.

Added: 
    llvm/test/Transforms/SLPVectorizer/crash_extractelement_poison.ll

Modified: 
    llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index b7bcfba7f86df6..993fd6ab1b0b41 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
+}