Repository: llvm/llvm-project
File: llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp (C++)

Code before:

  APInt DemandHi(2, 2);
  EXPECT_FALSE(DAG->isKnownNeverZero(VSelect444Big, DemandHi));
  EXPECT_TRUE(DAG->isKnownNeverZero(VSelect0444, DemandHi));

  APInt DemandAll(2, 3);
  EXPECT_FALSE(DAG->isKnownNeverZero(VSelect0444, DemandAll));
  EXPECT_FALSE(DAG->isKnownNeverZero(VSelect040Big, DemandAll));
  EXPECT_FALSE(DAG->isKnownNeverZero(VSelect444Big, DemandAll));
  EXPECT_TRUE(DAG->isKnownNeverZero(VSelect4444, DemandAll));
}

TEST_F(AArch64SelectionDAGTest, isKnownNeverZero_OR_DemandedElts) {
  SDLoc Loc;
  
  EVT VT = EVT::getVectorVT(Context, MVT::i32, 4);

  // Operand A: <0, 1, 0, 0>  (Lane 1 is Non-Zero!)
  SDValue A = DAG->getConstant(0, Loc, MVT::i32);
  SDValue NonZero = DAG->getConstant(1, Loc, MVT::i32);
  SDValue VecA = DAG->getBuildVector(VT, Loc, {A, NonZero, A, A});

  // Operand B: <0, 0, 0, 0> (All Zero)
  SDValue VecB = DAG->getConstant(0, Loc, VT);

  SDValue Op = DAG->getNode(ISD::OR, Loc, VT, VecA, VecB);

  // Lane 1 Only : Safe
  APInt DemandedLane1(4, 2);    
  EXPECT_TRUE(DAG->isKnownNeverZero(Op, DemandedLane1));

  // Lane 0 and 1 : Unsafe
  APInt DemandedLane2(4, 3);    
  EXPECT_FALSE(DAG->isKnownNeverZero(Op, DemandedLane2));
}
} // end namespace llvm


Review comment (bug):
I suspect this will constant fold and return a new build_vector rather than an OR node. Does this test fail without your other changes?

Code after:

  APInt DemandLo(2, 1);
  EXPECT_FALSE(DAG->isKnownNeverZero(Vec04, DemandLo));

  APInt DemandHi(2, 2);
  EXPECT_TRUE(DAG->isKnownNeverZero(Vec04, DemandHi));

  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);
  EXPECT_FALSE(DAG->isKnownNeverZero(Splat0));
  EXPECT_TRUE(DAG->isKnownNeverZero(Splat4));
  EXPECT_FALSE(DAG->isKnownNeverZero(SplatBig));
}

TEST_F(AArch64SelectionDAGTest, KnownNeverZero_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,
                                  Register::index2VirtReg(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->isKnownNeverZero(Select40));
  EXPECT_FALSE(DAG->isKnownNeverZero(Select40));
  EXPECT_TRUE(DAG->isKnownNeverZero(Select43));
  EXPECT_TRUE(DAG->isKnownNeverZero(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,
                                     Register::index2VirtReg(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);