[llvm] [VPlan] Support tailfolded loops in multi-use-reductions (PR #214455)

[Luke Lau via llvm-commits]

================
@@ -2035,32 +2043,58 @@ bool VPlanTransforms::handleMultiUseReductions(VPlan &Plan,
 
     // MinOrMaxPhiR has users outside the reduction cycle in the loop. Check if
     // the only other user is a FindLastIV reduction. MinOrMaxPhiR must have
-    // exactly 2 users:
+    // exactly 2 users if not tailfolded:
     // 1) the min/max operation of the reduction cycle, and
     // 2) the compare of a FindLastIV reduction cycle. This compare must match
     // the min/max operation - comparing MinOrMaxPhiR with the operand of the
     // min/max operation, and be used only by the select of the FindLastIV
     // reduction cycle.
+    // There is an additional user if the loop was tailfolded:
+    // 3) the select operation of the vector.latch block. This select uses the
+    // the original MinOrMaxPhiR if the mask is zero.
     RecurKind RdxKind = MinOrMaxPhiR->getRecurrenceKind();
     assert(
         RecurrenceDescriptor::isIntMinMaxRecurrenceKind(RdxKind) &&
         "only min/max recurrences support users outside the reduction chain");
 
-    auto *MinOrMaxOp =
+    auto *MinOrMaxBackedgeR =
         dyn_cast<VPRecipeWithIRFlags>(MinOrMaxPhiR->getBackedgeValue());
-    if (!MinOrMaxOp)
+    if (!MinOrMaxBackedgeR)
       return false;
 
-    // Check that MinOrMaxOp is a VPWidenIntrinsicRecipe or VPReplicateRecipe
-    // with an intrinsic that matches the reduction kind.
+    // If the loop is tailfolded then the backedge won't be the
+    // reduction-intrinsic but the select in the vector.latch block that wraps
+    // the reduction-intrinsic.
+    auto *MinOrMaxOp = MinOrMaxBackedgeR;
+    VPValue *HeaderMask = Plan.getVectorLoopRegion()->getHeaderMask();
+    if (VPValue *MinOrMaxTailfold;
+        HeaderMask &&
+        match(MinOrMaxBackedgeR,
+              m_SelectLike(m_Specific(HeaderMask), m_VPValue(MinOrMaxTailfold),
+                           m_Specific(MinOrMaxPhiR)))) {
+      MinOrMaxOp =
+          dyn_cast<VPRecipeWithIRFlags>(MinOrMaxTailfold->getDefiningRecipe());
+      if (!MinOrMaxOp)
+        return false;
+    }
+
+    // Check that MinOrMaxOp is a VPWidenIntrinsicRecipe or
+    // VPReplicateRecipe with an intrinsic that matches the reduction kind.
     Intrinsic::ID ExpectedIntrinsicID = getMinMaxReductionIntrinsicOp(RdxKind);
     if (!match(MinOrMaxOp, m_Intrinsic(ExpectedIntrinsicID)))
       return false;
 
-    // MinOrMaxOp must have 2 users: 1) MinOrMaxPhiR and 2)
+    // If the loop is tailfolded MinOrMaxOp should only feed into the predicated
+    // select MinOrMaxBackedgeR. This is ensured by the MinOrMaxBackedgeR
+    // select-matching.
+    if ((!HeaderMask && MinOrMaxOp->getNumUsers() != 2) ||
+        (HeaderMask && MinOrMaxOp->getNumUsers() != 1))
+      return false;
----------------
lukel97 wrote:

Is this `return false` ever hit? I think it might be dead, and the `MinOrMaxBackedgeR` assert below is the important one to check. If it's dead then I would just omit it

https://github.com/llvm/llvm-project/pull/214455