[AArch64] Generate usdot instruction with multiple zext users in loop #129718
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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.
sushgokh
requested review from
madhur13490,
davemgreen,
sjoerdmeijer and
david-arm
|
@llvm/pr-subscribers-backend-aarch64 Author: Sushant Gokhale (sushgokh) ChangesCurrently, If there are two partial reduce equations in loop body such as: and This patch checks if multiple users of Full diff: https://github.com/llvm/llvm-project/pull/129718.diff 2 Files Affected:
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index b4bebc61f5dbe..c5620558178c8 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -16898,14 +16898,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
|
|
I had started to look at a similar usdot issue. This PR looks more complete and looks good to me. |
simonwallis2
approved these changes
Mar 4, 2025
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/10/builds/715 Here is the relevant piece of the build log for the reference |
jph-13
pushed a commit
to jph-13/llvm-project
that referenced
this pull request
Mar 21, 2025
…llvm#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.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Currently,
partial_reduce(acc,mul(sext, zext))is reduced tousdotin loop only ifzexthas single user i.e.mulIf there are two partial reduce equations in loop body such as:
and
zexthas no other users other thanmul1/mul2, then this won't result inusdotinstructions.This patch checks if multiple users of
zext, like above, satisfy the same set of conditions as for a single user so thatusdotinstructions are generated.