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