[SelectionDAG][AArch64] Add dot product lowering in NEON for PARTIAL_REDUCE_*MLA ISD nodes

JamesChesterman · NickGuy-Arm · commit 1ed6ae3610d5 · 2025-05-13T17:58:27.000+01:00
Lowering for fixed width vectors added to tablegen.
There is also custom lowering to ensure that the USDOT patterns are
still lowered for fixed width vectors. It also ensures that the
v16i8 -&gt; v4i64 partial reduction case is lowered here instead of
being split (as there is not a v2i64 dot product instruction).
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -7744,7 +7744,11 @@ SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
     return LowerVECTOR_HISTOGRAM(Op, DAG);
   case ISD::PARTIAL_REDUCE_SMLA:
   case ISD::PARTIAL_REDUCE_UMLA:
-    return LowerPARTIAL_REDUCE_MLA(Op, DAG);
+  case ISD::PARTIAL_REDUCE_SMLA: {
+    if (SDValue Result = LowerPARTIAL_REDUCE_MLA(Op, DAG))
+      return Result;
+    return expandPartialReduceMLA(Op.getNode(), DAG);
+  }
   }
 }
 
@@ -27569,6 +27573,15 @@ void AArch64TargetLowering::ReplaceNodeResults(
     if (SDValue Res = LowerVECTOR_COMPRESS(SDValue(N, 0), DAG))
       Results.push_back(Res);
     return;
+  case ISD::PARTIAL_REDUCE_UMLA:
+  case ISD::PARTIAL_REDUCE_SMLA: {
+    SDValue Res;
+    if (Res = LowerPARTIAL_REDUCE_MLA(SDValue(N, 0), DAG))
+      Results.push_back(Res);
+    else
+      Results.push_back(expandPartialReduceMLA(N, DAG));
+    return;
+  }
   case ISD::ADD:
   case ISD::FADD:
     ReplaceAddWithADDP(N, Results, DAG, Subtarget);
@@ -29524,6 +29537,7 @@ SDValue AArch64TargetLowering::LowerVECTOR_HISTOGRAM(SDValue Op,
 SDValue
 AArch64TargetLowering::LowerPARTIAL_REDUCE_MLA(SDValue Op,
                                                SelectionDAG &DAG) const {
+
   SDLoc DL(Op);
 
   SDValue Acc = Op.getOperand(0);
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -1474,6 +1474,17 @@ defm SDOTlane : SIMDThreeSameVectorDotIndex<0, 0, 0b10, "sdot", AArch64sdot>;
 defm UDOTlane : SIMDThreeSameVectorDotIndex<1, 0, 0b10, "udot", AArch64udot>;
 }
 
+let Predicates = [HasNEON, HasDotProd] in {
+  def : Pat<(v4i32 (partial_reduce_umla (v4i32 V128:$Acc), (v16i8 V128:$MulLHS), (v16i8 V128:$MulRHS))),
+            (v4i32 (UDOTv16i8 V128:$Acc, V128:$MulLHS, V128:$MulRHS))>;
+  def : Pat<(v4i32 (partial_reduce_smla (v4i32 V128:$Acc), (v16i8 V128:$MulLHS), (v16i8 V128:$MulRHS))),
+            (v4i32 (SDOTv16i8 V128:$Acc, V128:$MulLHS, V128:$MulRHS))>;
+  def : Pat<(v2i32 (partial_reduce_umla (v2i32 V64:$Acc), (v8i8 V64:$MulLHS), (v8i8 V64:$MulRHS))),
+            (v2i32 (UDOTv8i8 V64:$Acc, V64:$MulLHS, V64:$MulRHS))>;
+  def : Pat<(v2i32 (partial_reduce_smla (v2i32 V64:$Acc), (v8i8 V64:$MulLHS), (v8i8 V64:$MulRHS))),
+            (v2i32 (SDOTv8i8 V64:$Acc, V64:$MulLHS, V64:$MulRHS))>; 
+} // End HasNEON, HasDotProd
+
 // ARMv8.6-A BFloat
 let Predicates = [HasNEON, HasBF16] in {
 defm BFDOT       : SIMDThreeSameVectorBFDot<1, "bfdot">;
diff --git a/llvm/test/CodeGen/AArch64/neon-partial-reduce-dot-product.ll b/llvm/test/CodeGen/AArch64/neon-partial-reduce-dot-product.ll
@@ -2,7 +2,8 @@
 ; RUN: llc -mtriple aarch64 -mattr=+neon,+dotprod < %s | FileCheck %s --check-prefixes=CHECK,CHECK-DOT,CHECK-NOI8MM
 ; 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
-; RUN: llc -mtriple aarch64 -mattr=+neon,+dotprod,+i8mm -aarch64-enable-partial-reduce-nodes < %s | FileCheck %s --check-prefixes=CHECK,CHECK-NOI8MM
+; 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:
@@ -49,17 +50,20 @@ define <4 x i32> @udot_in_loop(ptr %p1, ptr %p2){
 ; CHECK-NODOT-NEXT:    mov x8, xzr
 ; CHECK-NODOT-NEXT:  .LBB1_1: // %vector.body
 ; CHECK-NODOT-NEXT:    // =>This Inner Loop Header: Depth=1
-; CHECK-NODOT-NEXT:    ldr q2, [x0, x8]
-; CHECK-NODOT-NEXT:    ldr q3, [x1, x8]
-; CHECK-NODOT-NEXT:    mov v0.16b, v1.16b
+; CHECK-NODOT-NEXT:    ldr q0, [x0, x8]
+; CHECK-NODOT-NEXT:    ldr q2, [x1, x8]
 ; CHECK-NODOT-NEXT:    add x8, x8, #16
 ; CHECK-NODOT-NEXT:    umull v4.8h, v2.8b, v3.8b
 ; CHECK-NODOT-NEXT:    umull2 v2.8h, v2.16b, v3.16b
 ; CHECK-NODOT-NEXT:    cmp x8, #16
-; CHECK-NODOT-NEXT:    uaddw v1.4s, v1.4s, v4.4h
-; CHECK-NODOT-NEXT:    uaddw2 v1.4s, v1.4s, v4.8h
-; CHECK-NODOT-NEXT:    uaddw v1.4s, v1.4s, v2.4h
-; CHECK-NODOT-NEXT:    uaddw2 v1.4s, v1.4s, v2.8h
+; CHECK-NODOT-NEXT:    umull v3.8h, v0.8b, v2.8b
+; CHECK-NODOT-NEXT:    umull2 v2.8h, v0.16b, v2.16b
+; CHECK-NODOT-NEXT:    mov v0.16b, v1.16b
+; CHECK-NODOT-NEXT:    ushll v1.4s, v2.4h, #0
+; CHECK-NODOT-NEXT:    uaddw v4.4s, v0.4s, v3.4h
+; CHECK-NODOT-NEXT:    uaddw2 v1.4s, v1.4s, v3.8h
+; CHECK-NODOT-NEXT:    uaddw2 v2.4s, v4.4s, v2.8h
+; CHECK-NODOT-NEXT:    add v1.4s, v1.4s, v2.4s
 ; CHECK-NODOT-NEXT:    b.ne .LBB1_1
 ; CHECK-NODOT-NEXT:  // %bb.2: // %end
 ; CHECK-NODOT-NEXT:    ret
@@ -563,22 +567,6 @@ define <4 x i32> @udot_no_bin_op(<4 x i32> %acc, <16 x i8> %a){
 }
 
 define <4 x i32> @udot_no_bin_op_in_loop(ptr %p){
-; CHECK-DOT-LABEL: udot_no_bin_op_in_loop:
-; CHECK-DOT:       // %bb.0: // %entry
-; CHECK-DOT-NEXT:    movi v1.2d, #0000000000000000
-; CHECK-DOT-NEXT:    movi v2.16b, #1
-; CHECK-DOT-NEXT:    mov x8, xzr
-; CHECK-DOT-NEXT:  .LBB16_1: // %vector.body
-; CHECK-DOT-NEXT:    // =>This Inner Loop Header: Depth=1
-; CHECK-DOT-NEXT:    ldr q3, [x0, x8]
-; CHECK-DOT-NEXT:    mov v0.16b, v1.16b
-; CHECK-DOT-NEXT:    add x8, x8, #16
-; CHECK-DOT-NEXT:    cmp x8, #16
-; CHECK-DOT-NEXT:    udot v1.4s, v3.16b, v2.16b
-; CHECK-DOT-NEXT:    b.ne .LBB16_1
-; CHECK-DOT-NEXT:  // %bb.2: // %end
-; CHECK-DOT-NEXT:    ret
-;
 ; CHECK-NODOT-LABEL: udot_no_bin_op_in_loop:
 ; CHECK-NODOT:       // %bb.0: // %entry
 ; CHECK-NODOT-NEXT:    movi v1.2d, #0000000000000000
@@ -598,6 +586,66 @@ define <4 x i32> @udot_no_bin_op_in_loop(ptr %p){
 ; CHECK-NODOT-NEXT:    b.ne .LBB16_1
 ; CHECK-NODOT-NEXT:  // %bb.2: // %end
 ; CHECK-NODOT-NEXT:    ret
+;
+; CHECK-NEWLOWERING-I8MM-LABEL: udot_no_bin_op_in_loop:
+; CHECK-NEWLOWERING-I8MM:       // %bb.0: // %entry
+; CHECK-NEWLOWERING-I8MM-NEXT:    adrp x8, .LCPI16_0
+; CHECK-NEWLOWERING-I8MM-NEXT:    movi v4.2d, #0000000000000000
+; CHECK-NEWLOWERING-I8MM-NEXT:    adrp x9, .LCPI16_2
+; CHECK-NEWLOWERING-I8MM-NEXT:    ldr q1, [x8, :lo12:.LCPI16_0]
+; CHECK-NEWLOWERING-I8MM-NEXT:    adrp x8, .LCPI16_1
+; CHECK-NEWLOWERING-I8MM-NEXT:    adrp x10, .LCPI16_3
+; CHECK-NEWLOWERING-I8MM-NEXT:    ldr q2, [x8, :lo12:.LCPI16_1]
+; CHECK-NEWLOWERING-I8MM-NEXT:    ldr q3, [x9, :lo12:.LCPI16_2]
+; CHECK-NEWLOWERING-I8MM-NEXT:    ldr q5, [x10, :lo12:.LCPI16_3]
+; CHECK-NEWLOWERING-I8MM-NEXT:    mov x8, xzr
+; CHECK-NEWLOWERING-I8MM-NEXT:  .LBB16_1: // %vector.body
+; CHECK-NEWLOWERING-I8MM-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEWLOWERING-I8MM-NEXT:    ldr q6, [x0, x8]
+; CHECK-NEWLOWERING-I8MM-NEXT:    mov v0.16b, v4.16b
+; CHECK-NEWLOWERING-I8MM-NEXT:    add x8, x8, #16
+; CHECK-NEWLOWERING-I8MM-NEXT:    cmp x8, #16
+; CHECK-NEWLOWERING-I8MM-NEXT:    tbl v7.16b, { v6.16b }, v2.16b
+; CHECK-NEWLOWERING-I8MM-NEXT:    tbl v4.16b, { v6.16b }, v1.16b
+; CHECK-NEWLOWERING-I8MM-NEXT:    tbl v16.16b, { v6.16b }, v3.16b
+; CHECK-NEWLOWERING-I8MM-NEXT:    tbl v6.16b, { v6.16b }, v5.16b
+; CHECK-NEWLOWERING-I8MM-NEXT:    add v7.4s, v0.4s, v7.4s
+; CHECK-NEWLOWERING-I8MM-NEXT:    add v6.4s, v6.4s, v16.4s
+; CHECK-NEWLOWERING-I8MM-NEXT:    add v4.4s, v4.4s, v7.4s
+; CHECK-NEWLOWERING-I8MM-NEXT:    add v4.4s, v6.4s, v4.4s
+; CHECK-NEWLOWERING-I8MM-NEXT:    b.ne .LBB16_1
+; CHECK-NEWLOWERING-I8MM-NEXT:  // %bb.2: // %end
+; CHECK-NEWLOWERING-I8MM-NEXT:    ret
+;
+; CHECK-NEWLOWERING-NOI8MM-LABEL: udot_no_bin_op_in_loop:
+; CHECK-NEWLOWERING-NOI8MM:       // %bb.0: // %entry
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    adrp x8, .LCPI16_0
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    movi v4.2d, #0000000000000000
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    adrp x9, .LCPI16_2
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    ldr q1, [x8, :lo12:.LCPI16_0]
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    adrp x8, .LCPI16_1
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    adrp x10, .LCPI16_3
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    ldr q2, [x8, :lo12:.LCPI16_1]
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    ldr q3, [x9, :lo12:.LCPI16_2]
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    ldr q5, [x10, :lo12:.LCPI16_3]
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    mov x8, xzr
+; CHECK-NEWLOWERING-NOI8MM-NEXT:  .LBB16_1: // %vector.body
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    ldr q6, [x0, x8]
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    mov v0.16b, v4.16b
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    add x8, x8, #16
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    cmp x8, #16
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    tbl v7.16b, { v6.16b }, v2.16b
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    tbl v4.16b, { v6.16b }, v1.16b
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    tbl v16.16b, { v6.16b }, v3.16b
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    tbl v6.16b, { v6.16b }, v5.16b
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    add v7.4s, v0.4s, v7.4s
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    add v6.4s, v6.4s, v16.4s
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    add v4.4s, v4.4s, v7.4s
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    add v4.4s, v6.4s, v4.4s
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    b.ne .LBB16_1
+; CHECK-NEWLOWERING-NOI8MM-NEXT:  // %bb.2: // %end
+; CHECK-NEWLOWERING-NOI8MM-NEXT:    ret
 
 entry:
   br label %vector.body
