[AArch64][SVE] Fix selection failure during lowering of shuffle_vector

[bsmith]

Summary

The lowering code for shuffle_vector has a code path that looks through
extract_subvector, this code path did not properly account for the
potential presense of larger than Neon vector types and could produce
unselectable DAG nodes.

[david-arm]

llvm/test/CodeGen/AArch64/sve-shuffle-crash.ll

1 ↗

(On Diff #406836)

Hi @bsmith, there don't seem to be any CHECK lines here. Could you add some please using the llvm/utils/update_llc_test_checks.py script?

[bsmith]

llvm/test/CodeGen/AArch64/sve-shuffle-crash.ll

1 ↗

(On Diff #406836)

I intentionally didn't put any in here, the resulting codegen is large and uninteresting in this case, the failure mode is a crash which will be caught regardless.

I can put some in though if you think it would be valuable?

[david-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9732

You're introducing a new TypeSize -> uint64_t implicit typecast here. Can you use this instead?

if (Extract.getOperand(0).getValueSizeInBits().getFixedSizeInBits() > 128)

llvm/test/CodeGen/AArch64/sve-shuffle-crash.ll

1 ↗

(On Diff #406836)

Hmm, I do realise the CHECK lines probably look awful, but maybe it's still worth it? I imagine if the code quality is really that terrible, we'll probably want to improve that at some point and at least we'll be defending it. If we improve codegen in future we'll see reduced CHECK lines here.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9732

Looking at the following code can this be tightened to if (!V.getOperand(0).getValueType().is128BitVector())? It's only used post type legalisation so I think there's no other type that would be acceptable by this point.

llvm/test/CodeGen/AArch64/sve-shuffle-crash.ll

3 ↗

(On Diff #406836)

Not really against this file but there already exists sve-fixed-length-shuffles.ll for a similar "ensure the compiler doesn't crash on this shuffle" test and so perhaps this test can be added to that file?

[paulwalker-arm]

llvm/test/CodeGen/AArch64/sve-shuffle-crash.ll

1 ↗

(On Diff #406836)

@david-arm Given the test relies on <32 x i32> when the target vector length is 256-bit, I'm not sure we'll ever care about the code quality. If work is done to specifically improve it I would expect there to be a more focused test. Not saying your request is wrong, just something to consider.

[david-arm]

llvm/test/CodeGen/AArch64/sve-shuffle-crash.ll

1 ↗

(On Diff #406836)

OK fair enough. I guess I'm more worried about us missing something by not having any CHECK lines at all. Currently this test is just making sure we don't crash. I guess if @bsmith moves the test to another file it will at least need a CHECK-LABEL. :) Perhaps it would be good to have one CHECK line for a load and one for a store just to show we've generated something?

3 ↗

(On Diff #406836)

+1

[david-arm]

Hi @paulwalker-arm, this bug was actually found in user-written code in Gromacs, although only there was only one instance of this I think. So it is something users may see, just not very often!

[bsmith]

• Tighten type size check
• Move test to already existing file
• Add rudimentary check lines to test

[sdesmalen]

llvm/test/CodeGen/AArch64/sve-fixed-length-shuffles.ll

18

nit: For tests that rely on vscale_range I personally find the tests more readable if instead of referencing some attribute, it just has vscale_range(2, 2) here directly, and the RUN line has the +sve attribute. For this test it doesn't really matter too much I guess, but if the test file is long, the #<some number> make it really obscure what vscale_range it refers to.

[cameron.mcinally]

In D119252#3307242, @david-arm wrote:

Hi @paulwalker-arm, this bug was actually found in user-written code in Gromacs, although only there was only one instance of this I think. So it is something users may see, just not very often!

I found this separately (D119285) in a SPEC code, on a mixed width loop of 8b/64b vectors. But I agree, this isn't a common case.

[Matt]

FWIW, I can confirm the presence of ICE using the test case from https://reviews.llvm.org/D119285, https://llvm.godbolt.org/z/f1r6he545, as well as that this patch fixes the issue; thank you!

[paulwalker-arm]

llvm/test/CodeGen/AArch64/sve-fixed-length-shuffles.ll

18

I tend to agree with the first point when it comes to a single file when several different vscale_range values are being tested, although in this case all the tests are using the same value. The second point I personally disagree with because I prefer the convenience of generally being able to do llc < test.ll without having to always dig into the test file in order to run it.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

auto getScaledOffsetDup = [](SDValue BitCast, int &LaneC, MVT &CastVT) {
     SDValue Extract = BitCast.getOperand(0);
     unsigned ExtIdx = Extract.getConstantOperandVal(1);
     unsigned SrcEltBitWidth = Extract.getScalarValueSizeInBits();
     unsigned ExtIdxInBits = ExtIdx * SrcEltBitWidth;
     unsigned CastedEltBitWidth = BitCast.getScalarValueSizeInBits();
     if (ExtIdxInBits % CastedEltBitWidth != 0)
       return false;
 
+    // Can't handle cases where vector size is not 128-bit
+    if (!Extract.getOperand(0).getValueType().is128BitVector())
+      return false;
+
     // Update the lane value by offsetting with the scaled extract index.
     LaneC += ExtIdxInBits / CastedEltBitWidth;
 
     // Determine the casted vector type of the wide vector input.
     // dup (bitcast (extract_subv X, C)), LaneC --> dup (bitcast X), LaneC'
     // Examples:
     // dup (bitcast (extract_subv v2f64 X, 1) to v2f32), 1 --> dup v4f32 X, 3
     // dup (bitcast (extract_subv v16i8 X, 8) to v4i16), 1 --> dup v8i16 X, 5

llvm/test/CodeGen/AArch64/sve-fixed-length-shuffles.ll

-; RUN: llc -aarch64-sve-vector-bits-min=256 < %s | FileCheck %s
+; RUN: llc < %s | FileCheck %s
 
 target triple = "aarch64-unknown-linux-gnu"
 
 ; Currently there is no custom lowering for vector shuffles operating on types
 ; bigger than NEON. However, having no support opens us up to a code generator
 ; hang when expanding BUILD_VECTOR. Here we just validate the promblematic case
 ; successfully exits code generation.
 define void @hang_when_merging_stores_after_legalisation(<8 x i32>* %a, <2 x i32> %b) #0 {
 ; CHECK-LABEL: hang_when_merging_stores_after_legalisation:
   %splat = shufflevector <2 x i32> %b, <2 x i32> undef, <8 x i32> zeroinitializer
   %interleaved.vec = shufflevector <8 x i32> %splat, <8 x i32> undef, <8 x i32> <i32 0, i32 4, i32 1, i32 5, i32 2, i32 6, i32 3, i32 7>
   store <8 x i32> %interleaved.vec, <8 x i32>* %a, align 4
   ret void
 }
 
-attributes #0 = { nounwind "target-features"="+sve" }
+; Ensure we don't crash when trying to lower a shuffle via and extract
+define void @crash_when_lowering_extract_shuffle(<32 x i32>* %dst, i1 %cond) #0 {
+; CHECK-LABEL: crash_when_lowering_extract_shuffle:
+; CHECK:       ld1w { z3.s }, p0/z, [x0]
+; CHECK:       st1w { z3.s }, p0, [x0]
+  %broadcast.splat = shufflevector <32 x i1> zeroinitializer, <32 x i1> zeroinitializer, <32 x i32> zeroinitializer
+  br i1 %cond, label %exit, label %vector.body
+
+vector.body:
+  %1 = load <32 x i32>, <32 x i32>* %dst, align 16
+  %predphi = select <32 x i1> %broadcast.splat, <32 x i32> zeroinitializer, <32 x i32> %1
+  store <32 x i32> %predphi, <32 x i32>* %dst, align 16
+  br label %exit
+
+exit:
+  ret void
+}
+
+attributes #0 = { vscale_range(2,2) "target-features"="+sve" }