@@ -628,3 +628,163 @@ void VPlanTransforms::attachCheckBlock(VPlan &Plan, Value *Cond,
628628 Term->addMetadata (LLVMContext::MD_prof, BranchWeights);
629629 }
630630}
631+
632+ bool VPlanTransforms::handleMaxMinNumReductions (VPlan &Plan) {
633+ auto GetMinMaxCompareValue = [](VPReductionPHIRecipe *RedPhiR) -> VPValue * {
634+ auto *MinMaxR = dyn_cast<VPRecipeWithIRFlags>(
635+ RedPhiR->getBackedgeValue ()->getDefiningRecipe ());
636+ if (!MinMaxR)
637+ return nullptr ;
638+
639+ auto *RepR = dyn_cast<VPReplicateRecipe>(MinMaxR);
640+ if (!isa<VPWidenIntrinsicRecipe>(MinMaxR) &&
641+ !(RepR && isa<IntrinsicInst>(RepR->getUnderlyingInstr ())))
642+ return nullptr ;
643+
644+ #ifndef NDEBUG
645+ Intrinsic::ID RdxIntrinsicId =
646+ RedPhiR->getRecurrenceKind () == RecurKind::FMaxNum ? Intrinsic::maxnum
647+ : Intrinsic::minnum;
648+ assert ((isa<VPWidenIntrinsicRecipe>(MinMaxR) &&
649+ cast<VPWidenIntrinsicRecipe>(MinMaxR)->getVectorIntrinsicID () ==
650+ RdxIntrinsicId) ||
651+ (RepR &&
652+ cast<IntrinsicInst>(RepR->getUnderlyingInstr ())->getIntrinsicID () ==
653+ RdxIntrinsicId) &&
654+ " Intrinsic did not match recurrence kind"
655+ #endif
656+
657+ if (MinMaxR->getOperand (0 ) == RedPhiR)
658+ return MinMaxR->getOperand (1 );
659+
660+ assert (MinMaxR->getOperand (1 ) == RedPhiR &&
661+ " Reduction phi operand expected"
662+ return MinMaxR->getOperand (0 );
663+ };
664+
665+ VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion ();
666+ VPReductionPHIRecipe *RedPhiR = nullptr ;
667+ bool HasUnsupportedPhi = false ;
668+ for (auto &R : LoopRegion->getEntryBasicBlock ()->phis ()) {
669+ if (isa<VPCanonicalIVPHIRecipe, VPWidenIntOrFpInductionRecipe>(&R))
670+ continue ;
671+ auto *Cur = dyn_cast<VPReductionPHIRecipe>(&R);
672+ if (!Cur) {
673+ // TODO: Also support fixed-order recurrence phis.
674+ HasUnsupportedPhi = true ;
675+ continue ;
676+ }
677+ // For now, only a single reduction is supported.
678+ // TODO: Support multiple MaxNum/MinNum reductions and other reductions.
679+ if (RedPhiR)
680+ return false ;
681+ if (Cur->getRecurrenceKind () != RecurKind::FMaxNum &&
682+ Cur->getRecurrenceKind () != RecurKind::FMinNum) {
683+ HasUnsupportedPhi = true ;
684+ continue ;
685+ }
686+ RedPhiR = Cur;
687+ }
688+
689+ if (!RedPhiR)
690+ return true ;
691+
692+ // We won't be able to resume execution in the scalar tail, if there are
693+ // unsupported header phis or there is no scalar tail at all, due to
694+ // tail-folding.
695+ if (HasUnsupportedPhi || !Plan.hasScalarTail ())
696+ return false ;
697+
698+ VPValue *MinMaxOp = GetMinMaxCompareValue (RedPhiR);
699+ if (!MinMaxOp)
700+ return false ;
701+
702+ RecurKind RedPhiRK = RedPhiR->getRecurrenceKind ();
703+ assert ((RedPhiRK == RecurKind::FMaxNum || RedPhiRK == RecurKind::FMinNum) &&
704+ " unsupported reduction"
705+
706+ // / Check if the vector loop of \p Plan can early exit and restart
707+ // / execution of last vector iteration in the scalar loop. This requires all
708+ // / recipes up to early exit point be side-effect free as they are
709+ // / re-executed. Currently we check that the loop is free of any recipe that
710+ // / may write to memory. Expected to operate on an early VPlan w/o nested
711+ // / regions.
712+ for (VPBlockBase *VPB : vp_depth_first_shallow (
713+ Plan.getVectorLoopRegion ()->getEntryBasicBlock ())) {
714+ auto *VPBB = cast<VPBasicBlock>(VPB);
715+ for (auto &R : *VPBB) {
716+ if (R.mayWriteToMemory () &&
717+ !match (&R, m_BranchOnCount (m_VPValue (), m_VPValue ())))
718+ return false ;
719+ }
720+ }
721+
722+ VPBasicBlock *LatchVPBB = LoopRegion->getExitingBasicBlock ();
723+ VPBuilder Builder (LatchVPBB->getTerminator ());
724+ auto *LatchExitingBranch = cast<VPInstruction>(LatchVPBB->getTerminator ());
725+ assert (LatchExitingBranch->getOpcode () == VPInstruction::BranchOnCount &&
726+ " Unexpected terminator"
727+ auto *IsLatchExitTaken =
728+ Builder.createICmp (CmpInst::ICMP_EQ, LatchExitingBranch->getOperand (0 ),
729+ LatchExitingBranch->getOperand (1 ));
730+
731+ VPValue *IsNaN = Builder.createFCmp (CmpInst::FCMP_UNO, MinMaxOp, MinMaxOp);
732+ VPValue *AnyNaN = Builder.createNaryOp (VPInstruction::AnyOf, {IsNaN});
733+ auto *AnyExitTaken =
734+ Builder.createNaryOp (Instruction::Or, {AnyNaN, IsLatchExitTaken});
735+ Builder.createNaryOp (VPInstruction::BranchOnCond, AnyExitTaken);
736+ LatchExitingBranch->eraseFromParent ();
737+
738+ // If we exit early due to NaNs, compute the final reduction result based on
739+ // the reduction phi at the beginning of the last vector iteration.
740+ auto *RdxResult = find_singleton<VPSingleDefRecipe>(
741+ RedPhiR->users (), [](VPUser *U, bool ) -> VPSingleDefRecipe * {
742+ auto *VPI = dyn_cast<VPInstruction>(U);
743+ if (VPI && VPI->getOpcode () == VPInstruction::ComputeReductionResult)
744+ return VPI;
745+ return nullptr ;
746+ });
747+
748+ auto *MiddleVPBB = Plan.getMiddleBlock ();
749+ Builder.setInsertPoint (MiddleVPBB, MiddleVPBB->begin ());
750+ auto *NewSel =
751+ Builder.createSelect (AnyNaN, RedPhiR, RdxResult->getOperand (1 ));
752+ RdxResult->setOperand (1 , NewSel);
753+
754+ auto *ScalarPH = Plan.getScalarPreheader ();
755+ // Update resume phis for inductions in the scalar preheader. If AnyNaN is
756+ // true, the resume from the start of the last vector iteration via the
757+ // canonical IV, otherwise from the original value.
758+ for (auto &R : ScalarPH->phis ()) {
759+ auto *ResumeR = cast<VPPhi>(&R);
760+ VPValue *VecV = ResumeR->getOperand (0 );
761+ if (VecV == RdxResult)
762+ continue ;
763+ if (auto *DerivedIV = dyn_cast<VPDerivedIVRecipe>(VecV)) {
764+ if (DerivedIV->getNumUsers () == 1 &&
765+ DerivedIV->getOperand (1 ) == &Plan.getVectorTripCount ()) {
766+ auto *NewSel = Builder.createSelect (AnyNaN, Plan.getCanonicalIV (),
767+ &Plan.getVectorTripCount ());
768+ DerivedIV->moveAfter (&*Builder.getInsertPoint ());
769+ DerivedIV->setOperand (1 , NewSel);
770+ continue ;
771+ }
772+ }
773+ // Bail out and abandon the current, partially modified, VPlan if we
774+ // encounter resume phi that cannot be updated yet.
775+ if (VecV != &Plan.getVectorTripCount ()) {
776+ LLVM_DEBUG (dbgs () << " Found resume phi we cannot update for VPlan with "
777+ " FMaxNum/FMinNum reduction.\n "
778+ return false ;
779+ }
780+ auto *NewSel = Builder.createSelect (AnyNaN, Plan.getCanonicalIV (), VecV);
781+ ResumeR->setOperand (0 , NewSel);
782+ }
783+
784+ auto *MiddleTerm = MiddleVPBB->getTerminator ();
785+ Builder.setInsertPoint (MiddleTerm);
786+ VPValue *MiddleCond = MiddleTerm->getOperand (0 );
787+ VPValue *NewCond = Builder.createAnd (MiddleCond, Builder.createNot (AnyNaN));
788+ MiddleTerm->setOperand (0 , NewCond);
789+ return true ;
790+ }
0 commit comments