[llvm] [DAG] Simplify isKnownToBeAPowerOfTwo for ZERO_EXTEND (PR #182226)

via llvm-commits llvm-commits at lists.llvm.org
Tue Mar 3 13:09:55 PST 2026


https://github.com/fadedreamz updated https://github.com/llvm/llvm-project/pull/182226

>From 775f920a4e654b8fe604a867d213c35f1a219214 Mon Sep 17 00:00:00 2001
From: Mohammad Nazmul Alam <fadedreamz at gmail.com>
Date: Thu, 19 Feb 2026 00:47:43 -0500
Subject: [PATCH 1/3] [DAG] Simplify isKnownToBeAPowerOfTwo for ZERO_EXTEND

This commit utilizes already computed DemandedElts & OrZero flag
to call isKnownToBeAPowerOfTwo for ZERO_EXTEND.

Fixes: #181648
---
 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 581553d41cb6d..d2aa850225b76 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -4753,7 +4753,7 @@ bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val,
                                   Depth + 1);
 
   case ISD::ZERO_EXTEND:
-    return isKnownToBeAPowerOfTwo(Val.getOperand(0), /*OrZero=*/false,
+    return isKnownToBeAPowerOfTwo(Val.getOperand(0), DemandedElts, OrZero,
                                   Depth + 1);
 
   case ISD::VSCALE:

>From 230e2f37ffdb2b9ebf26d8a4926546af410f3fc7 Mon Sep 17 00:00:00 2001
From: Mohammad Nazmul Alam <fadedreamz at gmail.com>
Date: Sun, 1 Mar 2026 02:21:09 -0500
Subject: [PATCH 2/3] add test case

---
 .../AArch64/AArch64SelectionDAGTest.cpp       | 29 +++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
index b2b8dfb0c21fc..11292a0ada40a 100644
--- a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
+++ b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
@@ -914,6 +914,35 @@ TEST_F(AArch64SelectionDAGTest, KnownToBeAPowerOfTwo_Constants) {
   EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SplatBig, /*OrZero=*/true));
 }
 
+TEST_F(AArch64SelectionDAGTest, KnownToBeAPowerOfTwo_ZeroExtend) {
+  SDLoc Loc;
+  auto Cst0 = DAG->getConstant(0, Loc, MVT::i16);
+  auto Cst4 = DAG->getConstant(4, Loc, MVT::i16);
+  auto CstBig = DAG->getConstant(2 << 14, Loc, MVT::i16);
+  auto node0 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::i32, Cst0);
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(node0)); // this gets transformed
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(node0, /*OrZero=*/true)); // so need to pass OrZero=true
+  auto node4 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::i32, Cst4);
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(node4));
+  auto nodeBig = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::i32, CstBig);
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(nodeBig));
+  auto VecVT = MVT::v2i16;
+  auto Vec04 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::v2i32, DAG->getBuildVector(VecVT, Loc, {Cst0, Cst4}));
+  auto Vec4Big = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::v2i32, DAG->getBuildVector(VecVT, Loc, {Cst4, CstBig}));
+  auto Vec0Big = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::v2i32, DAG->getBuildVector(VecVT, Loc, {Cst0, CstBig}));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Vec04));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec4Big));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec0Big));
+
+  auto SplatVT = MVT::nxv2i16;
+  auto Splat0 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::nxv2i32, DAG->getSplat(SplatVT, Loc, Cst0));
+  auto Splat4 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::nxv2i32, DAG->getSplat(SplatVT, Loc, Cst4));
+  auto SplatBig = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::nxv2i32, DAG->getSplat(SplatVT, Loc, CstBig));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Splat0));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Splat4));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SplatBig));
+}
+
 TEST_F(AArch64SelectionDAGTest, KnownToBeAPowerOfTwo_Select) {
   SDLoc Loc;
   auto Cst0 = DAG->getConstant(0, Loc, MVT::i32);

>From c6d8e7553cebd71d1e651ad4101f312d447dca28 Mon Sep 17 00:00:00 2001
From: Mohammad Nazmul Alam <fadedreamz at gmail.com>
Date: Tue, 3 Mar 2026 16:09:37 -0500
Subject: [PATCH 3/3] update test case

---
 .../AArch64/AArch64SelectionDAGTest.cpp       | 64 ++++++++++++++-----
 1 file changed, 47 insertions(+), 17 deletions(-)

diff --git a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
index 11292a0ada40a..b0e4fd4d55ff1 100644
--- a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
+++ b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
@@ -918,29 +918,59 @@ TEST_F(AArch64SelectionDAGTest, KnownToBeAPowerOfTwo_ZeroExtend) {
   SDLoc Loc;
   auto Cst0 = DAG->getConstant(0, Loc, MVT::i16);
   auto Cst4 = DAG->getConstant(4, Loc, MVT::i16);
-  auto CstBig = DAG->getConstant(2 << 14, Loc, MVT::i16);
-  auto node0 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::i32, Cst0);
-  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(node0)); // this gets transformed
-  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(node0, /*OrZero=*/true)); // so need to pass OrZero=true
+  auto Cst6 = DAG->getConstant(6, Loc, MVT::i16);
+
+  auto node0 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::i32, Cst0); // this will be folded, and not pass through ZERO_EXTEND case.
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(node0));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(node0, /*OrZero=*/true));
   auto node4 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::i32, Cst4);
   EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(node4));
-  auto nodeBig = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::i32, CstBig);
-  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(nodeBig));
+  auto node6 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::i32, Cst6);
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(node6));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(node6, /*OrZero=*/true));
+
   auto VecVT = MVT::v2i16;
-  auto Vec04 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::v2i32, DAG->getBuildVector(VecVT, Loc, {Cst0, Cst4}));
-  auto Vec4Big = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::v2i32, DAG->getBuildVector(VecVT, Loc, {Cst4, CstBig}));
-  auto Vec0Big = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::v2i32, DAG->getBuildVector(VecVT, Loc, {Cst0, CstBig}));
-  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Vec04));
-  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec4Big));
-  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec0Big));
+  auto Vec04 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::v2i32,
+    DAG->getBuildVector(VecVT, Loc, {Cst0, Cst4}));
+  auto Vec46 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::v2i32,
+    DAG->getBuildVector(VecVT, Loc, {Cst4, Cst6}));
+  auto Vec06 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::v2i32,
+      DAG->getBuildVector(VecVT, Loc, {Cst0, Cst6}));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec04));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Vec04, /*OrZero=*/true));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec46));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec46, /*OrZero=*/true));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec06));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Vec06, /*OrZero=*/true));
 
+  // Scalable vector tests: ZERO_EXTEND of SPLAT_VECTOR of constants folds.
   auto SplatVT = MVT::nxv2i16;
-  auto Splat0 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::nxv2i32, DAG->getSplat(SplatVT, Loc, Cst0));
-  auto Splat4 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::nxv2i32, DAG->getSplat(SplatVT, Loc, Cst4));
-  auto SplatBig = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::nxv2i32, DAG->getSplat(SplatVT, Loc, CstBig));
-  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Splat0));
+  auto Splat0 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::nxv2i32,
+                              DAG->getSplat(SplatVT, Loc, Cst0));
+  auto Splat4 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::nxv2i32,
+                              DAG->getSplat(SplatVT, Loc, Cst4));
+  auto Splat6 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::nxv2i32,
+                                    DAG->getSplat(SplatVT, Loc, Cst6));
+
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Splat0));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Splat0, /*OrZero=*/true));
   EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Splat4));
-  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SplatBig));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(Splat4, /*OrZero=*/true));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Splat6));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(Splat6, /*OrZero=*/true));
+
+  auto Cst2_16 = DAG->getConstant(2, Loc, MVT::i16);
+  auto Cond = DAG->getCopyFromReg(DAG->getEntryNode(), Loc, 1, MVT::i1);
+  auto selectNode = DAG->getNode(ISD::SELECT, Loc, MVT::i16, Cond, Cst2_16, Cst4);
+  // zext of select(2,4) should still be power of 2
+  auto zextNode = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::i32, selectNode);
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(zextNode));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(zextNode, /*OrZero=*/true));
+
+  auto selectNode2 = DAG->getNode(ISD::SELECT, Loc, MVT::i16, Cond, Cst6, Cst4); // select(6,4)
+  auto zextNode2 = DAG->getNode(ISD::ZERO_EXTEND, Loc, MVT::i32, selectNode2);
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(zextNode2));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(zextNode2, /*OrZero=*/true));
 }
 
 TEST_F(AArch64SelectionDAGTest, KnownToBeAPowerOfTwo_Select) {



More information about the llvm-commits mailing list