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

Philip Reames via llvm-commits llvm-commits at lists.llvm.org
Fri Mar 15 07:56:57 PDT 2024


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

>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 1/6] [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;
 }
 

>From cb5ced00f0c6a7d02223b7f294ce6996ccdd882e Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Thu, 14 Mar 2024 16:30:41 -0700
Subject: [PATCH 2/6] 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 502c7a3d6bae39..b5b03bef6f8545 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -10516,7 +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);
+  SmallVector<int> ShuffleMask(VL.size(), PoisonMaskElem);
   for (unsigned I = 0, E = VL.size(); I < E; ++I) {
     Value *V = VL[I];
     // No need to shuffle duplicates for constants.

>From b86bc2de0d1ec3caa70b199672525845362d01d2 Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Fri, 15 Mar 2024 07:34:07 -0700
Subject: [PATCH 3/6] compute the correct shuffle mask

---
 llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index b5b03bef6f8545..c8efcbe7cafbde 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -10507,7 +10507,7 @@ InstructionCost BoUpSLP::getGatherCost(ArrayRef<Value *> VL,
   // Check if the same elements are inserted several times and count them as
   // shuffle candidates.
   APInt ShuffledElements = APInt::getZero(VL.size());
-  DenseSet<Value *> UniqueElements;
+  DenseMap<Value *, unsigned> UniqueElements;
   constexpr TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput;
   InstructionCost Cost;
   auto EstimateInsertCost = [&](unsigned I, Value *V) {
@@ -10522,15 +10522,20 @@ InstructionCost BoUpSLP::getGatherCost(ArrayRef<Value *> VL,
     // No need to shuffle duplicates for constants.
     if ((ForPoisonSrc && isConstant(V)) || isa<UndefValue>(V)) {
       ShuffledElements.setBit(I);
+      ShuffleMask[I] = I;
       continue;
     }
-    if (!UniqueElements.insert(V).second) {
-      DuplicateNonConst = true;
-      ShuffledElements.setBit(I);
+
+    if (!UniqueElements.count(V)) {
+      EstimateInsertCost(I, V);
+      UniqueElements[V] = I;
       ShuffleMask[I] = I;
       continue;
     }
-    EstimateInsertCost(I, V);
+
+    DuplicateNonConst = true;
+    ShuffledElements.setBit(I);
+    ShuffleMask[I] = UniqueElements[V];
   }
   if (ForPoisonSrc)
     Cost =

>From 836ef4c2debea36404fe608cbba035ace0f43f7d Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Fri, 15 Mar 2024 07:56:20 -0700
Subject: [PATCH 4/6] 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 c8efcbe7cafbde..fcc4468ced2b66 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -10522,7 +10522,7 @@ InstructionCost BoUpSLP::getGatherCost(ArrayRef<Value *> VL,
     // No need to shuffle duplicates for constants.
     if ((ForPoisonSrc && isConstant(V)) || isa<UndefValue>(V)) {
       ShuffledElements.setBit(I);
-      ShuffleMask[I] = I;
+      ShuffleMask[I] = isa<PoisonValue>(V) ? PoisonMaskElem : I;
       continue;
     }
 

>From 7aa4fd265c5513a1d76fe0c5c123e8eef200c6ae Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Fri, 15 Mar 2024 07:56:41 -0700
Subject: [PATCH 5/6] Update llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

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

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index fcc4468ced2b66..9a0161a776fcd9 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -10526,9 +10526,9 @@ InstructionCost BoUpSLP::getGatherCost(ArrayRef<Value *> VL,
       continue;
     }
 
-    if (!UniqueElements.count(V)) {
+    auto Res = UniqueElements.try_emplace(V, I);
+    if (Res.second) {
       EstimateInsertCost(I, V);
-      UniqueElements[V] = I;
       ShuffleMask[I] = I;
       continue;
     }

>From 41152f6f711e82883625e0042d79d34d516d8db8 Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Fri, 15 Mar 2024 07:56:49 -0700
Subject: [PATCH 6/6] 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 9a0161a776fcd9..10f382f2480d2c 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -10535,7 +10535,7 @@ InstructionCost BoUpSLP::getGatherCost(ArrayRef<Value *> VL,
 
     DuplicateNonConst = true;
     ShuffledElements.setBit(I);
-    ShuffleMask[I] = UniqueElements[V];
+    ShuffleMask[I] = Res.first->second;
   }
   if (ForPoisonSrc)
     Cost =



More information about the llvm-commits mailing list