[llvm] 58c8696 - [AArch64] Generate usdot instruction with multiple zext users in loop (#129718)
via llvm-commits
llvm-commits at lists.llvm.org
Tue Mar 4 20:29:17 PST 2025
Author: Sushant Gokhale
Date: 2025-03-04T20:29:14-08:00
New Revision: 58c869682ac260cbd4e1818b08c4557fd6025aed
URL: https://github.com/llvm/llvm-project/commit/58c869682ac260cbd4e1818b08c4557fd6025aed
DIFF: https://github.com/llvm/llvm-project/commit/58c869682ac260cbd4e1818b08c4557fd6025aed.diff
LOG: [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.
Added:
Modified:
llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
llvm/test/CodeGen/AArch64/neon-partial-reduce-dot-product.ll
Removed:
################################################################################
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 2dca8c0da4756..52ec4753ec4c1 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/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;
}
diff --git a/llvm/test/CodeGen/AArch64/neon-partial-reduce-dot-product.ll b/llvm/test/CodeGen/AArch64/neon-partial-reduce-dot-product.ll
index 1c9849bdaed3c..c48ebbad4fe21 100644
--- a/llvm/test/CodeGen/AArch64/neon-partial-reduce-dot-product.ll
+++ b/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
More information about the llvm-commits
mailing list