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

[Philip Reames via llvm-commits]

https://github.com/preames created https://github.com/llvm/llvm-project/pull/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.

>From 6734e37e873263eecd9b2f871f1c4a0cbda85f06 Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Thu, 14 Mar 2024 15:52:22 -0700
Subject: [PATCH] [SLP] Compute a shuffle mask for SK_Broadcast shuffle

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.
---
 .../Transforms/Vectorize/SLPVectorizer.cpp    | 20 ++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index b4cce680e2876f..885015df1f3b81 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -7744,16 +7744,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(), -1);
+      transform(VL, ShuffleMask.begin(), [](Value *V) {
+        return !isa<UndefValue>(V);
+      });
       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)