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

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$

Test Plan: New regression test included.

Reviewers: RKSimon

Reviewed By: RKSimon

Subscribers: RKSimon, llvm-commits

Differential Revision: http://reviews.llvm.org/D10607

llvm-svn: 240991

2 files changed, 103 insertions, 13 deletions

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 29f9df6425e..631f5852334 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -5579,12 +5579,12 @@ SDValue DAGCombiner::visitSETCC(SDNode *N) {
                        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) {
@@ -5614,7 +5614,6 @@ static SDNode *tryToFoldExtendOfConstant(SDNode *N, const TargetLowering &TLI,
   // 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);
@@ -5627,14 +5626,13 @@ static SDNode *tryToFoldExtendOfConstant(SDNode *N, const TargetLowering &TLI,
     }
 
     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();
diff --git a/llvm/test/CodeGen/X86/fold-vector-sext-crash2.ll b/llvm/test/CodeGen/X86/fold-vector-sext-crash2.ll
new file mode 100644
index 00000000000..44c836195ab
--- /dev/null
+++ b/llvm/test/CodeGen/X86/fold-vector-sext-crash2.ll
@@ -0,0 +1,92 @@
+; 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
+}