diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp index 80bdc67a49fe5..ee92e102baf87 100644 --- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp +++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp @@ -3760,40 +3760,7 @@ BoUpSLP::findReusedOrderedScalars(const BoUpSLP::TreeEntry &TE) { OrdersType CurrentOrder(NumScalars, NumScalars); SmallVector Positions; SmallBitVector UsedPositions(NumScalars); - DenseMap UsedEntries; - DenseMap> ValueToEntryPos; - for (Value *V : TE.Scalars) { - if (!isa(V)) - continue; - const auto *LocalSTE = getTreeEntry(V); - if (!LocalSTE) - continue; - unsigned Lane = - std::distance(LocalSTE->Scalars.begin(), find(LocalSTE->Scalars, V)); - if (Lane >= NumScalars) - continue; - ++UsedEntries.try_emplace(LocalSTE, 0).first->getSecond(); - ValueToEntryPos.try_emplace(V, LocalSTE, Lane); - } - if (UsedEntries.empty()) - return std::nullopt; - const TreeEntry &BestSTE = - *std::max_element(UsedEntries.begin(), UsedEntries.end(), - [](const std::pair &P1, - const std::pair &P2) { - return P1.second < P2.second; - }) - ->first; - UsedEntries.erase(&BestSTE); - const TreeEntry *SecondBestSTE = nullptr; - if (!UsedEntries.empty()) - SecondBestSTE = - std::max_element(UsedEntries.begin(), UsedEntries.end(), - [](const std::pair &P1, - const std::pair &P2) { - return P1.second < P2.second; - }) - ->first; + const TreeEntry *STE = nullptr; // Try to find all gathered scalars that are gets vectorized in other // vectorize node. Here we can have only one single tree vector node to // correctly identify order of the gathered scalars. @@ -3801,46 +3768,53 @@ BoUpSLP::findReusedOrderedScalars(const BoUpSLP::TreeEntry &TE) { Value *V = TE.Scalars[I]; if (!isa(V)) continue; - const auto [LocalSTE, Lane] = ValueToEntryPos.lookup(V); - if (!LocalSTE || (LocalSTE != &BestSTE && LocalSTE != SecondBestSTE)) - continue; - if (CurrentOrder[Lane] != NumScalars) { - if ((CurrentOrder[Lane] >= BestSTE.Scalars.size() || - BestSTE.Scalars[CurrentOrder[Lane]] == V) && - (Lane != I || LocalSTE == SecondBestSTE)) - continue; - UsedPositions.reset(CurrentOrder[Lane]); + if (const auto *LocalSTE = getTreeEntry(V)) { + if (!STE) + STE = LocalSTE; + else if (STE != LocalSTE) + // Take the order only from the single vector node. + return std::nullopt; + unsigned Lane = + std::distance(STE->Scalars.begin(), find(STE->Scalars, V)); + if (Lane >= NumScalars) + return std::nullopt; + if (CurrentOrder[Lane] != NumScalars) { + if (Lane != I) + continue; + UsedPositions.reset(CurrentOrder[Lane]); + } + // The partial identity (where only some elements of the gather node are + // in the identity order) is good. + CurrentOrder[Lane] = I; + UsedPositions.set(I); } - // The partial identity (where only some elements of the gather node are - // in the identity order) is good. - CurrentOrder[Lane] = I; - UsedPositions.set(I); } // Need to keep the order if we have a vector entry and at least 2 scalars or // the vectorized entry has just 2 scalars. - if (BestSTE.Scalars.size() != 2 && UsedPositions.count() <= 1) - return std::nullopt; - auto IsIdentityOrder = [&](ArrayRef CurrentOrder) { - for (unsigned I = 0; I < NumScalars; ++I) - if (CurrentOrder[I] != I && CurrentOrder[I] != NumScalars) - return false; - return true; - }; - if (IsIdentityOrder(CurrentOrder)) - return OrdersType(); - auto *It = CurrentOrder.begin(); - for (unsigned I = 0; I < NumScalars;) { - if (UsedPositions.test(I)) { - ++I; - continue; - } - if (*It == NumScalars) { - *It = I; - ++I; + if (STE && (UsedPositions.count() > 1 || STE->Scalars.size() == 2)) { + auto &&IsIdentityOrder = [NumScalars](ArrayRef CurrentOrder) { + for (unsigned I = 0; I < NumScalars; ++I) + if (CurrentOrder[I] != I && CurrentOrder[I] != NumScalars) + return false; + return true; + }; + if (IsIdentityOrder(CurrentOrder)) + return OrdersType(); + auto *It = CurrentOrder.begin(); + for (unsigned I = 0; I < NumScalars;) { + if (UsedPositions.test(I)) { + ++I; + continue; + } + if (*It == NumScalars) { + *It = I; + ++I; + } + ++It; } - ++It; + return std::move(CurrentOrder); } - return std::move(CurrentOrder); + return std::nullopt; } namespace {