[llvm] f337525 - [SLP] Compute a shuffle mask for SK_Broadcast shuffle (#85327)

[via llvm-commits]

Author: Philip Reames
Date: 2024-03-15T07:41:26-07:00
New Revision: f337525ee8b44da0f1e98eecd7f51bf5805c3760

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

LOG: [SLP] Compute a shuffle mask for SK_Broadcast shuffle (#85327)

This is the first of a couple of small patches to compute shuffle masks
for the couple of cases where we call getShuffleCost without one. My
goal is to add an invariant that all calls to getShuffleCost for fixed
length vectors have a mask.

---------

Co-authored-by: Alexey Bataev <a.bataev at gmx.com>

Added: 
    

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 cbddf29bb27d34..58dbff143929db 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -7691,16 +7691,22 @@ class BoUpSLP::ShuffleCostEstimator : public BaseShuffleAnalysis {
       bool NeedShuffle =
           count(VL, *It) > 1 &&
           (VL.front() != *It || !all_of(VL.drop_front(), UndefValue::classof));
+      if (!NeedShuffle)
+        return TTI.getVectorInstrCost(Instruction::InsertElement, VecTy,
+                                      CostKind, std::distance(VL.begin(), It),
+                                      PoisonValue::get(VecTy), *It);
+
+      SmallVector<int> ShuffleMask(VL.size(), PoisonMaskElem);
+      transform(VL, ShuffleMask.begin(), [](Value *V) {
+        return isa<PoisonValue>(V) ? PoisonMaskElem : 0;   
+      });
       InstructionCost InsertCost = TTI.getVectorInstrCost(
-          Instruction::InsertElement, VecTy, CostKind,
-          NeedShuffle ? 0 : std::distance(VL.begin(), It),
+          Instruction::InsertElement, VecTy, CostKind, 0,
           PoisonValue::get(VecTy), *It);
       return InsertCost +
-             (NeedShuffle ? TTI.getShuffleCost(
-                                TargetTransformInfo::SK_Broadcast, VecTy,
-                                /*Mask=*/std::nullopt, CostKind, /*Index=*/0,
-                                /*SubTp=*/nullptr, /*Args=*/*It)
-                          : TTI::TCC_Free);
+             TTI.getShuffleCost(TargetTransformInfo::SK_Broadcast, VecTy,
+                                ShuffleMask, CostKind, /*Index=*/0,
+                                /*SubTp=*/nullptr, /*Args=*/*It);
     }
     return GatherCost +
            (all_of(Gathers, UndefValue::classof)