[llvm] Enable known bits for predicated shifts (PR #200347)
Harry Ramsey via llvm-commits
llvm-commits at lists.llvm.org
Fri May 29 01:28:52 PDT 2026
https://github.com/Harry-Ramsey created https://github.com/llvm/llvm-project/pull/200347
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.
>From 08161b2c1cc3b825eb5665402c02695063257461 Mon Sep 17 00:00:00 2001
From: Harry Ramsey <harry.ramsey at arm.com>
Date: Thu, 28 May 2026 10:44:45 +0000
Subject: [PATCH 1/2] Enable known bits for predicated shifts
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.
---
llvm/test/CodeGen/AArch64/sve2-sra.ll | 112 ++++++++++++++++++++++++++
1 file changed, 112 insertions(+)
diff --git a/llvm/test/CodeGen/AArch64/sve2-sra.ll b/llvm/test/CodeGen/AArch64/sve2-sra.ll
index eafcd60bc1605..1ec6bd7158705 100644
--- a/llvm/test/CodeGen/AArch64/sve2-sra.ll
+++ b/llvm/test/CodeGen/AArch64/sve2-sra.ll
@@ -255,6 +255,118 @@ 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: lsr z1.b, z1.b, #1
+; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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: lsr z1.h, z1.h, #9
+; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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: lsr z1.s, z1.s, #25
+; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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: lsr z1.d, z1.d, #57
+; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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: lsl z1.b, z1.b, #7
+; CHECK-NEXT: lsr z0.b, z0.b, #2
+; CHECK-NEXT: asr z1.b, z1.b, #1
+; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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: lsl z1.h, z1.h, #15
+; CHECK-NEXT: lsr z0.h, z0.h, #10
+; CHECK-NEXT: asr z1.h, z1.h, #9
+; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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: lsl z1.s, z1.s, #31
+; CHECK-NEXT: lsr z0.s, z0.s, #26
+; CHECK-NEXT: asr z1.s, z1.s, #25
+; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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: lsl z1.d, z1.d, #63
+; CHECK-NEXT: lsr z0.d, z0.d, #58
+; CHECK-NEXT: asr z1.d, z1.d, #57
+; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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)
>From 2316ea97e101e4d2f84b445623e95d37a3a23c18 Mon Sep 17 00:00:00 2001
From: Harry Ramsey <harry.ramsey at arm.com>
Date: Thu, 28 May 2026 11:22:24 +0000
Subject: [PATCH 2/2] fixup! [AArch64][SVE] Enable known bits for predicated
shifts
---
.../lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 5 ---
.../Target/AArch64/AArch64ISelLowering.cpp | 16 ++++++++++
llvm/test/CodeGen/AArch64/sve2-sra.ll | 32 +++++++------------
3 files changed, 28 insertions(+), 25 deletions(-)
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 1ec6bd7158705..571d9a362dea4 100644
--- a/llvm/test/CodeGen/AArch64/sve2-sra.ll
+++ b/llvm/test/CodeGen/AArch64/sve2-sra.ll
@@ -259,8 +259,7 @@ define <vscale x 16 x i8> @usra_disjoint_shift_or16xi8(<vscale x 16 x i8> %a, <v
; CHECK-LABEL: usra_disjoint_shift_or16xi8:
; CHECK: // %bb.0:
; CHECK-NEXT: lsl z0.b, z0.b, #7
-; CHECK-NEXT: lsr z1.b, z1.b, #1
-; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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)
@@ -272,8 +271,7 @@ define <vscale x 8 x i16> @usra_disjoint_shift_or8xi16(<vscale x 8 x i16> %a, <v
; CHECK-LABEL: usra_disjoint_shift_or8xi16:
; CHECK: // %bb.0:
; CHECK-NEXT: lsl z0.h, z0.h, #7
-; CHECK-NEXT: lsr z1.h, z1.h, #9
-; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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)
@@ -285,8 +283,7 @@ define <vscale x 4 x i32> @usra_disjoint_shift_or4xi32(<vscale x 4 x i32> %a, <v
; CHECK-LABEL: usra_disjoint_shift_or4xi32:
; CHECK: // %bb.0:
; CHECK-NEXT: lsl z0.s, z0.s, #7
-; CHECK-NEXT: lsr z1.s, z1.s, #25
-; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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)
@@ -298,8 +295,7 @@ define <vscale x 2 x i64> @usra_disjoint_shift_or2xi64(<vscale x 2 x i64> %a, <v
; CHECK-LABEL: usra_disjoint_shift_or2xi64:
; CHECK: // %bb.0:
; CHECK-NEXT: lsl z0.d, z0.d, #7
-; CHECK-NEXT: lsr z1.d, z1.d, #57
-; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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)
@@ -310,10 +306,9 @@ define <vscale x 2 x i64> @usra_disjoint_shift_or2xi64(<vscale x 2 x i64> %a, <v
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: lsl z1.b, z1.b, #7
; CHECK-NEXT: lsr z0.b, z0.b, #2
-; CHECK-NEXT: asr z1.b, z1.b, #1
-; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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)
@@ -325,10 +320,9 @@ define <vscale x 16 x i8> @ssra_disjoint_shift_or16xi8(<vscale x 16 x i8> %a, <v
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: lsl z1.h, z1.h, #15
; CHECK-NEXT: lsr z0.h, z0.h, #10
-; CHECK-NEXT: asr z1.h, z1.h, #9
-; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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)
@@ -340,10 +334,9 @@ define <vscale x 8 x i16> @ssra_disjoint_shift_or8xi16(<vscale x 8 x i16> %a, <v
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: lsl z1.s, z1.s, #31
; CHECK-NEXT: lsr z0.s, z0.s, #26
-; CHECK-NEXT: asr z1.s, z1.s, #25
-; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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)
@@ -355,10 +348,9 @@ define <vscale x 4 x i32> @ssra_disjoint_shift_or4xi32(<vscale x 4 x i32> %a, <v
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: lsl z1.d, z1.d, #63
; CHECK-NEXT: lsr z0.d, z0.d, #58
-; CHECK-NEXT: asr z1.d, z1.d, #57
-; CHECK-NEXT: orr z0.d, z0.d, z1.d
+; 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)
More information about the llvm-commits
mailing list