[llvm] Support EXTRACT_SUBVECTOR in computeKnownBits for scalable vectors (PR #177163)
Sander de Smalen via llvm-commits
llvm-commits at lists.llvm.org
Fri Jan 23 09:04:51 PST 2026
https://github.com/sdesmalen-arm updated https://github.com/llvm/llvm-project/pull/177163
>From 31b9189771bb9622704fcabe2ee501b808bccca8 Mon Sep 17 00:00:00 2001
From: Sander de Smalen <sander.desmalen at arm.com>
Date: Wed, 21 Jan 2026 11:14:20 +0000
Subject: [PATCH 1/2] Support EXTRACT_SUBVECTOR in computeKnownBits for
scalable vectors
Rather than not supporting this case it would just be more conservative
as it will need to prove known bits for all elements.
---
.../lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 17 ++++---
.../AArch64/sve-fixed-length-shuffles.ll | 47 -------------------
.../AArch64/AArch64SelectionDAGTest.cpp | 17 ++++++-
3 files changed, 27 insertions(+), 54 deletions(-)
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 92ce9e3a299ec..ad1133a43c970 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -3333,6 +3333,8 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
KnownBits Known2;
unsigned NumElts = DemandedElts.getBitWidth();
+ assert((!Op.getValueType().isScalableVector() || NumElts == 1) &&
+ "DemandedElts for scalable vectors must be 1 to represent all lanes");
assert((!Op.getValueType().isFixedLengthVector() ||
NumElts == Op.getValueType().getVectorNumElements()) &&
"Unexpected vector size");
@@ -3511,12 +3513,15 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
case ISD::EXTRACT_SUBVECTOR: {
// Offset the demanded elts by the subvector index.
SDValue Src = Op.getOperand(0);
- // Bail until we can represent demanded elements for scalable vectors.
- if (Op.getValueType().isScalableVector() || Src.getValueType().isScalableVector())
- break;
- uint64_t Idx = Op.getConstantOperandVal(1);
- unsigned NumSrcElts = Src.getValueType().getVectorNumElements();
- APInt DemandedSrcElts = DemandedElts.zext(NumSrcElts).shl(Idx);
+
+ APInt DemandedSrcElts;
+ if (Src.getValueType().isScalableVector())
+ DemandedSrcElts = APInt(1, 1);
+ else {
+ uint64_t Idx = Op.getConstantOperandVal(1);
+ unsigned NumSrcElts = Src.getValueType().getVectorNumElements();
+ DemandedSrcElts = DemandedElts.zext(NumSrcElts).shl(Idx);
+ }
Known = computeKnownBits(Src, DemandedSrcElts, Depth + 1);
break;
}
diff --git a/llvm/test/CodeGen/AArch64/sve-fixed-length-shuffles.ll b/llvm/test/CodeGen/AArch64/sve-fixed-length-shuffles.ll
index 4fd94f7b9a4cb..59329624e75bb 100644
--- a/llvm/test/CodeGen/AArch64/sve-fixed-length-shuffles.ll
+++ b/llvm/test/CodeGen/AArch64/sve-fixed-length-shuffles.ll
@@ -26,53 +26,6 @@ define void @hang_when_merging_stores_after_legalisation(ptr %a, <2 x i32> %b) v
define void @crash_when_lowering_extract_shuffle(ptr %dst, i1 %cond) vscale_range(2,2) #0 {
; CHECK-LABEL: crash_when_lowering_extract_shuffle:
; CHECK: // %bb.0:
-; CHECK-NEXT: tbnz w1, #0, .LBB1_2
-; CHECK-NEXT: // %bb.1: // %vector.body
-; CHECK-NEXT: movi v2.2d, #0000000000000000
-; CHECK-NEXT: movi v0.2d, #0000000000000000
-; CHECK-NEXT: movi v3.2d, #0000000000000000
-; CHECK-NEXT: ptrue p0.s
-; CHECK-NEXT: umov w8, v2.b[8]
-; CHECK-NEXT: mov v0.b[1], v2.b[1]
-; CHECK-NEXT: fmov s1, w8
-; CHECK-NEXT: mov v0.b[2], v2.b[2]
-; CHECK-NEXT: mov v1.b[1], v2.b[9]
-; CHECK-NEXT: mov v0.b[3], v2.b[3]
-; CHECK-NEXT: mov v1.b[2], v2.b[10]
-; CHECK-NEXT: mov v0.b[4], v2.b[4]
-; CHECK-NEXT: mov v1.b[3], v2.b[11]
-; CHECK-NEXT: mov v0.b[5], v2.b[5]
-; CHECK-NEXT: mov v1.b[4], v2.b[12]
-; CHECK-NEXT: mov v0.b[6], v2.b[6]
-; CHECK-NEXT: mov v1.b[5], v2.b[13]
-; CHECK-NEXT: mov v0.b[7], v2.b[7]
-; CHECK-NEXT: mov v1.b[6], v2.b[14]
-; CHECK-NEXT: uunpklo z0.h, z0.b
-; CHECK-NEXT: mov v1.b[7], v2.b[15]
-; CHECK-NEXT: movprfx z2, z3
-; CHECK-NEXT: ext z2.b, z2.b, z3.b, #16
-; CHECK-NEXT: uunpklo z0.s, z0.h
-; CHECK-NEXT: ext v3.16b, v2.16b, v2.16b, #8
-; CHECK-NEXT: sunpklo z2.h, z2.b
-; CHECK-NEXT: uunpklo z1.h, z1.b
-; CHECK-NEXT: lsl z0.s, z0.s, #31
-; CHECK-NEXT: sunpklo z3.h, z3.b
-; CHECK-NEXT: sunpklo z2.s, z2.h
-; CHECK-NEXT: uunpklo z1.s, z1.h
-; CHECK-NEXT: asr z0.s, z0.s, #31
-; CHECK-NEXT: sunpklo z3.s, z3.h
-; CHECK-NEXT: cmpne p1.s, p0/z, z2.s, #0
-; CHECK-NEXT: movi v2.2d, #0000000000000000
-; CHECK-NEXT: lsl z1.s, z1.s, #31
-; CHECK-NEXT: cmpne p3.s, p0/z, z0.s, #0
-; CHECK-NEXT: asr z1.s, z1.s, #31
-; CHECK-NEXT: cmpne p2.s, p0/z, z3.s, #0
-; CHECK-NEXT: st1w { z2.s }, p1, [x0, #2, mul vl]
-; CHECK-NEXT: st1w { z2.s }, p3, [x0]
-; CHECK-NEXT: cmpne p0.s, p0/z, z1.s, #0
-; CHECK-NEXT: st1w { z2.s }, p2, [x0, #3, mul vl]
-; CHECK-NEXT: st1w { z2.s }, p0, [x0, #1, mul vl]
-; CHECK-NEXT: .LBB1_2: // %exit
; CHECK-NEXT: ret
%broadcast.splat = shufflevector <32 x i1> zeroinitializer, <32 x i1> zeroinitializer, <32 x i32> zeroinitializer
br i1 %cond, label %exit, label %vector.body
diff --git a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
index 6e6fc06e94ff0..5ee83b6091406 100644
--- a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
+++ b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
@@ -105,7 +105,7 @@ TEST_F(AArch64SelectionDAGTest, computeKnownBitsSVE_ZERO_EXTEND_VECTOR_INREG) {
auto OutVecVT = EVT::getVectorVT(Context, Int16VT, 2, true);
auto InVec = DAG->getConstant(0, Loc, InVecVT);
auto Op = DAG->getNode(ISD::ZERO_EXTEND_VECTOR_INREG, Loc, OutVecVT, InVec);
- auto DemandedElts = APInt(2, 3);
+ auto DemandedElts = APInt(1, 1);
KnownBits Known = DAG->computeKnownBits(Op, DemandedElts);
// We don't know anything for SVE at the moment.
@@ -721,6 +721,21 @@ TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_SUB) {
EXPECT_EQ(Known.One, APInt(8, 0x1));
}
+TEST_F(AArch64SelectionDAGTest, ComputeKnownBitsSVE_EXTRACT_SUBVECTOR) {
+ SDLoc Loc;
+ auto IntVT = EVT::getIntegerVT(Context, 8);
+ auto ScalableVecVT =
+ EVT::getVectorVT(Context, IntVT, 16, /*IsScalable=*/true);
+ auto FixedVecVT = EVT::getVectorVT(Context, IntVT, 16, /*IsScalable=*/false);
+ auto IdxVT = EVT::getIntegerVT(Context, 64);
+ auto Vec = DAG->getConstant(1, Loc, ScalableVecVT);
+ auto ZeroIdx = DAG->getConstant(0, Loc, IdxVT);
+ auto Op = DAG->getNode(ISD::EXTRACT_SUBVECTOR, Loc, FixedVecVT, Vec, ZeroIdx);
+ KnownBits Known = DAG->computeKnownBits(Op);
+ EXPECT_EQ(Known.Zero, APInt(8, 0xFE));
+ EXPECT_EQ(Known.One, APInt(8, 0x1));
+}
+
// Piggy-backing on the AArch64 tests to verify SelectionDAG::computeKnownBits.
TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_USUBO_CARRY) {
SDLoc Loc;
>From 332cbdfaf7fd2d932ab75dbab08967c9af60831b Mon Sep 17 00:00:00 2001
From: Sander de Smalen <sander.desmalen at arm.com>
Date: Fri, 23 Jan 2026 17:03:53 +0000
Subject: [PATCH 2/2] Add comment
---
llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index ad1133a43c970..bac8811755a5e 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -3516,7 +3516,7 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
APInt DemandedSrcElts;
if (Src.getValueType().isScalableVector())
- DemandedSrcElts = APInt(1, 1);
+ DemandedSrcElts = APInt(1, 1); // <=> 'demand all elements'
else {
uint64_t Idx = Op.getConstantOperandVal(1);
unsigned NumSrcElts = Src.getValueType().getVectorNumElements();
More information about the llvm-commits
mailing list