[llvm] [LoopVectorizer] Add support for partial reductions (PR #92418)

[Graham Hunter via llvm-commits]

================
@@ -341,6 +341,15 @@ class AArch64TTIImpl : public BasicTTIImplBase<AArch64TTIImpl> {
     return BaseT::isLegalNTLoad(DataType, Alignment);
   }
 
+  bool isPartialReductionSupported(const Instruction *ReductionInstr,
+                                   Type *InputType, unsigned ScaleFactor,
+                                   bool IsInputASignExtended,
+                                   bool IsInputBSignExtended,
+                                   const Instruction *BinOp) const {
+    return ScaleFactor == 4 && (ST->isSVEorStreamingSVEAvailable() ||
----------------
huntergr-arm wrote:

I think we need to be stricter here -- there's nothing implemented to support FP types, or for non-power-of-two integers, or something like 4 i32s being combined into an i128. Similarly there's no checks here for ReductionInstr being an add, or BinOp being non-null or a mul.

I believe we would reject at least some of those already due to other checks in the vectorizer, but this function should be properly checking the inputs.

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