Skip to content

[AArch64] match TRN starting from undef elements #167955

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Nov 19, 2025
Merged

[AArch64] match TRN starting from undef elements #167955

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e30eaeb
[AArch64] match TRN starting from undef elements (#167955)
ginsbach Nov 9, 2025
2d0763d
use 'poison' rather than 'undef' in new tests
ginsbach Nov 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions llvm/lib/Target/AArch64/AArch64PerfectShuffle.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6685,15 +6685,30 @@ inline bool isUZPMask(ArrayRef<int> M, unsigned NumElts,
/// <0, 8, 2, 10, 4, 12, 6, 14> or
/// <1, 9, 3, 11, 5, 13, 7, 15>
inline bool isTRNMask(ArrayRef<int> M, unsigned NumElts,
unsigned &WhichResult) {
unsigned &WhichResultOut) {
if (NumElts % 2 != 0)
return false;
WhichResult = (M[0] == 0 ? 0 : 1);
// Check the first non-undef element for trn1 vs trn2.
unsigned WhichResult = 2;
for (unsigned i = 0; i != NumElts; i += 2) {
if (M[i] >= 0) {
WhichResult = ((unsigned)M[i] == i ? 0 : 1);
break;
}
if (M[i + 1] >= 0) {
WhichResult = ((unsigned)M[i + 1] == i + NumElts ? 0 : 1);
break;
}
}
if (WhichResult == 2)
return false;

for (unsigned i = 0; i < NumElts; i += 2) {
if ((M[i] >= 0 && (unsigned)M[i] != i + WhichResult) ||
(M[i + 1] >= 0 && (unsigned)M[i + 1] != i + NumElts + WhichResult))
return false;
}
WhichResultOut = WhichResult;
return true;
}

Expand Down
114 changes: 113 additions & 1 deletion llvm/test/CodeGen/AArch64/arm64-trn.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -246,7 +246,7 @@ define <4 x float> @vtrnQf(ptr %A, ptr %B) nounwind {
ret <4 x float> %tmp5
}

; Undef shuffle indices should not prevent matching to VTRN:
; Undef shuffle indices (even at the start of the shuffle mask) should not prevent matching to VTRN:

define <8 x i8> @vtrni8_undef(ptr %A, ptr %B) nounwind {
; CHECKLE-LABEL: vtrni8_undef:
Expand Down Expand Up @@ -302,3 +302,115 @@ define <8 x i16> @vtrnQi16_undef(ptr %A, ptr %B) nounwind {
%tmp5 = add <8 x i16> %tmp3, %tmp4
ret <8 x i16> %tmp5
}

define <8 x i16> @vtrnQi16_undef_01(ptr %A, ptr %B) nounwind {
; CHECKLE-LABEL: vtrnQi16_undef_01:
; CHECKLE: // %bb.0:
; CHECKLE-NEXT: ldr q0, [x0]
; CHECKLE-NEXT: ldr q1, [x1]
; CHECKLE-NEXT: trn1 v2.8h, v0.8h, v1.8h
; CHECKLE-NEXT: trn2 v0.8h, v0.8h, v1.8h
; CHECKLE-NEXT: add v0.8h, v2.8h, v0.8h
; CHECKLE-NEXT: ret
;
; CHECKBE-LABEL: vtrnQi16_undef_01:
; CHECKBE: // %bb.0:
; CHECKBE-NEXT: ld1 { v0.8h }, [x0]
; CHECKBE-NEXT: ld1 { v1.8h }, [x1]
; CHECKBE-NEXT: trn1 v2.8h, v0.8h, v1.8h
; CHECKBE-NEXT: trn2 v0.8h, v0.8h, v1.8h
; CHECKBE-NEXT: add v0.8h, v2.8h, v0.8h
; CHECKBE-NEXT: rev64 v0.8h, v0.8h
; CHECKBE-NEXT: ext v0.16b, v0.16b, v0.16b, #8
; CHECKBE-NEXT: ret
%tmp1 = load <8 x i16>, ptr %A
%tmp2 = load <8 x i16>, ptr %B
%tmp3 = shufflevector <8 x i16> %tmp1, <8 x i16> %tmp2, <8 x i32> <i32 poison, i32 poison, i32 2, i32 poison, i32 4, i32 12, i32 6, i32 14>
%tmp4 = shufflevector <8 x i16> %tmp1, <8 x i16> %tmp2, <8 x i32> <i32 poison, i32 poison, i32 3, i32 11, i32 poison, i32 13, i32 poison, i32 poison>
%tmp5 = add <8 x i16> %tmp3, %tmp4
ret <8 x i16> %tmp5
}

define <8 x i16> @vtrnQi16_undef_0(ptr %A, ptr %B) nounwind {
; CHECKLE-LABEL: vtrnQi16_undef_0:
; CHECKLE: // %bb.0:
; CHECKLE-NEXT: ldr q0, [x0]
; CHECKLE-NEXT: ldr q1, [x1]
; CHECKLE-NEXT: trn1 v2.8h, v0.8h, v1.8h
; CHECKLE-NEXT: trn2 v0.8h, v0.8h, v1.8h
; CHECKLE-NEXT: add v0.8h, v2.8h, v0.8h
; CHECKLE-NEXT: ret
;
; CHECKBE-LABEL: vtrnQi16_undef_0:
; CHECKBE: // %bb.0:
; CHECKBE-NEXT: ld1 { v0.8h }, [x0]
; CHECKBE-NEXT: ld1 { v1.8h }, [x1]
; CHECKBE-NEXT: trn1 v2.8h, v0.8h, v1.8h
; CHECKBE-NEXT: trn2 v0.8h, v0.8h, v1.8h
; CHECKBE-NEXT: add v0.8h, v2.8h, v0.8h
; CHECKBE-NEXT: rev64 v0.8h, v0.8h
; CHECKBE-NEXT: ext v0.16b, v0.16b, v0.16b, #8
; CHECKBE-NEXT: ret
%tmp1 = load <8 x i16>, ptr %A
%tmp2 = load <8 x i16>, ptr %B
%tmp3 = shufflevector <8 x i16> %tmp1, <8 x i16> %tmp2, <8 x i32> <i32 poison, i32 8, i32 poison, i32 poison, i32 4, i32 12, i32 6, i32 14>
%tmp4 = shufflevector <8 x i16> %tmp1, <8 x i16> %tmp2, <8 x i32> <i32 poison, i32 9, i32 3, i32 11, i32 5, i32 13, i32 poison, i32 poison>
%tmp5 = add <8 x i16> %tmp3, %tmp4
ret <8 x i16> %tmp5
}

define <4 x i32> @vtrnQi32_undef_1(ptr %A, ptr %B) nounwind {
; CHECKLE-LABEL: vtrnQi32_undef_1:
; CHECKLE: // %bb.0:
; CHECKLE-NEXT: ldr q0, [x0]
; CHECKLE-NEXT: ldr q1, [x1]
; CHECKLE-NEXT: trn1 v2.4s, v0.4s, v1.4s
; CHECKLE-NEXT: trn2 v0.4s, v0.4s, v1.4s
; CHECKLE-NEXT: add v0.4s, v2.4s, v0.4s
; CHECKLE-NEXT: ret
;
; CHECKBE-LABEL: vtrnQi32_undef_1:
; CHECKBE: // %bb.0:
; CHECKBE-NEXT: ld1 { v0.4s }, [x0]
; CHECKBE-NEXT: ld1 { v1.4s }, [x1]
; CHECKBE-NEXT: trn1 v2.4s, v0.4s, v1.4s
; CHECKBE-NEXT: trn2 v0.4s, v0.4s, v1.4s
; CHECKBE-NEXT: add v0.4s, v2.4s, v0.4s
; CHECKBE-NEXT: rev64 v0.4s, v0.4s
; CHECKBE-NEXT: ext v0.16b, v0.16b, v0.16b, #8
; CHECKBE-NEXT: ret
%tmp1 = load <4 x i32>, ptr %A
%tmp2 = load <4 x i32>, ptr %B
%tmp3 = shufflevector <4 x i32> %tmp1, <4 x i32> %tmp2, <4 x i32> <i32 0, i32 poison, i32 2, i32 6>
%tmp4 = shufflevector <4 x i32> %tmp1, <4 x i32> %tmp2, <4 x i32> <i32 1, i32 poison, i32 3, i32 7>
%tmp5 = add <4 x i32> %tmp3, %tmp4
ret <4 x i32> %tmp5
}

define <16 x i8> @vtrnQi8_undef_012(ptr %A, ptr %B) nounwind {
; CHECKLE-LABEL: vtrnQi8_undef_012:
; CHECKLE: // %bb.0:
; CHECKLE-NEXT: ldr q0, [x0]
; CHECKLE-NEXT: ldr q1, [x1]
; CHECKLE-NEXT: trn1 v2.16b, v0.16b, v1.16b
; CHECKLE-NEXT: trn2 v0.16b, v0.16b, v1.16b
; CHECKLE-NEXT: add v0.16b, v2.16b, v0.16b
; CHECKLE-NEXT: ret
;
; CHECKBE-LABEL: vtrnQi8_undef_012:
; CHECKBE: // %bb.0:
; CHECKBE-NEXT: ld1 { v0.16b }, [x0]
; CHECKBE-NEXT: ld1 { v1.16b }, [x1]
; CHECKBE-NEXT: trn1 v2.16b, v0.16b, v1.16b
; CHECKBE-NEXT: trn2 v0.16b, v0.16b, v1.16b
; CHECKBE-NEXT: add v0.16b, v2.16b, v0.16b
; CHECKBE-NEXT: rev64 v0.16b, v0.16b
; CHECKBE-NEXT: ext v0.16b, v0.16b, v0.16b, #8
; CHECKBE-NEXT: ret
%tmp1 = load <16 x i8>, ptr %A
%tmp2 = load <16 x i8>, ptr %B
%tmp3 = shufflevector <16 x i8> %tmp1, <16 x i8> %tmp2, <16 x i32> <i32 poison, i32 poison, i32 poison, i32 18, i32 4, i32 poison, i32 6, i32 22, i32 poison, i32 24, i32 10, i32 26, i32 12, i32 28, i32 14, i32 30>
%tmp4 = shufflevector <16 x i8> %tmp1, <16 x i8> %tmp2, <16 x i32> <i32 poison, i32 poison, i32 poison, i32 19, i32 5, i32 21, i32 7, i32 poison, i32 9, i32 25, i32 11, i32 27, i32 poison, i32 29, i32 15, i32 31>
%tmp5 = add <16 x i8> %tmp3, %tmp4
ret <16 x i8> %tmp5
}
20 changes: 10 additions & 10 deletions llvm/test/CodeGen/AArch64/insert-extend.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,24 +85,24 @@ define i32 @large(ptr nocapture noundef readonly %p1, i32 noundef %st1, ptr noca
; CHECK-NEXT: addp v2.4s, v3.4s, v2.4s
Copy link
Member

@MacDue MacDue Nov 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it may be worth adding some smaller test cases for this PR (possibly in a new file), as the current test updates are fairly large and do not directly test this change.

Copy link
Contributor Author

@ginsbach ginsbach Nov 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, great suggestion. In the second commit, I have added additional tests to arm64-trn.ll, loosely based on equivalent tests in arm64-zip.ll.

; CHECK-NEXT: zip1 v16.4s, v5.4s, v4.4s
; CHECK-NEXT: sub v7.4s, v3.4s, v7.4s
; CHECK-NEXT: trn1 v4.4s, v5.4s, v4.4s
; CHECK-NEXT: zip2 v3.4s, v6.4s, v7.4s
; CHECK-NEXT: mov v6.s[1], v7.s[0]
; CHECK-NEXT: ext v7.16b, v5.16b, v16.16b, #8
; CHECK-NEXT: mov v5.s[3], v4.s[2]
; CHECK-NEXT: ext v4.16b, v2.16b, v2.16b, #8
; CHECK-NEXT: mov v6.d[1], v7.d[1]
; CHECK-NEXT: mov v3.d[1], v5.d[1]
; CHECK-NEXT: uzp1 v1.4s, v4.4s, v0.4s
; CHECK-NEXT: uzp2 v4.4s, v4.4s, v0.4s
; CHECK-NEXT: ext v7.16b, v2.16b, v2.16b, #8
; CHECK-NEXT: ext v5.16b, v5.16b, v16.16b, #8
; CHECK-NEXT: mov v3.d[1], v4.d[1]
; CHECK-NEXT: uzp1 v1.4s, v7.4s, v0.4s
; CHECK-NEXT: uzp2 v4.4s, v7.4s, v0.4s
; CHECK-NEXT: mov v6.d[1], v5.d[1]
; CHECK-NEXT: addp v0.4s, v2.4s, v0.4s
; CHECK-NEXT: sub v1.4s, v1.4s, v4.4s
; CHECK-NEXT: rev64 v7.4s, v0.4s
; CHECK-NEXT: add v5.4s, v3.4s, v6.4s
; CHECK-NEXT: sub v3.4s, v6.4s, v3.4s
; CHECK-NEXT: rev64 v7.4s, v0.4s
; CHECK-NEXT: sub v1.4s, v1.4s, v4.4s
; CHECK-NEXT: rev64 v2.4s, v1.4s
; CHECK-NEXT: rev64 v4.4s, v5.4s
; CHECK-NEXT: rev64 v6.4s, v3.4s
; CHECK-NEXT: addp v16.4s, v0.4s, v5.4s
; CHECK-NEXT: rev64 v2.4s, v1.4s
; CHECK-NEXT: sub v0.4s, v0.4s, v7.4s
; CHECK-NEXT: zip1 v21.4s, v16.4s, v16.4s
; CHECK-NEXT: sub v4.4s, v5.4s, v4.4s
Expand Down
Loading