[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

1 ↗

(On Diff #156379)

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

1488 ↗

(On Diff #156508)

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

llvm/trunk/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::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();

llvm/trunk/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:    v_mov_b32_e32 v0, 0
+; CHECK-NEXT:    v_mov_b32_e32 v1, 1
+; CHECK-NEXT:    v_mov_b32_e32 v2, 1
+; 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 }