[llvm] [DAG] isKnownToBeAPowerOfTwo - add DemandedElts, OrZero arguments and support for zero constants (PR #181485)

Simon Pilgrim via llvm-commits llvm-commits at lists.llvm.org
Sun Feb 15 02:57:17 PST 2026


https://github.com/RKSimon updated https://github.com/llvm/llvm-project/pull/181485

>From b4d36fa4196c94b498105e208dd13b905eb59593 Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev at redking.me.uk>
Date: Sat, 14 Feb 2026 16:05:10 +0000
Subject: [PATCH 1/3] [DAG] isKnownToBeAPowerOfTwo - add DemandedElts, OrZero
 arguments and support for zero constants

Update the signature for isKnownToBeAPowerOfTwo methods to take DemandedElts and OrZero arguments.

So far all I've done is add OrZero support for constants (scalars/buildvector/splat), with the intention to incrementally extend support in the future (or spread the work as beginner GFI patches).

Similarly the DemandedElts argument is currently only used for constant build vector tests but support can be extended in future patches.

Fixes #178938
---
 llvm/include/llvm/CodeGen/SelectionDAG.h      | 16 +++-
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 83 ++++++++++++++-----
 .../AArch64/AArch64SelectionDAGTest.cpp       | 46 ++++++++++
 3 files changed, 120 insertions(+), 25 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/SelectionDAG.h b/llvm/include/llvm/CodeGen/SelectionDAG.h
index 9b4b57099889a..d3c7f6285bd36 100644
--- a/llvm/include/llvm/CodeGen/SelectionDAG.h
+++ b/llvm/include/llvm/CodeGen/SelectionDAG.h
@@ -2243,8 +2243,20 @@ class SelectionDAG {
 
   /// Test if the given value is known to have exactly one bit set. This differs
   /// from computeKnownBits in that it doesn't necessarily determine which bit
-  /// is set.
-  LLVM_ABI bool isKnownToBeAPowerOfTwo(SDValue Val, unsigned Depth = 0) const;
+  /// is set. If 'OrZero' is set, then return true if the given value is either
+  /// a power of two or zero.
+  LLVM_ABI bool isKnownToBeAPowerOfTwo(SDValue Val, bool OrZero = false,
+                                       unsigned Depth = 0) const;
+
+  /// Test if the given value is known to have exactly one bit set. This differs
+  /// from computeKnownBits in that it doesn't necessarily determine which bit
+  /// is set. The DemandedElts argument allows
+  /// us to only collect the minimum sign bits of the requested vector
+  /// elements.If 'OrZero' is set, then return true if the given value is either
+  /// a power of two or zero.
+  LLVM_ABI bool isKnownToBeAPowerOfTwo(SDValue Val, const APInt &DemandedElts,
+                                       bool OrZero = false,
+                                       unsigned Depth = 0) const;
 
   /// Test if the given _fp_ value is known to be an integer power-of-2, either
   /// positive or negative.
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 2771e96931f19..8f1a3c4cfd2ca 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -4640,17 +4640,42 @@ SelectionDAG::computeOverflowForSignedMul(SDValue N0, SDValue N1) const {
   return OFK_Sometime;
 }
 
-bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val, unsigned Depth) const {
+bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val, bool OrZero,
+                                          unsigned Depth) const {
+  EVT VT = Val.getValueType();
+
+  // Since the number of lanes in a scalable vector is unknown at compile time,
+  // we track one bit which is implicitly broadcast to all lanes.  This means
+  // that all lanes in a scalable vector are considered demanded.
+  APInt DemandedElts = VT.isFixedLengthVector()
+                           ? APInt::getAllOnes(VT.getVectorNumElements())
+                           : APInt(1, 1);
+
+  return isKnownToBeAPowerOfTwo(Val, DemandedElts, OrZero, Depth);
+}
+
+bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val,
+                                          const APInt &DemandedElts,
+                                          bool OrZero, unsigned Depth) const {
   if (Depth >= MaxRecursionDepth)
     return false; // Limit search depth.
 
   EVT OpVT = Val.getValueType();
   unsigned BitWidth = OpVT.getScalarSizeInBits();
+  unsigned NumElts = DemandedElts.getBitWidth();
+  assert((!OpVT.isScalableVector() || NumElts == 1) &&
+         "DemandedElts for scalable vectors must be 1 to represent all lanes");
+  assert(
+      (!OpVT.isFixedLengthVector() || NumElts == OpVT.getVectorNumElements()) &&
+      "Unexpected vector size");
 
-  // Is the constant a known power of 2?
-  if (ISD::matchUnaryPredicate(Val, [BitWidth](ConstantSDNode *C) {
-        return C->getAPIntValue().zextOrTrunc(BitWidth).isPowerOf2();
-      }))
+  auto IsPowerOfTwoOrZero = [BitWidth, OrZero](ConstantSDNode *C) {
+    APInt V = C->getAPIntValue().zextOrTrunc(BitWidth);
+    return (OrZero && V.isZero()) || V.isPowerOf2();
+  };
+
+  // Is the constant a known power of 2 or zero?
+  if (ISD::matchUnaryPredicate(Val, IsPowerOfTwoOrZero))
     return true;
 
   // A left-shift of a constant one will have exactly one bit set because
@@ -4659,7 +4684,8 @@ bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val, unsigned Depth) const {
     auto *C = isConstOrConstSplat(Val.getOperand(0));
     if (C && C->getAPIntValue() == 1)
       return true;
-    return isKnownToBeAPowerOfTwo(Val.getOperand(0), Depth + 1) &&
+    return isKnownToBeAPowerOfTwo(Val.getOperand(0), /*OrZero=*/false,
+                                  Depth + 1) &&
            isKnownNeverZero(Val, Depth);
   }
 
@@ -4669,42 +4695,51 @@ bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val, unsigned Depth) const {
     auto *C = isConstOrConstSplat(Val.getOperand(0));
     if (C && C->getAPIntValue().isSignMask())
       return true;
-    return isKnownToBeAPowerOfTwo(Val.getOperand(0), Depth + 1) &&
+    return isKnownToBeAPowerOfTwo(Val.getOperand(0), /*OrZero=*/false,
+                                  Depth + 1) &&
            isKnownNeverZero(Val, Depth);
   }
 
   if (Val.getOpcode() == ISD::ROTL || Val.getOpcode() == ISD::ROTR)
-    return isKnownToBeAPowerOfTwo(Val.getOperand(0), Depth + 1);
-
-  // Are all operands of a build vector constant powers of two?
-  if (Val.getOpcode() == ISD::BUILD_VECTOR)
-    if (llvm::all_of(Val->ops(), [BitWidth](SDValue E) {
-          if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(E))
-            return C->getAPIntValue().zextOrTrunc(BitWidth).isPowerOf2();
-          return false;
-        }))
+    return isKnownToBeAPowerOfTwo(Val.getOperand(0), /*OrZero=*/false,
+                                  Depth + 1);
+
+  // Are all operands of a build vector constant powers of two or zero?
+  if (Val.getOpcode() == ISD::BUILD_VECTOR) {
+    bool AllPowersOfTwoOrZero = true;
+    for (unsigned I = 0; I != NumElts && AllPowersOfTwoOrZero; ++I)
+      if (DemandedElts[I]) {
+        auto *C = dyn_cast<ConstantSDNode>(Val.getOperand(I));
+        AllPowersOfTwoOrZero &= C && IsPowerOfTwoOrZero(C);
+      }
+    if (AllPowersOfTwoOrZero)
       return true;
+  }
 
   // Is the operand of a splat vector a constant power of two?
   if (Val.getOpcode() == ISD::SPLAT_VECTOR)
     if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(Val->getOperand(0)))
-      if (C->getAPIntValue().zextOrTrunc(BitWidth).isPowerOf2())
+      if (IsPowerOfTwoOrZero(C))
         return true;
 
   // vscale(power-of-two) is a power-of-two for some targets
   if (Val.getOpcode() == ISD::VSCALE &&
       getTargetLoweringInfo().isVScaleKnownToBeAPowerOfTwo() &&
-      isKnownToBeAPowerOfTwo(Val.getOperand(0), Depth + 1))
+      isKnownToBeAPowerOfTwo(Val.getOperand(0), /*OrZero=*/false, Depth + 1))
     return true;
 
   if (Val.getOpcode() == ISD::SMIN || Val.getOpcode() == ISD::SMAX ||
       Val.getOpcode() == ISD::UMIN || Val.getOpcode() == ISD::UMAX)
-    return isKnownToBeAPowerOfTwo(Val.getOperand(1), Depth + 1) &&
-           isKnownToBeAPowerOfTwo(Val.getOperand(0), Depth + 1);
+    return isKnownToBeAPowerOfTwo(Val.getOperand(1), /*OrZero=*/false,
+                                  Depth + 1) &&
+           isKnownToBeAPowerOfTwo(Val.getOperand(0), /*OrZero=*/false,
+                                  Depth + 1);
 
   if (Val.getOpcode() == ISD::SELECT || Val.getOpcode() == ISD::VSELECT)
-    return isKnownToBeAPowerOfTwo(Val.getOperand(2), Depth + 1) &&
-           isKnownToBeAPowerOfTwo(Val.getOperand(1), Depth + 1);
+    return isKnownToBeAPowerOfTwo(Val.getOperand(2), /*OrZero=*/false,
+                                  Depth + 1) &&
+           isKnownToBeAPowerOfTwo(Val.getOperand(1), /*OrZero=*/false,
+                                  Depth + 1);
 
   // Looking for `x & -x` pattern:
   // If x == 0:
@@ -4712,12 +4747,14 @@ bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val, unsigned Depth) const {
   // If x != 0:
   //    x & -x -> non-zero pow2
   // so if we find the pattern return whether we know `x` is non-zero.
+  // TODO OrZero handling
   SDValue X;
   if (sd_match(Val, m_And(m_Value(X), m_Neg(m_Deferred(X)))))
     return isKnownNeverZero(X, Depth);
 
   if (Val.getOpcode() == ISD::ZERO_EXTEND)
-    return isKnownToBeAPowerOfTwo(Val.getOperand(0), Depth + 1);
+    return isKnownToBeAPowerOfTwo(Val.getOperand(0), /*OrZero=*/false,
+                                  Depth + 1);
 
   // More could be done here, though the above checks are enough
   // to handle some common cases.
diff --git a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
index e8c46a21d8e72..df73ef3534e7b 100644
--- a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
+++ b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
@@ -868,6 +868,52 @@ TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_VSHL) {
   EXPECT_EQ(Known.One, APInt(8, 0x80));
 }
 
+// Piggy-backing on the AArch64 tests to verify
+// SelectionDAG::KnownToBeAPowerOfTwo.
+TEST_F(AArch64SelectionDAGTest, KnownToBeAPowerOfTwo_Constants) {
+  SDLoc Loc;
+  auto Cst0 = DAG->getConstant(0, Loc, MVT::i32);
+  auto Cst4 = DAG->getConstant(4, Loc, MVT::i32);
+  auto CstBig = DAG->getConstant(2 << 17, Loc, MVT::i32);
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Cst0));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Cst0, /*OrZero=*/true));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Cst4));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(CstBig));
+
+  auto VecVT = MVT::v2i16;
+  auto Vec04 = DAG->getBuildVector(VecVT, Loc, {Cst0, Cst4});
+  auto Vec44 = DAG->getBuildVector(VecVT, Loc, {Cst4, Cst4});
+  auto Vec4Big = DAG->getBuildVector(VecVT, Loc, {Cst4, CstBig});
+  auto Vec0Big = DAG->getBuildVector(VecVT, Loc, {Cst0, CstBig});
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec04));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Vec04, /*OrZero=*/true));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Vec44));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Vec44, /*OrZero=*/true));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec4Big));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Vec4Big, /*OrZero=*/true));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec0Big));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Vec0Big, /*OrZero=*/true));
+
+  APInt DemandLo(2, 1);
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec04, DemandLo));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Vec04, DemandLo, /*OrZero=*/true));
+
+  APInt DemandHi(2, 2);
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Vec04, DemandHi));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Vec04, DemandHi, /*OrZero=*/true));
+
+  auto SplatVT = MVT::nxv2i16;
+  auto Splat0 = DAG->getSplat(SplatVT, Loc, Cst0);
+  auto Splat4 = DAG->getSplat(SplatVT, Loc, Cst4);
+  auto SplatBig = DAG->getSplat(SplatVT, Loc, CstBig);
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Splat0));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Splat4, /*OrZero=*/true));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Splat4));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Splat4, /*OrZero=*/true));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(SplatBig));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SplatBig, /*OrZero=*/true));
+}
+
 TEST_F(AArch64SelectionDAGTest, isSplatValue_Fixed_BUILD_VECTOR) {
   TargetLowering TL(*TM, *STI);
 

>From 198cff58c2a268b4492a6fe3c5dfc6a1e3229454 Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev at redking.me.uk>
Date: Sun, 15 Feb 2026 09:42:01 +0000
Subject: [PATCH 2/3] Fix comment whitespace

---
 llvm/include/llvm/CodeGen/SelectionDAG.h | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/SelectionDAG.h b/llvm/include/llvm/CodeGen/SelectionDAG.h
index d3c7f6285bd36..c45dec89b4b11 100644
--- a/llvm/include/llvm/CodeGen/SelectionDAG.h
+++ b/llvm/include/llvm/CodeGen/SelectionDAG.h
@@ -2250,10 +2250,9 @@ class SelectionDAG {
 
   /// Test if the given value is known to have exactly one bit set. This differs
   /// from computeKnownBits in that it doesn't necessarily determine which bit
-  /// is set. The DemandedElts argument allows
-  /// us to only collect the minimum sign bits of the requested vector
-  /// elements.If 'OrZero' is set, then return true if the given value is either
-  /// a power of two or zero.
+  /// is set. The DemandedElts argument allows us to only collect the minimum
+  /// sign bits of the requested vector elements. If 'OrZero' is set, then
+  /// return true if the given value is either a power of two or zero.
   LLVM_ABI bool isKnownToBeAPowerOfTwo(SDValue Val, const APInt &DemandedElts,
                                        bool OrZero = false,
                                        unsigned Depth = 0) const;

>From 3428e0d1c2a5314b7a27ef892fb267ad7e37200e Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev at redking.me.uk>
Date: Sun, 15 Feb 2026 10:27:34 +0000
Subject: [PATCH 3/3] Use all_of/enumerate for BUILD_VECTOR iteration

---
 llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 8f1a3c4cfd2ca..b87bc4b18ebc9 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -4669,7 +4669,7 @@ bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val,
       (!OpVT.isFixedLengthVector() || NumElts == OpVT.getVectorNumElements()) &&
       "Unexpected vector size");
 
-  auto IsPowerOfTwoOrZero = [BitWidth, OrZero](ConstantSDNode *C) {
+  auto IsPowerOfTwoOrZero = [BitWidth, OrZero](const ConstantSDNode *C) {
     APInt V = C->getAPIntValue().zextOrTrunc(BitWidth);
     return (OrZero && V.isZero()) || V.isPowerOf2();
   };
@@ -4705,16 +4705,12 @@ bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val,
                                   Depth + 1);
 
   // Are all operands of a build vector constant powers of two or zero?
-  if (Val.getOpcode() == ISD::BUILD_VECTOR) {
-    bool AllPowersOfTwoOrZero = true;
-    for (unsigned I = 0; I != NumElts && AllPowersOfTwoOrZero; ++I)
-      if (DemandedElts[I]) {
-        auto *C = dyn_cast<ConstantSDNode>(Val.getOperand(I));
-        AllPowersOfTwoOrZero &= C && IsPowerOfTwoOrZero(C);
-      }
-    if (AllPowersOfTwoOrZero)
-      return true;
-  }
+  if (Val.getOpcode() == ISD::BUILD_VECTOR &&
+      all_of(enumerate(Val->ops()), [&](auto P) {
+        auto *C = dyn_cast<ConstantSDNode>(P.value());
+        return !DemandedElts[P.index()] || (C && IsPowerOfTwoOrZero(C));
+      }))
+    return true;
 
   // Is the operand of a splat vector a constant power of two?
   if (Val.getOpcode() == ISD::SPLAT_VECTOR)



More information about the llvm-commits mailing list