[CodeGen] Fix ICE in SelectionDAG::computeKnownBits

[scott.linder]

Summary

SelectionDAG::computeKnownBits asserts handling EXTRACT_SUBVECTOR when zero extending the demanded elements mask if it is already as long as the source vector.

One assumption with the fix is that it is valid for EXTRACT_SUBVECTOR to ask for the entire original vector, which appears valid from reading the code. The test is just for the absence of the assert, as there is nothing in the IR/DAG to check for.

[RKSimon]

I think SelectionDAG::ComputeNumSignBits and TargetLowering::SimplifyDemandedVectorElts suffer from the same issue - can you deal with those as well ?

test/CodeGen/AMDGPU/extract-subvector-equal-length.ll

2

FileCheck ?

[scott.linder]

Addressed feedback; updated other places where DemandedElts is zero extended to the source vector width.

I do not have a FileCheck in the test because the passing condition is just for there to be no assert during CodeGen. There is nothing wrong with the IR/DAG/ISA at any point when the bug is or isn't present, so I don't know what to check for.

[RKSimon]

Any chance that you can add tests for SelectionDAG::ComputeNumSignBits and TargetLowering::SimplifyDemandedVectorElts ?

[arsenm]

In D49574#1170008, @scott.linder wrote:

Addressed feedback; updated other places where DemandedElts is zero extended to the source vector width.

I do not have a FileCheck in the test because the passing condition is just for there to be no assert during CodeGen. There is nothing wrong with the IR/DAG/ISA at any point when the bug is or isn't present, so I don't know what to check for.

The operation ultimately produces something . In this case I would probably just use update_llc_checks

[scott.linder]

I can update the test with update_llc_checks, but I'm not sure about testing the other functions. This test case is a minimized version of a test which happened to trigger the assert; to test the others would it be reasonable to use unit tests instead of LLVM IR?

[scott.linder]

lib/CodeGen/SelectionDAG/TargetLowering.cpp

1486–1487

In writing tests for this version of the code in TargetLowering, I noticed that the condition here is actually more strict than in the equivalent code in SelectionDAG. In this function the subsequent zext can never assert. I still believe my original patch is correct, and I believe this condition can be relaxed to just ugt, but the fact that the code differs makes me wonder why.

Can someone who is more familiar with the code see a reason for the different condition?

[scott.linder]

I have updated the lit test with update_llc_test_checks, updated the condition in TargetLowering to allow for the case where DemandedElts is already wide enough, and added unittests for the cases which trigger the ICE in each function. This is still dependent on my assumptions about the exact semantics of EXTRACT_SUBVECTOR being correct.

[arsenm]

In D49574#1170244, @scott.linder wrote:

I can update the test with update_llc_checks, but I'm not sure about testing the other functions. This test case is a minimized version of a test which happened to trigger the assert; to test the others would it be reasonable to use unit tests instead of LLVM IR?

I'm not sure it's possible with SelectionDAG. I don't think we have any other way

[scott.linder]

I noticed I had assumed AArch64 was available in the unittests, so I updated those to be skipped in the case where it isn't.

@arsenm what do you mean is impossible? These appear to be some of the first unittests for SelectionDAG, but I don't see how to test these cases without adding them. Even if there is a reason not to add them, I still can't add more IR tests for the remaining cases.

[RKSimon]

https://bugs.llvm.org/show_bug.cgi?id=38459 looks similar

[RKSimon]

I feel guilty now, but it might make sense to keep just the original computeKnownBits fix you did in this patch, get that reviewed/committed so we can merge into 7.00 to fix PR38459. And move the other fixes/tests into a new patch?

[scott.linder]

That sounds fine to me; here is the original patch, with the updated test.

I've confirmed this still resolves https://bugs.llvm.org/show_bug.cgi?id=38459

[RKSimon]

LGTM - cheers

lib/CodeGen/SelectionDAG/SelectionDAG.cpp

case ISD::EXTRACT_SUBVECTOR: {
     // If we know the element index, just demand that subvector elements,
     // otherwise demand them all.
     SDValue Src = Op.getOperand(0);
     ConstantSDNode *SubIdx = dyn_cast<ConstantSDNode>(Op.getOperand(1));
     unsigned NumSrcElts = Src.getValueType().getVectorNumElements();
     if (SubIdx && SubIdx->getAPIntValue().ule(NumSrcElts - NumElts)) {
       // Offset the demanded elts by the subvector index.
       uint64_t Idx = SubIdx->getZExtValue();
-      APInt DemandedSrc = DemandedElts.zext(NumSrcElts).shl(Idx);
+      APInt DemandedSrc = DemandedElts.zextOrSelf(NumSrcElts).shl(Idx);
       computeKnownBits(Src, Known, DemandedSrc, Depth + 1);
     } else {
       computeKnownBits(Src, Known, Depth + 1);
     }
     break;
   }
   case ISD::BITCAST: {
     SDValue N0 = Op.getOperand(0);
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());
     computeKnownBits(Op.getOperand(0), Known, InDemandedElts, Depth + 1);
     Known = Known.zext(BitWidth);
     Known.Zero.setBitsFrom(InVT.getScalarSizeInBits());
     break;
   }
   case ISD::ZERO_EXTEND: {
     EVT InVT = Op.getOperand(0).getValueType();
     computeKnownBits(Op.getOperand(0), Known, 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 =
case ISD::EXTRACT_SUBVECTOR: {
     // If we know the element index, just demand that subvector elements,
     // otherwise demand them all.
     SDValue Src = Op.getOperand(0);
     ConstantSDNode *SubIdx = dyn_cast<ConstantSDNode>(Op.getOperand(1));
     unsigned NumSrcElts = Src.getValueType().getVectorNumElements();
     if (SubIdx && SubIdx->getAPIntValue().ule(NumSrcElts - NumElts)) {
       // Offset the demanded elts by the subvector index.
       uint64_t Idx = SubIdx->getZExtValue();
-      APInt DemandedSrc = DemandedElts.zext(NumSrcElts).shl(Idx);
+      APInt DemandedSrc = DemandedElts.zextOrSelf(NumSrcElts).shl(Idx);
       return ComputeNumSignBits(Src, DemandedSrc, Depth + 1);
     }
     return ComputeNumSignBits(Src, Depth + 1);
   }
   case ISD::CONCAT_VECTORS:
     // Determine the minimum number of sign bits across all demanded
     // elts of the input vectors. Early out if the result is already 1.
     Tmp = std::numeric_limits<unsigned>::max();

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

test/CodeGen/AMDGPU/extract-subvector-equal-length.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -march=amdgcn -mtriple=amdgcn-- -verify-machineinstrs < %s | FileCheck %s
+
+; Test for ICE in SelectionDAG::computeKnownBits when visiting EXTRACT_SUBVECTOR
+; with DemandedElts already as wide as the source vector.
+
+define <3 x i32> @quux() #0 {
+; CHECK-LABEL: quux:
+; CHECK:       ; %bb.0: ; %bb
+; CHECK-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; CHECK-NEXT:    s_mov_b32 s9, 1
+; CHECK-NEXT:    s_mov_b32 s8, 0
+; CHECK-NEXT:    s_mov_b32 s10, s9
+; CHECK-NEXT:    v_mov_b32_e32 v0, s8
+; CHECK-NEXT:    v_mov_b32_e32 v1, s9
+; CHECK-NEXT:    v_mov_b32_e32 v2, s10
+; CHECK-NEXT:    v_mov_b32_e32 v3, s11
+; CHECK-NEXT:    s_setpc_b64 s[30:31]
+bb:
+  %tmp = shufflevector <4 x i8> <i8 1, i8 2, i8 3, i8 4>, <4 x i8> undef, <3 x i32> <i32 0, i32 1, i32 2>
+  %tmp1 = extractelement <3 x i8> %tmp, i64 0
+  %tmp2 = zext i8 %tmp1 to i32
+  %tmp3 = insertelement <3 x i32> undef, i32 %tmp2, i32 0
+  %tmp4 = extractelement <3 x i8> %tmp, i64 1
+  %tmp5 = zext i8 %tmp4 to i32
+  %tmp6 = insertelement <3 x i32> %tmp3, i32 %tmp5, i32 1
+  %tmp7 = extractelement <3 x i8> %tmp, i64 2
+  %tmp8 = zext i8 %tmp7 to i32
+  %tmp9 = insertelement <3 x i32> %tmp6, i32 %tmp8, i32 2
+  %tmp10 = lshr <3 x i32> %tmp9, <i32 1, i32 1, i32 1>
+  ret <3 x i32> %tmp10
+}
+
+attributes #0 = { noinline optnone }

unittests/CodeGen/CMakeLists.txt

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

unittests/CodeGen/SelectionDAGTest.cpp

+//===- llvm/unittest/CodeGen/SelectionDAGTest.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 SelectionDAGTest : 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);
+    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(SelectionDAGTest, 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(SelectionDAGTest, 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(SelectionDAGTest, 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(SelectionDAGTest, 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(SelectionDAGTest, 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();
+}