diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp index 893c02a81d4b7..1081f1303d727 100644 --- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp +++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp @@ -14078,16 +14078,43 @@ static Instruction *tryGetSecondaryReductionRoot(PHINode *Phi, return nullptr; } +/// \p Returns the first operand of \p I that does not match \p Phi. If +/// operand is not an instruction it returns nullptr. +static Instruction *getNonPhiOperand(Instruction *I, PHINode *Phi) { + Value *Op0 = nullptr; + Value *Op1 = nullptr; + if (!matchRdxBop(I, Op0, Op1)) + return nullptr; + return dyn_cast(Op0 == Phi ? Op1 : Op0); +} + +/// \Returns true if \p I is a candidate instruction for reduction vectorization. +static bool isReductionCandidate(Instruction *I) { + bool IsSelect = match(I, m_Select(m_Value(), m_Value(), m_Value())); + Value *B0 = nullptr, *B1 = nullptr; + bool IsBinop = matchRdxBop(I, B0, B1); + return IsBinop || IsSelect; +} + bool SLPVectorizerPass::vectorizeHorReduction( PHINode *P, Instruction *Root, BasicBlock *BB, BoUpSLP &R, TargetTransformInfo *TTI, SmallVectorImpl &PostponedInsts) { if (!ShouldVectorizeHor) return false; - if (!isa(Root)) - P = nullptr; + bool TryOperandsAsNewSeeds = P && isa(Root); if (Root->getParent() != BB || isa(Root)) return false; + + // If we can find a secondary reduction root, use that instead. + auto SelectRoot = [&]() { + if (TryOperandsAsNewSeeds && isReductionCandidate(Root) && + HorizontalReduction::getRdxKind(Root) != RecurKind::None) + if (Instruction *NewRoot = tryGetSecondaryReductionRoot(P, Root)) + return NewRoot; + return Root; + }; + // Start analysis starting from Root instruction. If horizontal reduction is // found, try to vectorize it. If it is not a horizontal reduction or // vectorization is not possible or not effective, and currently analyzed @@ -14100,28 +14127,32 @@ bool SLPVectorizerPass::vectorizeHorReduction( // If a horizintal reduction was not matched or vectorized we collect // instructions for possible later attempts for vectorization. std::queue> Stack; - Stack.emplace(Root, 0); + Stack.emplace(SelectRoot(), 0); SmallPtrSet VisitedInstrs; bool Res = false; - auto &&TryToReduce = [this, TTI, &P, &R](Instruction *Inst, Value *&B0, - Value *&B1) -> Value * { + auto &&TryToReduce = [this, TTI, &R](Instruction *Inst) -> Value * { if (R.isAnalyzedReductionRoot(Inst)) return nullptr; - bool IsBinop = matchRdxBop(Inst, B0, B1); - bool IsSelect = match(Inst, m_Select(m_Value(), m_Value(), m_Value())); - if (IsBinop || IsSelect) { - assert((!P || is_contained(P->operands(), Inst)) && - "Phi needs to use the binary operator"); - if (P && HorizontalReduction::getRdxKind(Inst) != RecurKind::None) - if (Instruction *NewRoot = tryGetSecondaryReductionRoot(P, Inst)) - Inst = NewRoot; - - HorizontalReduction HorRdx; - if (HorRdx.matchAssociativeReduction(Inst, *SE, *DL, *TLI)) - return HorRdx.tryToReduce(R, TTI, *TLI); + if (!isReductionCandidate(Inst)) + return nullptr; + HorizontalReduction HorRdx; + if (!HorRdx.matchAssociativeReduction(Inst, *SE, *DL, *TLI)) + return nullptr; + return HorRdx.tryToReduce(R, TTI, *TLI); + }; + auto TryAppendToPostponedInsts = [&](Instruction *FutureSeed) { + if (TryOperandsAsNewSeeds && FutureSeed == Root) { + FutureSeed = getNonPhiOperand(Root, P); + if (!FutureSeed) + return false; } - return nullptr; + // Do not collect CmpInst or InsertElementInst/InsertValueInst as their + // analysis is done separately. + if (!isa(FutureSeed)) + PostponedInsts.push_back(FutureSeed); + return true; }; + while (!Stack.empty()) { Instruction *Inst; unsigned Level; @@ -14132,37 +14163,19 @@ bool SLPVectorizerPass::vectorizeHorReduction( // iteration while stack was populated before that happened. if (R.isDeleted(Inst)) continue; - Value *B0 = nullptr, *B1 = nullptr; - if (Value *V = TryToReduce(Inst, B0, B1)) { + if (Value *VectorizedV = TryToReduce(Inst)) { Res = true; - // Set P to nullptr to avoid re-analysis of phi node in - // matchAssociativeReduction function unless this is the root node. - P = nullptr; - if (auto *I = dyn_cast(V)) { + if (auto *I = dyn_cast(VectorizedV)) { // Try to find another reduction. Stack.emplace(I, Level); continue; } } else { - bool IsBinop = B0 && B1; - if (P && IsBinop) { - Inst = dyn_cast(B0); - if (Inst == P) - Inst = dyn_cast(B1); - if (!Inst) { - // Set P to nullptr to avoid re-analysis of phi node in - // matchAssociativeReduction function unless this is the root node. - P = nullptr; - continue; - } + // We could not vectorize `Inst` so try to use it as a future seed. + if (!TryAppendToPostponedInsts(Inst)) { + assert(Stack.empty() && "Expected empty stack"); + break; } - // Set P to nullptr to avoid re-analysis of phi node in - // matchAssociativeReduction function unless this is the root node. - P = nullptr; - // Do not collect CmpInst or InsertElementInst/InsertValueInst as their - // analysis is done separately. - if (!isa(Inst)) - PostponedInsts.push_back(Inst); } // Try to vectorize operands.