[llvm] 8ea8dd9 - [SLP] Fix crash on trying to reshuffle a scalar that was vectorized.

[Alexey Bataev via llvm-commits]

Author: Alexey Bataev
Date: 2023-11-15T11:01:45-08:00
New Revision: 8ea8dd9a017182d167f39f521ef397afba5a0fd5

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

LOG: [SLP] Fix crash on trying to reshuffle a scalar that was vectorized.

If the buildvector node contains extractelement, which vector operand
depends on vector node, need to check if the node is ready and use
vectorized value instead of the original vector operation.

Added: 
    llvm/test/Transforms/SLPVectorizer/X86/gather_extract_from_vectorbuild.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 da18ea8e2ce5d58..d8208cabe6dde0d 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -10149,6 +10149,9 @@ class BoUpSLP::ShuffleInstructionBuilder final : public BaseShuffleAnalysis {
         continue;
       auto *EI = cast<ExtractElementInst>(E->Scalars[I]);
       VecBase = EI->getVectorOperand();
+      if (const TreeEntry *TE = R.getTreeEntry(VecBase))
+        VecBase = TE->VectorizedValue;
+      assert(VecBase && "Expected vectorized value.");
       UniqueBases.insert(VecBase);
       // If the only one use is vectorized - can delete the extractelement
       // itself.
@@ -10186,6 +10189,9 @@ class BoUpSLP::ShuffleInstructionBuilder final : public BaseShuffleAnalysis {
         if (SubMask[I] == PoisonMaskElem)
           continue;
         Value *VecOp = cast<ExtractElementInst>(V)->getVectorOperand();
+        if (const TreeEntry *TE = R.getTreeEntry(VecOp))
+          VecOp = TE->VectorizedValue;
+        assert(VecOp && "Expected vectorized value.");
         const int Size =
             cast<FixedVectorType>(VecOp->getType())->getNumElements();
 #ifndef NDEBUG
@@ -10593,6 +10599,21 @@ ResTy BoUpSLP::processBuildVector(const TreeEntry *E, Args &...Params) {
     ExtractShuffles =
         tryToGatherExtractElements(GatheredScalars, ExtractMask, NumParts);
     if (!ExtractShuffles.empty()) {
+      SmallVector<const TreeEntry *> ExtractEntries;
+      for (auto [Idx, I] : enumerate(ExtractMask)) {
+        if (I == PoisonMaskElem)
+          continue;
+        if (const auto *TE = getTreeEntry(
+                cast<ExtractElementInst>(E->Scalars[Idx])->getVectorOperand()))
+          ExtractEntries.push_back(TE);
+      }
+      if (Value *Delayed = ShuffleBuilder.needToDelay(E, ExtractEntries)) {
+        // Delay emission of gathers which are not ready yet.
+        PostponedGathers.insert(E);
+        // Postpone gather emission, will be emitted after the end of the
+        // process to keep correct order.
+        return Delayed;
+      }
       if (Value *VecBase = ShuffleBuilder.adjustExtracts(
               E, ExtractMask, NumParts, UseVecBaseAsInput)) {
         ExtractVecBase = VecBase;
@@ -10779,12 +10800,16 @@ ResTy BoUpSLP::processBuildVector(const TreeEntry *E, Args &...Params) {
           if (isa<UndefValue>(E->Scalars[I]))
             continue;
           auto *EI = cast<ExtractElementInst>(E->Scalars[I]);
+          Value *VecOp = EI->getVectorOperand();
+          if (const auto *TE = getTreeEntry(VecOp))
+            if (TE->VectorizedValue)
+              VecOp = TE->VectorizedValue;
           if (!Vec1) {
-            Vec1 = EI->getVectorOperand();
+            Vec1 = VecOp;
           } else if (Vec1 != EI->getVectorOperand()) {
             assert((!Vec2 || Vec2 == EI->getVectorOperand()) &&
                    "Expected only 1 or 2 vectors shuffle.");
-            Vec2 = EI->getVectorOperand();
+            Vec2 = VecOp;
           }
         }
       }

diff  --git a/llvm/test/Transforms/SLPVectorizer/X86/gather_extract_from_vectorbuild.ll b/llvm/test/Transforms/SLPVectorizer/X86/gather_extract_from_vectorbuild.ll
new file mode 100644
index 000000000000000..3bccfac8566de85
--- /dev/null
+++ b/llvm/test/Transforms/SLPVectorizer/X86/gather_extract_from_vectorbuild.ll
@@ -0,0 +1,33 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt -passes=slp-vectorizer -S -mtriple=x86_64 < %s | FileCheck %s
+; Vectorization tree roots at vector build sequence (insertelement),
+; SLP crashed on generating vector code for pair {%i4, 0.0} trying to produce
+; a shuffle with %ins1 as a source because it was marked deleted
+; due to vectorization.
+
+define void @test() {
+; CHECK-LABEL: define void @test() {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[TMP0:%.*]] = phi <2 x float> [ zeroinitializer, [[ENTRY:%.*]] ], [ [[TMP3:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[TMP1:%.*]] = fadd <2 x float> zeroinitializer, [[TMP0]]
+; CHECK-NEXT:    [[TMP2:%.*]] = select <2 x i1> zeroinitializer, <2 x float> [[TMP1]], <2 x float> zeroinitializer
+; CHECK-NEXT:    [[TMP3]] = shufflevector <2 x float> [[TMP2]], <2 x float> <float poison, float 0.000000e+00>, <2 x i32> <i32 0, i32 3>
+; CHECK-NEXT:    br label [[LOOP]]
+;
+entry:
+  br label %loop
+
+loop:
+  %ph0 = phi float [ 0.000000e+00, %entry ], [ %i4, %loop ]
+  %ph1 = phi float [ 0.000000e+00, %entry ], [ 0.000000e+00, %loop ]
+  %i = fadd float 0.000000e+00, %ph0
+  %i1 = fadd float 0.000000e+00, %ph1
+  %i2 = select i1 false, float %i, float 0.000000e+00
+  %i3 = select i1 false, float %i1, float 0.000000e+00
+  %ins0 = insertelement <2 x float> zeroinitializer, float %i2, i64 0
+  %ins1 = insertelement <2 x float> %ins0, float %i3, i64 1
+  %i4 = extractelement <2 x float> %ins1, i64 0
+  br label %loop
+}