[llvm] f42072e - [Analysis] Add `KnownBits` optimization for `pdep` and `pext` (#204223)
via llvm-commits
llvm-commits at lists.llvm.org
Sat Jun 20 21:19:33 PDT 2026
Author: Jan Schultke
Date: 2026-06-21T12:19:29+08:00
New Revision: f42072ed8de552ce25b7a978c052acca6ab13246
URL: https://github.com/llvm/llvm-project/commit/f42072ed8de552ce25b7a978c052acca6ab13246
DIFF: https://github.com/llvm/llvm-project/commit/f42072ed8de552ce25b7a978c052acca6ab13246.diff
LOG: [Analysis] Add `KnownBits` optimization for `pdep` and `pext` (#204223)
Fixes #204136
Added:
llvm/test/Analysis/ValueTracking/knownbits-pdep.ll
llvm/test/Analysis/ValueTracking/knownbits-pext.ll
Modified:
llvm/include/llvm/Support/KnownBits.h
llvm/lib/Analysis/ValueTracking.cpp
llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
llvm/lib/Support/KnownBits.cpp
llvm/unittests/Support/KnownBitsTest.cpp
Removed:
################################################################################
diff --git a/llvm/include/llvm/Support/KnownBits.h b/llvm/include/llvm/Support/KnownBits.h
index e47383dc6b2d1..dc80a3f10b087 100644
--- a/llvm/include/llvm/Support/KnownBits.h
+++ b/llvm/include/llvm/Support/KnownBits.h
@@ -479,6 +479,12 @@ struct KnownBits {
/// Compute known bits for clmul(LHS, RHS).
LLVM_ABI static KnownBits clmul(const KnownBits &LHS, const KnownBits &RHS);
+ /// Compute known bits for pext(LHS, RHS).
+ LLVM_ABI static KnownBits pext(const KnownBits &LHS, const KnownBits &RHS);
+
+ /// Compute known bits for pdep(LHS, RHS).
+ LLVM_ABI static KnownBits pdep(const KnownBits &LHS, const KnownBits &RHS);
+
/// Determine if these known bits always give the same ICMP_EQ result.
LLVM_ABI static std::optional<bool> eq(const KnownBits &LHS,
const KnownBits &RHS);
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 9cb41a0b6689e..9334551005224 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2089,6 +2089,16 @@ static void computeKnownBitsFromOperator(const Operator *I,
computeKnownBits(I->getOperand(1), DemandedElts, Known2, Q, Depth + 1);
Known = KnownBits::clmul(Known, Known2);
break;
+ case Intrinsic::pext:
+ computeKnownBits(I->getOperand(0), DemandedElts, Known, Q, Depth + 1);
+ computeKnownBits(I->getOperand(1), DemandedElts, Known2, Q, Depth + 1);
+ Known = KnownBits::pext(Known, Known2);
+ break;
+ case Intrinsic::pdep:
+ computeKnownBits(I->getOperand(0), DemandedElts, Known, Q, Depth + 1);
+ computeKnownBits(I->getOperand(1), DemandedElts, Known2, Q, Depth + 1);
+ Known = KnownBits::pdep(Known, Known2);
+ break;
case Intrinsic::uadd_sat:
computeKnownBits(I->getOperand(0), DemandedElts, Known, Q, Depth + 1);
computeKnownBits(I->getOperand(1), DemandedElts, Known2, Q, Depth + 1);
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 44120cceed2a3..76556531ada42 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -3955,19 +3955,13 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
case ISD::PDEP: {
Known = computeKnownBits(Op.getOperand(1), DemandedElts, Depth + 1);
Known2 = computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
- // Zeros are retained from the mask operand. But not ones.
- Known.One.clearAllBits();
- // The result will have at least as many trailing zeros as the non-mask
- // operand since bits can only map to the same or higher bit position.
- Known.Zero.setLowBits(Known2.countMinTrailingZeros());
+ Known = KnownBits::pdep(Known2, Known);
break;
}
case ISD::PEXT: {
Known = computeKnownBits(Op.getOperand(1), DemandedElts, Depth + 1);
- // The result has as many leading zeros as the number of zeroes in the mask.
- unsigned Count = Known.Zero.popcount();
- Known.Zero = APInt::getHighBitsSet(BitWidth, Count);
- Known.One.clearAllBits();
+ Known2 = computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
+ Known = KnownBits::pext(Known2, Known);
break;
}
case ISD::CLMUL: {
diff --git a/llvm/lib/Support/KnownBits.cpp b/llvm/lib/Support/KnownBits.cpp
index 5becdac89fcaa..a2f71061fe0ad 100644
--- a/llvm/lib/Support/KnownBits.cpp
+++ b/llvm/lib/Support/KnownBits.cpp
@@ -645,6 +645,71 @@ KnownBits KnownBits::clmul(const KnownBits &LHS, const KnownBits &RHS) {
return Res;
}
+KnownBits KnownBits::pext(const KnownBits &LHS, const KnownBits &RHS) {
+ // For each source position I where mask[I] could be set, the output index j
+ // lies in [M0, M1] where these track the range of possible set-bit counts
+ // seen so far in mask.
+ //
+ // The output bit j
+ // - can be 0 if any candidate LHS[I] could be zero or popcount(mask) could
+ // be <= j, and
+ // - can be 1 only if some candidate LHS[I] could be one and popcount(mask)
+ // is known > j.
+ unsigned BitWidth = LHS.getBitWidth();
+ KnownBits Res(BitWidth);
+ Res.setAllConflict();
+
+ unsigned M0 = 0, M1 = 0;
+ for (unsigned I = 0; I < BitWidth; ++I) {
+ if (!RHS.Zero[I]) {
+ APInt Range = APInt::getBitsSet(BitWidth, M0, M1 + 1);
+ if (!LHS.Zero[I])
+ Res.Zero &= ~Range; // some position in Range could be 1
+ if (!LHS.One[I])
+ Res.One &= ~Range; // some position in Range could be 0
+ }
+ if (RHS.One[I])
+ ++M0, ++M1;
+ else if (!RHS.Zero[I])
+ ++M1;
+ }
+
+ // Output positions j >= M0 may have no source (popcount(mask) <= j), in
+ // which case they default to zero.
+ Res.One &= APInt::getLowBitsSet(BitWidth, M0);
+ return Res;
+}
+
+KnownBits KnownBits::pdep(const KnownBits &LHS, const KnownBits &RHS) {
+ // For each output position I where mask[I] could be set, the source index j
+ // lies in [M0, M1] where these track possible counts of set mask bits < I.
+ //
+ // The output bit
+ // - can be 0 if mask[I] or any candidate LHS[j] could be zero, and
+ // - can be 1 only if both mask[I] and some candidate LHS[j] could be one.
+ unsigned BitWidth = LHS.getBitWidth();
+ KnownBits Res(BitWidth);
+ Res.setAllConflict();
+
+ unsigned M0 = 0, M1 = 0;
+ for (unsigned I = 0; I < BitWidth; ++I) {
+ if (!RHS.One[I])
+ Res.One.clearBit(I); // mask[I] could be 0 -> output[I] could be 0
+ if (!RHS.Zero[I]) {
+ APInt Range = APInt::getBitsSet(BitWidth, M0, M1 + 1);
+ if (!Range.isSubsetOf(LHS.One))
+ Res.One.clearBit(I); // some candidate could be 0
+ if (!Range.isSubsetOf(LHS.Zero))
+ Res.Zero.clearBit(I); // some candidate could be 1
+ }
+ if (RHS.One[I])
+ ++M0, ++M1;
+ else if (!RHS.Zero[I])
+ ++M1;
+ }
+ return Res;
+}
+
std::optional<bool> KnownBits::eq(const KnownBits &LHS, const KnownBits &RHS) {
if (LHS.isConstant() && RHS.isConstant())
return LHS.getConstant() == RHS.getConstant();
diff --git a/llvm/test/Analysis/ValueTracking/knownbits-pdep.ll b/llvm/test/Analysis/ValueTracking/knownbits-pdep.ll
new file mode 100644
index 0000000000000..ba18549f4514f
--- /dev/null
+++ b/llvm/test/Analysis/ValueTracking/knownbits-pdep.ll
@@ -0,0 +1,89 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+
+; pdep(x, 0b11110000): bits 0..3 of result are always zero.
+define i1 @pdep_low_bits_zero(i8 %x) nounwind {
+; CHECK-LABEL: @pdep_low_bits_zero(
+; CHECK-NEXT: ret i1 true
+;
+ %pdep = call i8 @llvm.pdep.i8(i8 %x, i8 240)
+ %and = and i8 %pdep, 15
+ %r = icmp eq i8 %and, 0
+ ret i1 %r
+}
+
+define <2 x i1> @pdep_low_bits_zero_vector(<2 x i8> %x) nounwind {
+; CHECK-LABEL: @pdep_low_bits_zero_vector(
+; CHECK-NEXT: ret <2 x i1> splat (i1 true)
+;
+ %pdep = call <2 x i8> @llvm.pdep.v2i8(<2 x i8> %x, <2 x i8> splat (i8 -16))
+ %and = and <2 x i8> %pdep, splat (i8 15)
+ %r = icmp eq <2 x i8> %and, zeroinitializer
+ ret <2 x i1> %r
+}
+
+; pdep(x, 0b00001111): bits 4..7 of result are always zero.
+define i1 @pdep_high_bits_zero(i8 %x) nounwind {
+; CHECK-LABEL: @pdep_high_bits_zero(
+; CHECK-NEXT: ret i1 true
+;
+ %pdep = call i8 @llvm.pdep.i8(i8 %x, i8 15)
+ %and = and i8 %pdep, 240
+ %r = icmp eq i8 %and, 0
+ ret i1 %r
+}
+
+define <2 x i1> @pdep_high_bits_zero_vector(<2 x i8> %x) nounwind {
+; CHECK-LABEL: @pdep_high_bits_zero_vector(
+; CHECK-NEXT: ret <2 x i1> splat (i1 true)
+;
+ %pdep = call <2 x i8> @llvm.pdep.v2i8(<2 x i8> %x, <2 x i8> splat (i8 15))
+ %and = and <2 x i8> %pdep, splat (i8 240)
+ %r = icmp eq <2 x i8> %and, zeroinitializer
+ ret <2 x i1> %r
+}
+
+; pdep(x, 0) -> all bits zero.
+define i1 @pdep_zero_mask_known_zero(i8 %x) nounwind {
+; CHECK-LABEL: @pdep_zero_mask_known_zero(
+; CHECK-NEXT: ret i1 true
+;
+ %pdep = call i8 @llvm.pdep.i8(i8 %x, i8 0)
+ %r = icmp eq i8 %pdep, 0
+ ret i1 %r
+}
+
+define <2 x i1> @pdep_zero_mask_known_zero_vector(<2 x i8> %x) nounwind {
+; CHECK-LABEL: @pdep_zero_mask_known_zero_vector(
+; CHECK-NEXT: ret <2 x i1> splat (i1 true)
+;
+ %pdep = call <2 x i8> @llvm.pdep.v2i8(<2 x i8> %x, <2 x i8> zeroinitializer)
+ %r = icmp eq <2 x i8> %pdep, zeroinitializer
+ ret <2 x i1> %r
+}
+
+; Negative: pdep(x, 0b11001100) can have bits 2,3,6,7 set.
+; Sign bit may be set, but known bits analysis can narrow the range.
+define i1 @pdep_mask_bits_not_known(i8 %x) nounwind {
+; CHECK-LABEL: @pdep_mask_bits_not_known(
+; CHECK-NEXT: [[PDEP:%.*]] = call i8 @llvm.pdep.i8(i8 [[X:%.*]], i8 -52)
+; CHECK-NEXT: [[R:%.*]] = icmp sgt i8 [[PDEP]], -1
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %pdep = call i8 @llvm.pdep.i8(i8 %x, i8 204)
+ %and = and i8 %pdep, 128
+ %r = icmp eq i8 %and, 0
+ ret i1 %r
+}
+
+define <2 x i1> @pdep_mask_bits_not_known_vector(<2 x i8> %x) nounwind {
+; CHECK-LABEL: @pdep_mask_bits_not_known_vector(
+; CHECK-NEXT: [[PDEP:%.*]] = call <2 x i8> @llvm.pdep.v2i8(<2 x i8> [[X:%.*]], <2 x i8> splat (i8 -52))
+; CHECK-NEXT: [[R:%.*]] = icmp sgt <2 x i8> [[PDEP]], splat (i8 -1)
+; CHECK-NEXT: ret <2 x i1> [[R]]
+;
+ %pdep = call <2 x i8> @llvm.pdep.v2i8(<2 x i8> %x, <2 x i8> splat (i8 -52))
+ %and = and <2 x i8> %pdep, splat (i8 128)
+ %r = icmp eq <2 x i8> %and, zeroinitializer
+ ret <2 x i1> %r
+}
diff --git a/llvm/test/Analysis/ValueTracking/knownbits-pext.ll b/llvm/test/Analysis/ValueTracking/knownbits-pext.ll
new file mode 100644
index 0000000000000..8ad80199c18de
--- /dev/null
+++ b/llvm/test/Analysis/ValueTracking/knownbits-pext.ll
@@ -0,0 +1,91 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+
+; pext(x, 0b00001111) produces a 4-bit result; bit 4 and above are always zero.
+define i1 @pext_high_bits_zero(i8 %x) nounwind {
+; CHECK-LABEL: @pext_high_bits_zero(
+; CHECK-NEXT: ret i1 true
+;
+ %pext = call i8 @llvm.pext.i8(i8 %x, i8 15)
+ %and = and i8 %pext, 240
+ %r = icmp eq i8 %and, 0
+ ret i1 %r
+}
+
+; Vector: pext(x, 0b00001111) produces a 4-bit result; bit 4 and above are zero.
+define <2 x i1> @pext_high_bits_zero_vector(<2 x i8> %x) nounwind {
+; CHECK-LABEL: @pext_high_bits_zero_vector(
+; CHECK-NEXT: ret <2 x i1> splat (i1 true)
+;
+ %pext = call <2 x i8> @llvm.pext.v2i8(<2 x i8> %x, <2 x i8> splat (i8 15))
+ %and = and <2 x i8> %pext, splat (i8 240)
+ %r = icmp eq <2 x i8> %and, zeroinitializer
+ ret <2 x i1> %r
+}
+
+; pext(x, 0b11001100) can produce up to 4 bits; bits 4..7 are zero.
+define i1 @pext_sparse_mask_high_zero(i8 %x) nounwind {
+; CHECK-LABEL: @pext_sparse_mask_high_zero(
+; CHECK-NEXT: ret i1 true
+;
+ %pext = call i8 @llvm.pext.i8(i8 %x, i8 204)
+ %and = and i8 %pext, 240
+ %r = icmp eq i8 %and, 0
+ ret i1 %r
+}
+
+define <2 x i1> @pext_sparse_mask_high_zero_vector(<2 x i8> %x) nounwind {
+; CHECK-LABEL: @pext_sparse_mask_high_zero_vector(
+; CHECK-NEXT: ret <2 x i1> splat (i1 true)
+;
+ %pext = call <2 x i8> @llvm.pext.v2i8(<2 x i8> %x, <2 x i8> splat (i8 -52))
+ %and = and <2 x i8> %pext, splat (i8 240)
+ %r = icmp eq <2 x i8> %and, zeroinitializer
+ ret <2 x i1> %r
+}
+
+; pext(x, 0) -> all bits zero: the result is always zero.
+define i1 @pext_zero_mask_known_zero(i8 %x) nounwind {
+; CHECK-LABEL: @pext_zero_mask_known_zero(
+; CHECK-NEXT: ret i1 true
+;
+ %pext = call i8 @llvm.pext.i8(i8 %x, i8 0)
+ %r = icmp eq i8 %pext, 0
+ ret i1 %r
+}
+
+define <2 x i1> @pext_zero_mask_known_zero_vector(<2 x i8> %x) nounwind {
+; CHECK-LABEL: @pext_zero_mask_known_zero_vector(
+; CHECK-NEXT: ret <2 x i1> splat (i1 true)
+;
+ %pext = call <2 x i8> @llvm.pext.v2i8(<2 x i8> %x, <2 x i8> zeroinitializer)
+ %r = icmp eq <2 x i8> %pext, zeroinitializer
+ ret <2 x i1> %r
+}
+
+; Negative: pext(x, 0b11001100) can have bits 0-3 set; can't prove bit 0 is zero.
+define i1 @pext_low_bits_not_known(i8 %x) nounwind {
+; CHECK-LABEL: @pext_low_bits_not_known(
+; CHECK-NEXT: [[PEXT:%.*]] = call i8 @llvm.pext.i8(i8 [[X:%.*]], i8 -52)
+; CHECK-NEXT: [[AND:%.*]] = and i8 [[PEXT]], 1
+; CHECK-NEXT: [[R:%.*]] = icmp eq i8 [[AND]], 0
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %pext = call i8 @llvm.pext.i8(i8 %x, i8 204)
+ %and = and i8 %pext, 1
+ %r = icmp eq i8 %and, 0
+ ret i1 %r
+}
+
+define <2 x i1> @pext_low_bits_not_known_vector(<2 x i8> %x) nounwind {
+; CHECK-LABEL: @pext_low_bits_not_known_vector(
+; CHECK-NEXT: [[PEXT:%.*]] = call <2 x i8> @llvm.pext.v2i8(<2 x i8> [[X:%.*]], <2 x i8> splat (i8 -52))
+; CHECK-NEXT: [[AND:%.*]] = and <2 x i8> [[PEXT]], splat (i8 1)
+; CHECK-NEXT: [[R:%.*]] = icmp eq <2 x i8> [[AND]], zeroinitializer
+; CHECK-NEXT: ret <2 x i1> [[R]]
+;
+ %pext = call <2 x i8> @llvm.pext.v2i8(<2 x i8> %x, <2 x i8> splat (i8 -52))
+ %and = and <2 x i8> %pext, splat (i8 1)
+ %r = icmp eq <2 x i8> %and, zeroinitializer
+ ret <2 x i1> %r
+}
diff --git a/llvm/unittests/Support/KnownBitsTest.cpp b/llvm/unittests/Support/KnownBitsTest.cpp
index 581efad2aee8b..f6f5505c67e68 100644
--- a/llvm/unittests/Support/KnownBitsTest.cpp
+++ b/llvm/unittests/Support/KnownBitsTest.cpp
@@ -654,6 +654,8 @@ TEST(KnownBitsTest, BinaryExhaustive) {
testBinaryOpExhaustive("avgCeilS", KnownBits::avgCeilS, APIntOps::avgCeilS);
testBinaryOpExhaustive("clmul", KnownBits::clmul, APIntOps::clmul);
+ testBinaryOpExhaustive("pext", KnownBits::pext, APIntOps::compressBits);
+ testBinaryOpExhaustive("pdep", KnownBits::pdep, APIntOps::expandBits);
}
TEST(KnownBitsTest, UnaryExhaustive) {
More information about the llvm-commits
mailing list