[llvm] 4676805 - [AArch64][DAGCombine] Fix a bug in performBuildVectorCombine where it could produce an invalid EXTRACT_SUBVECTOR

[Usman Nadeem via llvm-commits]

Author: Usman Nadeem
Date: 2022-08-24T16:24:19-07:00
New Revision: 46768052e0c9e21a7fcaa2881e29e6b4237b200f

URL: https://github.com/llvm/llvm-project/commit/46768052e0c9e21a7fcaa2881e29e6b4237b200f
DIFF: https://github.com/llvm/llvm-project/commit/46768052e0c9e21a7fcaa2881e29e6b4237b200f.diff

LOG: [AArch64][DAGCombine] Fix a bug in performBuildVectorCombine where it could produce an invalid EXTRACT_SUBVECTOR

EXTRACT_SUBVECTOR requires that Idx be a constant multiple of ResultType's
known minimum vector length.

Something like this will produce an invalid extract_subvector:

t1: v4i16 = .....
t2: i32 = extract_vector_elt t1, Constant:i64<1>
t3: i32 = extract_vector_elt t1, Constant:i64<2>
t4: v2i32 = BUILD_VECTOR t2, t3
// produces
t5: v2i32 = extract_subvector t...., Constant:i64<1>

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

Change-Id: I7a5acf054edee3e89c0f85a28d8869256403ce08

Added: 
    llvm/test/CodeGen/AArch64/aarch64-vectorcombine-invalid-extract-index-crash.ll

Modified: 
    llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 1fa724d34b15..3eb07a069502 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -16145,6 +16145,7 @@ static SDValue performBuildVectorCombine(SDNode *N,
                                          TargetLowering::DAGCombinerInfo &DCI,
                                          SelectionDAG &DAG) {
   SDLoc DL(N);
+  EVT VT = N->getValueType(0);
 
   // A build vector of two extracted elements is equivalent to an
   // extract subvector where the inner vector is any-extended to the
@@ -16155,7 +16156,7 @@ static SDValue performBuildVectorCombine(SDNode *N,
 
   // For now, only consider the v2i32 case, which arises as a result of
   // legalization.
-  if (N->getValueType(0) != MVT::v2i32)
+  if (VT != MVT::v2i32)
     return SDValue();
 
   SDValue Elt0 = N->getOperand(0), Elt1 = N->getOperand(1);
@@ -16168,7 +16169,10 @@ static SDValue performBuildVectorCombine(SDNode *N,
       // Both EXTRACT_VECTOR_ELT from same vector...
       Elt0->getOperand(0) == Elt1->getOperand(0) &&
       // ... and contiguous. First element's index +1 == second element's index.
-      Elt0->getConstantOperandVal(1) + 1 == Elt1->getConstantOperandVal(1)) {
+      Elt0->getConstantOperandVal(1) + 1 == Elt1->getConstantOperandVal(1) &&
+      // EXTRACT_SUBVECTOR requires that Idx be a constant multiple of
+      // ResultType's known minimum vector length.
+      Elt0->getConstantOperandVal(1) % VT.getVectorMinNumElements() == 0) {
     SDValue VecToExtend = Elt0->getOperand(0);
     EVT ExtVT = VecToExtend.getValueType().changeVectorElementType(MVT::i32);
     if (!DAG.getTargetLoweringInfo().isTypeLegal(ExtVT))

diff  --git a/llvm/test/CodeGen/AArch64/aarch64-vectorcombine-invalid-extract-index-crash.ll b/llvm/test/CodeGen/AArch64/aarch64-vectorcombine-invalid-extract-index-crash.ll
new file mode 100644
index 000000000000..e17add462ff8
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/aarch64-vectorcombine-invalid-extract-index-crash.ll
@@ -0,0 +1,28 @@
+; RUN: llc < %s -o /dev/null
+
+;
+; Test that we dont get a crash after an invalid build_vector combine.
+;
+
+target triple = "aarch64-unknown-linux-gnu"
+define void @test_crash(i8* %dst_ptr) {
+entry:
+  %vec_load = load <4 x i16>, <4 x i16>* undef, align 8
+  %0 = sext <4 x i16> %vec_load to <4 x i32>
+  %add71vec = add nsw <4 x i32> %0, <i32 32, i32 32, i32 32, i32 32>
+  %add104vec = add nsw <4 x i32> %add71vec, zeroinitializer
+  %add105vec = add nsw <4 x i32> zeroinitializer, %add104vec
+  %vec = lshr <4 x i32> %add105vec, <i32 6, i32 6, i32 6, i32 6>
+  %1 = trunc <4 x i32> %vec to <4 x i16>
+  %2 = shufflevector <4 x i16> %1, <4 x i16> undef, <2 x i32> <i32 1, i32 2>
+  %3 = sext <2 x i16> %2 to <2 x i32>
+  %4 = bitcast i8* %dst_ptr to <4 x i8>*
+  %5 = shufflevector <2 x i32> %3, <2 x i32> poison, <4 x i32> <i32 0, i32 1, i32 undef, i32 undef>
+  %6 = shufflevector <4 x i32> undef, <4 x i32> %5, <4 x i32> <i32 0, i32 4, i32 5, i32 undef>
+  %7 = insertelement <4 x i32> %6, i32 undef, i64 3
+  %8 = add nsw <4 x i32> %7, zeroinitializer
+  %9 = select <4 x i1> zeroinitializer, <4 x i32> %8, <4 x i32> undef
+  %10 = trunc <4 x i32> %9 to <4 x i8>
+  store <4 x i8> %10, <4 x i8>* %4, align 1
+  ret void
+}