[llvm] [AArch64] Fold shifts of guarded vector inputs into USHL and SSHL (PR #210880)
via llvm-commits
llvm-commits at lists.llvm.org
Mon Jul 20 23:50:48 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-aarch64
Author: Adam Scott (as4230)
<details>
<summary>Changes</summary>
Follow-up to #<!-- -->207628, which folded `select(icmp ult(amt, EltSize), shl(x, amt), 0)` into USHL. This handles the other way to write the same guard, where the select zeroes the shift input instead of the result:
shl(select(icmp ult(amt, EltSize), x, 0), amt) -> ushl(x, amt)
srl(select(icmp ult(amt, EltSize), x, 0), amt) -> ushl(x, -amt)
sra(select(icmp ult(amt, EltSize), x, 0), amt) -> sshl(x, -amt)
For v4i32:
Before:
```
movi v2.4s, #<!-- -->63
movi v3.4s, #<!-- -->32
and v1.16b, v1.16b, v2.16b
cmhi v2.4s, v3.4s, v1.4s
and v0.16b, v2.16b, v0.16b
ushl v0.4s, v0.4s, v1.4s
```
After:
```
movi v2.4s, #<!-- -->63
and v1.16b, v1.16b, v2.16b
ushl v0.4s, v0.4s, v1.4s
```
This form needs no umin clamp. The lanes where the select chooses zero are shifted by EltSize or more which is poison. That is also why sra folds here when it could not in #<!-- -->207628, since sshl's sign fill only differs from the select's zero in those poison lanes.
This is the AArch64 counterpart of the shift(select) half of #<!-- -->86922
---
Full diff: https://github.com/llvm/llvm-project/pull/210880.diff
2 Files Affected:
- (modified) llvm/lib/Target/AArch64/AArch64ISelLowering.cpp (+47-12)
- (modified) llvm/test/CodeGen/AArch64/vselect-masked-shift.ll (+135)
``````````diff
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 9872ec9a4a1fe..f7e4a129f31f4 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -1247,7 +1247,7 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
setTargetDAGCombine(ISD::SCALAR_TO_VECTOR);
- setTargetDAGCombine(ISD::SHL);
+ setTargetDAGCombine({ISD::SHL, ISD::SRL, ISD::SRA});
setTargetDAGCombine(ISD::VECTOR_DEINTERLEAVE);
setTargetDAGCombine(ISD::CTPOP);
@@ -29419,12 +29419,19 @@ static SDValue matchZeroSelectArm(SDValue TVal, SDValue FVal,
// 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.
+// clamped to EltSize with umin. sra uses sshl and is only valid when the
+// out-of-range lanes are poison.
static SDValue foldMaskedShiftToUSHL(SelectionDAG &DAG,
const AArch64Subtarget *Subtarget,
SDNode *N, SDValue X, SDValue Amt,
SDValue Cond, ISD::CondCode RequiredCC,
- bool IsSRL) {
+ unsigned ShiftOpcode,
+ bool AmtOutOfRangeIsPoison) {
+ assert((ShiftOpcode == ISD::SHL || ShiftOpcode == ISD::SRL ||
+ ShiftOpcode == ISD::SRA) &&
+ "Unexpected shift opcode");
+ assert((ShiftOpcode != ISD::SRA || AmtOutOfRangeIsPoison) &&
+ "sshl sign fills where the select needs zero");
using namespace llvm::SDPatternMatch;
EVT VT = N->getValueType(0);
// ushl only exists for 64 and 128 bit vectors.
@@ -29441,21 +29448,22 @@ static SDValue foldMaskedShiftToUSHL(SelectionDAG &DAG,
SDLoc DL(N);
// Amounts that might exceed 127 need the umin clamp.
- if (!DAG.computeKnownBits(Amt).getMaxValue().ule(127)) {
+ 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)
+ // There is no shift right register instruction but ushl and sshl shift
+ // right when the amount is negative.
+ if (ShiftOpcode != ISD::SHL)
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);
+ unsigned IID = ShiftOpcode == ISD::SRA ? Intrinsic::aarch64_neon_sshl
+ : Intrinsic::aarch64_neon_ushl;
+ return DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, VT,
+ DAG.getTargetConstant(IID, DL, MVT::i32), X, Amt);
}
// vselect(setcc_ult(amt, EltSize), shl(x, amt), zeros) -> ushl(x, amt)
@@ -29472,7 +29480,29 @@ performVSelectMaskedShiftCombine(SDNode *N, SelectionDAG &DAG,
return foldMaskedShiftToUSHL(DAG, Subtarget, N, Shift.getOperand(0),
Shift.getOperand(1), N->getOperand(0),
- RequiredCC, Shift.getOpcode() == ISD::SRL);
+ RequiredCC, Shift.getOpcode(),
+ /*AmtOutOfRangeIsPoison=*/false);
+}
+
+// shl/srl(vselect(setcc_ult(amt, EltSize), x, zeros), amt) -> ushl(x, amt)
+// sra(vselect(setcc_ult(amt, EltSize), x, zeros), amt) -> sshl(x, -amt)
+// The zeroed lanes are shifted by EltSize or more which is poison.
+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(),
+ /*AmtOutOfRangeIsPoison=*/true);
}
// vselect (v1i1 setcc) ->
@@ -31334,7 +31364,12 @@ 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:
+ case ISD::SRA:
+ 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 2d5db9ad8b2f1..33b410f7b65b2 100644
--- a/llvm/test/CodeGen/AArch64/vselect-masked-shift.ll
+++ b/llvm/test/CodeGen/AArch64/vselect-masked-shift.ll
@@ -243,6 +243,141 @@ entry:
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: and v1.16b, v1.16b, v2.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: and v1.16b, v1.16b, v2.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: and v1.16b, v1.16b, v2.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)
+ %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: 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: and v1.16b, v1.16b, v2.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 <4 x i32> @masked_input_sra_v4i32(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: masked_input_sra_v4i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.4s, #63
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: neg v1.4s, v1.4s
+; CHECK-NEXT: sshl 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)
+ %v = select <4 x i1> %ok, <4 x i32> %x, <4 x i32> zeroinitializer
+ %res = ashr <4 x i32> %v, %m
+ ret <4 x i32> %res
+}
+
+define <8 x i16> @masked_input_sra_v8i16(<8 x i16> %x, <8 x i16> %amt) {
+; CHECK-LABEL: masked_input_sra_v8i16:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: movi v2.8h, #31
+; CHECK-NEXT: and v1.16b, v1.16b, v2.16b
+; CHECK-NEXT: neg v1.8h, v1.8h
+; CHECK-NEXT: sshl v0.8h, v0.8h, v1.8h
+; CHECK-NEXT: ret
+entry:
+ %m = and <8 x i16> %amt, splat (i16 31)
+ %ok = icmp ult <8 x i16> %m, splat (i16 16)
+ %v = select <8 x i1> %ok, <8 x i16> %x, <8 x i16> zeroinitializer
+ %res = ashr <8 x i16> %v, %m
+ ret <8 x i16> %res
+}
+
+define <2 x i64> @unbounded_input_shl_v2i64(<2 x i64> %x, <2 x i64> %amt) {
+; CHECK-LABEL: unbounded_input_shl_v2i64:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: ushl v0.2d, v0.2d, v1.2d
+; CHECK-NEXT: ret
+entry:
+ %ok = icmp ult <2 x i64> %amt, splat (i64 64)
+ %v = select <2 x i1> %ok, <2 x i64> %x, <2 x i64> zeroinitializer
+ %res = shl <2 x i64> %v, %amt
+ ret <2 x i64> %res
+}
+
+define <4 x i32> @unbounded_input_sra_v4i32(<4 x i32> %x, <4 x i32> %amt) {
+; CHECK-LABEL: unbounded_input_sra_v4i32:
+; CHECK: // %bb.0: // %entry
+; CHECK-NEXT: neg v1.4s, v1.4s
+; CHECK-NEXT: sshl v0.4s, v0.4s, v1.4s
+; CHECK-NEXT: ret
+entry:
+ %ok = icmp ult <4 x i32> %amt, splat (i32 32)
+ %v = select <4 x i1> %ok, <4 x i32> %x, <4 x i32> zeroinitializer
+ %res = ashr <4 x i32> %v, %amt
+ ret <4 x i32> %res
+}
+
attributes #0 = { "target-features"="+sve" }
define <8 x i32> @neg_masked_shl_v8i32_sve_vls(<8 x i32> %x, <8 x i32> %amt) #1 {
``````````
</details>
https://github.com/llvm/llvm-project/pull/210880
More information about the llvm-commits
mailing list