[llvm] [AArch64] Fold vector shifts guarded against oversized amounts into USHL (PR #207628)
Adam Scott via llvm-commits
llvm-commits at lists.llvm.org
Sun Jul 5 18:01:14 PDT 2026
https://github.com/as4230 created https://github.com/llvm/llvm-project/pull/207628
`select(icmp ult(amt, EltSize), shl(x, amt), 0)`, where EltSize is the {8, 16, 32, 64} lane size, is the usual way to guard a variable vector shift against shl poison. On AArch64 the guard is unnecessary because USHL already returns zero once the shift amount reaches the lane size.
For v4i32:
Before:
```
movi v2.4s, #63
movi v3.4s, #32
and v1.16b, v1.16b, v2.16b
ushl v0.4s, v0.4s, v1.4s
cmhi v1.4s, v3.4s, v1.4s
and v0.16b, v1.16b, v0.16b
```
After:
```
movi v2.4s, #63
and v1.16b, v1.16b, v2.16b
ushl v0.4s, v0.4s, v1.4s
```
USHL reads each lane's shift amount as a signed value from its low byte so amounts above 127 would be misread. The fold applies directly when known bits can prove the amounts are at most 127 and otherwise the amounts are clamped to EltSize with umin first, which is still one instruction cheaper than the select.
NEON has no umin for 64-bit lanes so unbounded v2i64 amounts use the SVE umin when available and otherwise keep the select.
The guard can also zero the input before the shift instead of the result afterward. That form never needs the clamp since the lanes the select zeroes are shifted by EltSize or more and the shl is already poison there.
lshr gets the same fold since it already lowers to ushl with a negated amount.
This is the AArch64 counterpart of #86922, which added the same folds for AVX2's variable shifts.
Fixes #200698.
>From b761c4523f21844c0992dfdfa800cb9f15122d3e Mon Sep 17 00:00:00 2001
From: Adam Scott <adamscott200322 at gmail.com>
Date: Sun, 5 Jul 2026 16:40:22 +0000
Subject: [PATCH 1/2] [AArch64] Add tests for masked vector shift selects. NFC
---
.../CodeGen/AArch64/vselect-masked-shift.ll | 343 ++++++++++++++++++
1 file changed, 343 insertions(+)
create mode 100644 llvm/test/CodeGen/AArch64/vselect-masked-shift.ll
diff --git a/llvm/test/CodeGen/AArch64/vselect-masked-shift.ll b/llvm/test/CodeGen/AArch64/vselect-masked-shift.ll
new file mode 100644
index 0000000000000..1c066a888ef34
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/vselect-masked-shift.ll
@@ -0,0 +1,343 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc -mtriple=aarch64-none-elf < %s | FileCheck %s
+
+define <4 x i32> @masked_shl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: masked_shl_v4i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #63
+; CHECK-NEXT: movi v3.4s, #32
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
+; CHECK-NEXT: cmhi v1.4s, v3.4s, v1.4s
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %m = and <4 x i32> %amt, splat (i32 63)
+ %shl = shl <4 x i32> %x, %m
+ %ok = icmp ult <4 x i32> %m, splat (i32 32)
+ %res = select <4 x i1> %ok, <4 x i32> %shl, <4 x i32> zeroinitializer
+ ret <4 x i32> %res
+}
+
+define <4 x i32> @masked_shl_v4i32_swapped(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: masked_shl_v4i32_swapped:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #63
+; CHECK-NEXT: movi v3.4s, #32
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
+; CHECK-NEXT: cmhi v1.4s, v3.4s, v1.4s
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %m = and <4 x i32> %amt, splat (i32 63)
+ %shl = shl <4 x i32> %x, %m
+ %oob = icmp uge <4 x i32> %m, splat (i32 32)
+ %res = select <4 x i1> %oob, <4 x i32> zeroinitializer, <4 x i32> %shl
+ ret <4 x i32> %res
+}
+
+define <4 x i32> @masked_srl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: masked_srl_v4i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #63
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: movi v2.4s, #32
+; CHECK-NEXT: neg v3.4s, v1.4s
+; CHECK-NEXT: cmhi v1.4s, v2.4s, v1.4s
+; CHECK-NEXT: ushl v0.4s, v0.4s, v3.4s
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %m = and <4 x i32> %amt, splat (i32 63)
+ %srl = lshr <4 x i32> %x, %m
+ %ok = icmp ult <4 x i32> %m, splat (i32 32)
+ %res = select <4 x i1> %ok, <4 x i32> %srl, <4 x i32> zeroinitializer
+ ret <4 x i32> %res
+}
+
+define <16 x i8> @masked_shl_v16i8(<16 x i8> %x, <16 x i8> %amt) {
+; CHECK-LABEL: masked_shl_v16i8:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.16b, #15
+; CHECK-NEXT: movi v3.16b, #8
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: ushl v0.16b, v0.16b, v1.16b
+; CHECK-NEXT: cmhi v1.16b, v3.16b, v1.16b
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %m = and <16 x i8> %amt, splat (i8 15)
+ %shl = shl <16 x i8> %x, %m
+ %ok = icmp ult <16 x i8> %m, splat (i8 8)
+ %res = select <16 x i1> %ok, <16 x i8> %shl, <16 x i8> zeroinitializer
+ ret <16 x i8> %res
+}
+
+define <8 x i16> @masked_shl_v8i16(<8 x i16> %x, <8 x i16> %amt) {
+; CHECK-LABEL: masked_shl_v8i16:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.8h, #31
+; CHECK-NEXT: movi v3.8h, #16
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: ushl v0.8h, v0.8h, v1.8h
+; CHECK-NEXT: cmhi v1.8h, v3.8h, v1.8h
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %m = and <8 x i16> %amt, splat (i16 31)
+ %shl = shl <8 x i16> %x, %m
+ %ok = icmp ult <8 x i16> %m, splat (i16 16)
+ %res = select <8 x i1> %ok, <8 x i16> %shl, <8 x i16> zeroinitializer
+ ret <8 x i16> %res
+}
+
+define <2 x i64> @masked_shl_v2i64(<2 x i64> %x, <2 x i64> %amt) {
+; CHECK-LABEL: masked_shl_v2i64:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: mov w8, #127 // =0x7f
+; CHECK-NEXT: dup v2.2d, x8
+; CHECK-NEXT: mov w8, #64 // =0x40
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: dup v2.2d, x8
+; CHECK-NEXT: ushl v0.2d, v0.2d, v1.2d
+; CHECK-NEXT: cmhi v1.2d, v2.2d, v1.2d
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %m = and <2 x i64> %amt, splat (i64 127)
+ %shl = shl <2 x i64> %x, %m
+ %ok = icmp ult <2 x i64> %m, splat (i64 64)
+ %res = select <2 x i1> %ok, <2 x i64> %shl, <2 x i64> zeroinitializer
+ ret <2 x i64> %res
+}
+
+define <2 x i32> @masked_shl_v2i32(<2 x i32> %x, <2 x i32> %amt) {
+; CHECK-LABEL: masked_shl_v2i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.2s, #63
+; CHECK-NEXT: movi v3.2s, #32
+; CHECK-NEXT: and v1.8b, v1.8b, v2.8b
+; CHECK-NEXT: ushl v0.2s, v0.2s, v1.2s
+; CHECK-NEXT: cmhi v1.2s, v3.2s, v1.2s
+; CHECK-NEXT: and v0.8b, v1.8b, v0.8b
+; CHECK-NEXT: ret
+entry:
+ %m = and <2 x i32> %amt, splat (i32 63)
+ %shl = shl <2 x i32> %x, %m
+ %ok = icmp ult <2 x i32> %m, splat (i32 32)
+ %res = select <2 x i1> %ok, <2 x i32> %shl, <2 x i32> zeroinitializer
+ ret <2 x i32> %res
+}
+
+define <4 x i32> @unbounded_shl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: unbounded_shl_v4i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #32
+; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
+; CHECK-NEXT: cmhi v1.4s, v2.4s, v1.4s
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %shl = shl <4 x i32> %x, %amt
+ %ok = icmp ult <4 x i32> %amt, splat (i32 32)
+ %res = select <4 x i1> %ok, <4 x i32> %shl, <4 x i32> zeroinitializer
+ ret <4 x i32> %res
+}
+
+define <4 x i32> @unbounded_srl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: unbounded_srl_v4i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #32
+; CHECK-NEXT: neg v3.4s, v1.4s
+; CHECK-NEXT: ushl v0.4s, v0.4s, v3.4s
+; CHECK-NEXT: cmhi v1.4s, v2.4s, v1.4s
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %srl = lshr <4 x i32> %x, %amt
+ %ok = icmp ult <4 x i32> %amt, splat (i32 32)
+ %res = select <4 x i1> %ok, <4 x i32> %srl, <4 x i32> zeroinitializer
+ ret <4 x i32> %res
+}
+
+define <4 x i32> @masked_shl_v4i32_mask255(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: masked_shl_v4i32_mask255:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.2d, #0x0000ff000000ff
+; CHECK-NEXT: movi v3.4s, #32
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
+; CHECK-NEXT: cmhi v1.4s, v3.4s, v1.4s
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %m = and <4 x i32> %amt, splat (i32 255)
+ %shl = shl <4 x i32> %x, %m
+ %ok = icmp ult <4 x i32> %m, splat (i32 32)
+ %res = select <4 x i1> %ok, <4 x i32> %shl, <4 x i32> zeroinitializer
+ ret <4 x i32> %res
+}
+
+define <4 x i32> @masked_input_shl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: masked_input_shl_v4i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #63
+; CHECK-NEXT: movi v3.4s, #32
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: cmhi v2.4s, v3.4s, v1.4s
+; CHECK-NEXT: and v0.16b, v2.16b, v0.16b
+; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
+; CHECK-NEXT: ret
+entry:
+ %m = and <4 x i32> %amt, splat (i32 63)
+ %ok = icmp ult <4 x i32> %m, splat (i32 32)
+ %zx = select <4 x i1> %ok, <4 x i32> %x, <4 x i32> zeroinitializer
+ %res = shl <4 x i32> %zx, %m
+ ret <4 x i32> %res
+}
+
+define <4 x i32> @masked_input_shl_v4i32_swapped(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: masked_input_shl_v4i32_swapped:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #63
+; CHECK-NEXT: movi v3.4s, #32
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: cmhi v2.4s, v3.4s, v1.4s
+; CHECK-NEXT: and v0.16b, v2.16b, v0.16b
+; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
+; CHECK-NEXT: ret
+entry:
+ %m = and <4 x i32> %amt, splat (i32 63)
+ %oob = icmp uge <4 x i32> %m, splat (i32 32)
+ %zx = select <4 x i1> %oob, <4 x i32> zeroinitializer, <4 x i32> %x
+ %res = shl <4 x i32> %zx, %m
+ ret <4 x i32> %res
+}
+
+define <4 x i32> @masked_input_srl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: masked_input_srl_v4i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #63
+; CHECK-NEXT: movi v3.4s, #32
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: cmhi v2.4s, v3.4s, v1.4s
+; CHECK-NEXT: neg v1.4s, v1.4s
+; CHECK-NEXT: and v0.16b, v2.16b, v0.16b
+; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
+; CHECK-NEXT: ret
+entry:
+ %m = and <4 x i32> %amt, splat (i32 63)
+ %ok = icmp ult <4 x i32> %m, splat (i32 32)
+ %zx = select <4 x i1> %ok, <4 x i32> %x, <4 x i32> zeroinitializer
+ %res = lshr <4 x i32> %zx, %m
+ ret <4 x i32> %res
+}
+
+define <4 x i32> @unbounded_input_shl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: unbounded_input_shl_v4i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #32
+; CHECK-NEXT: cmhi v2.4s, v2.4s, v1.4s
+; CHECK-NEXT: and v0.16b, v2.16b, v0.16b
+; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
+; CHECK-NEXT: ret
+entry:
+ %ok = icmp ult <4 x i32> %amt, splat (i32 32)
+ %zx = select <4 x i1> %ok, <4 x i32> %x, <4 x i32> zeroinitializer
+ %res = shl <4 x i32> %zx, %amt
+ ret <4 x i32> %res
+}
+
+define <4 x i32> @masked_input_shl_srl_multiuse_v4i32(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: masked_input_shl_srl_multiuse_v4i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #63
+; CHECK-NEXT: movi v3.4s, #32
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: cmhi v2.4s, v3.4s, v1.4s
+; CHECK-NEXT: and v0.16b, v2.16b, v0.16b
+; CHECK-NEXT: neg v2.4s, v1.4s
+; CHECK-NEXT: ushl v1.4s, v0.4s, v1.4s
+; CHECK-NEXT: ushl v0.4s, v0.4s, v2.4s
+; CHECK-NEXT: orr v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %m = and <4 x i32> %amt, splat (i32 63)
+ %ok = icmp ult <4 x i32> %m, splat (i32 32)
+ %zx = select <4 x i1> %ok, <4 x i32> %x, <4 x i32> zeroinitializer
+ %shl = shl <4 x i32> %zx, %m
+ %shr = lshr <4 x i32> %zx, %m
+ %res = or <4 x i32> %shl, %shr
+ ret <4 x i32> %res
+}
+
+define <2 x i64> @unbounded_shl_v2i64_sve(<2 x i64> %x, <2 x i64> %amt) #0 {
+; CHECK-LABEL: unbounded_shl_v2i64_sve:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: mov z2.d, #64 // =0x40
+; CHECK-NEXT: ushl v0.2d, v0.2d, v1.2d
+; CHECK-NEXT: cmhi v1.2d, v2.2d, v1.2d
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %shl = shl <2 x i64> %x, %amt
+ %ok = icmp ult <2 x i64> %amt, splat (i64 64)
+ %res = select <2 x i1> %ok, <2 x i64> %shl, <2 x i64> zeroinitializer
+ ret <2 x i64> %res
+}
+
+define <2 x i64> @neg_unbounded_shl_v2i64(<2 x i64> %x, <2 x i64> %amt) {
+; CHECK-LABEL: neg_unbounded_shl_v2i64:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: mov w8, #64 // =0x40
+; CHECK-NEXT: ushl v0.2d, v0.2d, v1.2d
+; CHECK-NEXT: dup v2.2d, x8
+; CHECK-NEXT: cmhi v1.2d, v2.2d, v1.2d
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %shl = shl <2 x i64> %x, %amt
+ %ok = icmp ult <2 x i64> %amt, splat (i64 64)
+ %res = select <2 x i1> %ok, <2 x i64> %shl, <2 x i64> zeroinitializer
+ ret <2 x i64> %res
+}
+
+define <4 x i32> @neg_masked_shl_v4i32_wrong_bound(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: neg_masked_shl_v4i32_wrong_bound:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #63
+; CHECK-NEXT: movi v3.4s, #16
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
+; CHECK-NEXT: cmhi v1.4s, v3.4s, v1.4s
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %m = and <4 x i32> %amt, splat (i32 63)
+ %shl = shl <4 x i32> %x, %m
+ %ok = icmp ult <4 x i32> %m, splat (i32 16)
+ %res = select <4 x i1> %ok, <4 x i32> %shl, <4 x i32> zeroinitializer
+ ret <4 x i32> %res
+}
+
+define <4 x i32> @neg_masked_ashr_v4i32(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: neg_masked_ashr_v4i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #63
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: movi v2.4s, #32
+; CHECK-NEXT: neg v3.4s, v1.4s
+; CHECK-NEXT: cmhi v1.4s, v2.4s, v1.4s
+; CHECK-NEXT: sshl v0.4s, v0.4s, v3.4s
+; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: ret
+entry:
+ %m = and <4 x i32> %amt, splat (i32 63)
+ %ashr = ashr <4 x i32> %x, %m
+ %ok = icmp ult <4 x i32> %m, splat (i32 32)
+ %res = select <4 x i1> %ok, <4 x i32> %ashr, <4 x i32> zeroinitializer
+ ret <4 x i32> %res
+}
+
+attributes #0 = { "target-features"="+sve" }
>From 22ff828904f9b8fe4f1299e5236e6c4b8f1610c0 Mon Sep 17 00:00:00 2001
From: Adam Scott <adamscott200322 at gmail.com>
Date: Sun, 5 Jul 2026 23:42:49 +0000
Subject: [PATCH 2/2] [AArch64] Fold vector shifts guarded against oversized
amounts into USHL
---
.../Target/AArch64/AArch64ISelLowering.cpp | 111 +++++++++++++++++-
.../CodeGen/AArch64/vselect-masked-shift.ll | 81 +++++--------
2 files changed, 141 insertions(+), 51 deletions(-)
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 9ef66ecad411c..6805c21c20cf1 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -1246,7 +1246,7 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
setTargetDAGCombine(ISD::SCALAR_TO_VECTOR);
- setTargetDAGCombine(ISD::SHL);
+ setTargetDAGCombine({ISD::SHL, ISD::SRL});
setTargetDAGCombine(ISD::VECTOR_DEINTERLEAVE);
setTargetDAGCombine(ISD::CTPOP);
@@ -28932,6 +28932,108 @@ static SDValue performVselectPowCombine(SDNode *N,
return DAG.getNode(ISD::VSELECT, DL, VT, Cond, TrueVal, NewPow);
}
+// A vselect can zero the lanes with an out-of-range shift amount with either:
+// vselect(setcc_ult(amt, EltSize), val, zeros)
+// vselect(setcc_uge(amt, EltSize), zeros, val)
+// Returns val for whichever polarity is present and sets RequiredCC to the
+// condition code the setcc has to be using.
+static SDValue matchZeroSelectArm(SDValue TVal, SDValue FVal,
+ ISD::CondCode &RequiredCC) {
+ if (ISD::isConstantSplatVectorAllZeros(FVal.getNode())) {
+ RequiredCC = ISD::SETULT;
+ return TVal;
+ }
+ if (ISD::isConstantSplatVectorAllZeros(TVal.getNode())) {
+ RequiredCC = ISD::SETUGE;
+ return FVal;
+ }
+ return SDValue();
+}
+
+// ushl already produces zero for shift amounts of EltSize or more, so a
+// select zeroing those lanes is redundant. However, ushl reads each lane's
+// shift amount as a signed value from its low byte and would misread amounts
+// above 127, so the amounts must either be provably at most 127 or get
+// clamped to EltSize with umin.
+static SDValue foldMaskedShiftToUSHL(SelectionDAG &DAG,
+ const AArch64Subtarget *Subtarget,
+ SDNode *N, SDValue X, SDValue Amt,
+ SDValue Cond, ISD::CondCode RequiredCC,
+ bool IsSRL, bool AmtOutOfRangeIsPoison) {
+ using namespace llvm::SDPatternMatch;
+ EVT VT = N->getValueType(0);
+ if (!Subtarget->isNeonAvailable() || !VT.isFixedLengthVector() ||
+ !VT.isInteger() || !DAG.getTargetLoweringInfo().isTypeLegal(VT))
+ return SDValue();
+
+ unsigned EltSize = VT.getScalarSizeInBits();
+ if (!sd_match(Cond, m_SetCC(m_Specific(Amt), m_SpecificInt(EltSize),
+ m_SpecificCondCode(RequiredCC))))
+ return SDValue();
+
+ SDLoc DL(N);
+
+ // Amounts that might exceed 127 need the umin clamp, unless the
+ // out-of-range lanes are poison and any result is acceptable for them.
+ if (!AmtOutOfRangeIsPoison &&
+ !DAG.computeKnownBits(Amt).getMaxValue().ule(127)) {
+ // Only SVE has a umin for 64-bit lanes.
+ if (EltSize == 64 && !Subtarget->isSVEAvailable())
+ return SDValue();
+ Amt = DAG.getNode(ISD::UMIN, DL, VT, Amt, DAG.getConstant(EltSize, DL, VT));
+ }
+
+ // There is no shift right register instruction but ushl shifts right when
+ // the amount is negative.
+ if (IsSRL)
+ Amt = DAG.getNode(ISD::SUB, DL, VT, DAG.getConstant(0, DL, VT), Amt);
+ return DAG.getNode(
+ ISD::INTRINSIC_WO_CHAIN, DL, VT,
+ DAG.getTargetConstant(Intrinsic::aarch64_neon_ushl, DL, MVT::i32), X,
+ Amt);
+}
+
+// vselect(setcc_ult(amt, EltSize), shl(x, amt), zeros) -> ushl(x, amt)
+// vselect(setcc_ult(amt, EltSize), srl(x, amt), zeros) -> ushl(x, -amt)
+static SDValue
+performVSelectMaskedShiftCombine(SDNode *N, SelectionDAG &DAG,
+ const AArch64Subtarget *Subtarget) {
+ ISD::CondCode RequiredCC;
+ SDValue Shift =
+ matchZeroSelectArm(N->getOperand(1), N->getOperand(2), RequiredCC);
+ if (!Shift ||
+ (Shift.getOpcode() != ISD::SHL && Shift.getOpcode() != ISD::SRL))
+ return SDValue();
+
+ return foldMaskedShiftToUSHL(DAG, Subtarget, N, Shift.getOperand(0),
+ Shift.getOperand(1), N->getOperand(0),
+ RequiredCC, Shift.getOpcode() == ISD::SRL,
+ /*AmtOutOfRangeIsPoison=*/false);
+}
+
+// shl/srl(vselect(setcc_ult(amt, EltSize), x, zeros), amt) -> ushl(x, amt)
+//
+// The lanes where the vselect chooses zero are shifted by EltSize or more,
+// which is poison, so they need no clamping.
+static SDValue
+performShiftOfZeroSelectCombine(SDNode *N, SelectionDAG &DAG,
+ const AArch64Subtarget *Subtarget) {
+ SDValue Sel = N->getOperand(0);
+ if (Sel.getOpcode() != ISD::VSELECT)
+ return SDValue();
+
+ ISD::CondCode RequiredCC;
+ SDValue X =
+ matchZeroSelectArm(Sel.getOperand(1), Sel.getOperand(2), RequiredCC);
+ if (!X)
+ return SDValue();
+
+ return foldMaskedShiftToUSHL(DAG, Subtarget, N, X, N->getOperand(1),
+ Sel.getOperand(0), RequiredCC,
+ N->getOpcode() == ISD::SRL,
+ /*AmtOutOfRangeIsPoison=*/true);
+}
+
// vselect (v1i1 setcc) ->
// vselect (v1iXX setcc) (XX is the size of the compared operand type)
// FIXME: Currently the type legalizer can't handle VSELECT having v1i1 as
@@ -28945,6 +29047,9 @@ static SDValue performVSelectCombine(SDNode *N,
if (auto SwapResult = trySwapVSelectOperands(N, DAG))
return SwapResult;
+ if (SDValue Shift = performVSelectMaskedShiftCombine(N, DAG, Subtarget))
+ return Shift;
+
SDValue N0 = N->getOperand(0);
SDValue IfTrue = N->getOperand(1);
SDValue IfFalse = N->getOperand(2);
@@ -30734,7 +30839,11 @@ SDValue AArch64TargetLowering::PerformDAGCombine(SDNode *N,
case ISD::SCALAR_TO_VECTOR:
return performScalarToVectorCombine(N, DCI, DAG);
case ISD::SHL:
+ if (SDValue R = performShiftOfZeroSelectCombine(N, DAG, Subtarget))
+ return R;
return performSHLCombine(N, DCI, DAG);
+ case ISD::SRL:
+ return performShiftOfZeroSelectCombine(N, DAG, Subtarget);
case ISD::CTPOP:
return performCTPOPCombine(N, DCI, DAG);
case ISD::BITCAST:
diff --git a/llvm/test/CodeGen/AArch64/vselect-masked-shift.ll b/llvm/test/CodeGen/AArch64/vselect-masked-shift.ll
index 1c066a888ef34..7e68d48586842 100644
--- a/llvm/test/CodeGen/AArch64/vselect-masked-shift.ll
+++ b/llvm/test/CodeGen/AArch64/vselect-masked-shift.ll
@@ -5,11 +5,8 @@ define <4 x i32> @masked_shl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
; CHECK-LABEL: masked_shl_v4i32:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.4s, #63
-; CHECK-NEXT: movi v3.4s, #32
; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
-; CHECK-NEXT: cmhi v1.4s, v3.4s, v1.4s
-; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
; CHECK-NEXT: ret
entry:
%m = and <4 x i32> %amt, splat (i32 63)
@@ -23,11 +20,8 @@ define <4 x i32> @masked_shl_v4i32_swapped(<4 x i32> %x, <4 x i32> %amt) {
; CHECK-LABEL: masked_shl_v4i32_swapped:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.4s, #63
-; CHECK-NEXT: movi v3.4s, #32
; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
-; CHECK-NEXT: cmhi v1.4s, v3.4s, v1.4s
-; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
; CHECK-NEXT: ret
entry:
%m = and <4 x i32> %amt, splat (i32 63)
@@ -42,11 +36,8 @@ define <4 x i32> @masked_srl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.4s, #63
; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
-; CHECK-NEXT: movi v2.4s, #32
-; CHECK-NEXT: neg v3.4s, v1.4s
-; CHECK-NEXT: cmhi v1.4s, v2.4s, v1.4s
-; CHECK-NEXT: ushl v0.4s, v0.4s, v3.4s
-; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: neg v1.4s, v1.4s
+; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
; CHECK-NEXT: ret
entry:
%m = and <4 x i32> %amt, splat (i32 63)
@@ -60,11 +51,8 @@ define <16 x i8> @masked_shl_v16i8(<16 x i8> %x, <16 x i8> %amt) {
; CHECK-LABEL: masked_shl_v16i8:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.16b, #15
-; CHECK-NEXT: movi v3.16b, #8
; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
; CHECK-NEXT: ushl v0.16b, v0.16b, v1.16b
-; CHECK-NEXT: cmhi v1.16b, v3.16b, v1.16b
-; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
; CHECK-NEXT: ret
entry:
%m = and <16 x i8> %amt, splat (i8 15)
@@ -78,11 +66,8 @@ define <8 x i16> @masked_shl_v8i16(<8 x i16> %x, <8 x i16> %amt) {
; CHECK-LABEL: masked_shl_v8i16:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.8h, #31
-; CHECK-NEXT: movi v3.8h, #16
; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
; CHECK-NEXT: ushl v0.8h, v0.8h, v1.8h
-; CHECK-NEXT: cmhi v1.8h, v3.8h, v1.8h
-; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
; CHECK-NEXT: ret
entry:
%m = and <8 x i16> %amt, splat (i16 31)
@@ -97,12 +82,8 @@ define <2 x i64> @masked_shl_v2i64(<2 x i64> %x, <2 x i64> %amt) {
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: mov w8, #127 // =0x7f
; CHECK-NEXT: dup v2.2d, x8
-; CHECK-NEXT: mov w8, #64 // =0x40
; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
-; CHECK-NEXT: dup v2.2d, x8
; CHECK-NEXT: ushl v0.2d, v0.2d, v1.2d
-; CHECK-NEXT: cmhi v1.2d, v2.2d, v1.2d
-; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
; CHECK-NEXT: ret
entry:
%m = and <2 x i64> %amt, splat (i64 127)
@@ -116,11 +97,8 @@ define <2 x i32> @masked_shl_v2i32(<2 x i32> %x, <2 x i32> %amt) {
; CHECK-LABEL: masked_shl_v2i32:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.2s, #63
-; CHECK-NEXT: movi v3.2s, #32
; CHECK-NEXT: and v1.8b, v1.8b, v2.8b
; CHECK-NEXT: ushl v0.2s, v0.2s, v1.2s
-; CHECK-NEXT: cmhi v1.2s, v3.2s, v1.2s
-; CHECK-NEXT: and v0.8b, v1.8b, v0.8b
; CHECK-NEXT: ret
entry:
%m = and <2 x i32> %amt, splat (i32 63)
@@ -134,9 +112,8 @@ define <4 x i32> @unbounded_shl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
; CHECK-LABEL: unbounded_shl_v4i32:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.4s, #32
+; CHECK-NEXT: umin v1.4s, v1.4s, v2.4s
; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
-; CHECK-NEXT: cmhi v1.4s, v2.4s, v1.4s
-; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
; CHECK-NEXT: ret
entry:
%shl = shl <4 x i32> %x, %amt
@@ -149,10 +126,9 @@ define <4 x i32> @unbounded_srl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
; CHECK-LABEL: unbounded_srl_v4i32:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.4s, #32
-; CHECK-NEXT: neg v3.4s, v1.4s
-; CHECK-NEXT: ushl v0.4s, v0.4s, v3.4s
-; CHECK-NEXT: cmhi v1.4s, v2.4s, v1.4s
-; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
+; CHECK-NEXT: umin v1.4s, v1.4s, v2.4s
+; CHECK-NEXT: neg v1.4s, v1.4s
+; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
; CHECK-NEXT: ret
entry:
%srl = lshr <4 x i32> %x, %amt
@@ -167,9 +143,8 @@ define <4 x i32> @masked_shl_v4i32_mask255(<4 x i32> %x, <4 x i32> %amt) {
; CHECK-NEXT: movi v2.2d, #0x0000ff000000ff
; CHECK-NEXT: movi v3.4s, #32
; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: umin v1.4s, v1.4s, v3.4s
; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
-; CHECK-NEXT: cmhi v1.4s, v3.4s, v1.4s
-; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
; CHECK-NEXT: ret
entry:
%m = and <4 x i32> %amt, splat (i32 255)
@@ -183,10 +158,7 @@ define <4 x i32> @masked_input_shl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
; CHECK-LABEL: masked_input_shl_v4i32:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.4s, #63
-; CHECK-NEXT: movi v3.4s, #32
; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
-; CHECK-NEXT: cmhi v2.4s, v3.4s, v1.4s
-; CHECK-NEXT: and v0.16b, v2.16b, v0.16b
; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
; CHECK-NEXT: ret
entry:
@@ -201,10 +173,7 @@ define <4 x i32> @masked_input_shl_v4i32_swapped(<4 x i32> %x, <4 x i32> %amt) {
; CHECK-LABEL: masked_input_shl_v4i32_swapped:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.4s, #63
-; CHECK-NEXT: movi v3.4s, #32
; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
-; CHECK-NEXT: cmhi v2.4s, v3.4s, v1.4s
-; CHECK-NEXT: and v0.16b, v2.16b, v0.16b
; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
; CHECK-NEXT: ret
entry:
@@ -219,11 +188,8 @@ define <4 x i32> @masked_input_srl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
; CHECK-LABEL: masked_input_srl_v4i32:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.4s, #63
-; CHECK-NEXT: movi v3.4s, #32
; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
-; CHECK-NEXT: cmhi v2.4s, v3.4s, v1.4s
; CHECK-NEXT: neg v1.4s, v1.4s
-; CHECK-NEXT: and v0.16b, v2.16b, v0.16b
; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
; CHECK-NEXT: ret
entry:
@@ -237,9 +203,6 @@ entry:
define <4 x i32> @unbounded_input_shl_v4i32(<4 x i32> %x, <4 x i32> %amt) {
; CHECK-LABEL: unbounded_input_shl_v4i32:
; CHECK: // %bb.0: // %entry
-; CHECK-NEXT: movi v2.4s, #32
-; CHECK-NEXT: cmhi v2.4s, v2.4s, v1.4s
-; CHECK-NEXT: and v0.16b, v2.16b, v0.16b
; CHECK-NEXT: ushl v0.4s, v0.4s, v1.4s
; CHECK-NEXT: ret
entry:
@@ -253,10 +216,7 @@ define <4 x i32> @masked_input_shl_srl_multiuse_v4i32(<4 x i32> %x, <4 x i32> %a
; CHECK-LABEL: masked_input_shl_srl_multiuse_v4i32:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.4s, #63
-; CHECK-NEXT: movi v3.4s, #32
; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
-; CHECK-NEXT: cmhi v2.4s, v3.4s, v1.4s
-; CHECK-NEXT: and v0.16b, v2.16b, v0.16b
; CHECK-NEXT: neg v2.4s, v1.4s
; CHECK-NEXT: ushl v1.4s, v0.4s, v1.4s
; CHECK-NEXT: ushl v0.4s, v0.4s, v2.4s
@@ -275,10 +235,9 @@ entry:
define <2 x i64> @unbounded_shl_v2i64_sve(<2 x i64> %x, <2 x i64> %amt) #0 {
; CHECK-LABEL: unbounded_shl_v2i64_sve:
; CHECK: // %bb.0: // %entry
-; CHECK-NEXT: mov z2.d, #64 // =0x40
+; CHECK-NEXT: // kill: def $q1 killed $q1 def $z1
+; CHECK-NEXT: umin z1.d, z1.d, #64
; CHECK-NEXT: ushl v0.2d, v0.2d, v1.2d
-; CHECK-NEXT: cmhi v1.2d, v2.2d, v1.2d
-; CHECK-NEXT: and v0.16b, v1.16b, v0.16b
; CHECK-NEXT: ret
entry:
%shl = shl <2 x i64> %x, %amt
@@ -287,6 +246,28 @@ entry:
ret <2 x i64> %res
}
+define <1 x i128> @neg_masked_shl_v1i128(<1 x i128> %x, <1 x i128> %amt) {
+; CHECK-LABEL: neg_masked_shl_v1i128:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: and x8, x2, #0x3f
+; CHECK-NEXT: lsr x9, x0, #1
+; CHECK-NEXT: lsl x11, x1, x2
+; CHECK-NEXT: eor x10, x8, #0x3f
+; CHECK-NEXT: cmp x8, #128
+; CHECK-NEXT: lsr x9, x9, x10
+; CHECK-NEXT: lsl x10, x0, x2
+; CHECK-NEXT: orr x9, x11, x9
+; CHECK-NEXT: csel x0, x10, xzr, lo
+; CHECK-NEXT: csel x1, x9, xzr, lo
+; CHECK-NEXT: ret
+entry:
+ %m = and <1 x i128> %amt, splat (i128 63)
+ %shl = shl <1 x i128> %x, %m
+ %ok = icmp ult <1 x i128> %m, splat (i128 128)
+ %res = select <1 x i1> %ok, <1 x i128> %shl, <1 x i128> zeroinitializer
+ ret <1 x i128> %res
+}
+
define <2 x i64> @neg_unbounded_shl_v2i64(<2 x i64> %x, <2 x i64> %amt) {
; CHECK-LABEL: neg_unbounded_shl_v2i64:
; CHECK: // %bb.0: // %entry
More information about the llvm-commits
mailing list