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

[Philip Reames via llvm-commits]

https://github.com/preames updated https://github.com/llvm/llvm-project/pull/85327

>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 1/3] [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)

>From 254303ba599669eaabfcefa436b956c482090321 Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Thu, 14 Mar 2024 16:28:46 -0700
Subject: [PATCH 2/3] Update llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

Co-authored-by: Alexey Bataev <a.bataev at gmx.com>
---
 llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 885015df1f3b81..75630b500e2a1d 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -7749,7 +7749,7 @@ class BoUpSLP::ShuffleCostEstimator : public BaseShuffleAnalysis {
                                       CostKind, std::distance(VL.begin(), It),
                                       PoisonValue::get(VecTy), *It);
 
-      SmallVector<int> ShuffleMask(VL.size(), -1);
+      SmallVector<int> ShuffleMask(VL.size(), PoisonMaskElem);
       transform(VL, ShuffleMask.begin(), [](Value *V) {
         return !isa<UndefValue>(V);
       });

>From c0cd3151de14e736af1048c298d6968514c7a6e0 Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Thu, 14 Mar 2024 16:29:14 -0700
Subject: [PATCH 3/3] Update llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

Co-authored-by: Alexey Bataev <a.bataev at gmx.com>
---
 llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 75630b500e2a1d..0abd56b32ceabc 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -7751,7 +7751,7 @@ class BoUpSLP::ShuffleCostEstimator : public BaseShuffleAnalysis {
 
       SmallVector<int> ShuffleMask(VL.size(), PoisonMaskElem);
       transform(VL, ShuffleMask.begin(), [](Value *V) {
-        return !isa<UndefValue>(V);
+        return isa<PoisonValue>(V) ? PoisonMaskElem : 0;   
       });
       InstructionCost InsertCost = TTI.getVectorInstrCost(
           Instruction::InsertElement, VecTy, CostKind, 0,