[DAGCombiner] Fix & simplify constant folding of sext/zext.

[chfast]

Summary

This patch fixes the cases of sext/zext constant folding in DAG combiner where constans do not fit 64 bits. The fix simply removes un$

[RKSimon]

Thanks for finding this! I think the problems you're seeing revolve around the fact that integer build vectors scalar input operands can be of a wider type than the vector type - typically for legalization reasons. So your pcmpeqd regression is because the comparison result (v4i1) actually has scalar inputs of i32 that get implicitly truncated by BUILD_VECTOR.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5609 ↗

(On Diff #28126)

The operands of an integer build vector may be wider than the vector scalar type. You need to do something like
const APInt C = cast<ConstantSDNode>(Op)->getAPIntValue().trunc(EVTBits);

5611 ↗

(On Diff #28126)

I think you need to just drop the 'getZExtValue()' and keep as a APInt:
Elts.push_back(DAG.getConstant(C.shl(ShAmt).ashr(ShAmt), DL, SVT);

5613 ↗

(On Diff #28126)

Elts.push_back(DAG.getConstant(C.shl(ShAmt).lshr(ShAmt), DL, SVT);

test/CodeGen/X86/fold-vector-sext-crash2.ll

2 ↗

(On Diff #28126)

Missing -mtriple args

[chfast]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5611 ↗

(On Diff #28126)

That's definitely safer change. But isn't it the same?

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5609 ↗

(On Diff #28126)

Sorry that should have been
const APInt &C = cast<ConstantSDNode>(Op)->getAPIntValue().zextOrTrunc(EVTBits);

The aim is to ensure that the constant starts out as the bitwidth of input vector's scalar size so that the sext/zext below can work correctly.

5611 ↗

(On Diff #28126)

Sorry you're right - if you make the zextOrTrunc change above this should work.

[chfast]

Trunc the constant value first as it might be wider than the scalar type. Thanks @RKSimon for the tip.

[chfast]

Add -march option to the regression test.

[chfast]

Ping.

[chfast]

Simon, can you check it again?

[RKSimon]

Please can you add a zext test - you've updated that code path as well.

[chfast]

Extend the test to check zext variant and constant values in output.

[RKSimon]

Thanks, LGTM

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

 }
 
 SDValue DAGCombiner::visitSETCC(SDNode *N) {
   return SimplifySetCC(N->getValueType(0), N->getOperand(0), N->getOperand(1),
                        cast<CondCodeSDNode>(N->getOperand(2))->get(),
                        SDLoc(N));
 }
 
-// tryToFoldExtendOfConstant - Try to fold a sext/zext/aext
-// dag node into a ConstantSDNode or a build_vector of constants.
-// This function is called by the DAGCombiner when visiting sext/zext/aext
-// dag nodes (see for example method DAGCombiner::visitSIGN_EXTEND).
-// Vector extends are not folded if operations are legal; this is to
-// avoid introducing illegal build_vector dag nodes.
+/// Try to fold a sext/zext/aext dag node into a ConstantSDNode or 
+/// a build_vector of constants.
+/// This function is called by the DAGCombiner when visiting sext/zext/aext
+/// dag nodes (see for example method DAGCombiner::visitSIGN_EXTEND).
+/// Vector extends are not folded if operations are legal; this is to
+/// avoid introducing illegal build_vector dag nodes.
 static SDNode *tryToFoldExtendOfConstant(SDNode *N, const TargetLowering &TLI,
                                          SelectionDAG &DAG, bool LegalTypes,
                                          bool LegalOperations) {
   unsigned Opcode = N->getOpcode();
   SDValue N0 = N->getOperand(0);
   EVT VT = N->getValueType(0);
 
   assert((Opcode == ISD::SIGN_EXTEND || Opcode == ISD::ZERO_EXTEND ||
static SDNode *tryToFoldExtendOfConstant(SDNode *N, const TargetLowering &TLI,
   if (!(VT.isVector() &&
       (!LegalTypes || (!LegalOperations && TLI.isTypeLegal(SVT))) &&
       ISD::isBuildVectorOfConstantSDNodes(N0.getNode())))
     return nullptr;
 
   // We can fold this node into a build_vector.
   unsigned VTBits = SVT.getSizeInBits();
   unsigned EVTBits = N0->getValueType(0).getScalarType().getSizeInBits();
-  unsigned ShAmt = VTBits - EVTBits;
   SmallVector<SDValue, 8> Elts;
   unsigned NumElts = VT.getVectorNumElements();
   SDLoc DL(N);
 
   for (unsigned i=0; i != NumElts; ++i) {
     SDValue Op = N0->getOperand(i);
     if (Op->getOpcode() == ISD::UNDEF) {
       Elts.push_back(DAG.getUNDEF(SVT));
       continue;
     }
 
     SDLoc DL(Op);
-    ConstantSDNode *CurrentND = cast<ConstantSDNode>(Op);
-    const APInt &C = APInt(VTBits, CurrentND->getAPIntValue().getZExtValue());
+    // Get the constant value and if needed trunc it to the size of the type.
+    // Nodes like build_vector might have constants wider than the scalar type.
+    APInt C = cast<ConstantSDNode>(Op)->getAPIntValue().zextOrTrunc(EVTBits);
     if (Opcode == ISD::SIGN_EXTEND || Opcode == ISD::SIGN_EXTEND_VECTOR_INREG)
-      Elts.push_back(DAG.getConstant(C.shl(ShAmt).ashr(ShAmt).getZExtValue(),
-                                     DL, SVT));
+      Elts.push_back(DAG.getConstant(C.sext(VTBits), DL, SVT));
     else
-      Elts.push_back(DAG.getConstant(C.shl(ShAmt).lshr(ShAmt).getZExtValue(),
-                                     DL, SVT));
+      Elts.push_back(DAG.getConstant(C.zext(VTBits), DL, SVT));
   }
 
   return DAG.getNode(ISD::BUILD_VECTOR, DL, VT, Elts).getNode();
 }
 
 // ExtendUsesToFormExtLoad - Trying to extend uses of a load to enable this:
 // "fold ({s|z|a}ext (load x)) -> ({s|z|a}ext (truncate ({s|z|a}extload x)))"
 // transformation. Returns true if extension are possible and the above

llvm/trunk/test/CodeGen/X86/fold-vector-sext-crash2.ll

+; RUN: llc < %s -march=x86    | FileCheck %s -check-prefix=X32
+; RUN: llc < %s -march=x86-64 | FileCheck %s -check-prefix=X64
+
+; DAGCombiner crashes during sext folding
+
+define <2 x i256> @test_sext1() {
+  %Se = sext <2 x i8> <i8 -100, i8 -99> to <2 x i256>
+  %Shuff = shufflevector <2 x i256> zeroinitializer, <2 x i256> %Se, <2 x i32> <i32 1, i32 3>
+  ret <2 x i256> %Shuff
+
+  ; X64-LABEL: test_sext1
+  ; X64:       movq $-1
+  ; X64-NEXT:  movq $-1
+  ; X64-NEXT:  movq $-1
+  ; X64-NEXT:  movq $-99
+
+  ; X32-LABEL: test_sext1
+  ; X32:       movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-99
+}
+
+define <2 x i256> @test_sext2() {
+  %Se = sext <2 x i128> <i128 -2000, i128 -1999> to <2 x i256>
+  %Shuff = shufflevector <2 x i256> zeroinitializer, <2 x i256> %Se, <2 x i32> <i32 1, i32 3>
+  ret <2 x i256> %Shuff
+
+  ; X64-LABEL: test_sext2
+  ; X64:       movq $-1
+  ; X64-NEXT:  movq $-1
+  ; X64-NEXT:  movq $-1
+  ; X64-NEXT:  movq $-1999
+
+  ; X32-LABEL: test_sext2
+  ; X32:       movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1999
+}
+
+define <2 x i256> @test_zext1() {
+  %Se = zext <2 x i8> <i8 -1, i8 -2> to <2 x i256>
+  %Shuff = shufflevector <2 x i256> zeroinitializer, <2 x i256> %Se, <2 x i32> <i32 1, i32 3>
+  ret <2 x i256> %Shuff
+
+  ; X64-LABEL: test_zext1
+  ; X64:       movq $0
+  ; X64-NEXT:  movq $0
+  ; X64-NEXT:  movq $0
+  ; X64-NEXT:  movq $254
+
+  ; X32-LABEL: test_zext1
+  ; X32:       movl $0
+  ; X32-NEXT:  movl $0
+  ; X32-NEXT:  movl $0
+  ; X32-NEXT:  movl $0
+  ; X32-NEXT:  movl $0
+  ; X32-NEXT:  movl $0
+  ; X32-NEXT:  movl $0
+  ; X32-NEXT:  movl $254
+}
+
+define <2 x i256> @test_zext2() {
+  %Se = zext <2 x i128> <i128 -1, i128 -2> to <2 x i256>
+  %Shuff = shufflevector <2 x i256> zeroinitializer, <2 x i256> %Se, <2 x i32> <i32 1, i32 3>
+  ret <2 x i256> %Shuff
+
+  ; X64-LABEL: test_zext2
+  ; X64:       movq $0
+  ; X64-NEXT:  movq $0
+  ; X64-NEXT:  movq $-1
+  ; X64-NEXT:  movq $-2
+
+  ; X32-LABEL: test_zext2
+  ; X32:       movl $0
+  ; X32-NEXT:  movl $0
+  ; X32-NEXT:  movl $0
+  ; X32-NEXT:  movl $0
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-1
+  ; X32-NEXT:  movl $-2
+}