[AArch64] Improve codegen for partial.reduce.add v16i8 -> v2i32

Rather than expanding, we can handle this case natively by
widening the accumulator.

Author: sdesmalen-arm

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -1458,6 +1458,7 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
 
       setPartialReduceMLAAction(MLAOps, MVT::v4i32, MVT::v16i8, Legal);
       setPartialReduceMLAAction(MLAOps, MVT::v2i32, MVT::v8i8, Legal);
+      setPartialReduceMLAAction(MLAOps, MVT::v2i32, MVT::v16i8, Custom);
       setPartialReduceMLAAction(MLAOps, MVT::v2i64, MVT::v16i8, Custom);
 
       if (Subtarget->hasMatMulInt8()) {
@@ -30768,6 +30769,18 @@ AArch64TargetLowering::LowerPARTIAL_REDUCE_MLA(SDValue Op,
       ResultVT.isFixedLengthVector() &&
       useSVEForFixedLengthVectorVT(ResultVT, /*OverrideNEON=*/true);
 
+  // We can handle this case natively by accumulating into a wider
+  // zero-padded vector.
+  if (!ConvertToScalable && ResultVT == MVT::v2i32 && OpVT == MVT::v16i8) {
+    SDValue ZeroVec = DAG.getConstant(0, DL, MVT::v4i32);
+    SDValue WideAcc = DAG.getInsertSubvector(DL, ZeroVec, Acc, 0);
+    SDValue Wide = DAG.getNode(Op.getOpcode(), DL, MVT::v4i32,
+                               WideAcc, LHS, RHS);
+    SDValue Lo = DAG.getExtractSubvector(DL, MVT::v2i32, Wide, 0);
+    SDValue Hi = DAG.getExtractSubvector(DL, MVT::v2i32, Wide, 2);
+    return DAG.getNode(ISD::ADD, DL, ResultVT, Lo, Hi);
+  }
+
   if (ConvertToScalable) {
     ResultVT = getContainerForFixedLengthVector(DAG, ResultVT);
     OpVT = getContainerForFixedLengthVector(DAG, LHS.getValueType());

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

@@ -1453,30 +1453,48 @@ define <4 x i32> @partial_reduce_shl_zext_non_const_rhs(<16 x i8> %l, <4 x i32>
 }
 
 define <2 x i32> @udot_v16i8tov2i32(<2 x i32> %acc, <16 x i8> %input) {
-; CHECK-COMMON-LABEL: udot_v16i8tov2i32:
-; CHECK-COMMON:       // %bb.0: // %entry
-; CHECK-COMMON-NEXT:    ushll v2.8h, v1.8b, #0
-; CHECK-COMMON-NEXT:    // kill: def $d0 killed $d0 def $q0
-; CHECK-COMMON-NEXT:    ushll2 v1.8h, v1.16b, #0
-; CHECK-COMMON-NEXT:    ushll v3.4s, v2.4h, #0
-; CHECK-COMMON-NEXT:    uaddw v0.4s, v0.4s, v2.4h
-; CHECK-COMMON-NEXT:    ushll2 v4.4s, v2.8h, #0
-; CHECK-COMMON-NEXT:    ext v2.16b, v2.16b, v2.16b, #8
-; CHECK-COMMON-NEXT:    ext v3.16b, v3.16b, v3.16b, #8
-; CHECK-COMMON-NEXT:    add v0.2s, v3.2s, v0.2s
-; CHECK-COMMON-NEXT:    ext v3.16b, v4.16b, v4.16b, #8
-; CHECK-COMMON-NEXT:    uaddw v0.4s, v0.4s, v2.4h
-; CHECK-COMMON-NEXT:    ushll v2.4s, v1.4h, #0
-; CHECK-COMMON-NEXT:    add v0.2s, v3.2s, v0.2s
-; CHECK-COMMON-NEXT:    ext v2.16b, v2.16b, v2.16b, #8
-; CHECK-COMMON-NEXT:    ushll2 v3.4s, v1.8h, #0
-; CHECK-COMMON-NEXT:    uaddw v0.4s, v0.4s, v1.4h
-; CHECK-COMMON-NEXT:    ext v1.16b, v1.16b, v1.16b, #8
-; CHECK-COMMON-NEXT:    add v0.2s, v2.2s, v0.2s
-; CHECK-COMMON-NEXT:    ext v2.16b, v3.16b, v3.16b, #8
-; CHECK-COMMON-NEXT:    uaddw v0.4s, v0.4s, v1.4h
-; CHECK-COMMON-NEXT:    add v0.2s, v2.2s, v0.2s
-; CHECK-COMMON-NEXT:    ret
+; CHECK-NODOT-LABEL: udot_v16i8tov2i32:
+; CHECK-NODOT:       // %bb.0: // %entry
+; CHECK-NODOT-NEXT:    ushll v2.8h, v1.8b, #0
+; CHECK-NODOT-NEXT:    // kill: def $d0 killed $d0 def $q0
+; CHECK-NODOT-NEXT:    ushll2 v1.8h, v1.16b, #0
+; CHECK-NODOT-NEXT:    ushll v3.4s, v2.4h, #0
+; CHECK-NODOT-NEXT:    uaddw v0.4s, v0.4s, v2.4h
+; CHECK-NODOT-NEXT:    ushll2 v4.4s, v2.8h, #0
+; CHECK-NODOT-NEXT:    ext v2.16b, v2.16b, v2.16b, #8
+; CHECK-NODOT-NEXT:    ext v3.16b, v3.16b, v3.16b, #8
+; CHECK-NODOT-NEXT:    add v0.2s, v3.2s, v0.2s
+; CHECK-NODOT-NEXT:    ext v3.16b, v4.16b, v4.16b, #8
+; CHECK-NODOT-NEXT:    uaddw v0.4s, v0.4s, v2.4h
+; CHECK-NODOT-NEXT:    ushll v2.4s, v1.4h, #0
+; CHECK-NODOT-NEXT:    add v0.2s, v3.2s, v0.2s
+; CHECK-NODOT-NEXT:    ext v2.16b, v2.16b, v2.16b, #8
+; CHECK-NODOT-NEXT:    ushll2 v3.4s, v1.8h, #0
+; CHECK-NODOT-NEXT:    uaddw v0.4s, v0.4s, v1.4h
+; CHECK-NODOT-NEXT:    ext v1.16b, v1.16b, v1.16b, #8
+; CHECK-NODOT-NEXT:    add v0.2s, v2.2s, v0.2s
+; CHECK-NODOT-NEXT:    ext v2.16b, v3.16b, v3.16b, #8
+; CHECK-NODOT-NEXT:    uaddw v0.4s, v0.4s, v1.4h
+; CHECK-NODOT-NEXT:    add v0.2s, v2.2s, v0.2s
+; CHECK-NODOT-NEXT:    ret
+;
+; CHECK-DOT-LABEL: udot_v16i8tov2i32:
+; CHECK-DOT:       // %bb.0: // %entry
+; CHECK-DOT-NEXT:    movi v2.16b, #1
+; CHECK-DOT-NEXT:    fmov d0, d0
+; CHECK-DOT-NEXT:    udot v0.4s, v1.16b, v2.16b
+; CHECK-DOT-NEXT:    ext v1.16b, v0.16b, v0.16b, #8
+; CHECK-DOT-NEXT:    add v0.2s, v0.2s, v1.2s
+; CHECK-DOT-NEXT:    ret
+;
+; CHECK-DOT-I8MM-LABEL: udot_v16i8tov2i32:
+; CHECK-DOT-I8MM:       // %bb.0: // %entry
+; CHECK-DOT-I8MM-NEXT:    movi v2.16b, #1
+; CHECK-DOT-I8MM-NEXT:    fmov d0, d0
+; CHECK-DOT-I8MM-NEXT:    udot v0.4s, v1.16b, v2.16b
+; CHECK-DOT-I8MM-NEXT:    ext v1.16b, v0.16b, v0.16b, #8
+; CHECK-DOT-I8MM-NEXT:    add v0.2s, v0.2s, v1.2s
+; CHECK-DOT-I8MM-NEXT:    ret
 entry:
     %input.wide = zext <16 x i8> %input to <16 x i32>
     %partial.reduce = tail call <2 x i32> @llvm.vector.partial.reduce.add(<2 x i32> %acc, <16 x i32> %input.wide)