[COST] Improve shuffle kind detection if shuffle mask is provided.

[ABataev]

Summary

Added an extra analysis for better choosing of shuffle kind in
getShuffleCost functions for better cost estimation if mask was
provided.

[CarolineConcatto]

Hello @ABataev,
Thank you for the patch and for adding me as a reviewer.
I did a comment on your patch, hope you don't mind.
Carol

llvm/include/llvm/CodeGen/BasicTTIImpl.h

798

Why are you changing this to broadcast I don't see that happening in SLPVectorizer.cpp.
Before your implementation

Kind == TTI::SK_PermuteSingleSrc

would call:

return getPermuteShuffleOverhead(cast<FixedVectorType>(Tp));

But now it calls:

return getBroadcastShuffleOverhead(cast<FixedVectorType>(Tp));

Is that correct?

In my view, some tests would fail because now it has another cost model being called.

Oddly, there is not test failure, despite, in my view and with my knowledge, it would make some tests fail because you are changing the code path for the cost model.

Do you mind adding some tests for that as well?

[ABataev]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

798

Yes, there was no such thing in SLP vectorizer, I just extended functionality. Will try to add the tests for it.

[ABataev]

Added detection for broadcast case.

[sdesmalen]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

791

nit: maybe improveShuffleKindFromMask is more descriptive?

793–807

nit: This is probably just personal preference, but can this be implemented with a switch? i.e.

if (Mask.empty())
  return Kind;

switch (Kind) {
default:
  break;
case Kind == TTI::SK_PermuteSingleSrc:
  if (ShuffleVectorInst::isReverseMask(Mask))
    return TargetTransformInfo::SK_Reverse;
  if (ShuffleVectorInst::isZeroEltSplatMask(Mask))
    return TargetTransformInfo::SK_Broadcast;
  break;
case TTI::SK_PermuteTwoSrc:
  if (ShuffleVectorInst::isSelectMask(Mask))
    return TargetTransformInfo::SK_Select;
  if (ShuffleVectorInst::isTransposeMask(Mask))
    return TargetTransformInfo::SK_Transpose;
  break;
}
return Kind;

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

4311

Is this change supposed to be in this patch?

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

4311

Ye, without it SLP is unable to detect broadcast case. I can prepare a separate patch for it.

[ABataev]

Address comments

[ABataev]

Rebase

[dmgreen]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

803

Some of these Mask functions do not handle all sizes of mask. Do we need to be careful about shuffles that changesLength?

[ABataev]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

803

Here we don't have info about length changes, looks like all these functions use the size of the Mask array as vector factor.

[dmgreen]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

803

I may have the wrong impression of what is going on exactly, but I think it's about the elements being out of range for the mask size.
(Hopefully) Try this example:
https://godbolt.org/z/WWrPbPxcW

[ABataev]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

803

This function does not use 2 different length, it has just one, deduced out of Mask size. It is just impossible to have a change of size situation here, I assume. We do not work with actual shuffles, functions are static.
Also, could you provide more details about your example? What shall I expect as a result!

[dmgreen]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

803

Also, could you provide more details about your example? What shall I expect as a result!

It should not crash with this patch :-)

It says for me:

SLP: Adding cost 0 for bundle that starts with   %add.i33 = fadd fast half %21, %22.
SLP: Current total cost = 0
SLP: Adding cost 6 for bundle that starts with   %21 = extractelement <8 x half> %20, i32 0.
SLP: Current total cost = 6
clang: /work/llvm-project/llvm/lib/IR/Instructions.cpp:2122: bool isSingleSourceMaskImpl(llvm::ArrayRef<int>, int): Assertion `I
>= 0 && I < (NumOpElts * 2) && "Out-of-bounds shuffle mask element"' failed.
...
0.      Program arguments: bin/clang -target arm-arm-none-eabi -mcpu=cortex-m55 -Ofast tmp.c -S -o - -mllvm -debug -mllvm -print-
after-all
...
 #9 0x0000558afa5bc6de llvm::ShuffleVectorInst::isSingleSourceMask(llvm::ArrayRef<int>) (bin/clang+0x728e6de)
#10 0x0000558afa5bc911 llvm::ShuffleVectorInst::isSelectMask(llvm::ArrayRef<int>) (bin/clang+0x728e911)
#11 0x0000558af90f9c40 llvm::ARMTTIImpl::getShuffleCost(llvm::TargetTransformInfo::ShuffleKind, llvm::VectorType*, llvm::ArrayRef
<int>, int, llvm::VectorType*) ARMTargetTransformInfo.cpp:0:0
#12 0x0000558af9fcac6a llvm::TargetTransformInfo::getShuffleCost(llvm::TargetTransformInfo::ShuffleKind, llvm::VectorType*, llvm:
:ArrayRef<int>, int, llvm::VectorType*) const (bin/clang+0x6c9cc6a)
#13 0x0000558afb0d4f50 llvm::slpvectorizer::BoUpSLP::getEntryCost(llvm::slpvectorizer::BoUpSLP::TreeEntry*) (bin/clang+0x7da6f50)

[ABataev]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

803

Also, could you provide more details about your example? What shall I expect as a result!

It should not crash with this patch :-)

It says for me:

SLP: Adding cost 0 for bundle that starts with   %add.i33 = fadd fast half %21, %22.
SLP: Current total cost = 0
SLP: Adding cost 6 for bundle that starts with   %21 = extractelement <8 x half> %20, i32 0.
SLP: Current total cost = 6
clang: /work/llvm-project/llvm/lib/IR/Instructions.cpp:2122: bool isSingleSourceMaskImpl(llvm::ArrayRef<int>, int): Assertion `I
>= 0 && I < (NumOpElts * 2) && "Out-of-bounds shuffle mask element"' failed.
...
0.      Program arguments: bin/clang -target arm-arm-none-eabi -mcpu=cortex-m55 -Ofast tmp.c -S -o - -mllvm -debug -mllvm -print-
after-all
...
 #9 0x0000558afa5bc6de llvm::ShuffleVectorInst::isSingleSourceMask(llvm::ArrayRef<int>) (bin/clang+0x728e6de)
#10 0x0000558afa5bc911 llvm::ShuffleVectorInst::isSelectMask(llvm::ArrayRef<int>) (bin/clang+0x728e911)
#11 0x0000558af90f9c40 llvm::ARMTTIImpl::getShuffleCost(llvm::TargetTransformInfo::ShuffleKind, llvm::VectorType*, llvm::ArrayRef
<int>, int, llvm::VectorType*) ARMTargetTransformInfo.cpp:0:0
#12 0x0000558af9fcac6a llvm::TargetTransformInfo::getShuffleCost(llvm::TargetTransformInfo::ShuffleKind, llvm::VectorType*, llvm:
:ArrayRef<int>, int, llvm::VectorType*) const (bin/clang+0x6c9cc6a)
#13 0x0000558afb0d4f50 llvm::slpvectorizer::BoUpSLP::getEntryCost(llvm::slpvectorizer::BoUpSLP::TreeEntry*) (bin/clang+0x7da6f50)

803

Ok, I see. I added some checks in SLP vectorizer but looks like we need some extra checks in other places. Or relax asserts in isSingleSourceMaskImpl function (maybe, using an extra parameter)

[ABataev]

Rebase + generalized a check for mask overflow

[ABataev]

Rebase

[dmgreen]

Thanks. LGTM

llvm/include/llvm/CodeGen/BasicTTIImpl.h

if (auto *VTy = dyn_cast<VectorType>(Ty)) {
       SmallVector<Type *> Tys(Args.size(), Ty);
       return getScalarizationOverhead(VTy, Args, Tys) + Num * Cost;
     }
 
     // We don't know anything about this scalar instruction.
     return OpCost;
   }
 
+  TTI::ShuffleKind improveShuffleKindFromMask(TTI::ShuffleKind Kind,
+                                              ArrayRef<int> Mask) const {
+    int Limit = Mask.size() * 2;
+    if (Mask.empty() ||
+        // Extra check required by isSingleSourceMaskImpl function (called by
+        // ShuffleVectorInst::isSingleSourceMask).
+        any_of(Mask, [Limit](int I) { return I >= Limit; }))
+      return Kind;
+    switch (Kind) {
+    case TTI::SK_PermuteSingleSrc:
+      if (ShuffleVectorInst::isReverseMask(Mask))
+        return TTI::SK_Reverse;
+      if (ShuffleVectorInst::isZeroEltSplatMask(Mask))
+        return TTI::SK_Broadcast;
+      break;
+    case TTI::SK_PermuteTwoSrc:
+      if (ShuffleVectorInst::isSelectMask(Mask))
+        return TTI::SK_Select;
+      if (ShuffleVectorInst::isTransposeMask(Mask))
+        return TTI::SK_Transpose;
+      break;
+    case TTI::SK_Select:
+    case TTI::SK_Reverse:
+    case TTI::SK_Broadcast:
+    case TTI::SK_Transpose:
+    case TTI::SK_InsertSubvector:
+    case TTI::SK_ExtractSubvector:
+      break;
+    }
+    return Kind;
+  }
+
   InstructionCost getShuffleCost(TTI::ShuffleKind Kind, VectorType *Tp,
                                  ArrayRef<int> Mask, int Index,
                                  VectorType *SubTp) {
 
-    switch (Kind) {
+    switch (improveShuffleKindFromMask(Kind, Mask)) {
     case TTI::SK_Broadcast:
       return getBroadcastShuffleOverhead(cast<FixedVectorType>(Tp));
     case TTI::SK_Select:
     case TTI::SK_Reverse:
     case TTI::SK_Transpose:
     case TTI::SK_PermuteSingleSrc:
     case TTI::SK_PermuteTwoSrc:
       return getPermuteShuffleOverhead(cast<FixedVectorType>(Tp));

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

AArch64TTIImpl::getArithmeticReductionCost(unsigned Opcode, VectorType *ValTy,
   return BaseT::getArithmeticReductionCost(Opcode, ValTy, IsPairwiseForm,
                                            CostKind);
 }
 
 InstructionCost AArch64TTIImpl::getShuffleCost(TTI::ShuffleKind Kind,
                                                VectorType *Tp,
                                                ArrayRef<int> Mask, int Index,
                                                VectorType *SubTp) {
+  Kind = improveShuffleKindFromMask(Kind, Mask);
   if (Kind == TTI::SK_Broadcast || Kind == TTI::SK_Transpose ||
       Kind == TTI::SK_Select || Kind == TTI::SK_PermuteSingleSrc ||
       Kind == TTI::SK_Reverse) {
     static const CostTblEntry ShuffleTbl[] = {
       // Broadcast shuffle kinds can be performed with 'dup'.
       { TTI::SK_Broadcast, MVT::v8i8,  1 },
       { TTI::SK_Broadcast, MVT::v16i8, 1 },
       { TTI::SK_Broadcast, MVT::v4i16, 1 },

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

Value *GCNTTIImpl::rewriteIntrinsicWithAddressSpace(IntrinsicInst *II,
   default:
     return nullptr;
   }
 }
 
 InstructionCost GCNTTIImpl::getShuffleCost(TTI::ShuffleKind Kind,
                                            VectorType *VT, ArrayRef<int> Mask,
                                            int Index, VectorType *SubTp) {
+  Kind = improveShuffleKindFromMask(Kind, Mask);
   if (ST->hasVOP3PInsts()) {
     if (cast<FixedVectorType>(VT)->getNumElements() == 2 &&
         DL.getTypeSizeInBits(VT->getElementType()) == 16) {
       // With op_sel VOP3P instructions freely can access the low half or high
       // half of a register, so any swizzle is free.
 
       switch (Kind) {
       case TTI::SK_Broadcast:

llvm/lib/Target/ARM/ARMTargetTransformInfo.cpp

InstructionCost ARMTTIImpl::getMemcpyCost(const Instruction *I) {
   if (NumOps == -1)
     return 4;
   return NumOps;
 }
 
 InstructionCost ARMTTIImpl::getShuffleCost(TTI::ShuffleKind Kind,
                                            VectorType *Tp, ArrayRef<int> Mask,
                                            int Index, VectorType *SubTp) {
+  Kind = improveShuffleKindFromMask(Kind, Mask);
   if (ST->hasNEON()) {
     if (Kind == TTI::SK_Broadcast) {
       static const CostTblEntry NEONDupTbl[] = {
           // VDUP handles these cases.
           {ISD::VECTOR_SHUFFLE, MVT::v2i32, 1},
           {ISD::VECTOR_SHUFFLE, MVT::v2f32, 1},
           {ISD::VECTOR_SHUFFLE, MVT::v2i64, 1},
           {ISD::VECTOR_SHUFFLE, MVT::v2f64, 1},

llvm/lib/Target/SystemZ/SystemZTargetTransformInfo.cpp

InstructionCost SystemZTTIImpl::getArithmeticInstrCost(
   return BaseT::getArithmeticInstrCost(Opcode, Ty, CostKind, Op1Info, Op2Info,
                                        Opd1PropInfo, Opd2PropInfo, Args, CxtI);
 }
 
 InstructionCost SystemZTTIImpl::getShuffleCost(TTI::ShuffleKind Kind,
                                                VectorType *Tp,
                                                ArrayRef<int> Mask, int Index,
                                                VectorType *SubTp) {
+  Kind = improveShuffleKindFromMask(Kind, Mask);
   if (ST->hasVector()) {
     unsigned NumVectors = getNumVectorRegs(Tp);
 
     // TODO: Since fp32 is expanded, the shuffle cost should always be 0.
 
     // FP128 values are always in scalar registers, so there is no work
     // involved with a shuffle, except for broadcast. In that case register
     // moves are done with a single instruction per element.

llvm/lib/Target/X86/X86TargetTransformInfo.cpp

 InstructionCost X86TTIImpl::getShuffleCost(TTI::ShuffleKind Kind,
                                            VectorType *BaseTp,
                                            ArrayRef<int> Mask, int Index,
                                            VectorType *SubTp) {
   // 64-bit packed float vectors (v2f32) are widened to type v4f32.
   // 64-bit packed integer vectors (v2i32) are widened to type v4i32.
   std::pair<int, MVT> LT = TLI->getTypeLegalizationCost(DL, BaseTp);
 
+  Kind = improveShuffleKindFromMask(Kind, Mask);
   // Treat Transpose as 2-op shuffles - there's no difference in lowering.
   if (Kind == TTI::SK_Transpose)
     Kind = TTI::SK_PermuteTwoSrc;
 
   // For Broadcasts we are splatting the first element from the first input
   // register, so only need to reference that input and all the output
   // registers are the same.
   if (Kind == TTI::SK_Broadcast)

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

BoUpSLP::isGatherShuffledEntry(const TreeEntry *TE, SmallVectorImpl<int> &Mask,
   assert(VLIt != VectorizableTree.end() &&
          "Gathered values should be in the tree.");
   Mask.assign(TE->Scalars.size(), UndefMaskElem);
   Entries.clear();
   DenseMap<const TreeEntry *, int> Used;
   int NumShuffles = 0;
   for (int I = 0, E = TE->Scalars.size(); I < E; ++I) {
     Value *V = TE->Scalars[I];
     if (isa<UndefValue>(V))
       continue;
     const TreeEntry *VTE = getTreeEntry(V);
     if (!VTE) {
       // Check if it is used in one of the gathered entries.
       const auto *It =
           find_if(make_range(VectorizableTree.begin(), VLIt),
                   [V](const std::unique_ptr<TreeEntry> &EntryPtr) {
                     return EntryPtr->State == TreeEntry::NeedToGather &&
if (VTE) {
         }
       }
       int FoundLane =
           findLaneForValue(VTE->Scalars, VTE->ReuseShuffleIndices, V);
       unsigned Sz = VTE->Scalars.size();
       if (!VTE->ReuseShuffleIndices.empty())
         Sz = VTE->ReuseShuffleIndices.size();
       Mask[I] = Res.first->second * Sz + FoundLane;
-      // Extra check required by isSingleSourceMaskImpl function (called by
-      // ShuffleVectorInst::isSingleSourceMask).
-      if (Mask[I] >= 2 * E)
-        return None;
       continue;
     }
     return None;
   }
-  if (NumShuffles == 1) {
-    if (ShuffleVectorInst::isReverseMask(Mask))
-      return TargetTransformInfo::SK_Reverse;
+  if (NumShuffles == 1)
     return TargetTransformInfo::SK_PermuteSingleSrc;
-  }
-  if (NumShuffles == 2) {
-    if (ShuffleVectorInst::isSelectMask(Mask))
-      return TargetTransformInfo::SK_Select;
-    if (ShuffleVectorInst::isTransposeMask(Mask))
-      return TargetTransformInfo::SK_Transpose;
+  if (NumShuffles == 2)
     return TargetTransformInfo::SK_PermuteTwoSrc;
-  }
   return None;
 }
 
 InstructionCost
 BoUpSLP::getGatherCost(FixedVectorType *Ty,
                        const DenseSet<unsigned> &ShuffledIndices) const {
   unsigned NumElts = Ty->getNumElements();
   APInt DemandedElts = APInt::getNullValue(NumElts);

llvm/test/Transforms/SLPVectorizer/X86/remark_extract_broadcast.ll

 ; CHECK-NEXT:    store <8 x i16> [[TMP10]], <8 x i16>* [[TMP11]], align 2
 ; CHECK-NEXT:    ret void
 ;
 ; YAML:      Pass:            slp-vectorizer
 ; YAML-NEXT: Name:            StoresVectorized
 ; YAML-NEXT: Function:        fextr
 ; YAML-NEXT: Args:
 ; YAML-NEXT:   - String:          'Stores SLP vectorized with cost '
-; YAML-NEXT:   - Cost:            '-1'
+; YAML-NEXT:   - Cost:            '-4'
 ; YAML-NEXT:   - String:          ' and with tree size '
 ; YAML-NEXT:   - TreeSize:        '4'
 
 entry:
   %LD = load <8 x i16>, <8 x i16>* undef
   %V0 = extractelement <8 x i16> %LD, i32 0
   br label %t