Address comments

Author: NickGuy-Arm

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -1451,8 +1451,16 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
     for (MVT VT : { MVT::v32i8, MVT::v16i16, MVT::v8i32, MVT::v4i64 })
       setOperationAction(ISD::ADD, VT, Custom);
     // FADDP custom lowering
-    for (MVT VT : { MVT::v16f16, MVT::v8f32, MVT::v4f64 })
+    for (MVT VT : {MVT::v16f16, MVT::v8f32, MVT::v4f64})
       setOperationAction(ISD::FADD, VT, Custom);
+
+    if (EnablePartialReduceNodes && Subtarget->hasDotProd()) {
+      setPartialReduceMLAAction(MVT::v2i64, MVT::v8i16, Legal);
+      setPartialReduceMLAAction(MVT::v4i32, MVT::v16i8, Legal);
+      setPartialReduceMLAAction(MVT::v2i32, MVT::v8i8, Legal);
+      setPartialReduceMLAAction(MVT::v2i64, MVT::v16i8, Custom);
+    }
+
   } else /* !isNeonAvailable */ {
     for (MVT VT : MVT::fixedlen_vector_valuetypes()) {
       for (unsigned Op = 0; Op < ISD::BUILTIN_OP_END; ++Op)
@@ -1872,15 +1880,6 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
     setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv16i8, Custom);
   }
 
-  if (EnablePartialReduceNodes && Subtarget->hasNEON() &&
-      Subtarget->hasDotProd()) {
-    setPartialReduceMLAAction(MVT::v2i64, MVT::v8i16, Legal);
-    setPartialReduceMLAAction(MVT::v4i32, MVT::v16i8, Legal);
-    setPartialReduceMLAAction(MVT::v4i32, MVT::v16i8, Legal);
-    setPartialReduceMLAAction(MVT::v2i32, MVT::v8i8, Legal);
-    setPartialReduceMLAAction(MVT::v2i64, MVT::v16i8, Custom);
-  }
-
   // Handle operations that are only available in non-streaming SVE mode.
   if (Subtarget->isSVEAvailable()) {
     for (auto VT : {MVT::nxv16i8,  MVT::nxv8i16, MVT::nxv4i32,  MVT::nxv2i64,
@@ -29544,10 +29543,12 @@ SDValue
 AArch64TargetLowering::LowerPARTIAL_REDUCE_MLA(SDValue Op,
                                                SelectionDAG &DAG) const {
   bool Scalable = Op.getValueType().isScalableVector();
-  if (Scalable && !Subtarget->isSVEorStreamingSVEAvailable())
-    return SDValue();
-  if (!Scalable && (!Subtarget->isNeonAvailable() || !Subtarget->hasDotProd()))
-    return SDValue();
+
+  assert((!Scalable || Subtarget->isSVEorStreamingSVEAvailable()) &&
+         "SVE or StreamingSVE must be available when using scalable vectors.");
+  assert(
+      (Scalable || (Subtarget->isNeonAvailable() || Subtarget->hasDotProd())) &&
+      "Neon or dotprod must be available when using fixed-width vectors.");
 
   SDLoc DL(Op);
 

llvm/test/CodeGen/AArch64/neon-partial-reduce-dot-product.ll

@@ -3,7 +3,6 @@
 ; RUN: llc -mtriple aarch64 -mattr=+neon < %s | FileCheck %s --check-prefixes=CHECK,CHECK-NOI8MM,CHECK-NODOT
 ; RUN: llc -mtriple aarch64 -mattr=+neon,+dotprod,+i8mm < %s | FileCheck %s --check-prefixes=CHECK,CHECK-DOT,CHECK-I8MM,CHECK-DOT-I8MM
 ; RUN: llc -mtriple aarch64 -mattr=+neon,+dotprod,+i8mm -aarch64-enable-partial-reduce-nodes < %s | FileCheck %s --check-prefixes=CHECK,CHECK-DOT,CHECK-I8MM,CHECK-NEWLOWERING-I8MM
-; RUN: llc -mtriple aarch64 -mattr=+neon,+dotprod -aarch64-enable-partial-reduce-nodes < %s | FileCheck %s --check-prefixes=CHECK,CHECK-DOT,CHECK-NOI8MM,CHECK-NEWLOWERING-NOI8MM
 
 define <4 x i32> @udot(<4 x i32> %acc, <16 x i8> %u, <16 x i8> %s) {
 ; CHECK-DOT-LABEL: udot: