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

[Fedor Nikolaev via llvm-commits]

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

>From f6acb34220f11f60e06f7e8f16767d5052b3bc72 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       | 57 +++++++++++++++++++
 2 files changed, 59 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..b2b8dfb0c21fc 100644
--- a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
+++ b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
@@ -914,6 +914,63 @@ TEST_F(AArch64SelectionDAGTest, KnownToBeAPowerOfTwo_Constants) {
   EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SplatBig, /*OrZero=*/true));
 }
 
+TEST_F(AArch64SelectionDAGTest, KnownToBeAPowerOfTwo_Select) {
+  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);
+
+  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 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});
+
+  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);
+
+  APInt DemandLo(2, 1);
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(VSelect0444, DemandLo));
+  EXPECT_TRUE(
+      DAG->isKnownToBeAPowerOfTwo(VSelect0444, DemandLo, /*OrZero=*/true));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(VSelect444Big, DemandLo));
+
+  APInt DemandHi(2, 2);
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(VSelect444Big, DemandHi));
+  EXPECT_TRUE(
+      DAG->isKnownToBeAPowerOfTwo(VSelect444Big, DemandHi, /*OrZero=*/true));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(VSelect0444, DemandHi));
+
+  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, /*OrZero=*/true));
+  EXPECT_TRUE(
+      DAG->isKnownToBeAPowerOfTwo(VSelect444Big, DemandAll, /*OrZero=*/true));
+}
+
 TEST_F(AArch64SelectionDAGTest, isSplatValue_Fixed_BUILD_VECTOR) {
   TargetLowering TL(*TM, *STI);