Skip to content

Commit

Permalink
[SLP]Further improvement of the cost model for scalars used in buildv…
Browse files Browse the repository at this point in the history 
…ectors.

Further improvement of the cost model for the scalars used in
buildvectors sequences. The main functionality is outlined into
a separate function.
The cost is calculated in the following way:
1. If the Base vector is not undef vector, resizing the very first mask to
have common VF and perform action for 2 input vectors (including non-undef
Base). Other shuffle masks are combined with the resulting after the 1 stage and processed as a shuffle of 2 elements.
2. If the Base is undef vector and have only 1 shuffle mask, perform the
action only for 1 vector with the given mask, if it is not the identity
mask.
3. If > 2 masks are used, perform serie of shuffle actions for 2 vectors,
combing the masks properly between the steps.

The original implementation misses the very first analysis for the Base
vector, so the cost might too optimistic in some cases. But it improves
the cost for the insertelements which are part of the current SLP graph.

Part of D107966.

Differential Revision: https://reviews.llvm.org/D115750
  • Loading branch information
@alexey-bataev
alexey-bataev committed May 11, 2022
1 parent 400587b commit f5d45d7
Show file tree
Hide file tree
Showing 4 changed files with 251 additions and 133 deletions.
320 changes: 218 additions & 102 deletions llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
View file
Expand Up @@ -6579,6 +6579,126 @@ static bool areTwoInsertFromSameBuildVector(InsertElementInst *VU,
return false;
}

/// Checks if the \p IE1 instructions is followed by \p IE2 instruction in the
/// buildvector sequence.
static bool isFirstInsertElement(const InsertElementInst *IE1,
const InsertElementInst *IE2) {
const auto *I1 = IE1;
const auto *I2 = IE2;
const InsertElementInst *PrevI1;
const InsertElementInst *PrevI2;
do {
if (I2 == IE1)
return true;
if (I1 == IE2)
return false;
PrevI1 = I1;
PrevI2 = I2;
if (I1 && (I1 == IE1 || I1->hasOneUse()))
I1 = dyn_cast<InsertElementInst>(I1->getOperand(0));
if (I2 && (I2 == IE2 || I2->hasOneUse()))
I2 = dyn_cast<InsertElementInst>(I2->getOperand(0));
} while ((I1 && PrevI1 != I1) || (I2 && PrevI2 != I2));
llvm_unreachable("Two different buildvectors not expected.");
}

/// Does the analysis of the provided shuffle masks and performs the requested
/// actions on the vectors with the given shuffle masks. It tries to do it in
/// several steps.
/// 1. If the Base vector is not undef vector, resizing the very first mask to
/// have common VF and perform action for 2 input vectors (including non-undef
/// Base). Other shuffle masks are combined with the resulting after the 1 stage
/// and processed as a shuffle of 2 elements.
/// 2. If the Base is undef vector and have only 1 shuffle mask, perform the
/// action only for 1 vector with the given mask, if it is not the identity
/// mask.
/// 3. If > 2 masks are used, perform the remaining shuffle actions for 2
/// vectors, combing the masks properly between the steps.
template <typename T>
static T *performExtractsShuffleAction(
MutableArrayRef<std::pair<T *, SmallVector<int>>> ShuffleMask, Value *Base,
function_ref<unsigned(T *)> GetVF,
function_ref<std::pair<T *, bool>(T *, ArrayRef<int>)> ResizeAction,
function_ref<T *(ArrayRef<int>, ArrayRef<T *>)> Action) {
assert(!ShuffleMask.empty() && "Empty list of shuffles for inserts.");
SmallVector<int> Mask(ShuffleMask.begin()->second);
auto VMIt = std::next(ShuffleMask.begin());
T *Prev = nullptr;
bool IsBaseNotUndef = !isUndefVector(Base);
if (IsBaseNotUndef) {
// Base is not undef, need to combine it with the next subvectors.
std::pair<T *, bool> Res = ResizeAction(ShuffleMask.begin()->first, Mask);
for (unsigned Idx = 0, VF = Mask.size(); Idx < VF; ++Idx) {
if (Mask[Idx] == UndefMaskElem)
Mask[Idx] = Idx;
else
Mask[Idx] = (Res.second ? Idx : Mask[Idx]) + VF;
}
Prev = Action(Mask, {nullptr, Res.first});
} else if (ShuffleMask.size() == 1) {
// Base is undef and only 1 vector is shuffled - perform the action only for
// single vector, if the mask is not the identity mask.
std::pair<T *, bool> Res = ResizeAction(ShuffleMask.begin()->first, Mask);
if (Res.second)
// Identity mask is found.
Prev = Res.first;
else
Prev = Action(Mask, {ShuffleMask.begin()->first});
} else {
// Base is undef and at least 2 input vectors shuffled - perform 2 vectors
// shuffles step by step, combining shuffle between the steps.
unsigned Vec1VF = GetVF(ShuffleMask.begin()->first);
unsigned Vec2VF = GetVF(VMIt->first);
if (Vec1VF == Vec2VF) {
// No need to resize the input vectors since they are of the same size, we
// can shuffle them directly.
ArrayRef<int> SecMask = VMIt->second;
for (unsigned I = 0, VF = Mask.size(); I < VF; ++I) {
if (SecMask[I] != UndefMaskElem) {
assert(Mask[I] == UndefMaskElem && "Multiple uses of scalars.");
Mask[I] = SecMask[I] + Vec1VF;
}
}
Prev = Action(Mask, {ShuffleMask.begin()->first, VMIt->first});
} else {
// Vectors of different sizes - resize and reshuffle.
std::pair<T *, bool> Res1 =
ResizeAction(ShuffleMask.begin()->first, Mask);
std::pair<T *, bool> Res2 = ResizeAction(VMIt->first, VMIt->second);
ArrayRef<int> SecMask = VMIt->second;
for (unsigned I = 0, VF = Mask.size(); I < VF; ++I) {
if (Mask[I] != UndefMaskElem) {
assert(SecMask[I] == UndefMaskElem && "Multiple uses of scalars.");
if (Res1.second)
Mask[I] = I;
} else if (SecMask[I] != UndefMaskElem) {
assert(Mask[I] == UndefMaskElem && "Multiple uses of scalars.");
Mask[I] = (Res2.second ? I : SecMask[I]) + VF;
}
}
Prev = Action(Mask, {Res1.first, Res2.first});
}
VMIt = std::next(VMIt);
}
// Perform requested actions for the remaining masks/vectors.
for (auto E = ShuffleMask.end(); VMIt != E; ++VMIt) {
// Shuffle other input vectors, if any.
std::pair<T *, bool> Res = ResizeAction(VMIt->first, VMIt->second);
ArrayRef<int> SecMask = VMIt->second;
for (unsigned I = 0, VF = Mask.size(); I < VF; ++I) {
if (SecMask[I] != UndefMaskElem) {
assert((Mask[I] == UndefMaskElem || IsBaseNotUndef) &&
"Multiple uses of scalars.");
Mask[I] = (Res.second ? I : SecMask[I]) + VF;
} else if (Mask[I] != UndefMaskElem) {
Mask[I] = I;
}
}
Prev = Action(Mask, {Prev, Res.first});
}
return Prev;
}

InstructionCost BoUpSLP::getTreeCost(ArrayRef<Value *> VectorizedVals) {
InstructionCost Cost = 0;
LLVM_DEBUG(dbgs() << "SLP: Calculating cost for tree of size "
Expand All @@ -6599,9 +6719,8 @@ InstructionCost BoUpSLP::getTreeCost(ArrayRef<Value *> VectorizedVals) {

SmallPtrSet<Value *, 16> ExtractCostCalculated;
InstructionCost ExtractCost = 0;
SmallVector<unsigned> VF;
SmallVector<SmallVector<int>> ShuffleMask;
SmallVector<Value *> FirstUsers;
SmallVector<MapVector<const TreeEntry *, SmallVector<int>>> ShuffleMasks;
SmallVector<std::pair<Value *, const TreeEntry *>> FirstUsers;
SmallVector<APInt> DemandedElts;
for (ExternalUser &EU : ExternalUses) {
// We only add extract cost once for the same scalar.
Expand Down Expand Up @@ -6630,37 +6749,52 @@ InstructionCost BoUpSLP::getTreeCost(ArrayRef<Value *> VectorizedVals) {
if (auto *FTy = dyn_cast<FixedVectorType>(VU->getType())) {
Optional<unsigned> InsertIdx = getInsertIndex(VU);
if (InsertIdx) {
auto *It = find_if(FirstUsers, [VU](Value *V) {
return areTwoInsertFromSameBuildVector(VU,
cast<InsertElementInst>(V));
});
const TreeEntry *ScalarTE = getTreeEntry(EU.Scalar);
auto *It =
find_if(FirstUsers,
[VU](const std::pair<Value *, const TreeEntry *> &Pair) {
return areTwoInsertFromSameBuildVector(
VU, cast<InsertElementInst>(Pair.first));
});
int VecId = -1;
if (It == FirstUsers.end()) {
VF.push_back(FTy->getNumElements());
ShuffleMask.emplace_back(VF.back(), UndefMaskElem);
(void)ShuffleMasks.emplace_back();
SmallVectorImpl<int> &Mask = ShuffleMasks.back()[ScalarTE];
if (Mask.empty())
Mask.assign(FTy->getNumElements(), UndefMaskElem);
// Find the insertvector, vectorized in tree, if any.
Value *Base = VU;
while (auto *IEBase = dyn_cast<InsertElementInst>(Base)) {
if (IEBase != EU.User && !IEBase->hasOneUse())
break;
// Build the mask for the vectorized insertelement instructions.
if (const TreeEntry *E = getTreeEntry(IEBase)) {
VU = IEBase;
do {
int Idx = E->findLaneForValue(Base);
ShuffleMask.back()[Idx] = Idx;
Base = cast<InsertElementInst>(Base)->getOperand(0);
IEBase = cast<InsertElementInst>(Base);
int Idx = *getInsertIndex(IEBase);
assert(Mask[Idx] == UndefMaskElem &&
"InsertElementInstruction used already.");
Mask[Idx] = Idx;
Base = IEBase->getOperand(0);
} while (E == getTreeEntry(Base));
break;
}
Base = cast<InsertElementInst>(Base)->getOperand(0);
}
FirstUsers.push_back(VU);
DemandedElts.push_back(APInt::getZero(VF.back()));
FirstUsers.emplace_back(VU, ScalarTE);
DemandedElts.push_back(APInt::getZero(FTy->getNumElements()));
VecId = FirstUsers.size() - 1;
} else {
if (isFirstInsertElement(VU, cast<InsertElementInst>(It->first)))
It->first = VU;
VecId = std::distance(FirstUsers.begin(), It);
}
int InIdx = *InsertIdx;
ShuffleMask[VecId][InIdx] = EU.Lane;
SmallVectorImpl<int> &Mask = ShuffleMasks[VecId][ScalarTE];
if (Mask.empty())
Mask.assign(FTy->getNumElements(), UndefMaskElem);
Mask[InIdx] = EU.Lane;
DemandedElts[VecId].setBit(InIdx);
continue;
}
Expand All @@ -6687,89 +6821,75 @@ InstructionCost BoUpSLP::getTreeCost(ArrayRef<Value *> VectorizedVals) {

InstructionCost SpillCost = getSpillCost();
Cost += SpillCost + ExtractCost;
if (FirstUsers.size() == 1) {
int Limit = ShuffleMask.front().size() * 2;
if (!all_of(ShuffleMask.front(),
[Limit](int Idx) { return Idx < Limit; }) ||
!ShuffleVectorInst::isIdentityMask(ShuffleMask.front())) {
InstructionCost C = TTI->getShuffleCost(
auto &&ResizeToVF = [this, &Cost](const TreeEntry *TE, ArrayRef<int> Mask) {
InstructionCost C = 0;
unsigned VF = Mask.size();
unsigned VecVF = TE->getVectorFactor();
if (VF != VecVF &&
(any_of(Mask, [VF](int Idx) { return Idx >= static_cast<int>(VF); }) ||
(all_of(Mask,
[VF](int Idx) { return Idx < 2 * static_cast<int>(VF); }) &&
!ShuffleVectorInst::isIdentityMask(Mask)))) {
SmallVector<int> OrigMask(VecVF, UndefMaskElem);
std::copy(Mask.begin(), std::next(Mask.begin(), std::min(VF, VecVF)),
OrigMask.begin());
C = TTI->getShuffleCost(
TTI::SK_PermuteSingleSrc,
cast<FixedVectorType>(FirstUsers.front()->getType()),
ShuffleMask.front());
LLVM_DEBUG(dbgs() << "SLP: Adding cost " << C
<< " for final shuffle of insertelement external users "
<< *VectorizableTree.front()->Scalars.front() << ".\n"
<< "SLP: Current total cost = " << Cost << "\n");
FixedVectorType::get(TE->getMainOp()->getType(), VecVF), OrigMask);
LLVM_DEBUG(
dbgs() << "SLP: Adding cost " << C
<< " for final shuffle of insertelement external users.\n";
TE->dump(); dbgs() << "SLP: Current total cost = " << Cost << "\n");
Cost += C;
return std::make_pair(TE, true);
}
return std::make_pair(TE, false);
};
// Calculate the cost of the reshuffled vectors, if any.
for (int I = 0, E = FirstUsers.size(); I < E; ++I) {
Value *Base = cast<Instruction>(FirstUsers[I].first)->getOperand(0);
unsigned VF = ShuffleMasks[I].begin()->second.size();
auto *FTy = FixedVectorType::get(
cast<VectorType>(FirstUsers[I].first->getType())->getElementType(), VF);
auto Vector = ShuffleMasks[I].takeVector();
auto &&EstimateShufflesCost = [this, FTy,
&Cost](ArrayRef<int> Mask,
ArrayRef<const TreeEntry *> TEs) {
assert((TEs.size() == 1 || TEs.size() == 2) &&
"Expected exactly 1 or 2 tree entries.");
if (TEs.size() == 1) {
int Limit = 2 * Mask.size();
if (!all_of(Mask, [Limit](int Idx) { return Idx < Limit; }) ||
!ShuffleVectorInst::isIdentityMask(Mask)) {
InstructionCost C =
TTI->getShuffleCost(TTI::SK_PermuteSingleSrc, FTy, Mask);
LLVM_DEBUG(dbgs() << "SLP: Adding cost " << C
<< " for final shuffle of insertelement "
"external users.\n";
TEs.front()->dump();
dbgs() << "SLP: Current total cost = " << Cost << "\n");
Cost += C;
}
} else {
InstructionCost C =
TTI->getShuffleCost(TTI::SK_PermuteTwoSrc, FTy, Mask);
LLVM_DEBUG(dbgs() << "SLP: Adding cost " << C
<< " for final shuffle of vector node and external "
"insertelement users.\n";
if (TEs.front()) { TEs.front()->dump(); } TEs.back()->dump();
dbgs() << "SLP: Current total cost = " << Cost << "\n");
Cost += C;
}
return TEs.back();
};
(void)performExtractsShuffleAction<const TreeEntry>(
makeMutableArrayRef(Vector.data(), Vector.size()), Base,
[](const TreeEntry *E) { return E->getVectorFactor(); }, ResizeToVF,
EstimateShufflesCost);
InstructionCost InsertCost = TTI->getScalarizationOverhead(
cast<FixedVectorType>(FirstUsers.front()->getType()),
DemandedElts.front(), /*Insert*/ true, /*Extract*/ false);
LLVM_DEBUG(dbgs() << "SLP: subtracting the cost " << InsertCost
<< " for insertelements gather.\n"
<< "SLP: Current total cost = " << Cost << "\n");
Cost -= InsertCost;
} else if (FirstUsers.size() >= 2) {
unsigned MaxVF = *std::max_element(VF.begin(), VF.end());
// Combined masks of the first 2 vectors.
SmallVector<int> CombinedMask(MaxVF, UndefMaskElem);
copy(ShuffleMask.front(), CombinedMask.begin());
APInt CombinedDemandedElts = DemandedElts.front().zextOrSelf(MaxVF);
auto *VecTy = FixedVectorType::get(
cast<VectorType>(FirstUsers.front()->getType())->getElementType(),
MaxVF);
for (int I = 0, E = ShuffleMask[1].size(); I < E; ++I) {
if (ShuffleMask[1][I] != UndefMaskElem) {
CombinedMask[I] = ShuffleMask[1][I] + MaxVF;
CombinedDemandedElts.setBit(I);
}
}
InstructionCost C =
TTI->getShuffleCost(TTI::SK_PermuteTwoSrc, VecTy, CombinedMask);
LLVM_DEBUG(dbgs() << "SLP: Adding cost " << C
<< " for final shuffle of vector node and external "
"insertelement users "
<< *VectorizableTree.front()->Scalars.front() << ".\n"
<< "SLP: Current total cost = " << Cost << "\n");
Cost += C;
InstructionCost InsertCost = TTI->getScalarizationOverhead(
VecTy, CombinedDemandedElts, /*Insert*/ true, /*Extract*/ false);
LLVM_DEBUG(dbgs() << "SLP: subtracting the cost " << InsertCost
<< " for insertelements gather.\n"
<< "SLP: Current total cost = " << Cost << "\n");
cast<FixedVectorType>(FirstUsers[I].first->getType()), DemandedElts[I],
/*Insert*/ true, /*Extract*/ false);
Cost -= InsertCost;
for (int I = 2, E = FirstUsers.size(); I < E; ++I) {
if (ShuffleMask[I].empty())
continue;
// Other elements - permutation of 2 vectors (the initial one and the
// next Ith incoming vector).
unsigned VF = ShuffleMask[I].size();
for (unsigned Idx = 0; Idx < VF; ++Idx) {
int Mask = ShuffleMask[I][Idx];
if (Mask != UndefMaskElem)
CombinedMask[Idx] = MaxVF + Mask;
else if (CombinedMask[Idx] != UndefMaskElem)
CombinedMask[Idx] = Idx;
}
for (unsigned Idx = VF; Idx < MaxVF; ++Idx)
if (CombinedMask[Idx] != UndefMaskElem)
CombinedMask[Idx] = Idx;
InstructionCost C =
TTI->getShuffleCost(TTI::SK_PermuteTwoSrc, VecTy, CombinedMask);
LLVM_DEBUG(dbgs() << "SLP: Adding cost " << C
<< " for final shuffle of vector node and external "
"insertelement users "
<< *VectorizableTree.front()->Scalars.front() << ".\n"
<< "SLP: Current total cost = " << Cost << "\n");
Cost += C;
InstructionCost InsertCost = TTI->getScalarizationOverhead(
cast<FixedVectorType>(FirstUsers[I]->getType()), DemandedElts[I],
/*Insert*/ true, /*Extract*/ false);
LLVM_DEBUG(dbgs() << "SLP: subtracting the cost " << InsertCost
<< " for insertelements gather.\n"
<< "SLP: Current total cost = " << Cost << "\n");
Cost -= InsertCost;
}
}

#ifndef NDEBUG
Expand Down Expand Up @@ -10376,9 +10496,10 @@ class HorizontalReduction {
for (unsigned Cnt = 0; Cnt < NumReducedVals; ++Cnt) {
if (Cnt >= Pos && Cnt < Pos + ReduxWidth)
continue;
if (VectorizedVals.count(Candidates[Cnt]))
continue;
LocalExternallyUsedValues[Candidates[Cnt]];
unsigned NumOps = VectorizedVals.lookup(Candidates[Cnt]) +
std::count(VL.begin(), VL.end(), Candidates[Cnt]);
if (NumOps != ReducedValsToOps.find(Candidates[Cnt])->second.size())
LocalExternallyUsedValues[Candidates[Cnt]];
}
V.buildExternalUses(LocalExternallyUsedValues);

Expand Down Expand Up @@ -10480,10 +10601,7 @@ class HorizontalReduction {
auto TVIt = TrackedVals.find(RdxVal);
if (TVIt != TrackedVals.end())
StableRdxVal = TVIt->second;
unsigned NumOps = 0;
auto It = VectorizedVals.find(RdxVal);
if (It != VectorizedVals.end())
NumOps = It->second;
unsigned NumOps = VectorizedVals.lookup(RdxVal);
for (Instruction *RedOp :
makeArrayRef(ReducedValsToOps.find(RdxVal)->second)
.drop_back(NumOps)) {
Expand Down Expand Up @@ -10550,8 +10668,6 @@ class HorizontalReduction {
return VectorizedTree;
}

unsigned numReductionValues() const { return ReducedVals.size(); }

private:
/// Calculate the cost of a reduction.
InstructionCost getReductionCost(TargetTransformInfo *TTI,
Expand Down

0 comments on commit f5d45d7

Please sign in to comment.