diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp --- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp +++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -594,7 +594,7 @@ /// Fix a reduction cross-iteration phi. This is the second phase of /// vectorizing this phi node. - void fixReduction(PHINode *Phi, VPTransformState &State); + void fixReduction(VPWidenPHIRecipe *Phi, VPTransformState &State); /// Clear NSW/NUW flags from reduction instructions if necessary. void clearReductionWrapFlags(RecurrenceDescriptor &RdxDesc, @@ -4063,12 +4063,19 @@ // the currently empty PHI nodes. At this point every instruction in the // original loop is widened to a vector form so we can use them to construct // the incoming edges. - for (PHINode &Phi : OrigLoop->getHeader()->phis()) { - // Handle first-order recurrences and reductions that need to be fixed. - if (Legal->isFirstOrderRecurrence(&Phi)) - fixFirstOrderRecurrence(&Phi, State); - else if (Legal->isReductionVariable(&Phi)) - fixReduction(&Phi, State); + VPBasicBlock *Header = State.Plan->getEntry()->getEntryBasicBlock(); + for (VPRecipeBase &R : *Header) { + if (isa(&R)) + continue; + VPWidenPHIRecipe *PhiR = dyn_cast(&R); + if (!PhiR) + break; + + auto *OrigPhi = cast(PhiR->getUnderlyingValue()); + if (PhiR->getRecurrenceDescriptor()) { + fixReduction(PhiR, State); + } else if (Legal->isFirstOrderRecurrence(OrigPhi)) + fixFirstOrderRecurrence(OrigPhi, State); } } @@ -4259,17 +4266,19 @@ LCSSAPhi.addIncoming(ExtractForPhiUsedOutsideLoop, LoopMiddleBlock); } -void InnerLoopVectorizer::fixReduction(PHINode *Phi, VPTransformState &State) { +void InnerLoopVectorizer::fixReduction(VPWidenPHIRecipe *PhiR, + VPTransformState &State) { + PHINode *OrigPhi = cast(PhiR->getUnderlyingValue()); // Get it's reduction variable descriptor. - assert(Legal->isReductionVariable(Phi) && + assert(Legal->isReductionVariable(OrigPhi) && "Unable to find the reduction variable"); - RecurrenceDescriptor RdxDesc = Legal->getReductionVars()[Phi]; + RecurrenceDescriptor RdxDesc = *PhiR->getRecurrenceDescriptor(); RecurKind RK = RdxDesc.getRecurrenceKind(); TrackingVH ReductionStartValue = RdxDesc.getRecurrenceStartValue(); Instruction *LoopExitInst = RdxDesc.getLoopExitInstr(); setDebugLocFromInst(Builder, ReductionStartValue); - bool IsInLoopReductionPhi = Cost->isInLoopReduction(Phi); + bool IsInLoopReductionPhi = Cost->isInLoopReduction(OrigPhi); VPValue *LoopExitInstDef = State.Plan->getVPValue(LoopExitInst); // This is the vector-clone of the value that leaves the loop. @@ -4283,10 +4292,10 @@ // Reductions do not have to start at zero. They can start with // any loop invariant values. BasicBlock *Latch = OrigLoop->getLoopLatch(); - Value *LoopVal = Phi->getIncomingValueForBlock(Latch); + Value *LoopVal = OrigPhi->getIncomingValueForBlock(Latch); for (unsigned Part = 0; Part < UF; ++Part) { - Value *VecRdxPhi = State.get(State.Plan->getVPValue(Phi), Part); + Value *VecRdxPhi = State.get(PhiR->getVPValue(), Part); Value *Val = State.get(State.Plan->getVPValue(LoopVal), Part); cast(VecRdxPhi) ->addIncoming(Val, LI->getLoopFor(LoopVectorBody)->getLoopLatch()); @@ -4300,7 +4309,7 @@ setDebugLocFromInst(Builder, LoopExitInst); - Type *PhiTy = Phi->getType(); + Type *PhiTy = OrigPhi->getType(); // If tail is folded by masking, the vector value to leave the loop should be // a Select choosing between the vectorized LoopExitInst and vectorized Phi, // instead of the former. For an inloop reduction the reduction will already @@ -4328,8 +4337,7 @@ TTI->preferPredicatedReductionSelect( RdxDesc.getOpcode(), PhiTy, TargetTransformInfo::ReductionFlags())) { - auto *VecRdxPhi = - cast(State.get(State.Plan->getVPValue(Phi), Part)); + auto *VecRdxPhi = cast(State.get(PhiR->getVPValue(), Part)); VecRdxPhi->setIncomingValueForBlock( LI->getLoopFor(LoopVectorBody)->getLoopLatch(), Sel); } @@ -4429,12 +4437,12 @@ // Fix the scalar loop reduction variable with the incoming reduction sum // from the vector body and from the backedge value. int IncomingEdgeBlockIdx = - Phi->getBasicBlockIndex(OrigLoop->getLoopLatch()); + OrigPhi->getBasicBlockIndex(OrigLoop->getLoopLatch()); assert(IncomingEdgeBlockIdx >= 0 && "Invalid block index"); // Pick the other block. int SelfEdgeBlockIdx = (IncomingEdgeBlockIdx ? 0 : 1); - Phi->setIncomingValue(SelfEdgeBlockIdx, BCBlockPhi); - Phi->setIncomingValue(IncomingEdgeBlockIdx, LoopExitInst); + OrigPhi->setIncomingValue(SelfEdgeBlockIdx, BCBlockPhi); + OrigPhi->setIncomingValue(IncomingEdgeBlockIdx, LoopExitInst); } void InnerLoopVectorizer::clearReductionWrapFlags(RecurrenceDescriptor &RdxDesc, diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h --- a/llvm/lib/Transforms/Vectorize/VPlan.h +++ b/llvm/lib/Transforms/Vectorize/VPlan.h @@ -1041,6 +1041,8 @@ /// Returns the \p I th incoming VPBasicBlock. VPBasicBlock *getIncomingBlock(unsigned I) { return IncomingBlocks[I]; } + + RecurrenceDescriptor *getRecurrenceDescriptor() { return RdxDesc; } }; /// A recipe for vectorizing a phi-node as a sequence of mask-based select