[AArch64] Generate usdot instruction with multiple zext users in loop (#129718)

Currently, `partial_reduce(acc,mul(sext, zext))` is reduced to `usdot`
in loop only if `zext` has single user i.e. `mul`

If there are two partial reduce equations in loop body such as: 
```
partial_reduce1(acc1,mul1(sext1, zext))
partial_reduce2(acc2,mul2(sext2, zext))
```
and `zext` has no other users other than `mul1`/`mul2`, then this won't
result in `usdot` instructions.

This patch checks if multiple users of `zext`, like above, satisfy the
same set of conditions as for a single user so that `usdot` instructions
are generated.

Author: sushgokh

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -16905,14 +16905,18 @@ bool AArch64TargetLowering::optimizeExtendOrTruncateConversion(
     // 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()))) ||
-          (match(SingleUser,
-                 m_Intrinsic<Intrinsic::experimental_vector_partial_reduce_add>(
-                     m_Value(), m_Specific(I))) &&
-           !shouldExpandPartialReductionIntrinsic(
-               cast<IntrinsicInst>(SingleUser))))
+    if (SrcWidth * 4 <= DstWidth) {
+      if (all_of(I->users(), [&](auto *U) {
+            auto *SingleUser = cast<Instruction>(&*U);
+            return (
+                match(SingleUser, m_c_Mul(m_Specific(I), m_SExt(m_Value()))) ||
+                (match(SingleUser,
+                       m_Intrinsic<
+                           Intrinsic::experimental_vector_partial_reduce_add>(
+                           m_Value(), m_Specific(I))) &&
+                 !shouldExpandPartialReductionIntrinsic(
+                     cast<IntrinsicInst>(SingleUser))));
+          }))
         return false;
     }
 

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

@@ -916,50 +916,57 @@ entry:
 }
 
 define <4 x i32> @usdot_multiple_zext_users(ptr %p1, ptr %p2, ptr %p3) {
-; CHECK-LABEL: usdot_multiple_zext_users:
-; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    adrp x8, .LCPI28_0
-; CHECK-NEXT:    movi v0.2d, #0000000000000000
-; CHECK-NEXT:    movi v2.2d, #0000000000000000
-; CHECK-NEXT:    ldr q1, [x8, :lo12:.LCPI28_0]
-; CHECK-NEXT:    adrp x8, .LCPI28_1
-; CHECK-NEXT:    adrp x9, .LCPI28_2
-; CHECK-NEXT:    adrp x10, .LCPI28_3
-; CHECK-NEXT:    ldr q3, [x8, :lo12:.LCPI28_1]
-; CHECK-NEXT:    ldr q4, [x9, :lo12:.LCPI28_2]
-; CHECK-NEXT:    ldr q5, [x10, :lo12:.LCPI28_3]
-; CHECK-NEXT:    mov x8, xzr
-; CHECK-NEXT:  .LBB28_1: // %vector.body
-; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:    ldr q6, [x2, x8]
-; CHECK-NEXT:    ldr q18, [x0, x8]
-; CHECK-NEXT:    ldr q19, [x1, x8]
-; CHECK-NEXT:    add x8, x8, #16
-; CHECK-NEXT:    tbl v7.16b, { v6.16b }, v1.16b
-; CHECK-NEXT:    tbl v16.16b, { v6.16b }, v3.16b
-; CHECK-NEXT:    tbl v17.16b, { v6.16b }, v4.16b
-; CHECK-NEXT:    tbl v6.16b, { v6.16b }, v5.16b
-; CHECK-NEXT:    cmp x8, #1024
-; CHECK-NEXT:    uzp1 v7.8h, v16.8h, v7.8h
-; CHECK-NEXT:    sshll v16.8h, v18.8b, #0
-; CHECK-NEXT:    uzp1 v6.8h, v6.8h, v17.8h
-; CHECK-NEXT:    sshll2 v17.8h, v18.16b, #0
-; CHECK-NEXT:    sshll v18.8h, v19.8b, #0
-; CHECK-NEXT:    sshll2 v19.8h, v19.16b, #0
-; CHECK-NEXT:    smlal v0.4s, v16.4h, v7.4h
-; CHECK-NEXT:    smlal v2.4s, v18.4h, v7.4h
-; CHECK-NEXT:    smull v20.4s, v17.4h, v6.4h
-; CHECK-NEXT:    smull v21.4s, v19.4h, v6.4h
-; CHECK-NEXT:    smlal2 v0.4s, v17.8h, v6.8h
-; CHECK-NEXT:    smlal2 v2.4s, v19.8h, v6.8h
-; CHECK-NEXT:    smlal2 v20.4s, v16.8h, v7.8h
-; CHECK-NEXT:    smlal2 v21.4s, v18.8h, v7.8h
-; CHECK-NEXT:    add v0.4s, v20.4s, v0.4s
-; CHECK-NEXT:    add v2.4s, v21.4s, v2.4s
-; CHECK-NEXT:    b.ne .LBB28_1
-; CHECK-NEXT:  // %bb.2: // %end
-; CHECK-NEXT:    add v0.4s, v2.4s, v0.4s
-; CHECK-NEXT:    ret
+; CHECK-NOI8MM-LABEL: usdot_multiple_zext_users:
+; CHECK-NOI8MM:       // %bb.0: // %entry
+; CHECK-NOI8MM-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NOI8MM-NEXT:    movi v1.2d, #0000000000000000
+; CHECK-NOI8MM-NEXT:    mov x8, xzr
+; CHECK-NOI8MM-NEXT:  .LBB28_1: // %vector.body
+; CHECK-NOI8MM-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NOI8MM-NEXT:    ldr q2, [x0, x8]
+; CHECK-NOI8MM-NEXT:    ldr q3, [x2, x8]
+; CHECK-NOI8MM-NEXT:    ldr q4, [x1, x8]
+; CHECK-NOI8MM-NEXT:    add x8, x8, #16
+; CHECK-NOI8MM-NEXT:    sshll v5.8h, v2.8b, #0
+; CHECK-NOI8MM-NEXT:    sshll2 v2.8h, v2.16b, #0
+; CHECK-NOI8MM-NEXT:    ushll2 v6.8h, v3.16b, #0
+; CHECK-NOI8MM-NEXT:    ushll v3.8h, v3.8b, #0
+; CHECK-NOI8MM-NEXT:    sshll v7.8h, v4.8b, #0
+; CHECK-NOI8MM-NEXT:    sshll2 v4.8h, v4.16b, #0
+; CHECK-NOI8MM-NEXT:    cmp x8, #1024
+; CHECK-NOI8MM-NEXT:    smull v16.4s, v2.4h, v6.4h
+; CHECK-NOI8MM-NEXT:    smlal v0.4s, v5.4h, v3.4h
+; CHECK-NOI8MM-NEXT:    smull v17.4s, v4.4h, v6.4h
+; CHECK-NOI8MM-NEXT:    smlal v1.4s, v7.4h, v3.4h
+; CHECK-NOI8MM-NEXT:    smlal2 v16.4s, v5.8h, v3.8h
+; CHECK-NOI8MM-NEXT:    smlal2 v0.4s, v2.8h, v6.8h
+; CHECK-NOI8MM-NEXT:    smlal2 v17.4s, v7.8h, v3.8h
+; CHECK-NOI8MM-NEXT:    smlal2 v1.4s, v4.8h, v6.8h
+; CHECK-NOI8MM-NEXT:    add v0.4s, v16.4s, v0.4s
+; CHECK-NOI8MM-NEXT:    add v1.4s, v17.4s, v1.4s
+; CHECK-NOI8MM-NEXT:    b.ne .LBB28_1
+; CHECK-NOI8MM-NEXT:  // %bb.2: // %end
+; CHECK-NOI8MM-NEXT:    add v0.4s, v1.4s, v0.4s
+; CHECK-NOI8MM-NEXT:    ret
+;
+; CHECK-I8MM-LABEL: usdot_multiple_zext_users:
+; CHECK-I8MM:       // %bb.0: // %entry
+; CHECK-I8MM-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-I8MM-NEXT:    movi v1.2d, #0000000000000000
+; CHECK-I8MM-NEXT:    mov x8, xzr
+; CHECK-I8MM-NEXT:  .LBB28_1: // %vector.body
+; CHECK-I8MM-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-I8MM-NEXT:    ldr q2, [x0, x8]
+; CHECK-I8MM-NEXT:    ldr q3, [x1, x8]
+; CHECK-I8MM-NEXT:    ldr q4, [x2, x8]
+; CHECK-I8MM-NEXT:    add x8, x8, #16
+; CHECK-I8MM-NEXT:    usdot v0.4s, v4.16b, v2.16b
+; CHECK-I8MM-NEXT:    usdot v1.4s, v4.16b, v3.16b
+; CHECK-I8MM-NEXT:    cmp x8, #1024
+; CHECK-I8MM-NEXT:    b.ne .LBB28_1
+; CHECK-I8MM-NEXT:  // %bb.2: // %end
+; CHECK-I8MM-NEXT:    add v0.4s, v1.4s, v0.4s
+; CHECK-I8MM-NEXT:    ret
 entry:
   br label %vector.body