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

[Fedor Nikolaev via llvm-commits]

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

Done

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