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

14477

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

14477

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

14477

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);
-    if (InVT.isSimple() && (NearestPow2 > 2) &&
-        ((NumElems * 2) < NearestPow2)) {
+    if (InVT.isSimple() && NearestPow2 > 2 && MaxIndex < NearestPow2 &&
+        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));
         SDValue VecIn1 = DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, SplitVT, Vec,
                                      DAG.getConstant(0, 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
+}