Adjust reduction phi in transformation stage

Author: Mel-Chen

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -1486,17 +1486,6 @@ class LoopVectorizationCostModel {
     return InLoopReductions.contains(Phi);
   }
 
-  /// Returns true if the predicated reduction select should be used.
-  bool usePredicatedReductionSelect(unsigned Opcode, Type *PhiTy) const {
-    // Force to use predicated reduction select since the EVL of the
-    // second-to-last iteration might not be VF*UF.
-    if (foldTailWithEVL())
-      return true;
-    return PreferPredicatedReductionSelect ||
-           TTI.preferPredicatedReductionSelect(
-               Opcode, PhiTy, TargetTransformInfo::ReductionFlags());
-  }
-
   /// Estimate cost of an intrinsic call instruction CI if it were vectorized
   /// with factor VF.  Return the cost of the instruction, including
   /// scalarization overhead if it's needed.
@@ -9461,8 +9450,10 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
                cast<VPInstruction>(&U)->getOpcode() ==
                    VPInstruction::ComputeReductionResult;
       });
-      if (CM.usePredicatedReductionSelect(
-              PhiR->getRecurrenceDescriptor().getOpcode(), PhiTy))
+      if (PreferPredicatedReductionSelect ||
+          TTI.preferPredicatedReductionSelect(
+              PhiR->getRecurrenceDescriptor().getOpcode(), PhiTy,
+              TargetTransformInfo::ReductionFlags()))
         PhiR->setOperand(1, NewExitingVPV);
     }
 

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -1357,6 +1357,30 @@ static void transformRecipestoEVLRecipes(VPlan &Plan, VPValue &EVL) {
   VPTypeAnalysis TypeInfo(CanonicalIVType);
   LLVMContext &Ctx = CanonicalIVType->getContext();
   SmallVector<VPValue *> HeaderMasks = collectAllHeaderMasks(Plan);
+
+  // Adjust reduction phi recipes
+  VPBasicBlock *Header = Plan.getVectorLoopRegion()->getEntryBasicBlock();
+  for (VPRecipeBase &R : Header->phis()) {
+    auto *PhiR = dyn_cast<VPReductionPHIRecipe>(&R);
+    if (!PhiR || PhiR->isInLoop())
+      continue;
+    VPValue *BackedgeVPV = PhiR->getBackedgeValue();
+    bool IsExiting = any_of(BackedgeVPV->users(), [](VPUser *U) {
+      return isa<VPInstruction>(U) && cast<VPInstruction>(U)->getOpcode() ==
+                                          VPInstruction::ComputeReductionResult;
+    });
+    if (IsExiting)
+      continue;
+
+    auto *FoundRdxSelect = find_if(BackedgeVPV->users(), [](VPUser *U) {
+      return match(U, m_Select(m_VPValue(), m_VPValue(), m_VPValue()));
+    });
+    assert(FoundRdxSelect && "Must have a reduction select when out-loop "
+                             "reduction with tail folding");
+    auto *RdxSelect = cast<VPRecipeBase>(*FoundRdxSelect);
+    PhiR->setOperand(1, RdxSelect->getVPSingleValue());
+  }
+
   for (VPValue *HeaderMask : collectAllHeaderMasks(Plan)) {
     for (VPUser *U : collectUsersRecursively(HeaderMask)) {
       auto *CurRecipe = dyn_cast<VPRecipeBase>(U);