[llvm] [LV][AArch64] Support partial reductions of extended compares (PR #212190)

[Graham Hunter via llvm-commits]

================
@@ -776,6 +776,43 @@ VPInstruction *vputils::findComputeReductionResult(VPReductionPHIRecipe *PhiR) {
       cast<VPSingleDefRecipe>(SelR));
 }
 
+/// Returns the type whose width \p Mask occupies or nullptr if unknown.
+static Type *getMaskMaterializationType(VPValue *Mask, unsigned Depth = 0) {
+  // Give up rather than recurse further. Falling back to i1 is always safe.
+  constexpr unsigned MaxMaskDepth = 3;
+  if (Depth > MaxMaskDepth)
+    return nullptr;
+
+  VPValue *A, *B;
+  if (match(Mask, m_Cmp(m_VPValue(A), m_VPValue()))) {
+    Type *CmpTy = A->getScalarType();
+    unsigned Bits = CmpTy->getPrimitiveSizeInBits();
+    // A pointer reports no size and an i1 is the width we came here to replace.
+    if (Bits <= 1)
+      return nullptr;
+    return IntegerType::get(CmpTy->getContext(), Bits);
+  }
+
+  if (match(Mask, m_Binary<Instruction::And>(m_VPValue(A), m_VPValue(B))) ||
+      match(Mask, m_Binary<Instruction::Or>(m_VPValue(A), m_VPValue(B))) ||
+      match(Mask, m_Binary<Instruction::Xor>(m_VPValue(A), m_VPValue(B)))) {
+    Type *TyA = getMaskMaterializationType(A, Depth + 1);
+    Type *TyB = getMaskMaterializationType(B, Depth + 1);
+    return TyA == TyB ? TyA : nullptr;
+  }
+
+  return nullptr;
+}
+
+Type *vputils::getExtendSrcTypeForPartialReduction(VPValue *ExtSrc) {
----------------
huntergr-arm wrote:

I wonder if it isn't better to leave it as i1, and make sure the cost functions for each target handle that appropriately.

For SVE at least we would want to reduce directly from a `<vscale x N x i1>` to a scalar accumulator using `incp`; we haven't used partial reduction intrinsics to reduce directly to a scalar inside the loop yet, but that is part of the intended design.

A little experimentation shows that allowing i1 in the TTI hook creates the appropriate recipes, but asserts when creating a type in calculateRegisterUsageForPlan -- that'll need a little intervention for a scaling factor greater than the VF.

@david-arm @lukel97  any preference?

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