[SLP]Fix reused extracts cost.

[ABataev]

Summary

If the extractelement instruction is used multiple times in the
different tree entries (either vectorized, or gathered), need to
compensate the scalar cost of such instructions. They are completely
removed if all users are part of the tree but we need to compensate the
cost only once for each instruction.

[RKSimon]

LGTM

[qiongsiwu1]

This patch causes a crash when compiling the IR below.

define dso_local void @foo() local_unnamed_addr {
entry:
  %0 = extractelement <2 x double> undef, i32 1
  br label %loop_do_

loop_do_:
  %a1 = phi double [ %a1, %loop_do_ ], [ %0, %entry ]
  %a2 = phi double [ undef, %loop_do_ ], [ undef, %entry ]
  br label %loop_do_
}

The patch moved the isSlat(VL) check after it processes ExtractValue, but this test case has a tree where one value is undef, and the other is an extractvalue. Should we match not only isa<ExtractElementInst, UndefValue>(V), but also an undef value itself as well?

Thanks a lot! FYI @RolandF

[ABataev]

In D114958#3176403, @qiongsiwu1 wrote:

This patch causes a crash when compiling the IR below.

define dso_local void @foo() local_unnamed_addr {
entry:
  %0 = extractelement <2 x double> undef, i32 1
  br label %loop_do_

loop_do_:
  %a1 = phi double [ %a1, %loop_do_ ], [ %0, %entry ]
  %a2 = phi double [ undef, %loop_do_ ], [ undef, %entry ]
  br label %loop_do_
}

The patch moved the isSlat(VL) check after we process ExtractValue, but this test case has a tree where one value is undef, and the other is an extractvalue. Should we match not only isa<ExtractElementInst, UndefValue>(V), but also an undef value itself as well?

Thanks a lot! FYI @RolandF

Hi, committed a101a9b64b31ed20e4d18c2e395b0a3742ad1836 yesterday to fix it, could you try to update your compiler?

[qiongsiwu1]

In D114958#3176411, @ABataev wrote:

In D114958#3176403, @qiongsiwu1 wrote:

This patch causes a crash when compiling the IR below.

define dso_local void @foo() local_unnamed_addr {
entry:
  %0 = extractelement <2 x double> undef, i32 1
  br label %loop_do_

loop_do_:
  %a1 = phi double [ %a1, %loop_do_ ], [ %0, %entry ]
  %a2 = phi double [ undef, %loop_do_ ], [ undef, %entry ]
  br label %loop_do_
}

The patch moved the isSlat(VL) check after we process ExtractValue, but this test case has a tree where one value is undef, and the other is an extractvalue. Should we match not only isa<ExtractElementInst, UndefValue>(V), but also an undef value itself as well?

Thanks a lot! FYI @RolandF

Hi, committed a101a9b64b31ed20e4d18c2e395b0a3742ad1836 yesterday to fix it, could you try to update your compiler?

Ah great! Thanks!

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

bool BoUpSLP::canReuseExtract(ArrayRef<Value *> VL, Value *OpValue,
   }
 
   return ShouldKeepOrder;
 }
 
 bool BoUpSLP::areAllUsersVectorized(Instruction *I,
                                     ArrayRef<Value *> VectorizedVals) const {
   return (I->hasOneUse() && is_contained(VectorizedVals, I)) ||
-         llvm::all_of(I->users(), [this](User *U) {
-           return ScalarToTreeEntry.count(U) > 0;
+         all_of(I->users(), [this](User *U) {
+           return ScalarToTreeEntry.count(U) > 0 || MustGather.contains(U);
          });
 }
 
 static std::pair<InstructionCost, InstructionCost>
 getVectorCallCosts(CallInst *CI, FixedVectorType *VecTy,
                    TargetTransformInfo *TTI, TargetLibraryInfo *TLI) {
   Intrinsic::ID ID = getVectorIntrinsicIDForCall(CI, TLI);
 
VecTy = FixedVectorType::get(
         IntegerType::get(F->getContext(), MinBWs[VL[0]].first), VL.size());
   unsigned EntryVF = E->getVectorFactor();
   auto *FinalVecTy = FixedVectorType::get(VecTy->getElementType(), EntryVF);
 
   bool NeedToShuffleReuses = !E->ReuseShuffleIndices.empty();
   // FIXME: it tries to fix a problem with MSVC buildbots.
   TargetTransformInfo &TTIRef = *TTI;
   auto &&AdjustExtractsCost = [this, &TTIRef, CostKind, VL, VecTy,
-                               VectorizedVals](InstructionCost &Cost,
-                                               bool IsGather) {
+                               VectorizedVals, E](InstructionCost &Cost) {
     DenseMap<Value *, int> ExtractVectorsTys;
+    SmallPtrSet<Value *, 4> CheckedExtracts;
     for (auto *V : VL) {
       if (isa<UndefValue>(V))
         continue;
       // If all users of instruction are going to be vectorized and this
       // instruction itself is not going to be vectorized, consider this
       // instruction as dead and remove its cost from the final cost of the
       // vectorized tree.
-      if (!areAllUsersVectorized(cast<Instruction>(V), VectorizedVals))
+      // Also, avoid adjusting the cost for extractelements with multiple uses
+      // in different graph entries.
+      const TreeEntry *VE = getTreeEntry(V);
+      if (!CheckedExtracts.insert(V).second ||
+          !areAllUsersVectorized(cast<Instruction>(V), VectorizedVals) ||
+          (VE && VE != E))
         continue;
       auto *EE = cast<ExtractElementInst>(V);
       Optional<unsigned> EEIdx = getExtractIndex(EE);
       if (!EEIdx)
         continue;
       unsigned Idx = *EEIdx;
       if (TTIRef.getNumberOfParts(VecTy) !=
           TTIRef.getNumberOfParts(EE->getVectorOperandType())) {
if (Shuffle.hasValue()) {
         // Detected that instead of gather we can emit a shuffle of single/two
         // previously vectorized nodes. Add the cost of the permutation rather
         // than gather.
         ::addMask(Mask, E->ReuseShuffleIndices);
         GatherCost = TTI->getShuffleCost(*Shuffle, FinalVecTy, Mask);
       }
       return GatherCost;
     }
-    if (isSplat(VL)) {
-      // Found the broadcasting of the single scalar, calculate the cost as the
-      // broadcast.
-      return TTI->getShuffleCost(TargetTransformInfo::SK_Broadcast, VecTy);
-    }
     if ((E->getOpcode() == Instruction::ExtractElement ||
          all_of(E->Scalars,
                 [](Value *V) {
                   return isa<ExtractElementInst, UndefValue>(V);
                 })) &&
         allSameType(VL)) {
       // Check that gather of extractelements can be represented as just a
       // shuffle of a single/two vectors the scalars are extracted from.
       SmallVector<int> Mask;
       Optional<TargetTransformInfo::ShuffleKind> ShuffleKind =
           isFixedVectorShuffle(VL, Mask);
       if (ShuffleKind.hasValue()) {
         // Found the bunch of extractelement instructions that must be gathered
         // into a vector and can be represented as a permutation elements in a
         // single input vector or of 2 input vectors.
         InstructionCost Cost =
             computeExtractCost(VL, VecTy, *ShuffleKind, Mask, *TTI);
-        AdjustExtractsCost(Cost, /*IsGather=*/true);
+        AdjustExtractsCost(Cost);
         if (NeedToShuffleReuses)
           Cost += TTI->getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc,
                                       FinalVecTy, E->ReuseShuffleIndices);
         return Cost;
       }
     }
+    if (isSplat(VL)) {
+      // Found the broadcasting of the single scalar, calculate the cost as the
+      // broadcast.
+      return TTI->getShuffleCost(TargetTransformInfo::SK_Broadcast, VecTy);
+    }
     InstructionCost ReuseShuffleCost = 0;
     if (NeedToShuffleReuses)
       ReuseShuffleCost = TTI->getShuffleCost(
           TTI::SK_PermuteSingleSrc, FinalVecTy, E->ReuseShuffleIndices);
     // Improve gather cost for gather of loads, if we can group some of the
     // loads into vector loads.
     if (VL.size() > 2 && E->getOpcode() == Instruction::Load &&
         !E->isAltShuffle()) {
case Instruction::ExtractElement: {
                   TTI::getCastContextHint(Ext), CostKind, Ext);
               continue;
             }
           }
           CommonCost -=
               TTI->getVectorInstrCost(Instruction::ExtractElement, VecTy, I);
         }
       } else {
-        AdjustExtractsCost(CommonCost, /*IsGather=*/false);
+        AdjustExtractsCost(CommonCost);
       }
       return CommonCost;
     }
     case Instruction::InsertElement: {
       assert(E->ReuseShuffleIndices.empty() &&
              "Unique insertelements only are expected.");
       auto *SrcVecTy = cast<FixedVectorType>(VL0->getType());
 

llvm/test/Transforms/SLPVectorizer/X86/extractelement-multiple-uses.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 -pass-remarks-output=%t | FileCheck %s
 ; RUN: FileCheck %s --input-file=%t --check-prefix=YAML
 
-; YAML: --- !Missed
+; YAML: --- !Passed
 ; YAML: Pass:            slp-vectorizer
-; YAML: Name:            NotBeneficial
+; YAML: Name:            VectorizedList
 ; YAML: Function:        multi_uses
 ; YAML: Args:
-; YAML:  - String:          'List vectorization was possible but not beneficial with cost '
-; YAML:  - Cost:            '0'
-; YAML:  - String:          ' >= '
-; YAML:  - Treshold:        '0'
+; YAML:  - String:          'SLP vectorized with cost '
+; YAML:  - Cost:            '-1'
+; YAML:  - String:          ' and with tree size '
+; YAML:  - TreeSize:        '3'
 
 define float @multi_uses(<2 x float> %x, <2 x float> %y) {
 ; CHECK-LABEL: @multi_uses(
-; CHECK-NEXT:    [[X0:%.*]] = extractelement <2 x float> [[X:%.*]], i32 0
-; CHECK-NEXT:    [[X1:%.*]] = extractelement <2 x float> [[X]], i32 1
 ; CHECK-NEXT:    [[Y1:%.*]] = extractelement <2 x float> [[Y:%.*]], i32 1
-; CHECK-NEXT:    [[X0X0:%.*]] = fmul float [[X0]], [[Y1]]
-; CHECK-NEXT:    [[X1X1:%.*]] = fmul float [[X1]], [[Y1]]
-; CHECK-NEXT:    [[ADD:%.*]] = fadd float [[X0X0]], [[X1X1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x float> poison, float [[Y1]], i32 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <2 x float> [[TMP1]], float [[Y1]], i32 1
+; CHECK-NEXT:    [[TMP3:%.*]] = fmul <2 x float> [[X:%.*]], [[TMP2]]
+; CHECK-NEXT:    [[TMP4:%.*]] = extractelement <2 x float> [[TMP3]], i32 0
+; CHECK-NEXT:    [[TMP5:%.*]] = extractelement <2 x float> [[TMP3]], i32 1
+; CHECK-NEXT:    [[ADD:%.*]] = fadd float [[TMP4]], [[TMP5]]
 ; CHECK-NEXT:    ret float [[ADD]]
 ;
   %x0 = extractelement <2 x float> %x, i32 0
   %x1 = extractelement <2 x float> %x, i32 1
   %y1 = extractelement <2 x float> %y, i32 1
   %x0x0 = fmul float %x0, %y1
   %x1x1 = fmul float %x1, %y1
   %add = fadd float %x0x0, %x1x1
   ret float %add
 }