[llvm] [LV] Handle partial reductions with predication (PR #194859)

[Sander de Smalen via llvm-commits]

================
@@ -6543,13 +6546,31 @@ static void transformToPartialReduction(const VPPartialReductionChain &Chain,
   }
 
   // Check if WidenRecipe is the final result of the reduction. If so look
-  // through selects for predicated reductions.
-  VPValue *Cond = nullptr;
-  VPValue *ExitValue = cast_or_null<VPInstruction>(
-      findUserOf(WidenRecipe, m_Select(m_VPValue(Cond), m_Specific(WidenRecipe),
-                                       m_Specific(RdxPhi))));
+  // through the select recipes introduced by tail-folding. When a Blend
+  // guards the reduction, the tail-fold select wraps the Blend,
+  // otherwise it wraps the WidenRecipe directly.
+  VPValue *ExitSearch =
+      Chain.Blend ? cast<VPValue>(Chain.Blend) : cast<VPValue>(WidenRecipe);
+
+  VPValue *TailFoldCond = nullptr;
+  auto *ExitValue = cast_or_null<VPInstruction>(findUserOf(
+      ExitSearch, m_Select(m_VPValue(TailFoldCond), m_Specific(ExitSearch),
+                           m_Specific(RdxPhi))));
+
+  VPValue *BlendCond = Chain.Blend ? Chain.Blend->getMask(0) : nullptr;
+  VPValue *Cond = TailFoldCond;
+
+  if (TailFoldCond && BlendCond) {
+    VPBuilder Builder(WidenRecipe);
+    Cond = Builder.createLogicalAnd(TailFoldCond, BlendCond,
+                                    WidenRecipe->getDebugLoc());
+  } else if (BlendCond) {
+    Cond = BlendCond;
+  }
----------------
sdesmalen-arm wrote:

nit: no need for `TailFoldCond` (can use `Cond` directly when testing for `ExitValue` instead of `TailFoldCond`), e.g.
```suggestion
  VPValue *Cond = nullptr;
  VPValue *ExitValue = cast_or_null<VPInstruction>(
      findUserOf(ExitSearch, m_Select(m_VPValue(Cond), m_Specific(ExitSearch),
                                      m_Specific(RdxPhi))));
  if (Chain.Blend) {
    VPValue *BlendCond = Chain.Blend->getMask(0);
    Cond = ExitValue ? VPBuilder(WidenRecipe)
                           .createLogicalAnd(Cond, BlendCond,
                                             WidenRecipe->getDebugLoc())
                     : BlendCond;
}
```

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