[CodeGen] Fix remaining zext() assertions in SelectionDAG

[scott.linder]

Summary

Includes remaining cases not committed in https://reviews.llvm.org/D49574

This introduces what appear to be the first unittests for SelectionDAG. Lit tests do not directly test the problem cases in these functions, so I think unittests are easier to write and audit here.

[RKSimon]

Can anyone comment if the unit test is the way to go for these? It's proving tricky to expose the potential issue with regular lit tests.

[efriedma]

I guess this approach is okay, specifically for testing the SelectionDAG analysis helpers.

unittests/CodeGen/SelectionDAGTest.cpp

39 ↗

(On Diff #160425)

This seems bad; the test silently fails to run if the AArch64 target isn't built?

[scott.linder]

unittests/CodeGen/SelectionDAGTest.cpp

39 ↗

(On Diff #160425)

This does seem bad, but as far as I can tell there is no appealing alternative. Not that it absolves me of this, but other tests (e.g. CodeGen/GlobalISel/LegalizerHelperTest.h) do the same. The core issue is not having a "generic target" always available for unittests.

Is there anything better to do in this case?

[efriedma]

unittests/CodeGen/SelectionDAGTest.cpp

39 ↗

(On Diff #160425)

How hard would it be to make a fake target for SelectionDAG unittests? A skeleton implementation of a Target, a TargetMachine, and a TargetSubtargetInfo is hopefully not that much code. But maybe it's not quite that simple...

That said, I'm not really that worried about the part where the test only runs if the aarch64 target is enabled; we have target-specific tests all over. Maybe enough to just rename SelectionDAGTest to AArch64SelectionDAGTest so someone doesn't accidentally add a new test without running it.

[scott.linder]

Update test name to be clear it only runs when aarch64 is available.

I added a FIXME about creating a fake target for unittests, but I do not have the experience right now to know how much work would be involved, and whether a fake target is acceptable to add to begin with.

[efriedma]

LGTM

[thakis]

llvm/trunk/unittests/CodeGen/AArch64SelectionDAGTest.cpp

164

Are you intentionally adding a main() here? add_llvm_unittests should provide one already, no?

[scott.linder]

llvm/trunk/unittests/CodeGen/AArch64SelectionDAGTest.cpp

164

It was intentional, but probably better to move the Initialize* calls into a SetUpTestCase instead. I will update in an NFC commit.

[scott.linder]

llvm/trunk/unittests/CodeGen/AArch64SelectionDAGTest.cpp

164

Fixed in r341574

llvm/trunk/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

case ISD::LOAD: {
     } else if (const MDNode *Ranges = LD->getRanges()) {
       if (LD->getExtensionType() == ISD::NON_EXTLOAD)
         computeKnownBitsFromRangeMetadata(*Ranges, Known);
     }
     break;
   }
   case ISD::ZERO_EXTEND_VECTOR_INREG: {
     EVT InVT = Op.getOperand(0).getValueType();
-    APInt InDemandedElts = DemandedElts.zext(InVT.getVectorNumElements());
+    APInt InDemandedElts = DemandedElts.zextOrSelf(InVT.getVectorNumElements());
     Known = computeKnownBits(Op.getOperand(0), InDemandedElts, Depth + 1);
     Known = Known.zext(BitWidth);
     Known.Zero.setBitsFrom(InVT.getScalarSizeInBits());
     break;
   }
   case ISD::ZERO_EXTEND: {
     EVT InVT = Op.getOperand(0).getValueType();
     Known = computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
case ISD::SIGN_EXTEND_INREG:
     // Max of the input and what this extends.
     Tmp = cast<VTSDNode>(Op.getOperand(1))->getVT().getScalarSizeInBits();
     Tmp = VTBits-Tmp+1;
     Tmp2 = ComputeNumSignBits(Op.getOperand(0), DemandedElts, Depth+1);
     return std::max(Tmp, Tmp2);
   case ISD::SIGN_EXTEND_VECTOR_INREG: {
     SDValue Src = Op.getOperand(0);
     EVT SrcVT = Src.getValueType();
-    APInt DemandedSrcElts = DemandedElts.zext(SrcVT.getVectorNumElements());
+    APInt DemandedSrcElts = DemandedElts.zextOrSelf(SrcVT.getVectorNumElements());
     Tmp = VTBits - SrcVT.getScalarSizeInBits();
     return ComputeNumSignBits(Src, DemandedSrcElts, Depth+1) + Tmp;
   }
 
   case ISD::SRA:
     Tmp = ComputeNumSignBits(Op.getOperand(0), DemandedElts, Depth+1);
     // SRA X, C   -> adds C sign bits.
     if (ConstantSDNode *C =

llvm/trunk/lib/CodeGen/SelectionDAG/TargetLowering.cpp

case ISD::INSERT_SUBVECTOR: {
     if (SimplifyDemandedVectorElts(Base, BaseElts, KnownUndef, KnownZero, TLO,
                                    Depth + 1))
       return true;
     KnownUndef.insertBits(SubUndef, SubIdx);
     KnownZero.insertBits(SubZero, SubIdx);
     break;
   }
   case ISD::EXTRACT_SUBVECTOR: {
-    if (!isa<ConstantSDNode>(Op.getOperand(1)))
-      break;
     SDValue Src = Op.getOperand(0);
+    ConstantSDNode *SubIdx = dyn_cast<ConstantSDNode>(Op.getOperand(1));
     unsigned NumSrcElts = Src.getValueType().getVectorNumElements();
-    const APInt& Idx = cast<ConstantSDNode>(Op.getOperand(1))->getAPIntValue();
-    if (Idx.uge(NumSrcElts - NumElts))
-      break;
+    if (SubIdx && SubIdx->getAPIntValue().ule(NumSrcElts - NumElts)) {
       // Offset the demanded elts by the subvector index.
-    uint64_t SubIdx = Idx.getZExtValue();
-    APInt SrcElts = DemandedElts.zext(NumSrcElts).shl(SubIdx);
+      uint64_t Idx = SubIdx->getZExtValue();
+      APInt SrcElts = DemandedElts.zextOrSelf(NumSrcElts).shl(Idx);
       APInt SrcUndef, SrcZero;
       if (SimplifyDemandedVectorElts(Src, SrcElts, SrcUndef, SrcZero, TLO,
                                      Depth + 1))
         return true;
-    KnownUndef = SrcUndef.extractBits(NumElts, SubIdx);
-    KnownZero = SrcZero.extractBits(NumElts, SubIdx);
+      KnownUndef = SrcUndef.extractBits(NumElts, Idx);
+      KnownZero = SrcZero.extractBits(NumElts, Idx);
+    }
     break;
   }
   case ISD::INSERT_VECTOR_ELT: {
     SDValue Vec = Op.getOperand(0);
     SDValue Scl = Op.getOperand(1);
     auto *CIdx = dyn_cast<ConstantSDNode>(Op.getOperand(2));
 
     // For a legal, constant insertion index, if we don't need this insertion

llvm/trunk/unittests/CodeGen/AArch64SelectionDAGTest.cpp

+//===- llvm/unittest/CodeGen/AArch64SelectionDAGTest.cpp -------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/SelectionDAG.h"
+#include "llvm/Analysis/OptimizationRemarkEmitter.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/CodeGen/TargetLowering.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/TargetRegistry.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Target/TargetMachine.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+void initLLVM() {
+  InitializeAllTargets();
+  InitializeAllTargetMCs();
+}
+
+class AArch64SelectionDAGTest : public testing::Test {
+protected:
+  void SetUp() override {
+    StringRef Assembly = "define void @f() { ret void }";
+
+    Triple TargetTriple("aarch64--");
+    std::string Error;
+    const Target *T = TargetRegistry::lookupTarget("", TargetTriple, Error);
+    // FIXME: These tests do not depend on AArch64 specifically, but we have to
+    // initialize a target. A skeleton Target for unittests would allow us to
+    // always run these tests.
+    if (!T)
+      return;
+
+    TargetOptions Options;
+    TM = std::unique_ptr<TargetMachine>(T->createTargetMachine(
+        "AArch64", "", "", Options, None, None, CodeGenOpt::Aggressive));
+    if (!TM)
+      return;
+
+    SMDiagnostic SMError;
+    M = parseAssemblyString(Assembly, SMError, Context);
+    if (!M)
+      report_fatal_error(SMError.getMessage());
+    M->setDataLayout(TM->createDataLayout());
+
+    F = M->getFunction("f");
+    if (!F)
+      report_fatal_error("F?");
+
+    MachineModuleInfo MMI(TM.get());
+
+    MF = make_unique<MachineFunction>(*F, *TM, *TM->getSubtargetImpl(*F), 0,
+                                      MMI);
+
+    DAG = make_unique<SelectionDAG>(*TM, CodeGenOpt::None);
+    if (!DAG)
+      report_fatal_error("DAG?");
+    OptimizationRemarkEmitter ORE(F);
+    DAG->init(*MF, ORE, nullptr, nullptr, nullptr);
+  }
+
+  LLVMContext Context;
+  std::unique_ptr<TargetMachine> TM = nullptr;
+  std::unique_ptr<Module> M;
+  Function *F;
+  std::unique_ptr<MachineFunction> MF;
+  std::unique_ptr<SelectionDAG> DAG;
+};
+
+TEST_F(AArch64SelectionDAGTest, computeKnownBits_ZERO_EXTEND_VECTOR_INREG) {
+  if (!TM)
+    return;
+  SDLoc Loc;
+  auto Int8VT = EVT::getIntegerVT(Context, 8);
+  auto Int16VT = EVT::getIntegerVT(Context, 16);
+  auto InVecVT = EVT::getVectorVT(Context, Int8VT, 4);
+  auto OutVecVT = EVT::getVectorVT(Context, Int16VT, 2);
+  auto InVec = DAG->getConstant(0, Loc, InVecVT);
+  auto Op = DAG->getZeroExtendVectorInReg(InVec, Loc, OutVecVT);
+  auto DemandedElts = APInt(4, 15);
+  KnownBits Known;
+  DAG->computeKnownBits(Op, Known, DemandedElts);
+  EXPECT_TRUE(Known.isZero());
+}
+
+TEST_F(AArch64SelectionDAGTest, computeKnownBits_EXTRACT_SUBVECTOR) {
+  if (!TM)
+    return;
+  SDLoc Loc;
+  auto IntVT = EVT::getIntegerVT(Context, 8);
+  auto VecVT = EVT::getVectorVT(Context, IntVT, 3);
+  auto IdxVT = EVT::getIntegerVT(Context, 64);
+  auto Vec = DAG->getConstant(0, Loc, VecVT);
+  auto ZeroIdx = DAG->getConstant(0, Loc, IdxVT);
+  auto Op = DAG->getNode(ISD::EXTRACT_SUBVECTOR, Loc, VecVT, Vec, ZeroIdx);
+  auto DemandedElts = APInt(3, 7);
+  KnownBits Known;
+  DAG->computeKnownBits(Op, Known, DemandedElts);
+  EXPECT_TRUE(Known.isZero());
+}
+
+TEST_F(AArch64SelectionDAGTest, ComputeNumSignBits_SIGN_EXTEND_VECTOR_INREG) {
+  if (!TM)
+    return;
+  SDLoc Loc;
+  auto Int8VT = EVT::getIntegerVT(Context, 8);
+  auto Int16VT = EVT::getIntegerVT(Context, 16);
+  auto InVecVT = EVT::getVectorVT(Context, Int8VT, 4);
+  auto OutVecVT = EVT::getVectorVT(Context, Int16VT, 2);
+  auto InVec = DAG->getConstant(1, Loc, InVecVT);
+  auto Op = DAG->getSignExtendVectorInReg(InVec, Loc, OutVecVT);
+  auto DemandedElts = APInt(4, 15);
+  EXPECT_EQ(DAG->ComputeNumSignBits(Op, DemandedElts), 15u);
+}
+
+TEST_F(AArch64SelectionDAGTest, ComputeNumSignBits_EXTRACT_SUBVECTOR) {
+  if (!TM)
+    return;
+  SDLoc Loc;
+  auto IntVT = EVT::getIntegerVT(Context, 8);
+  auto VecVT = EVT::getVectorVT(Context, IntVT, 3);
+  auto IdxVT = EVT::getIntegerVT(Context, 64);
+  auto Vec = DAG->getConstant(1, Loc, VecVT);
+  auto ZeroIdx = DAG->getConstant(0, Loc, IdxVT);
+  auto Op = DAG->getNode(ISD::EXTRACT_SUBVECTOR, Loc, VecVT, Vec, ZeroIdx);
+  auto DemandedElts = APInt(3, 7);
+  EXPECT_EQ(DAG->ComputeNumSignBits(Op, DemandedElts), 7u);
+}
+
+TEST_F(AArch64SelectionDAGTest, SimplifyDemandedVectorElts_EXTRACT_SUBVECTOR) {
+  if (!TM)
+    return;
+
+  TargetLowering TL(*TM);
+
+  SDLoc Loc;
+  auto IntVT = EVT::getIntegerVT(Context, 8);
+  auto VecVT = EVT::getVectorVT(Context, IntVT, 3);
+  auto IdxVT = EVT::getIntegerVT(Context, 64);
+  auto Vec = DAG->getConstant(1, Loc, VecVT);
+  auto ZeroIdx = DAG->getConstant(0, Loc, IdxVT);
+  auto Op = DAG->getNode(ISD::EXTRACT_SUBVECTOR, Loc, VecVT, Vec, ZeroIdx);
+  auto DemandedElts = APInt(3, 7);
+  auto KnownUndef = APInt(3, 0);
+  auto KnownZero = APInt(3, 0);
+  TargetLowering::TargetLoweringOpt TLO(*DAG, false, false);
+  EXPECT_EQ(TL.SimplifyDemandedVectorElts(Op, DemandedElts, KnownUndef,
+                                          KnownZero, TLO),
+            false);
+}
+
+} // end anonymous namespace
+
+int main(int argc, char **argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  initLLVM();
+  return RUN_ALL_TESTS();
+}

llvm/trunk/unittests/CodeGen/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
+  ${LLVM_TARGETS_TO_BUILD}
+  AsmParser
   AsmPrinter
   CodeGen
   Core
   MC
   SelectionDAG
   Support
   Target
   )
 
 add_llvm_unittest(CodeGenTests
+  AArch64SelectionDAGTest.cpp
   DIEHashTest.cpp
   LowLevelTypeTest.cpp
   MachineInstrBundleIteratorTest.cpp
   MachineInstrTest.cpp
   MachineOperandTest.cpp
   ScalableVectorMVTsTest.cpp
   )
 
 add_subdirectory(GlobalISel)