[AArch64] Improve codegen for some fixed-width partial reductions

This patch teaches optimizeExtendOrTruncateConversion to bail out
if the user of a zero-extend is a partial reduction intrinsic
that we know will get lowered efficiently to a udot instruction.

Author: david-arm

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -16866,9 +16866,14 @@ bool AArch64TargetLowering::optimizeExtendOrTruncateConversion(
     // mul(zext(i8), sext) can be transformed into smull(zext, sext) which
     // performs one extend implicitly. If DstWidth is at most 4 * SrcWidth, at
     // most one extra extend step is needed and using tbl is not profitable.
+    // Similarly, bail out if partial_reduce(acc, zext(i8)) can be lowered to a
+    // udot instruction.
     if (SrcWidth * 4 <= DstWidth && I->hasOneUser()) {
       auto *SingleUser = cast<Instruction>(*I->user_begin());
-      if (match(SingleUser, m_c_Mul(m_Specific(I), m_SExt(m_Value()))))
+      if (match(SingleUser, m_c_Mul(m_Specific(I), m_SExt(m_Value()))) ||
+          (isa<IntrinsicInst>(SingleUser) &&
+           !shouldExpandPartialReductionIntrinsic(
+               cast<IntrinsicInst>(SingleUser))))
         return false;
     }
 

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

@@ -574,35 +574,42 @@ 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-LABEL: udot_no_bin_op_in_loop:
-; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    adrp x8, .LCPI20_0
-; CHECK-NEXT:    movi v4.2d, #0000000000000000
-; CHECK-NEXT:    adrp x9, .LCPI20_2
-; CHECK-NEXT:    ldr q1, [x8, :lo12:.LCPI20_0]
-; CHECK-NEXT:    adrp x8, .LCPI20_1
-; CHECK-NEXT:    adrp x10, .LCPI20_3
-; CHECK-NEXT:    ldr q2, [x8, :lo12:.LCPI20_1]
-; CHECK-NEXT:    ldr q3, [x9, :lo12:.LCPI20_2]
-; CHECK-NEXT:    ldr q5, [x10, :lo12:.LCPI20_3]
-; CHECK-NEXT:    mov x8, xzr
-; CHECK-NEXT:  .LBB20_1: // %vector.body
-; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:    ldr q6, [x0, x8]
-; CHECK-NEXT:    mov v0.16b, v4.16b
-; CHECK-NEXT:    add x8, x8, #16
-; CHECK-NEXT:    cmp x8, #16
-; CHECK-NEXT:    tbl v7.16b, { v6.16b }, v2.16b
-; CHECK-NEXT:    tbl v4.16b, { v6.16b }, v1.16b
-; CHECK-NEXT:    tbl v16.16b, { v6.16b }, v3.16b
-; CHECK-NEXT:    tbl v6.16b, { v6.16b }, v5.16b
-; CHECK-NEXT:    add v7.4s, v0.4s, v7.4s
-; CHECK-NEXT:    add v6.4s, v6.4s, v16.4s
-; CHECK-NEXT:    add v4.4s, v4.4s, v7.4s
-; CHECK-NEXT:    add v4.4s, v6.4s, v4.4s
-; CHECK-NEXT:    b.ne .LBB20_1
-; CHECK-NEXT:  // %bb.2: // %end
-; CHECK-NEXT:    ret
+; 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:  .LBB20_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 .LBB20_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
+; CHECK-NODOT-NEXT:    mov x8, xzr
+; CHECK-NODOT-NEXT:  .LBB20_1: // %vector.body
+; CHECK-NODOT-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NODOT-NEXT:    ldr q0, [x0, x8]
+; CHECK-NODOT-NEXT:    add x8, x8, #16
+; CHECK-NODOT-NEXT:    cmp x8, #16
+; CHECK-NODOT-NEXT:    ushll v2.8h, v0.8b, #0
+; CHECK-NODOT-NEXT:    ushll2 v3.8h, v0.16b, #0
+; CHECK-NODOT-NEXT:    mov v0.16b, v1.16b
+; CHECK-NODOT-NEXT:    ushll v1.4s, v3.4h, #0
+; CHECK-NODOT-NEXT:    uaddw v4.4s, v0.4s, v2.4h
+; CHECK-NODOT-NEXT:    uaddw2 v1.4s, v1.4s, v2.8h
+; CHECK-NODOT-NEXT:    uaddw2 v2.4s, v4.4s, v3.8h
+; CHECK-NODOT-NEXT:    add v1.4s, v1.4s, v2.4s
+; CHECK-NODOT-NEXT:    b.ne .LBB20_1
+; CHECK-NODOT-NEXT:  // %bb.2: // %end
+; CHECK-NODOT-NEXT:    ret
 entry:
   br label %vector.body