From ef600cbd86db579f1dc490b2cb65808d6ff7067f Mon Sep 17 00:00:00 2001
From: Amara Emerson <aemerson@apple.com>
Date: Thu, 13 Sep 2018 21:28:58 +0000
Subject: [PATCH] [DAGCombine] Fix crash when store merging created an
 extract_subvector with invalid index.

Differential Revision: https://reviews.llvm.org/D51831

llvm-svn: 342183
---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp         |  9 ++++++++-
 .../X86/merge-vector-stores-scale-idx-crash.ll        | 19 +++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/X86/merge-vector-stores-scale-idx-crash.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 20e749b..71048b1 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -13843,17 +13843,24 @@ bool DAGCombiner::MergeStoresOfConstantsOrVecElts(
              Val.getOpcode() == ISD::EXTRACT_SUBVECTOR)) {
           SDValue Vec = Val.getOperand(0);
           EVT MemVTScalarTy = MemVT.getScalarType();
+          SDValue Idx = Val.getOperand(1);
           // We may need to add a bitcast here to get types to line up.
           if (MemVTScalarTy != Vec.getValueType()) {
             unsigned Elts = Vec.getValueType().getSizeInBits() /
                             MemVTScalarTy.getSizeInBits();
+            if (Val.getValueType().isVector()) {
+              unsigned IdxC = cast<ConstantSDNode>(Idx)->getZExtValue();
+              unsigned NewIdx =
+                  ((uint64_t)IdxC * MemVT.getVectorNumElements()) / Elts;
+              Idx = DAG.getConstant(NewIdx, SDLoc(Val), Idx.getValueType());
+            }
             EVT NewVecTy =
                 EVT::getVectorVT(*DAG.getContext(), MemVTScalarTy, Elts);
             Vec = DAG.getBitcast(NewVecTy, Vec);
           }
           auto OpC = (MemVT.isVector()) ? ISD::EXTRACT_SUBVECTOR
                                         : ISD::EXTRACT_VECTOR_ELT;
-          Val = DAG.getNode(OpC, SDLoc(Val), MemVT, Vec, Val.getOperand(1));
+          Val = DAG.getNode(OpC, SDLoc(Val), MemVT, Vec, Idx);
         }
         Ops.push_back(Val);
       }
diff --git a/llvm/test/CodeGen/X86/merge-vector-stores-scale-idx-crash.ll b/llvm/test/CodeGen/X86/merge-vector-stores-scale-idx-crash.ll
new file mode 100644
index 0000000..d85163a
--- /dev/null
+++ b/llvm/test/CodeGen/X86/merge-vector-stores-scale-idx-crash.ll
@@ -0,0 +1,19 @@
+; RUN: llc < %s  -mtriple=x86_64-apple-osx10.14 -mattr=+avx2 | FileCheck %s
+
+; Check that we don't crash due creating invalid extract_subvector indices in store merging.
+; CHECK-LABEL: testfn
+; CHECK: retq
+define void @testfn(i32* nocapture %p) {
+  %v0 = getelementptr i32, i32* %p, i64 12
+  %1 = bitcast i32* %v0 to <2 x i64>*
+  %2 = bitcast i32* %v0 to <4 x i32>*
+  %3 = getelementptr <2 x i64>, <2 x i64>* %1, i64 -3
+  store <2 x i64> undef, <2 x i64>* %3, align 16
+  %4 = shufflevector <4 x i64> zeroinitializer, <4 x i64> undef, <2 x i32> <i32 0, i32 1>
+  %5 = getelementptr <2 x i64>, <2 x i64>* %1, i64 -2
+  store <2 x i64> %4, <2 x i64>* %5, align 16
+  %6 = shufflevector <8 x i32> zeroinitializer, <8 x i32> undef, <4 x i32> <i32 4, i32 5, i32 6, i32 7>
+  %7 = getelementptr <4 x i32>, <4 x i32>* %2, i64 -1
+  store <4 x i32> %6, <4 x i32>* %7, align 16
+  ret void
+}
-- 
2.7.4