Fix an out-of-bounds shufflevector index bug

[george.burgess.iv]

Summary

I'm somewhat confident about my fix on its own, but I'm not at all familiar with vector instructions. Since I had to modify existing tests, please take a look. :)

Bug: MaxIndex represents the maximum index we're passing to ShuffleVector. If it's a power of two, we can end up splitting into two vectors which have a total size == MaxIndex. This caused an assertion failure in SelectionDAG::getVectorShuffle:

Assertion `llvm::all_of(Mask, [&](int M) { return M < (NElts * 2); }) && "Index out of range"' failed

(NElts == 4, and Mask == {0, 8, -1, -1})

[javed.absar]

test/CodeGen/ARM/crash-on-pow2-shufflevector.ll

20

Shouldn't this test have some sort of 'CHECK'?

[RKSimon]

test/CodeGen/ARM/crash-on-pow2-shufflevector.ll

20

Make the lit command something like:

; RUN: llc < %s -mtriple=armv7--linux-android | FileCheck %s

You should then be able to use utils/update_llc_test_checks.py to create the CHECKS

64

You should be able to reduce this test case a lot more, either by hand or with bugpoint

[jbhateja]

This comment was removed by jbhateja.

[george.burgess.iv]

Address all feedback

[george.burgess.iv]

Thanks for the comments!

I propose following as the fix , Simon and other reviews can comment.

That'll probably fix the assertion, but I don't think it'll do what we want here. AFAICT, in that function, VecIn1 contains elements [0..3] from the original vector, and VecIn2 contains elements [4..7]. However, we're trying to shuffle out elements 0 and 8.

Hopefully the new test-case is more clear about this.

test/CodeGen/ARM/crash-on-pow2-shufflevector.ll

20

Shouldn't this test have some sort of 'CHECK'?

I didn't do checks because goal of this test was to get an assertion to fail. If we'd rather have checks here, I'm happy to do that, too. ¯\_(ツ)_/¯

You should then be able to use utils/update_llc_test_checks.py to create the CHECKS

Thanks for the pointer! (Here and for bugpoint :) )

[jbhateja]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

14472

If MaxIndex is same as MaxSize of vector it will cause later split (line #14475) to be futile.

[RKSimon]

LGTM

[jbhateja]

In D37880#879678, @RKSimon wrote:

LGTM

Do you have any observations over alternate fixed proposed. I think it will make this full proof.

[george.burgess.iv]

Sorry for taking so long to get back to this, and thanks again for the comments/reviews!

Since -- as outlined above -- I still don't think the other proposed fix is correct, I plan to commit this as-is. I'll do so in two days, so people have time to voice objections if they have them. (And I'm always happy to revert if we do find a better way to fix this.)

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

14472

I tried adding InVT.getVectorNumElements() != NearestPow2 into the conditional below, and it seems to substantially change the output of some test cases (e.g. one of the tests in test/CodeGen/X86/avx512-shuffles/partial_permute.ll goes from 4 vector instructions to 18).

Since I don't know this code, and it looks like your comment applies to this code even without this patch, I'd rather not have that change as a part of this.

[jbhateja]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

14472

Hi George,

Your fix is fine, I looked at it again but for the boundary case it will be more appropriate to AND another condition for if block which is MaxIndex < NearestPow2. This shall save unnecessary splif of vector.

So if condition should be

if (InVT.isSimple() && (NearestPow2 > 2) &&
    ( MaxIndex < NearestPow2) &&
    ((NumElems * 2) < NearestPow2)) {

Thanks

[george.burgess.iv]

Address all feedback

[jbhateja]

LGTM with a minor comment, please run clang-format. over the changes for better inndentation.

Thanks

[george.burgess.iv]

Thanks!

please run clang-format. over the changes for better inndentation

Done; the formatting remained unchanged for me. ¯\_(ツ)_/¯

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (VecIn.size() == 2) {
     for (unsigned i = 0; i < NumElems; i++) {
       if (VectorMask[i] <= 0)
         continue;
       unsigned Index = N->getOperand(i).getConstantOperandVal(1);
       IndexVec[i] = Index;
       MaxIndex = std::max(MaxIndex, Index);
     }
 
-    NearestPow2 = PowerOf2Ceil(MaxIndex);
+    NearestPow2 = PowerOf2Ceil(MaxIndex + 1);
     if (InVT.isSimple() && (NearestPow2 > 2) &&
         ((NumElems * 2) < NearestPow2)) {
       unsigned SplitSize = NearestPow2 / 2;
       EVT SplitVT = EVT::getVectorVT(*DAG.getContext(),
                                      InVT.getVectorElementType(), SplitSize);
       if (TLI.isTypeLegal(SplitVT)) {
         SDValue VecIn2 = DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, SplitVT, Vec,
                                      DAG.getConstant(SplitSize, DL, IdxTy));

test/CodeGen/ARM/crash-on-pow2-shufflevector.ll

+; RUN: llc < %s -mtriple=armv7 | FileCheck %s
+;
+; Ensure that don't crash given a largeish power-of-two shufflevector index.
+
+%struct.desc = type { i32, [7 x i32] }
+
+define i32 @foo(%struct.desc* %descs, i32 %num, i32 %cw) local_unnamed_addr #0 {
+; CHECK-LABEL: foo:
+; CHECK:       @ BB#0: @ %entry
+; CHECK-NEXT:    mov r1, #32
+; CHECK-NEXT:    vld1.32 {d16, d17}, [r0], r1
+; CHECK-NEXT:    vld1.32 {d18, d19}, [r0]
+; CHECK-NEXT:    vtrn.32 q8, q9
+; CHECK-NEXT:    vadd.i32 d16, d16, d16
+; CHECK-NEXT:    vmov.32 r0, d16[1]
+; CHECK-NEXT:    bx lr
+entry:
+  %descs.vec = bitcast %struct.desc* %descs to <16 x i32>*
+  %wide.vec = load <16 x i32>, <16 x i32>* %descs.vec, align 4
+  %strided.vec = shufflevector <16 x i32> %wide.vec, <16 x i32> undef, <2 x i32> <i32 0, i32 8>
+  %bin.rdx20 = add <2 x i32> %strided.vec, %strided.vec
+  %0 = extractelement <2 x i32> %bin.rdx20, i32 1
+  ret i32 %0
+}

test/CodeGen/X86/avx512-shuffles/partial_permute.ll

; CHECK-NEXT:    retq
   %shuf = shufflevector <8 x i64> %vec, <8 x i64> undef, <4 x i32> <i32 7, i32 7, i32 5, i32 1>
   %res = select <4 x i1> <i1 1, i1 1, i1 0, i1 1>, <4 x i64> %shuf, <4 x i64> zeroinitializer
   ret <4 x i64> %res
 }
 
 define <2 x i64> @test_8xi64_to_2xi64_perm_mem_mask0(<8 x i64>* %vp) {
 ; CHECK-LABEL: test_8xi64_to_2xi64_perm_mem_mask0:
 ; CHECK:       # BB#0:
-; CHECK-NEXT:    vmovdqa64 (%rdi), %zmm0
-; CHECK-NEXT:    vextracti32x4 $2, %zmm0, %xmm1
-; CHECK-NEXT:    vpshufd {{.*#+}} xmm0 = xmm0[2,3,0,1]
-; CHECK-NEXT:    vpunpcklqdq {{.*#+}} xmm0 = xmm1[0],xmm0[0]
+; CHECK-NEXT:    vmovaps (%rdi), %zmm0
+; CHECK-NEXT:    vextractf64x4 $1, %zmm0, %ymm1
+; CHECK-NEXT:    vblendps {{.*#+}} xmm0 = xmm1[0,1],xmm0[2,3]
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    retq
   %vec = load <8 x i64>, <8 x i64>* %vp
   %res = shufflevector <8 x i64> %vec, <8 x i64> undef, <2 x i32> <i32 4, i32 1>
   ret <2 x i64> %res
 }
 define <2 x i64> @test_masked_8xi64_to_2xi64_perm_mem_mask0(<8 x i64>* %vp, <2 x i64> %vec2) {
 ; CHECK-LABEL: test_masked_8xi64_to_2xi64_perm_mem_mask0:
 ; CHECK:       # BB#0:
 ; CHECK-NEXT:    vmovdqa64 (%rdi), %zmm1
-; CHECK-NEXT:    vextracti32x4 $2, %zmm1, %xmm2
-; CHECK-NEXT:    vpshufd {{.*#+}} xmm1 = xmm1[2,3,0,1]
+; CHECK-NEXT:    vextracti64x4 $1, %zmm1, %ymm2
+; CHECK-NEXT:    vpblendd {{.*#+}} xmm1 = xmm2[0,1],xmm1[2,3]
 ; CHECK-NEXT:    movb $2, %al
 ; CHECK-NEXT:    kmovd %eax, %k1
-; CHECK-NEXT:    vpunpcklqdq {{.*#+}} xmm0 = xmm2[0],xmm1[0]
+; CHECK-NEXT:    vmovdqa64 %xmm1, %xmm0 {%k1}
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    retq
   %vec = load <8 x i64>, <8 x i64>* %vp
   %shuf = shufflevector <8 x i64> %vec, <8 x i64> undef, <2 x i32> <i32 4, i32 1>
   %res = select <2 x i1> <i1 0, i1 1>, <2 x i64> %shuf, <2 x i64> %vec2
   ret <2 x i64> %res
 }
 
 define <2 x i64> @test_masked_z_8xi64_to_2xi64_perm_mem_mask0(<8 x i64>* %vp) {
 ; CHECK-LABEL: test_masked_z_8xi64_to_2xi64_perm_mem_mask0:
 ; CHECK:       # BB#0:
 ; CHECK-NEXT:    vmovdqa64 (%rdi), %zmm0
-; CHECK-NEXT:    vextracti32x4 $2, %zmm0, %xmm1
-; CHECK-NEXT:    vpshufd {{.*#+}} xmm0 = xmm0[2,3,0,1]
+; CHECK-NEXT:    vextracti64x4 $1, %zmm0, %ymm1
+; CHECK-NEXT:    vpblendd {{.*#+}} xmm0 = xmm1[0,1],xmm0[2,3]
 ; CHECK-NEXT:    movb $2, %al
 ; CHECK-NEXT:    kmovd %eax, %k1
-; CHECK-NEXT:    vpunpcklqdq {{.*#+}} xmm0 = xmm1[0],xmm0[0]
+; CHECK-NEXT:    vmovdqa64 %xmm0, %xmm0 {%k1} {z}
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    retq
   %vec = load <8 x i64>, <8 x i64>* %vp
   %shuf = shufflevector <8 x i64> %vec, <8 x i64> undef, <2 x i32> <i32 4, i32 1>
   %res = select <2 x i1> <i1 0, i1 1>, <2 x i64> %shuf, <2 x i64> zeroinitializer
   ret <2 x i64> %res
 }
 
; CHECK-NEXT:    retq
   %res = select <2 x i1> <i1 0, i1 1>, <2 x double> %shuf, <2 x double> zeroinitializer
   ret <2 x double> %res
 }
 
 define <2 x double> @test_masked_8xdouble_to_2xdouble_perm_mem_mask1(<8 x double>* %vp, <2 x double> %vec2) {
 ; CHECK-LABEL: test_masked_8xdouble_to_2xdouble_perm_mem_mask1:
 ; CHECK:       # BB#0:
 ; CHECK-NEXT:    vmovapd (%rdi), %zmm1
-; CHECK-NEXT:    vextractf32x4 $2, %zmm1, %xmm2
+; CHECK-NEXT:    vextractf64x4 $1, %zmm1, %ymm2
 ; CHECK-NEXT:    movb $1, %al
 ; CHECK-NEXT:    kmovd %eax, %k1
 ; CHECK-NEXT:    vshufpd {{.*#+}} xmm0 {%k1} = xmm1[1],xmm2[0]
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    retq
   %vec = load <8 x double>, <8 x double>* %vp
   %shuf = shufflevector <8 x double> %vec, <8 x double> undef, <2 x i32> <i32 1, i32 4>
   %res = select <2 x i1> <i1 1, i1 0>, <2 x double> %shuf, <2 x double> %vec2
   ret <2 x double> %res
 }
 
 define <2 x double> @test_masked_z_8xdouble_to_2xdouble_perm_mem_mask1(<8 x double>* %vp) {
 ; CHECK-LABEL: test_masked_z_8xdouble_to_2xdouble_perm_mem_mask1:
 ; CHECK:       # BB#0:
 ; CHECK-NEXT:    vmovapd (%rdi), %zmm0
-; CHECK-NEXT:    vextractf32x4 $2, %zmm0, %xmm1
+; CHECK-NEXT:    vextractf64x4 $1, %zmm0, %ymm1
 ; CHECK-NEXT:    movb $1, %al
 ; CHECK-NEXT:    kmovd %eax, %k1
 ; CHECK-NEXT:    vshufpd {{.*#+}} xmm0 {%k1} {z} = xmm0[1],xmm1[0]
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    retq
   %vec = load <8 x double>, <8 x double>* %vp
   %shuf = shufflevector <8 x double> %vec, <8 x double> undef, <2 x i32> <i32 1, i32 4>
   %res = select <2 x i1> <i1 1, i1 0>, <2 x double> %shuf, <2 x double> zeroinitializer
   ret <2 x double> %res
 }
 

test/CodeGen/X86/vector-shuffle-512-v8.ll

; AVX512F-32-NEXT:    retl
   %res = shufflevector <8 x double> %v, <8 x double> undef, <4 x i32> <i32 2, i32 3, i32 4, i32 6>
   ret <4 x double> %res
 }
 
 ;FIXME: compressp
 define <2 x double> @test_v8f64_34 (<8 x double> %v) {
 ; AVX512F-LABEL: test_v8f64_34:
 ; AVX512F:       # BB#0:
-; AVX512F-NEXT:    vextractf32x4 $2, %zmm0, %xmm1
-; AVX512F-NEXT:    vextractf128 $1, %ymm0, %xmm0
-; AVX512F-NEXT:    vshufpd {{.*#+}} xmm0 = xmm0[1],xmm1[0]
+; AVX512F-NEXT:    vextractf64x4 $1, %zmm0, %ymm1
+; AVX512F-NEXT:    vshufpd {{.*#+}} ymm0 = ymm0[1],ymm1[0],ymm0[3],ymm1[2]
+; AVX512F-NEXT:    vpermpd {{.*#+}} ymm0 = ymm0[2,1,2,3]
+; AVX512F-NEXT:      # kill: %XMM0<def> %XMM0<kill> %YMM0<kill>
 ; AVX512F-NEXT:    vzeroupper
 ; AVX512F-NEXT:    retq
 ;
 ; AVX512F-32-LABEL: test_v8f64_34:
 ; AVX512F-32:       # BB#0:
-; AVX512F-32-NEXT:    vextractf32x4 $2, %zmm0, %xmm1
-; AVX512F-32-NEXT:    vextractf128 $1, %ymm0, %xmm0
-; AVX512F-32-NEXT:    vshufpd {{.*#+}} xmm0 = xmm0[1],xmm1[0]
+; AVX512F-32-NEXT:    vextractf64x4 $1, %zmm0, %ymm1
+; AVX512F-32-NEXT:    vshufpd {{.*#+}} ymm0 = ymm0[1],ymm1[0],ymm0[3],ymm1[2]
+; AVX512F-32-NEXT:    vpermpd {{.*#+}} ymm0 = ymm0[2,1,2,3]
+; AVX512F-32-NEXT:      # kill: %XMM0<def> %XMM0<kill> %YMM0<kill>
 ; AVX512F-32-NEXT:    vzeroupper
 ; AVX512F-32-NEXT:    retl
   %res = shufflevector <8 x double> %v, <8 x double> undef, <2 x i32> <i32 3, i32 4>
   ret <2 x double> %res
 }
 
 ; FIXME: vpcompress
 define <4 x i64> @test_v8i64_1257 (<8 x i64> %v) {