[llvm] Improved ISD::SRL handling in isKnownToBeAPowerOfTwo (PR #182562)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-selectiondag

Author: None (Iasonaskrpr)

<details>
<summary>Changes</summary>

Fixes #<!-- -->181651

Added DemandedElts argument to isConstOrConstSplat and to isKnowTobePowerOfTwo calls and OrZero || isKnownNeverZero(Val, Depth) is checked before isKnowTobePowerOfTwo. Also added unit tests.

---
Full diff: https://github.com/llvm/llvm-project/pull/182562.diff


2 Files Affected:

- (modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp (+4-4) 
- (modified) llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp (+74) 


``````````diff
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 581553d41cb6d..35a1b89448f59 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -4723,12 +4723,12 @@ bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val,
   case ISD::SRL: {
     // A logical right-shift of a constant sign-bit will have exactly
     // one bit set.
-    auto *C = isConstOrConstSplat(Val.getOperand(0));
+    auto *C = isConstOrConstSplat(Val.getOperand(0), DemandedElts);
     if (C && C->getAPIntValue().isSignMask())
       return true;
-    return isKnownToBeAPowerOfTwo(Val.getOperand(0), /*OrZero=*/false,
-                                  Depth + 1) &&
-           isKnownNeverZero(Val, Depth);
+    return (OrZero || isKnownNeverZero(Val, Depth)) &&
+           isKnownToBeAPowerOfTwo(Val.getOperand(0), DemandedElts, OrZero,
+                                  Depth + 1);
   }
 
   case ISD::ROTL:
diff --git a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
index b2b8dfb0c21fc..705b370206529 100644
--- a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
+++ b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
@@ -971,6 +971,80 @@ TEST_F(AArch64SelectionDAGTest, KnownToBeAPowerOfTwo_Select) {
       DAG->isKnownToBeAPowerOfTwo(VSelect444Big, DemandAll, /*OrZero=*/true));
 }
 
+TEST_F(AArch64SelectionDAGTest, KnownToBeAPowerOfTwo_ISD_SRL) {
+  SDLoc Loc;
+  auto Cst0 = DAG->getConstant(0, Loc, MVT::i32);
+  auto Cst1 = DAG->getConstant(1, Loc, MVT::i32);
+  auto Cst4 = DAG->getConstant(4, Loc, MVT::i32);
+  auto Cst7 = DAG->getConstant(7, Loc, MVT::i32);
+  auto CstBig = DAG->getConstant(2 << 17, Loc, MVT::i32);
+
+  auto SRL40 = DAG->getNode(ISD::SRL, Loc, MVT::i32, Cst4, Cst0);
+  auto SRL41 = DAG->getNode(ISD::SRL, Loc, MVT::i32, Cst4, Cst1);
+  auto SRL44 = DAG->getNode(ISD::SRL, Loc, MVT::i32, Cst4, Cst4);
+  auto SRL71 = DAG->getNode(ISD::SRL, Loc, MVT::i32, Cst7, Cst1);
+  auto SRLBig4 = DAG->getNode(ISD::SRL, Loc, MVT::i32, CstBig, Cst4);
+
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRL40));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRL41));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(SRL44));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRL44, /*OrZero=*/true));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(SRL71));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRLBig4));
+
+  auto VecVT = MVT::v2i16;
+  auto Vec00 = DAG->getBuildVector(VecVT, Loc, {Cst0, Cst0});
+  auto Vec04 = DAG->getBuildVector(VecVT, Loc, {Cst0, Cst4});
+  auto Vec11 = DAG->getBuildVector(VecVT, Loc, {Cst1, Cst1});
+  auto Vec44 = DAG->getBuildVector(VecVT, Loc, {Cst4, Cst4});
+  auto Vec74 = DAG->getBuildVector(VecVT, Loc, {Cst7, Cst4});
+  auto Vec4Big = DAG->getBuildVector(VecVT, Loc, {Cst4, CstBig});
+
+  auto SRL4404 = DAG->getNode(ISD::SRL, Loc, VecVT, Vec44, Vec04);
+  auto SRL7411 = DAG->getNode(ISD::SRL, Loc, VecVT, Vec74, Vec11);
+  auto SRL4Big04 = DAG->getNode(ISD::SRL, Loc, VecVT, Vec4Big, Vec04);
+  auto SRL4Big00 = DAG->getNode(ISD::SRL, Loc, VecVT, Vec4Big, Vec00);
+
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(SRL4404));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRL4404, /*OrZero=*/true));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(SRL7411));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(SRL7411, /*OrZero=*/true));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(SRL4Big00));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRL4Big00, /*OrZero=*/true));
+
+  APInt DemandLo(2, 1);
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(SRL7411, DemandLo));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(SRL7411, DemandLo, /*OrZero=*/true));
+
+  APInt DemandHi(2, 2);
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRL7411, DemandHi));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRL7411, DemandHi, /*OrZero=*/true));
+
+  APInt DemandAll(2, 3);
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRL4Big04, DemandHi));
+  EXPECT_TRUE(
+      DAG->isKnownToBeAPowerOfTwo(SRL4Big04, 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);
+
+  auto SRLSplatBig0 = DAG->getNode(ISD::SRL, Loc, VecVT, SplatBig, Splat0);
+  auto SRLSplatBig4 = DAG->getNode(ISD::SRL, Loc, VecVT, SplatBig, Splat4);
+  auto SRLSplat44 = DAG->getNode(ISD::SRL, Loc, VecVT, Splat4, Splat4);
+  auto SRLSplat40 = DAG->getNode(ISD::SRL, Loc, VecVT, Splat4, Splat0);
+
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(SRLSplatBig0));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRLSplatBig0, /*OrZero=*/true));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRLSplatBig4));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRLSplatBig4, /*OrZero=*/true));
+  EXPECT_FALSE(DAG->isKnownToBeAPowerOfTwo(SRLSplat44));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRLSplat44, /*OrZero=*/true));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRLSplat40));
+  EXPECT_TRUE(DAG->isKnownToBeAPowerOfTwo(SRLSplat40, /*OrZero=*/true));
+}
+
 TEST_F(AArch64SelectionDAGTest, isSplatValue_Fixed_BUILD_VECTOR) {
   TargetLowering TL(*TM, *STI);
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/182562