[AArch64] Add Neon USDOT support (#143525)

NickGuy-Arm · web-flow · commit 7c946e6e478e · 2025-06-10T16:54:19.000+01:00
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -1464,6 +1464,16 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
       setPartialReduceMLAAction(MLAOps, MVT::v4i32, MVT::v16i8, Legal);
       setPartialReduceMLAAction(MLAOps, MVT::v2i32, MVT::v8i8, Legal);
       setPartialReduceMLAAction(MLAOps, MVT::v2i64, MVT::v16i8, Custom);
+
+      if (Subtarget->hasMatMulInt8()) {
+        setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SUMLA, MVT::v4i32,
+                                  MVT::v16i8, Legal);
+        setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SUMLA, MVT::v2i64,
+                                  MVT::v16i8, Custom);
+
+        setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SUMLA, MVT::v2i32,
+                                  MVT::v8i8, Legal);
+      }
     }
 
   } else /* !isNeonAvailable */ {
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -1711,6 +1711,11 @@ multiclass SIMDSUDOTIndex {
 
 defm SUDOTlane : SIMDSUDOTIndex;
 
+def : Pat<(v2i32 (partial_reduce_sumla v2i32:$Acc, v8i8:$LHS, v8i8:$RHS)),
+              (USDOTv8i8 $Acc, $RHS, $LHS)>;
+def : Pat<(v4i32 (partial_reduce_sumla v4i32:$Acc, v16i8:$LHS, v16i8:$RHS)),
+              (USDOTv16i8 $Acc, $RHS, $LHS)>;
+
 }
 
 // ARMv8.2-A FP16 Fused Multiply-Add Long
diff --git a/llvm/test/CodeGen/AArch64/neon-partial-reduce-dot-product.ll b/llvm/test/CodeGen/AArch64/neon-partial-reduce-dot-product.ll
@@ -181,14 +181,7 @@ define <4 x i32> @usdot(<4 x i32> %acc, <16 x i8> %u, <16 x i8> %s) {
 ;
 ; CHECK-NEWLOWERING-I8MM-LABEL: usdot:
 ; CHECK-NEWLOWERING-I8MM:       // %bb.0:
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v3.8h, v1.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v4.8h, v2.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll2 v1.8h, v1.16b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll2 v2.8h, v2.16b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.4s, v4.4h, v3.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v0.4s, v4.8h, v3.8h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.4s, v2.4h, v1.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v0.4s, v2.8h, v1.8h
+; CHECK-NEWLOWERING-I8MM-NEXT:    usdot v0.4s, v1.16b, v2.16b
 ; CHECK-NEWLOWERING-I8MM-NEXT:    ret
   %u.wide = zext <16 x i8> %u to <16 x i32>
   %s.wide = sext <16 x i8> %s to <16 x i32>
@@ -247,15 +240,8 @@ define <4 x i32> @usdot_in_loop(ptr %p1, ptr %p2){
 ; CHECK-NEWLOWERING-I8MM-NEXT:    ldr q3, [x1, x8]
 ; CHECK-NEWLOWERING-I8MM-NEXT:    mov v0.16b, v1.16b
 ; CHECK-NEWLOWERING-I8MM-NEXT:    add x8, x8, #16
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v4.8h, v2.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v5.8h, v3.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll2 v2.8h, v2.16b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll2 v3.8h, v3.16b, #0
+; CHECK-NEWLOWERING-I8MM-NEXT:    usdot v1.4s, v3.16b, v2.16b
 ; CHECK-NEWLOWERING-I8MM-NEXT:    cmp x8, #16
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v1.4s, v4.4h, v5.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v1.4s, v4.8h, v5.8h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v1.4s, v2.4h, v3.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v1.4s, v2.8h, v3.8h
 ; CHECK-NEWLOWERING-I8MM-NEXT:    b.ne .LBB6_1
 ; CHECK-NEWLOWERING-I8MM-NEXT:  // %bb.2: // %end
 ; CHECK-NEWLOWERING-I8MM-NEXT:    ret
@@ -306,19 +292,7 @@ define <2 x i32> @usdot_narrow(<2 x i32> %acc, <8 x i8> %u, <8 x i8> %s) #0{
 ;
 ; CHECK-NEWLOWERING-I8MM-LABEL: usdot_narrow:
 ; CHECK-NEWLOWERING-I8MM:       // %bb.0:
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v1.8h, v1.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v2.8h, v2.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    // kill: def $d0 killed $d0 def $q0
-; CHECK-NEWLOWERING-I8MM-NEXT:    smull v3.4s, v2.4h, v1.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.4s, v2.4h, v1.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    ext v4.16b, v1.16b, v1.16b, #8
-; CHECK-NEWLOWERING-I8MM-NEXT:    ext v5.16b, v2.16b, v2.16b, #8
-; CHECK-NEWLOWERING-I8MM-NEXT:    smull2 v1.4s, v2.8h, v1.8h
-; CHECK-NEWLOWERING-I8MM-NEXT:    ext v3.16b, v3.16b, v3.16b, #8
-; CHECK-NEWLOWERING-I8MM-NEXT:    ext v1.16b, v1.16b, v1.16b, #8
-; CHECK-NEWLOWERING-I8MM-NEXT:    add v0.2s, v3.2s, v0.2s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.4s, v5.4h, v4.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    add v0.2s, v1.2s, v0.2s
+; CHECK-NEWLOWERING-I8MM-NEXT:    usdot v0.2s, v1.8b, v2.8b
 ; CHECK-NEWLOWERING-I8MM-NEXT:    ret
   %u.wide = zext <8 x i8> %u to <8 x i32>
   %s.wide = sext <8 x i8> %s to <8 x i32>
@@ -347,14 +321,7 @@ define <4 x i32> @sudot(<4 x i32> %acc, <16 x i8> %u, <16 x i8> %s) #0{
 ;
 ; CHECK-NEWLOWERING-I8MM-LABEL: sudot:
 ; CHECK-NEWLOWERING-I8MM:       // %bb.0:
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v3.8h, v1.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v4.8h, v2.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll2 v1.8h, v1.16b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll2 v2.8h, v2.16b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.4s, v4.4h, v3.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v0.4s, v4.8h, v3.8h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.4s, v2.4h, v1.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v0.4s, v2.8h, v1.8h
+; CHECK-NEWLOWERING-I8MM-NEXT:    usdot v0.4s, v2.16b, v1.16b
 ; CHECK-NEWLOWERING-I8MM-NEXT:    ret
   %s.wide = sext <16 x i8> %u to <16 x i32>
   %u.wide = zext <16 x i8> %s to <16 x i32>
@@ -413,15 +380,8 @@ define <4 x i32> @sudot_in_loop(ptr %p1, ptr %p2){
 ; CHECK-NEWLOWERING-I8MM-NEXT:    ldr q3, [x1, x8]
 ; CHECK-NEWLOWERING-I8MM-NEXT:    mov v0.16b, v1.16b
 ; CHECK-NEWLOWERING-I8MM-NEXT:    add x8, x8, #16
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v4.8h, v2.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v5.8h, v3.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll2 v2.8h, v2.16b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll2 v3.8h, v3.16b, #0
+; CHECK-NEWLOWERING-I8MM-NEXT:    usdot v1.4s, v2.16b, v3.16b
 ; CHECK-NEWLOWERING-I8MM-NEXT:    cmp x8, #16
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v1.4s, v4.4h, v5.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v1.4s, v4.8h, v5.8h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v1.4s, v2.4h, v3.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v1.4s, v2.8h, v3.8h
 ; CHECK-NEWLOWERING-I8MM-NEXT:    b.ne .LBB9_1
 ; CHECK-NEWLOWERING-I8MM-NEXT:  // %bb.2: // %end
 ; CHECK-NEWLOWERING-I8MM-NEXT:    ret
@@ -472,19 +432,7 @@ define <2 x i32> @sudot_narrow(<2 x i32> %acc, <8 x i8> %u, <8 x i8> %s) #0{
 ;
 ; CHECK-NEWLOWERING-I8MM-LABEL: sudot_narrow:
 ; CHECK-NEWLOWERING-I8MM:       // %bb.0:
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v1.8h, v1.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v2.8h, v2.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    // kill: def $d0 killed $d0 def $q0
-; CHECK-NEWLOWERING-I8MM-NEXT:    smull v3.4s, v2.4h, v1.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.4s, v2.4h, v1.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    ext v4.16b, v1.16b, v1.16b, #8
-; CHECK-NEWLOWERING-I8MM-NEXT:    ext v5.16b, v2.16b, v2.16b, #8
-; CHECK-NEWLOWERING-I8MM-NEXT:    smull2 v1.4s, v2.8h, v1.8h
-; CHECK-NEWLOWERING-I8MM-NEXT:    ext v3.16b, v3.16b, v3.16b, #8
-; CHECK-NEWLOWERING-I8MM-NEXT:    ext v1.16b, v1.16b, v1.16b, #8
-; CHECK-NEWLOWERING-I8MM-NEXT:    add v0.2s, v3.2s, v0.2s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.4s, v5.4h, v4.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    add v0.2s, v1.2s, v0.2s
+; CHECK-NEWLOWERING-I8MM-NEXT:    usdot v0.2s, v2.8b, v1.8b
 ; CHECK-NEWLOWERING-I8MM-NEXT:    ret
   %u.wide = sext <8 x i8> %u to <8 x i32>
   %s.wide = zext <8 x i8> %s to <8 x i32>
@@ -614,26 +562,10 @@ define <4 x i64> @usdot_8to64(<4 x i64> %acc, <16 x i8> %a, <16 x i8> %b){
 ;
 ; CHECK-NEWLOWERING-I8MM-LABEL: usdot_8to64:
 ; CHECK-NEWLOWERING-I8MM:       // %bb.0: // %entry
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v4.8h, v2.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll2 v2.8h, v2.16b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v5.8h, v3.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll2 v3.8h, v3.16b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v6.4s, v4.4h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v7.4s, v2.4h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v16.4s, v5.4h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v17.4s, v3.4h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll2 v4.4s, v4.8h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll2 v2.4s, v2.8h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll2 v5.4s, v5.8h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll2 v3.4s, v3.8h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.2d, v6.2s, v16.2s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v1.2d, v7.2s, v17.2s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v0.2d, v6.4s, v16.4s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v1.2d, v7.4s, v17.4s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.2d, v4.2s, v5.2s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v1.2d, v2.2s, v3.2s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v0.2d, v4.4s, v5.4s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v1.2d, v2.4s, v3.4s
+; CHECK-NEWLOWERING-I8MM-NEXT:    movi v4.2d, #0000000000000000
+; CHECK-NEWLOWERING-I8MM-NEXT:    usdot v4.4s, v2.16b, v3.16b
+; CHECK-NEWLOWERING-I8MM-NEXT:    saddw v0.2d, v0.2d, v4.2s
+; CHECK-NEWLOWERING-I8MM-NEXT:    saddw2 v0.2d, v0.2d, v4.4s
 ; CHECK-NEWLOWERING-I8MM-NEXT:    ret
 entry:
   %a.wide = zext <16 x i8> %a to <16 x i64>
@@ -679,26 +611,10 @@ define <4 x i64> @sudot_8to64(<4 x i64> %acc, <16 x i8> %a, <16 x i8> %b) {
 ;
 ; CHECK-NEWLOWERING-I8MM-LABEL: sudot_8to64:
 ; CHECK-NEWLOWERING-I8MM:       // %bb.0: // %entry
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v4.8h, v2.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll2 v2.8h, v2.16b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v5.8h, v3.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll2 v3.8h, v3.16b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v6.4s, v4.4h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v7.4s, v2.4h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v16.4s, v5.4h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v17.4s, v3.4h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll2 v4.4s, v4.8h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll2 v2.4s, v2.8h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll2 v5.4s, v5.8h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll2 v3.4s, v3.8h, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.2d, v6.2s, v16.2s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v1.2d, v7.2s, v17.2s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v0.2d, v6.4s, v16.4s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v1.2d, v7.4s, v17.4s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.2d, v4.2s, v5.2s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v1.2d, v2.2s, v3.2s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v0.2d, v4.4s, v5.4s
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v1.2d, v2.4s, v3.4s
+; CHECK-NEWLOWERING-I8MM-NEXT:    movi v4.2d, #0000000000000000
+; CHECK-NEWLOWERING-I8MM-NEXT:    usdot v4.4s, v3.16b, v2.16b
+; CHECK-NEWLOWERING-I8MM-NEXT:    saddw v0.2d, v0.2d, v4.2s
+; CHECK-NEWLOWERING-I8MM-NEXT:    saddw2 v0.2d, v0.2d, v4.4s
 ; CHECK-NEWLOWERING-I8MM-NEXT:    ret
 entry:
   %a.wide = sext <16 x i8> %a to <16 x i64>
@@ -1147,21 +1063,9 @@ define <4 x i32> @usdot_multiple_zext_users(ptr %p1, ptr %p2, ptr %p3) {
 ; CHECK-NEWLOWERING-I8MM-NEXT:    ldr q3, [x1, x8]
 ; CHECK-NEWLOWERING-I8MM-NEXT:    ldr q4, [x2, x8]
 ; CHECK-NEWLOWERING-I8MM-NEXT:    add x8, x8, #16
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v5.8h, v2.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll v6.8h, v4.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll v7.8h, v3.8b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll2 v2.8h, v2.16b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    ushll2 v4.8h, v4.16b, #0
-; CHECK-NEWLOWERING-I8MM-NEXT:    sshll2 v3.8h, v3.16b, #0
+; CHECK-NEWLOWERING-I8MM-NEXT:    usdot v0.4s, v4.16b, v2.16b
+; CHECK-NEWLOWERING-I8MM-NEXT:    usdot v1.4s, v4.16b, v3.16b
 ; CHECK-NEWLOWERING-I8MM-NEXT:    cmp x8, #1024
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.4s, v5.4h, v6.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v1.4s, v7.4h, v6.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v0.4s, v5.8h, v6.8h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v1.4s, v7.8h, v6.8h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v0.4s, v2.4h, v4.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal v1.4s, v3.4h, v4.4h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v0.4s, v2.8h, v4.8h
-; CHECK-NEWLOWERING-I8MM-NEXT:    smlal2 v1.4s, v3.8h, v4.8h
 ; CHECK-NEWLOWERING-I8MM-NEXT:    b.ne .LBB28_1
 ; CHECK-NEWLOWERING-I8MM-NEXT:  // %bb.2: // %end
 ; CHECK-NEWLOWERING-I8MM-NEXT:    add v0.4s, v1.4s, v0.4s

Original file line number	Diff line number	Diff line change
`@@ -1711,6 +1711,11 @@ multiclass SIMDSUDOTIndex {`
`1711`	`1711`
`1712`	`1712`	`defm SUDOTlane : SIMDSUDOTIndex;`
`1713`	`1713`
	`1714`	`+def : Pat<(v2i32 (partial_reduce_sumla v2i32:$Acc, v8i8:$LHS, v8i8:$RHS)),`
	`1715`	`+ (USDOTv8i8 $Acc, $RHS, $LHS)>;`
	`1716`	`+def : Pat<(v4i32 (partial_reduce_sumla v4i32:$Acc, v16i8:$LHS, v16i8:$RHS)),`
	`1717`	`+ (USDOTv16i8 $Acc, $RHS, $LHS)>;`
	`1718`	`+`
`1714`	`1719`	`}`
`1715`	`1720`
`1716`	`1721`	`// ARMv8.2-A FP16 Fused Multiply-Add Long`