[llvm] Enable known bits for predicated shifts (PR #200347)
via llvm-commits
llvm-commits at lists.llvm.org
Fri May 29 01:29:31 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-aarch64
Author: Harry Ramsey (Harry-Ramsey)
<details>
<summary>Changes</summary>
Allow SelectionDAG to query target known-bits information for scalable vector nodes, and known-bits cases for SVE predicated SHL, SRL and SRA nodes.
This enables DAG combines to prove disjointness for ORs involving scalable vector shifts, enabling USRA/SSRA instruction selection.
---
Full diff: https://github.com/llvm/llvm-project/pull/200347.diff
3 Files Affected:
- (modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp (-5)
- (modified) llvm/lib/Target/AArch64/AArch64ISelLowering.cpp (+16)
- (modified) llvm/test/CodeGen/AArch64/sve2-sra.ll (+104)
``````````diff
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index b92ba79564f56..a380bb789544f 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -4527,11 +4527,6 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
case ISD::INTRINSIC_WO_CHAIN:
case ISD::INTRINSIC_W_CHAIN:
case ISD::INTRINSIC_VOID:
- // TODO: Probably okay to remove after audit; here to reduce change size
- // in initial enablement patch for scalable vectors
- if (Op.getValueType().isScalableVector())
- break;
-
// Allow the target to implement this method for its nodes.
TLI->computeKnownBitsForTargetNode(Op, Known, DemandedElts, *this, Depth);
break;
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 710ca7fcba756..d74299a6013df 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -2926,6 +2926,22 @@ void AArch64TargetLowering::computeKnownBitsForTargetNode(
}
break;
}
+ case AArch64ISD::SHL_PRED:
+ case AArch64ISD::SRL_PRED:
+ case AArch64ISD::SRA_PRED: {
+ KnownBits KnownVal =
+ DAG.computeKnownBits(Op->getOperand(1), DemandedElts, Depth + 1);
+ KnownBits KnownAmt =
+ DAG.computeKnownBits(Op->getOperand(2), DemandedElts, Depth + 1);
+
+ if (Op.getOpcode() == AArch64ISD::SHL_PRED)
+ Known = KnownBits::shl(KnownVal, KnownAmt);
+ else if (Op.getOpcode() == AArch64ISD::SRL_PRED)
+ Known = KnownBits::lshr(KnownVal, KnownAmt);
+ else
+ Known = KnownBits::ashr(KnownVal, KnownAmt);
+ break;
+ }
case ISD::INTRINSIC_WO_CHAIN:
case ISD::INTRINSIC_VOID: {
unsigned IntNo = Op.getConstantOperandVal(0);
diff --git a/llvm/test/CodeGen/AArch64/sve2-sra.ll b/llvm/test/CodeGen/AArch64/sve2-sra.ll
index eafcd60bc1605..571d9a362dea4 100644
--- a/llvm/test/CodeGen/AArch64/sve2-sra.ll
+++ b/llvm/test/CodeGen/AArch64/sve2-sra.ll
@@ -255,6 +255,110 @@ define <vscale x 2 x i64> @ssra_intr_u_i64(<vscale x 2 x i1> %pg, <vscale x 2 x
ret <vscale x 2 x i64> %add
}
+define <vscale x 16 x i8> @usra_disjoint_shift_or16xi8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b) #0 {
+; CHECK-LABEL: usra_disjoint_shift_or16xi8:
+; CHECK: // %bb.0:
+; CHECK-NEXT: lsl z0.b, z0.b, #7
+; CHECK-NEXT: usra z0.b, z1.b, #1
+; CHECK-NEXT: ret
+ %shl = shl <vscale x 16 x i8> %a, splat (i8 7)
+ %srl = lshr <vscale x 16 x i8> %b, splat (i8 1)
+ %r = or <vscale x 16 x i8> %shl, %srl
+ ret <vscale x 16 x i8> %r
+}
+
+define <vscale x 8 x i16> @usra_disjoint_shift_or8xi16(<vscale x 8 x i16> %a, <vscale x 8 x i16> %b) #0 {
+; CHECK-LABEL: usra_disjoint_shift_or8xi16:
+; CHECK: // %bb.0:
+; CHECK-NEXT: lsl z0.h, z0.h, #7
+; CHECK-NEXT: usra z0.h, z1.h, #9
+; CHECK-NEXT: ret
+ %shl = shl <vscale x 8 x i16> %a, splat (i16 7)
+ %srl = lshr <vscale x 8 x i16> %b, splat (i16 9)
+ %r = or <vscale x 8 x i16> %shl, %srl
+ ret <vscale x 8 x i16> %r
+}
+
+define <vscale x 4 x i32> @usra_disjoint_shift_or4xi32(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b) #0 {
+; CHECK-LABEL: usra_disjoint_shift_or4xi32:
+; CHECK: // %bb.0:
+; CHECK-NEXT: lsl z0.s, z0.s, #7
+; CHECK-NEXT: usra z0.s, z1.s, #25
+; CHECK-NEXT: ret
+ %shl = shl <vscale x 4 x i32> %a, splat (i32 7)
+ %srl = lshr <vscale x 4 x i32> %b, splat (i32 25)
+ %r = or <vscale x 4 x i32> %shl, %srl
+ ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 2 x i64> @usra_disjoint_shift_or2xi64(<vscale x 2 x i64> %a, <vscale x 2 x i64> %b) #0 {
+; CHECK-LABEL: usra_disjoint_shift_or2xi64:
+; CHECK: // %bb.0:
+; CHECK-NEXT: lsl z0.d, z0.d, #7
+; CHECK-NEXT: usra z0.d, z1.d, #57
+; CHECK-NEXT: ret
+ %shl = shl <vscale x 2 x i64> %a, splat (i64 7)
+ %srl = lshr <vscale x 2 x i64> %b, splat (i64 57)
+ %r = or <vscale x 2 x i64> %shl, %srl
+ ret <vscale x 2 x i64> %r
+}
+
+define <vscale x 16 x i8> @ssra_disjoint_shift_or16xi8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b) #0 {
+; CHECK-LABEL: ssra_disjoint_shift_or16xi8:
+; CHECK: // %bb.0:
+; CHECK-NEXT: lsr z0.b, z0.b, #2
+; CHECK-NEXT: lsl z1.b, z1.b, #7
+; CHECK-NEXT: ssra z0.b, z1.b, #1
+; CHECK-NEXT: ret
+ %acc = lshr <vscale x 16 x i8> %a, splat (i8 2)
+ %sign = shl <vscale x 16 x i8> %b, splat (i8 7)
+ %sra = ashr <vscale x 16 x i8> %sign, splat (i8 1)
+ %r = or <vscale x 16 x i8> %acc, %sra
+ ret <vscale x 16 x i8> %r
+}
+
+define <vscale x 8 x i16> @ssra_disjoint_shift_or8xi16(<vscale x 8 x i16> %a, <vscale x 8 x i16> %b) #0 {
+; CHECK-LABEL: ssra_disjoint_shift_or8xi16:
+; CHECK: // %bb.0:
+; CHECK-NEXT: lsr z0.h, z0.h, #10
+; CHECK-NEXT: lsl z1.h, z1.h, #15
+; CHECK-NEXT: ssra z0.h, z1.h, #9
+; CHECK-NEXT: ret
+ %acc = lshr <vscale x 8 x i16> %a, splat (i16 10)
+ %sign = shl <vscale x 8 x i16> %b, splat (i16 15)
+ %sra = ashr <vscale x 8 x i16> %sign, splat (i16 9)
+ %r = or <vscale x 8 x i16> %acc, %sra
+ ret <vscale x 8 x i16> %r
+}
+
+define <vscale x 4 x i32> @ssra_disjoint_shift_or4xi32(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b) #0 {
+; CHECK-LABEL: ssra_disjoint_shift_or4xi32:
+; CHECK: // %bb.0:
+; CHECK-NEXT: lsr z0.s, z0.s, #26
+; CHECK-NEXT: lsl z1.s, z1.s, #31
+; CHECK-NEXT: ssra z0.s, z1.s, #25
+; CHECK-NEXT: ret
+ %acc = lshr <vscale x 4 x i32> %a, splat (i32 26)
+ %sign = shl <vscale x 4 x i32> %b, splat (i32 31)
+ %sra = ashr <vscale x 4 x i32> %sign, splat (i32 25)
+ %r = or <vscale x 4 x i32> %acc, %sra
+ ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 2 x i64> @ssra_disjoint_shift_or2xi64(<vscale x 2 x i64> %a, <vscale x 2 x i64> %b) #0 {
+; CHECK-LABEL: ssra_disjoint_shift_or2xi64:
+; CHECK: // %bb.0:
+; CHECK-NEXT: lsr z0.d, z0.d, #58
+; CHECK-NEXT: lsl z1.d, z1.d, #63
+; CHECK-NEXT: ssra z0.d, z1.d, #57
+; CHECK-NEXT: ret
+ %acc = lshr <vscale x 2 x i64> %a, splat (i64 58)
+ %sign = shl <vscale x 2 x i64> %b, splat (i64 63)
+ %sra = ashr <vscale x 2 x i64> %sign, splat (i64 57)
+ %r = or <vscale x 2 x i64> %acc, %sra
+ ret <vscale x 2 x i64> %r
+}
+
declare <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 immarg)
declare <vscale x 8 x i1> @llvm.aarch64.sve.ptrue.nxv8i1(i32 immarg)
declare <vscale x 4 x i1> @llvm.aarch64.sve.ptrue.nxv4i1(i32 immarg)
``````````
</details>
https://github.com/llvm/llvm-project/pull/200347
More information about the llvm-commits
mailing list