[llvm] [SLP] Compute a shuffle mask for getGatherCost (PR #85330)

[Philip Reames via llvm-commits]

https://github.com/preames created https://github.com/llvm/llvm-project/pull/85330

This is the second of a series 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 e33049503f10825d055da801e30511993ca1bdbd Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Thu, 14 Mar 2024 16:13:00 -0700
Subject: [PATCH] [SLP] Compute a shuffle mask for getGatherCost

This is the second of a series 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.
---
 llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index b4cce680e2876f..502c7a3d6bae39 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -10516,6 +10516,7 @@ InstructionCost BoUpSLP::getGatherCost(ArrayRef<Value *> VL,
           TTI->getVectorInstrCost(Instruction::InsertElement, VecTy, CostKind,
                                   I, Constant::getNullValue(VecTy), V);
   };
+  SmallVector<int> ShuffleMask(VL.size(), -1);
   for (unsigned I = 0, E = VL.size(); I < E; ++I) {
     Value *V = VL[I];
     // No need to shuffle duplicates for constants.
@@ -10526,6 +10527,7 @@ InstructionCost BoUpSLP::getGatherCost(ArrayRef<Value *> VL,
     if (!UniqueElements.insert(V).second) {
       DuplicateNonConst = true;
       ShuffledElements.setBit(I);
+      ShuffleMask[I] = I;
       continue;
     }
     EstimateInsertCost(I, V);
@@ -10535,8 +10537,8 @@ InstructionCost BoUpSLP::getGatherCost(ArrayRef<Value *> VL,
         TTI->getScalarizationOverhead(VecTy, ~ShuffledElements, /*Insert*/ true,
                                       /*Extract*/ false, CostKind);
   if (DuplicateNonConst)
-    Cost +=
-        TTI->getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc, VecTy);
+    Cost += TTI->getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc,
+                                VecTy, ShuffleMask);
   return Cost;
 }