[SLP]Improved isGatherShuffledEntry, NFC.

Reworked isGatherShuffledEntry function, simplified and moved
common code to the lambda (it shall go away when non-power-2 patch will
be landed).
This commit is contained in:
Alexey Bataev
2021-04-23 12:33:41 -07:00
parent 0cb75ec25f
commit 2b896fc0d5

View File

@@ -4296,70 +4296,68 @@ InstructionCost BoUpSLP::getTreeCost() {
Optional<TargetTransformInfo::ShuffleKind>
BoUpSLP::isGatherShuffledEntry(const TreeEntry *TE, SmallVectorImpl<int> &Mask,
SmallVectorImpl<const TreeEntry *> &Entries) {
auto *VLIt = find_if(VectorizableTree,
[TE](const std::unique_ptr<TreeEntry> &EntryPtr) {
return EntryPtr.get() == TE;
});
assert(VLIt != VectorizableTree.end() &&
"Gathered values should be in the tree.");
Mask.assign(TE->Scalars.size(), UndefMaskElem);
Entries.clear();
DenseMap<const TreeEntry *, int> Used;
int NumShuffles = 0;
DenseMap<Value *, const TreeEntry *> UsedValuesEntry;
unsigned VF = 0;
// FIXME: Shall be replaced by GetVF function once non-power-2 patch is
// landed.
auto &&GetVF = [](const TreeEntry *TE) {
if (!TE->ReuseShuffleIndices.empty())
return TE->ReuseShuffleIndices.size();
return TE->Scalars.size();
};
for (int I = 0, E = TE->Scalars.size(); I < E; ++I) {
Value *V = TE->Scalars[I];
if (isa<UndefValue>(V))
continue;
const TreeEntry *VTE = getTreeEntry(V);
const TreeEntry *VTE = UsedValuesEntry.lookup(V);
if (!VTE) {
// Check if it is used in one of the gathered entries.
const auto *It =
find_if(make_range(VectorizableTree.begin(), VLIt),
[V](const std::unique_ptr<TreeEntry> &EntryPtr) {
return EntryPtr->State == TreeEntry::NeedToGather &&
is_contained(EntryPtr->Scalars, V);
});
if (It != VLIt)
VTE = It->get();
}
if (VTE) {
auto Res = Used.try_emplace(VTE, NumShuffles);
if (Res.second) {
Entries.push_back(VTE);
++NumShuffles;
if (NumShuffles > 2)
return None;
if (NumShuffles == 2) {
unsigned FirstSz = Entries.front()->Scalars.size();
if (!Entries.front()->ReuseShuffleIndices.empty())
FirstSz = Entries.front()->ReuseShuffleIndices.size();
unsigned SecondSz = Entries.back()->Scalars.size();
if (!Entries.back()->ReuseShuffleIndices.empty())
SecondSz = Entries.back()->ReuseShuffleIndices.size();
if (FirstSz != SecondSz)
return None;
}
}
int FoundLane =
findLaneForValue(VTE->Scalars, VTE->ReuseShuffleIndices, V);
unsigned Sz = VTE->Scalars.size();
if (!VTE->ReuseShuffleIndices.empty())
Sz = VTE->ReuseShuffleIndices.size();
Mask[I] = Res.first->second * Sz + FoundLane;
// Extra check required by isSingleSourceMaskImpl function (called by
// ShuffleVectorInst::isSingleSourceMask).
if (Mask[I] >= 2 * E)
if (Entries.size() == 2)
return None;
continue;
VTE = getTreeEntry(V);
if (!VTE || find_if(
VectorizableTree,
[VTE, TE](const std::unique_ptr<TreeEntry> &EntryPtr) {
return EntryPtr.get() == VTE || EntryPtr.get() == TE;
})->get() == TE) {
// Check if it is used in one of the gathered entries.
const auto *It =
find_if(VectorizableTree,
[V, TE](const std::unique_ptr<TreeEntry> &EntryPtr) {
return EntryPtr.get() == TE ||
(EntryPtr->State == TreeEntry::NeedToGather &&
is_contained(EntryPtr->Scalars, V));
});
// The vector factor of shuffled entries must be the same.
if (It->get() == TE)
return None;
VTE = It->get();
}
Entries.push_back(VTE);
if (Entries.size() == 1) {
VF = GetVF(VTE);
} else if (VF != GetVF(VTE)) {
assert(Entries.size() == 2 && "Expected shuffle of 1 or 2 entries.");
assert(VF > 0 && "Expected non-zero vector factor.");
return None;
}
for (Value *SV : VTE->Scalars)
UsedValuesEntry.try_emplace(SV, VTE);
}
return None;
int FoundLane = findLaneForValue(VTE->Scalars, VTE->ReuseShuffleIndices, V);
Mask[I] = (Entries.front() == VTE ? 0 : VF) + FoundLane;
// Extra check required by isSingleSourceMaskImpl function (called by
// ShuffleVectorInst::isSingleSourceMask).
if (Mask[I] >= 2 * E)
return None;
}
if (NumShuffles == 1) {
if (Entries.size() == 1) {
if (ShuffleVectorInst::isReverseMask(Mask))
return TargetTransformInfo::SK_Reverse;
return TargetTransformInfo::SK_PermuteSingleSrc;
}
if (NumShuffles == 2) {
if (Entries.size() == 2) {
if (ShuffleVectorInst::isSelectMask(Mask))
return TargetTransformInfo::SK_Select;
if (ShuffleVectorInst::isTransposeMask(Mask))