[llvm] af15c46 - [SLP]Do not crash if number of vector registers does not feet the vector

[Alexey Bataev via llvm-commits]

Author: Alexey Bataev
Date: 2023-10-30T07:30:52-07:00
New Revision: af15c46777208a4cb4b276c4974a5b556608a415

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

LOG: [SLP]Do not crash if number of vector registers does not feet the vector
type.

Need to check, if the number of vector registers, returned by TTI, is
not greater than total number of mask element and not zero, before
trying to perform any operations. TTI still may return non-valid number
of registers.

Added: 
    llvm/test/Transforms/SLPVectorizer/X86/shuffle-multivector.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 bb4e743c1544a98..b6895c649f838c1 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -7580,8 +7580,8 @@ class BoUpSLP::ShuffleCostEstimator : public BaseShuffleAnalysis {
     auto *MaskVecTy =
         FixedVectorType::get(E1.Scalars.front()->getType(), Mask.size());
     unsigned NumParts = TTI.getNumberOfParts(MaskVecTy);
-    assert(NumParts > 0 && NumParts < Mask.size() &&
-           "Expected positive number of registers.");
+    if (NumParts == 0 || NumParts >= Mask.size())
+      NumParts = 1;
     unsigned SliceSize = Mask.size() / NumParts;
     const auto *It =
         find_if(Mask, [](int Idx) { return Idx != PoisonMaskElem; });
@@ -7598,8 +7598,8 @@ class BoUpSLP::ShuffleCostEstimator : public BaseShuffleAnalysis {
     auto *MaskVecTy =
         FixedVectorType::get(E1.Scalars.front()->getType(), Mask.size());
     unsigned NumParts = TTI.getNumberOfParts(MaskVecTy);
-    assert(NumParts > 0 && NumParts < Mask.size() &&
-           "Expected positive number of registers.");
+    if (NumParts == 0 || NumParts >= Mask.size())
+      NumParts = 1;
     unsigned SliceSize = Mask.size() / NumParts;
     const auto *It =
         find_if(Mask, [](int Idx) { return Idx != PoisonMaskElem; });

diff  --git a/llvm/test/Transforms/SLPVectorizer/X86/shuffle-multivector.ll b/llvm/test/Transforms/SLPVectorizer/X86/shuffle-multivector.ll
new file mode 100644
index 000000000000000..143052a3d9cd077
--- /dev/null
+++ b/llvm/test/Transforms/SLPVectorizer/X86/shuffle-multivector.ll
@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=slp-vectorizer -S < %s -mtriple=x86_64-unknown-linux -slp-threshold=-160 | FileCheck %s
+
+define void @test1(i128 %p0, i128 %p1, i128 %p2, i128 %p3, <4 x i128> %vec) {
+; CHECK-LABEL: @test1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = insertelement <2 x i128> poison, i128 [[P0:%.*]], i32 0
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x i128> [[TMP0]], i128 [[P1:%.*]], i32 1
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc <2 x i128> [[TMP1]] to <2 x i32>
+; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <2 x i32> [[TMP2]], <2 x i32> poison, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
+; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <2 x i128> poison, i128 [[P2:%.*]], i32 0
+; CHECK-NEXT:    [[TMP5:%.*]] = insertelement <2 x i128> [[TMP4]], i128 [[P3:%.*]], i32 1
+; CHECK-NEXT:    [[TMP6:%.*]] = trunc <2 x i128> [[TMP5]] to <2 x i32>
+; CHECK-NEXT:    [[TMP7:%.*]] = shufflevector <2 x i32> [[TMP6]], <2 x i32> poison, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
+; CHECK-NEXT:    [[T5:%.*]] = trunc i128 [[P1]] to i32
+; CHECK-NEXT:    [[TMP8:%.*]] = sdiv <4 x i32> [[TMP3]], [[TMP7]]
+; CHECK-NEXT:    [[TMP9:%.*]] = shufflevector <2 x i128> [[TMP1]], <2 x i128> [[TMP5]], <4 x i32> <i32 poison, i32 0, i32 3, i32 2>
+; CHECK-NEXT:    [[TMP10:%.*]] = shufflevector <4 x i128> [[VEC:%.*]], <4 x i128> [[TMP9]], <4 x i32> <i32 0, i32 5, i32 6, i32 7>
+; CHECK-NEXT:    [[TMP11:%.*]] = trunc <4 x i128> [[TMP10]] to <4 x i32>
+; CHECK-NEXT:    [[TMP12:%.*]] = sdiv <4 x i32> [[TMP8]], [[TMP11]]
+; CHECK-NEXT:    br label [[BB:%.*]]
+; CHECK:       bb:
+; CHECK-NEXT:    [[TMP13:%.*]] = phi <4 x i32> [ [[TMP12]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %t1 = trunc i128 %p0 to i32
+  %t2 = trunc i128 %p1 to i32
+  %t3 = trunc i128 %p2 to i32
+  %t4 = trunc i128 %p3 to i32
+  %t5 = trunc i128 %p1 to i32
+  %t6 = trunc i128 %p0 to i32
+  %t7 = trunc i128 %p3 to i32
+  %t8 = trunc i128 %p2 to i32
+  %m0 = sdiv i32 %t1, %t3
+  %m1 = sdiv i32 %t2, %t4
+  %m2 = sdiv i32 %t1, %t3
+  %m3 = sdiv i32 %t2, %t4
+  %e0 = extractelement <4 x i128> %vec, i32 0
+  %t9 = trunc i128 %e0 to i32
+  %d0 = sdiv i32 %m0, %t9
+  %d1 = sdiv i32 %m1, %t6
+  %d2 = sdiv i32 %m2, %t7
+  %d3 = sdiv i32 %m3, %t8
+  br label %bb
+
+bb:
+  %phi0 = phi i32 [ %d0, %entry ]
+  %phi1 = phi i32 [ %d1, %entry ]
+  %phi2 = phi i32 [ %d2, %entry ]
+  %phi3 = phi i32 [ %d3, %entry ]
+  ret void
+}