[llvm] [DAG] isKnownToBeAPowerOfTwo - add DemandedElts + OrZero handling to … (PR #181753)

[Fedor Nikolaev via llvm-commits]

https://github.com/felichita created https://github.com/llvm/llvm-project/pull/181753

[DAG] isKnownToBeAPowerOfTwo - add DemandedElts + OrZero handling to ISD::SELECT/VSELECT cases

Pass DemandedElts and OrZero arguments through SELECT/VSELECT nodes when checking if values are power-of-two.

Add unit tests for scalar SELECT (OrZero) and vector VSELECT (DemandedElts + OrZero).

Resolves [#181645](https://github.com/llvm/llvm-project/issues/181645)

>From e001b74d05bd181cd3e8c43e1058a23715df7aa6 Mon Sep 17 00:00:00 2001
From: Fedor Nikolaev <fridrixnm at gmail.com>
Date: Mon, 16 Feb 2026 22:29:06 +0100
Subject: [PATCH] [DAG] isKnownToBeAPowerOfTwo - add DemandedElts + OrZero
 handling to ISD::SELECT/VSELECT cases

Pass DemandedElts and OrZero arguments through SELECT/VSELECT nodes when
checking if values are power-of-two.

Add unit tests for scalar SELECT (OrZero) and vector VSELECT (DemandedElts + OrZero).

Resolves [#181645](https://github.com/llvm/llvm-project/issues/181645)
---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp |  4 +-
 .../AArch64/AArch64SelectionDAGTest.cpp       | 38 +++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 90a323166b76e..597036fa597de 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -4747,9 +4747,9 @@ bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val,
 
   case ISD::SELECT:
   case ISD::VSELECT:
-    return isKnownToBeAPowerOfTwo(Val.getOperand(2), /*OrZero=*/false,
+    return isKnownToBeAPowerOfTwo(Val.getOperand(2), DemandedElts, OrZero,
                                   Depth + 1) &&
-           isKnownToBeAPowerOfTwo(Val.getOperand(1), /*OrZero=*/false,
+           isKnownToBeAPowerOfTwo(Val.getOperand(1), DemandedElts, OrZero,
                                   Depth + 1);
 
   case ISD::ZERO_EXTEND:
diff --git a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
index df73ef3534e7b..1d1b21b912281 100644
--- a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
+++ b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
@@ -873,6 +873,7 @@ TEST_F(AArch64SelectionDAGTest, ComputeKnownBits_VSHL) {
 TEST_F(AArch64SelectionDAGTest, KnownToBeAPowerOfTwo_Constants) {
   SDLoc Loc;
   auto Cst0 = DAG->getConstant(0, Loc, MVT::i32);
+  auto Cst3 = DAG->getConstant(3, 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));
@@ -912,6 +913,43 @@ TEST_F(AArch64SelectionDAGTest, KnownToBeAPowerOfTwo_Constants) {
   EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Splat4, /*OrZero=*/true));
   EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(SplatBig));
   EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SplatBig, /*OrZero=*/true));
+
+  auto Cond = DAG->getCopyFromReg(DAG->getEntryNode(), Loc, 1, MVT::i1);
+  auto Select40 = DAG->getNode(ISD::SELECT, Loc, MVT::i32,
+                               Cond, Cst4, Cst0);
+  auto Select43 = DAG->getNode(ISD::SELECT, Loc, MVT::i32,
+                               Cond, Cst4, Cst3);
+  auto Select4Big = DAG->getNode(ISD::SELECT, Loc, MVT::i32,
+                                 Cond, Cst4, CstBig);
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Select40));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Select40, /*OrZero=*/true));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Select43));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Select4Big));
+
+  auto VecCond = DAG->getCopyFromReg(DAG->getEntryNode(), Loc, 2, MVT::v2i1);
+  auto VSelect0444 = DAG->getNode(ISD::VSELECT, Loc, VecVT,
+                                  VecCond, Vec04, Vec44);
+  auto VSelect4444 = DAG->getNode(ISD::VSELECT, Loc, VecVT,
+                                  VecCond, Vec44, Vec44);
+  auto VSelect040Big = DAG->getNode(ISD::VSELECT, Loc, VecVT,
+                                    VecCond, Vec04, Vec0Big);
+  auto VSelect444Big = DAG->getNode(ISD::VSELECT, Loc, VecVT,
+                                    VecCond, Vec44, Vec4Big);
+
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(VSelect444Big, DemandHi));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(VSelect444Big, DemandHi, true));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(VSelect0444, DemandHi));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(VSelect0444, DemandLo));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(VSelect0444, DemandLo, true));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(VSelect444Big, DemandLo));
+
+  APInt DemandAll(2, 3);
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(VSelect0444, DemandAll));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(VSelect040Big, DemandAll));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(VSelect444Big, DemandAll));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(VSelect4444, DemandAll));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(VSelect040Big, DemandAll, true));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(VSelect444Big, DemandAll, true));
 }
 
 TEST_F(AArch64SelectionDAGTest, isSplatValue_Fixed_BUILD_VECTOR) {