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

[Craig Topper via llvm-commits]

================
@@ -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);
----------------
topperc wrote:

The constant used for Vec4Big is too big for an i16 element. So it gets truncated to 0.

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