[SLP]Use revectorized value for extracts from buildvector, beeing

vectorized. When trying to reuse the extractelement instruction, emitted for the insertelement instruction, need to check, if the this insertelement instruction was vectorized. In this case, need to use vectorized value, not the original insertelement.
llvm · Jan 4, 2024 · 79e6231 · 79e6231
1 parent 96c23eb
commit 79e6231
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 2 deletions.
diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -11905,8 +11905,10 @@ Value *BoUpSLP::vectorizeTree(
         if (!Ex) {
           // "Reuse" the existing extract to improve final codegen.
           if (auto *ES = dyn_cast<ExtractElementInst>(Scalar)) {
-            Ex = Builder.CreateExtractElement(ES->getOperand(0),
-                                              ES->getOperand(1));
+            Value *V = ES->getVectorOperand();
+            if (const TreeEntry *ETE = getTreeEntry(V))
+              V = ETE->VectorizedValue;
+            Ex = Builder.CreateExtractElement(V, ES->getIndexOperand());
           } else {
             Ex = Builder.CreateExtractElement(Vec, Lane);
           }

diff --git a/llvm/test/Transforms/SLPVectorizer/X86/gather_extract_from_vectorbuild.ll b/llvm/test/Transforms/SLPVectorizer/X86/gather_extract_from_vectorbuild.ll
@@ -58,3 +58,48 @@ loop:
   %i5 = extractelement <2 x float> %ins1, i64 1
   br label %loop
 }
+
+define void @test2() {
+; CHECK-LABEL: define void @test2() {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[BB3:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[PH:%.*]] = phi float [ poison, [[BB2:%.*]] ], [ [[TMP3:%.*]], [[LOOP:%.*]] ]
+; CHECK-NEXT:    unreachable
+; CHECK:       bb2:
+; CHECK-NEXT:    br i1 poison, label [[BB3]], label [[BB1:%.*]]
+; CHECK:       bb3:
+; CHECK-NEXT:    br label [[LOOP]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[TMP0:%.*]] = phi <2 x float> [ zeroinitializer, [[BB3]] ], [ [[TMP2:%.*]], [[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]] = extractelement <2 x float> [[TMP2]], i64 1
+; CHECK-NEXT:    br i1 poison, label [[BB1]], label [[LOOP]]
+;
+entry:
+  br label %bb3
+
+bb1:
+  %ph = phi float [ poison, %bb2 ], [ %i5, %loop ]
+  unreachable
+
+bb2:
+  br i1 poison, label %bb3, label %bb1
+
+bb3:
+  br label %loop
+
+loop:
+  %ph0 = phi float [ 0.000000e+00, %bb3 ], [ %i4, %loop ]
+  %ph1 = phi float [ 0.000000e+00, %bb3 ], [ %i5, %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
+  %i5 = extractelement <2 x float> %ins1, i64 1
+  br i1 poison, label %bb1, label %loop
+}