[SLP] Fix for PR31879: vectorize repeated scalar ops that don't get put back into a vector

[ABataev]

Summary

Previously the cost of the existing ExtractElement/ExtractValue
instructions was considered as a dead cost only if it was detected that
they have only one use. But these instructions may be considered
dead also if users of the instructions are also going to be vectorized,
like:

%x0 = extractelement <2 x float> %x, i32 0
%x1 = extractelement <2 x float> %x, i32 1
%x0x0 = fmul float %x0, %x0
%x1x1 = fmul float %x1, %x1
%add = fadd float %x0x0, %x1x1

This can be transformed to

%1 = fmul <2 x float> %x, %x
%2 = extractelement <2 x float> %1, i32 0
%3 = extractelement <2 x float> %1, i32 1
%add = fadd float %2, %3

because though %x0 and %x1 have 2 users each other, these users are
part of the vectorized tree and we can consider these extractelement
instructions as dead.

[hfinkel]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1629 ↗

(On Diff #88234)

Maybe this should be called hasOneUser?

1672 ↗

(On Diff #88234)

Why are we checking this use(r) count at all? If canReuseExtract is true, then don't we just care that all users are part of the to-be-vectorized tree?

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1629 ↗

(On Diff #88234)

Ok, will do.

1672 ↗

(On Diff #88234)

It does not check if we can vectorize extract, it checks should we remove this extract from the code or there are other users of this extract. If we can remove this extract from the code, we can consider this instruction as dead and subtract from the cost, otherwise this instruction is alive and its cost should be considered during vectorization.

[ABataev]

Update after review

[hfinkel]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1672 ↗

(On Diff #88234)

Right, but what's special in this regard about all of the uses belonging to the same instruction? What if the test case, for example, looked like this:

%x0 = extractelement <2 x float> %x, i32 0
%x1 = extractelement <2 x float> %x, i32 1
%x0x1 = fmul float %x0, %x1
%x1x1 = fmul float %x1, %x1
%add = fadd float %x0x1, %x1x1
ret float %add

}

are the costs of all of these extracts still accounted for correctly? The final code presumably has a shuffle in place of one of the extracts?

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1672 ↗

(On Diff #88234)

Oh, now I see what are you talking about. But we still may run into something like this:

@a = global float 0.000000e+00, align 4

define float @f(<2 x float> %x) {
  %x0 = extractelement <2 x float> %x, i32 0
  %x1 = extractelement <2 x float> %x, i32 1
  %x0x0 = fmul float %x0, %x0
  %x1x1 = fmul float %x1, %x1
  %add = fadd float %x0x0, %x1x1
  store float %add, float* @a
  ret float %x0
}

where we have ret user of %x0, which is not a part of the vectorized tree. What we should do is to check that all users of extractelement instructions are going to be vectorized.

[ABataev]

Consider extractelement instructions dead if all users are going to be vectorized.

[hfinkel]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1663 ↗

(On Diff #88357)

Why is the hasOneUse check still needed if we're checking the ScalarToTreeEntry map? Or is this just a microoptimization?

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1663 ↗

(On Diff #88357)

Yes, it is a microoptimization. If we have only onlу user, this user will be vectorized/removed for sure

[hfinkel]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1663 ↗

(On Diff #88357)

Okay, please add a comment explaining that. Otherwise, LGTM.

llvm/trunk/lib/Transforms/Vectorize/SLPVectorizer.cpp

case Instruction::PHI: {
       return 0;
     }
     case Instruction::ExtractValue:
     case Instruction::ExtractElement: {
       if (canReuseExtract(VL, Opcode)) {
         int DeadCost = 0;
         for (unsigned i = 0, e = VL.size(); i < e; ++i) {
           Instruction *E = cast<Instruction>(VL[i]);
-          if (E->hasOneUse())
+          // If all users are going to be vectorized, instruction can be
+          // considered as dead.
+          // The same, if have only one user, it will be vectorized for sure.
+          if (E->hasOneUse() ||
+              std::all_of(E->user_begin(), E->user_end(), [this](User *U) {
+                return ScalarToTreeEntry.count(U) > 0;
+              }))
             // Take credit for instruction that will become dead.
             DeadCost +=
                 TTI->getVectorInstrCost(Instruction::ExtractElement, VecTy, i);
         }
         return -DeadCost;
       }
       return getGatherCost(VecTy);
     }

llvm/trunk/test/Transforms/SLPVectorizer/X86/extractelement.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -slp-vectorizer -S -mtriple=x86_64-unknown-linux -march=core-avx2 | FileCheck %s
 ; RUN: opt < %s -slp-vectorizer -S -mtriple=x86_64-unknown-linux -march=core-avx2 -slp-threshold=-1 -slp-vectorize-hor-store | FileCheck %s --check-prefix=THRESH1
 ; RUN: opt < %s -slp-vectorizer -S -mtriple=x86_64-unknown-linux -march=core-avx2 -slp-threshold=-2 -slp-vectorize-hor-store | FileCheck %s --check-prefix=THRESH2
 
 @a = global float 0.000000e+00, align 4
 
 define float @f(<2 x float> %x) {
 ; CHECK-LABEL: @f(
-; CHECK-NEXT:    [[X0:%.*]] = extractelement <2 x float> [[X:%.*]], i32 0
-; CHECK-NEXT:    [[X1:%.*]] = extractelement <2 x float> [[X]], i32 1
-; CHECK-NEXT:    [[X0X0:%.*]] = fmul float [[X0]], [[X0]]
-; CHECK-NEXT:    [[X1X1:%.*]] = fmul float [[X1]], [[X1]]
-; CHECK-NEXT:    [[ADD:%.*]] = fadd float [[X0X0]], [[X1X1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = fmul <2 x float> [[X:%.*]], [[X]]
+; CHECK-NEXT:    [[TMP2:%.*]] = extractelement <2 x float> [[TMP1]], i32 0
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <2 x float> [[TMP1]], i32 1
+; CHECK-NEXT:    [[ADD:%.*]] = fadd float [[TMP2]], [[TMP3]]
 ; CHECK-NEXT:    ret float [[ADD]]
 ;
   %x0 = extractelement <2 x float> %x, i32 0
   %x1 = extractelement <2 x float> %x, i32 1
   %x0x0 = fmul float %x0, %x0
   %x1x1 = fmul float %x1, %x1
   %add = fadd float %x0x0, %x1x1
   ret float %add
 }
 
 define float @f_used_out_of_tree(<2 x float> %x) {
 ; THRESH2-LABEL: @f_used_out_of_tree(
-; THRESH2-NEXT:    [[X0:%.*]] = extractelement <2 x float> [[X:%.*]], i32 0
-; THRESH2-NEXT:    [[X1:%.*]] = extractelement <2 x float> [[X]], i32 1
-; THRESH2-NEXT:    [[X0X0:%.*]] = fmul float [[X0]], [[X0]]
-; THRESH2-NEXT:    [[X1X1:%.*]] = fmul float [[X1]], [[X1]]
-; THRESH2-NEXT:    [[ADD:%.*]] = fadd float [[X0X0]], [[X1X1]]
+; THRESH2-NEXT:    [[TMP1:%.*]] = extractelement <2 x float> [[X:%.*]], i32 0
+; THRESH2-NEXT:    [[TMP2:%.*]] = fmul <2 x float> [[X]], [[X]]
+; THRESH2-NEXT:    [[TMP3:%.*]] = extractelement <2 x float> [[TMP2]], i32 0
+; THRESH2-NEXT:    [[TMP4:%.*]] = extractelement <2 x float> [[TMP2]], i32 1
+; THRESH2-NEXT:    [[ADD:%.*]] = fadd float [[TMP3]], [[TMP4]]
 ; THRESH2-NEXT:    store float [[ADD]], float* @a
-; THRESH2-NEXT:    ret float [[X0]]
+; THRESH2-NEXT:    ret float [[TMP1]]
 ;
   %x0 = extractelement <2 x float> %x, i32 0
   %x1 = extractelement <2 x float> %x, i32 1
   %x0x0 = fmul float %x0, %x0
   %x1x1 = fmul float %x1, %x1
   %add = fadd float %x0x0, %x1x1
   store float %add, float* @a
   ret float %x0
 }
 
 define float @f_used_twice_in_tree(<2 x float> %x) {
 ; THRESH1-LABEL: @f_used_twice_in_tree(
-; THRESH1-NEXT:    [[X0:%.*]] = extractelement <2 x float> [[X:%.*]], i32 0
-; THRESH1-NEXT:    [[X1:%.*]] = extractelement <2 x float> [[X]], i32 1
-; THRESH1-NEXT:    [[X0X0:%.*]] = fmul float [[X0]], [[X1]]
-; THRESH1-NEXT:    [[X1X1:%.*]] = fmul float [[X1]], [[X1]]
-; THRESH1-NEXT:    [[ADD:%.*]] = fadd float [[X0X0]], [[X1X1]]
+; THRESH1-NEXT:    [[TMP1:%.*]] = extractelement <2 x float> [[X:%.*]], i32 1
+; THRESH1-NEXT:    [[TMP2:%.*]] = insertelement <2 x float> undef, float [[TMP1]], i32 0
+; THRESH1-NEXT:    [[TMP3:%.*]] = insertelement <2 x float> [[TMP2]], float [[TMP1]], i32 1
+; THRESH1-NEXT:    [[TMP4:%.*]] = fmul <2 x float> [[X]], [[TMP3]]
+; THRESH1-NEXT:    [[TMP5:%.*]] = extractelement <2 x float> [[TMP4]], i32 0
+; THRESH1-NEXT:    [[TMP6:%.*]] = extractelement <2 x float> [[TMP4]], i32 1
+; THRESH1-NEXT:    [[ADD:%.*]] = fadd float [[TMP5]], [[TMP6]]
 ; THRESH1-NEXT:    ret float [[ADD]]
+;
   %x0 = extractelement <2 x float> %x, i32 0
   %x1 = extractelement <2 x float> %x, i32 1
   %x0x0 = fmul float %x0, %x1
   %x1x1 = fmul float %x1, %x1
   %add = fadd float %x0x0, %x1x1
   ret float %add
 }