[CostModel][AArch64] Add some initial costs for SK_Select and SK_PermuteSingleSrc

[RKSimon]

Summary

AArch64 was only setting costs for SK_Transpose, which meant that many of the simpler shuffles (e.g. SK_Select and SK_PermuteSingleSrc for larger vector elements) was being severely overestimated by the default shuffle expansion.

This patch adds costs to help improve SLP performance and avoid a regression in reductions in an upcoming patch that will allow SLP to recognise a lot more SK_Select shuffle patterns.

I'm not very knowledgable about AArch64 shuffle lowering so I've kept the extra costs to a minimum - someone who knows this code can add extra costs which should improve vectorization a lot more.

[evandro]

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

963

I'm not sure that the costs of 64 and 128 bits long vectors should be different.

978

Ditto.

[RKSimon]

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

963

I'm going off the codegen in the shuffle-select.ll tests below where 2X32 selects are a move and 4X32 selects are a rev+trn.

If you have better cost estimates please suggest them - I don't know much about the aarch64 microarch and I'm really just wanting to fix regressions in D48174 caused by the existing (dubious) shuffle costs.

978

I wasn't sure what the best form shuffle test were for these cases - do you know what the likely worst offending shuffle masks will be? If you can give me some hints I'd be happy to create a shuffle-single-src.ll costs test file to at least cover these types.

[evandro]

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

963

That's an artifact due to the mask values being different for both tests. Were they, say, <0, 2> in sel.v2i32() and <0, 2, 4, 6> in sel.v4i32(), both would result in a single instruction performing the shuffle.

[RKSimon]

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

963

But this is for SK_Select style patterns:

2x32 - <0,3> or <2,1> - this will always be 1 mov AFAICT
4x32 - <0,1,2,7>, <0,1,6,3>, <0,5,2,7>, etc. - I couldn't cause this to be anything more than 2 instructions, sometimes independent othertimes not.

That was the reasoning behind the costs - I'm a litttle more unsure about what would is the maximum for SK_PermuteSingleSrc but I reckoned it would be 2 instructions as well.

[RKSimon]

Can anyone recommend the best shuffle costs to put here for the 32/64 cases?

[evandro]

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

964

Independently of the costs, this lookup table could be simplified to a test on the number of vector elements.

979

Ditto.

[RKSimon]

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

964

That's trivial to do - the tables look easier to grok, but I'm happy to make the change.

[RKSimon]

Removed shuffle cost tables

[RKSimon]

Rebased now that D48174 has landed - shows the full vectorisation of the reduction_v4i32 test (prior to D48174 it only vectorised the second half).

[evandro]

Next time, when one patch causes a regression that another patch fixes, please, make the one a parent of the other.

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

967

Perhaps you should assert instead of testing here.

976

Ditto.

[evandro]

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

965

Please, comment on what you are trying to achieve here.

974

And here too.

[RKSimon]

I'm trying a different approach which should work better as additional costs are added for other types or shuffle kinds - this will be beneficial to aarch64 as SLP etc. starts using the ShuffleKind matches more directly (see D48236).

I've gone back to using the costs tables as they tend to be easier to understand at a glance and the number of entries/calls aren't really large enough to be a perf issue. I've also changed the key to use the ShuffleKind, which matches what we do on X86, again for clarity.

[evandro]

It looks reasonable to me. If no one else objects in a couple of days or so, it LGTM.

[efriedma]

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

972

I'm pretty sure the general case here is not rev+trn, but I guess it should be two instructions in general (it's always at most two vector-vector element insertions).

979

For a single source with two elements, a shuffle is either an identity shuffle, a splat, or a reversal; I guess those are all in fact 1 instruction.

980

egrep -c "[0-3],[0-3],[0-3],[0-3].*Cost [3-4]" lib/Target/AArch64/AArch64PerfectShuffle.h says there are 182 permutes which cost more than 2. But I guess this is a reasonable estimate not knowing anything else.

[RKSimon]

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

980

Out of interest, are those perfect shuffle costs based off throughput or number of instructions?

[RKSimon]

Updated SK_Select comment to make it clear that the shuffle could be a rev+trn, or something else....

[RKSimon]

Use worst case PermuteSingleSrc costs from perfect shuffle tool

[efriedma]

LGTM

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

980

All the shuffles it uses have the same throughput on an A57, so there isn't really any difference.

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

static const CostTblEntry TransposeTbl[] = {
         {ISD::VECTOR_SHUFFLE, MVT::v2f64, 1},
     };
     std::pair<int, MVT> LT = TLI->getTypeLegalizationCost(DL, Tp);
     if (const auto *Entry =
             CostTableLookup(TransposeTbl, ISD::VECTOR_SHUFFLE, LT.second))
       return LT.first * Entry->Cost;
   }
 
+  if (Kind == TTI::SK_Select) {
+    std::pair<int, MVT> LT = TLI->getTypeLegalizationCost(DL, Tp);
+    unsigned EltBits = LT.second.getScalarSizeInBits();
+    if (EltBits == 64)
+      return LT.first;
+    if (EltBits == 32)
+        return LT.first * LT.second.getSizeInBits() / 64;
+  }
+
+  if (Kind == TTI::SK_PermuteSingleSrc) {
+    std::pair<int, MVT> LT = TLI->getTypeLegalizationCost(DL, Tp);
+    unsigned EltBits = LT.second.getScalarSizeInBits();
+    if (EltBits == 64)
+      return LT.first;
+    if (EltBits == 32)
+      return LT.first * LT.second.getSizeInBits() / 64;
+  }
+
   return BaseT::getShuffleCost(Kind, Tp, Index, SubTp);
 }

test/Analysis/CostModel/AArch64/shuffle-select.ll

 ; CODE-LABEL: sel.v8i16
 ; CODE:       tbl v0.16b, { v0.16b, v1.16b }, v2.16b
 define <8 x i16> @sel.v8i16(<8 x i16> %v0, <8 x i16> %v1) {
   %tmp0 = shufflevector <8 x i16> %v0, <8 x i16> %v1, <8 x i32> <i32 0, i32 9, i32 2, i32 11, i32 4, i32 13, i32 6, i32 15>
   ret <8 x i16> %tmp0
 }
 
 ; COST-LABEL: sel.v2i32
-; COST:        Found an estimated cost of 6 for instruction: %tmp0 = shufflevector <2 x i32> %v0, <2 x i32> %v1, <2 x i32> <i32 0, i32 3>
+; COST:        Found an estimated cost of 1 for instruction: %tmp0 = shufflevector <2 x i32> %v0, <2 x i32> %v1, <2 x i32> <i32 0, i32 3>
 ; CODE-LABEL: sel.v2i32
 ; CODE:       mov v0.s[1], v1.s[1]
 define <2 x i32> @sel.v2i32(<2 x i32> %v0, <2 x i32> %v1) {
   %tmp0 = shufflevector <2 x i32> %v0, <2 x i32> %v1, <2 x i32> <i32 0, i32 3>
   ret <2 x i32> %tmp0
 }
 
 ; COST-LABEL: sel.v4i32
-; COST:       Found an estimated cost of 18 for instruction: %tmp0 = shufflevector <4 x i32> %v0, <4 x i32> %v1, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+; COST:       Found an estimated cost of 2 for instruction: %tmp0 = shufflevector <4 x i32> %v0, <4 x i32> %v1, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
 ; CODE-LABEL: sel.v4i32
 ; CODE:       rev64 v0.4s, v0.4s
 ; CODE:       trn2 v0.4s, v0.4s, v1.4s
 define <4 x i32> @sel.v4i32(<4 x i32> %v0, <4 x i32> %v1) {
   %tmp0 = shufflevector <4 x i32> %v0, <4 x i32> %v1, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
   ret <4 x i32> %tmp0
 }
 
 ; COST-LABEL: sel.v2i64
-; COST:       Found an estimated cost of 6 for instruction: %tmp0 = shufflevector <2 x i64> %v0, <2 x i64> %v1, <2 x i32> <i32 0, i32 3>
+; COST:       Found an estimated cost of 1 for instruction: %tmp0 = shufflevector <2 x i64> %v0, <2 x i64> %v1, <2 x i32> <i32 0, i32 3>
 ; CODE-LABEL: sel.v2i64
 ; CODE:       mov v0.d[1], v1.d[1]
 define <2 x i64> @sel.v2i64(<2 x i64> %v0, <2 x i64> %v1) {
   %tmp0 = shufflevector <2 x i64> %v0, <2 x i64> %v1, <2 x i32> <i32 0, i32 3>
   ret <2 x i64> %tmp0
 }
 
 ; COST-LABEL: sel.v2f32
-; COST:       Found an estimated cost of 6 for instruction: %tmp0 = shufflevector <2 x float> %v0, <2 x float> %v1, <2 x i32> <i32 0, i32 3>
+; COST:       Found an estimated cost of 1 for instruction: %tmp0 = shufflevector <2 x float> %v0, <2 x float> %v1, <2 x i32> <i32 0, i32 3>
 ; CODE-LABEL: sel.v2f32
 ; CODE:       mov v0.s[1], v1.s[1]
 define <2 x float> @sel.v2f32(<2 x float> %v0, <2 x float> %v1) {
   %tmp0 = shufflevector <2 x float> %v0, <2 x float> %v1, <2 x i32> <i32 0, i32 3>
   ret <2 x float> %tmp0
 }
 
 ; COST-LABEL: sel.v4f32
-; COST:       Found an estimated cost of 18 for instruction: %tmp0 = shufflevector <4 x float> %v0, <4 x float> %v1, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+; COST:       Found an estimated cost of 2 for instruction: %tmp0 = shufflevector <4 x float> %v0, <4 x float> %v1, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
 ; CODE-LABEL: sel.v4f32
 ; CODE:       rev64 v0.4s, v0.4s
 ; CODE:       trn2 v0.4s, v0.4s, v1.4s
 define <4 x float> @sel.v4f32(<4 x float> %v0, <4 x float> %v1) {
   %tmp0 = shufflevector <4 x float> %v0, <4 x float> %v1, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
   ret <4 x float> %tmp0
 }
 
 ; COST-LABEL: sel.v2f64
-; COST:       Found an estimated cost of 6 for instruction: %tmp0 = shufflevector <2 x double> %v0, <2 x double> %v1, <2 x i32> <i32 0, i32 3>
+; COST:       Found an estimated cost of 1 for instruction: %tmp0 = shufflevector <2 x double> %v0, <2 x double> %v1, <2 x i32> <i32 0, i32 3>
 ; CODE-LABEL: sel.v2f64
 ; CODE:       mov v0.d[1], v1.d[1]
 define <2 x double> @sel.v2f64(<2 x double> %v0, <2 x double> %v1) {
   %tmp0 = shufflevector <2 x double> %v0, <2 x double> %v1, <2 x i32> <i32 0, i32 3>
   ret <2 x double> %tmp0
 }

test/Transforms/SLPVectorizer/AArch64/transpose.ll

;
   %tmp2.1 = add i64 %tmp1.0, %tmp1.1
   store i64 %tmp2.0, i64* %c.0, align 8
   store i64 %tmp2.1, i64* %c.1, align 8
   ret void
 }
 
 define <4 x i32> @build_vec_v4i32(<4 x i32> %v0, <4 x i32> %v1) {
 ; CHECK-LABEL: @build_vec_v4i32(
-; CHECK-NEXT:    [[V0_0:%.*]] = extractelement <4 x i32> %v0, i32 0
-; CHECK-NEXT:    [[V0_1:%.*]] = extractelement <4 x i32> %v0, i32 1
-; CHECK-NEXT:    [[V0_2:%.*]] = extractelement <4 x i32> %v0, i32 2
-; CHECK-NEXT:    [[V0_3:%.*]] = extractelement <4 x i32> %v0, i32 3
-; CHECK-NEXT:    [[V1_0:%.*]] = extractelement <4 x i32> %v1, i32 0
-; CHECK-NEXT:    [[V1_1:%.*]] = extractelement <4 x i32> %v1, i32 1
-; CHECK-NEXT:    [[V1_2:%.*]] = extractelement <4 x i32> %v1, i32 2
-; CHECK-NEXT:    [[V1_3:%.*]] = extractelement <4 x i32> %v1, i32 3
-; CHECK-NEXT:    [[TMP0_0:%.*]] = add i32 [[V0_0]], [[V1_0]]
-; CHECK-NEXT:    [[TMP0_1:%.*]] = add i32 [[V0_1]], [[V1_1]]
-; CHECK-NEXT:    [[TMP0_2:%.*]] = add i32 [[V0_2]], [[V1_2]]
-; CHECK-NEXT:    [[TMP0_3:%.*]] = add i32 [[V0_3]], [[V1_3]]
-; CHECK-NEXT:    [[TMP1_0:%.*]] = sub i32 [[V0_0]], [[V1_0]]
-; CHECK-NEXT:    [[TMP1_1:%.*]] = sub i32 [[V0_1]], [[V1_1]]
-; CHECK-NEXT:    [[TMP1_2:%.*]] = sub i32 [[V0_2]], [[V1_2]]
-; CHECK-NEXT:    [[TMP1_3:%.*]] = sub i32 [[V0_3]], [[V1_3]]
-; CHECK-NEXT:    [[TMP2_0:%.*]] = add i32 [[TMP0_0]], [[TMP0_1]]
-; CHECK-NEXT:    [[TMP2_1:%.*]] = add i32 [[TMP1_0]], [[TMP1_1]]
-; CHECK-NEXT:    [[TMP2_2:%.*]] = add i32 [[TMP0_2]], [[TMP0_3]]
-; CHECK-NEXT:    [[TMP2_3:%.*]] = add i32 [[TMP1_2]], [[TMP1_3]]
-; CHECK-NEXT:    [[TMP3_0:%.*]] = insertelement <4 x i32> undef, i32 [[TMP2_0]], i32 0
-; CHECK-NEXT:    [[TMP3_1:%.*]] = insertelement <4 x i32> [[TMP3_0]], i32 [[TMP2_1]], i32 1
-; CHECK-NEXT:    [[TMP3_2:%.*]] = insertelement <4 x i32> [[TMP3_1]], i32 [[TMP2_2]], i32 2
-; CHECK-NEXT:    [[TMP3_3:%.*]] = insertelement <4 x i32> [[TMP3_2]], i32 [[TMP2_3]], i32 3
-; CHECK-NEXT:    ret <4 x i32> [[TMP3_3]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i32> %v0, <4 x i32> undef, <2 x i32> <i32 0, i32 2>
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <2 x i32> [[TMP1]], <2 x i32> undef, <4 x i32> <i32 0, i32 0, i32 1, i32 1>
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <4 x i32> %v1, <4 x i32> undef, <2 x i32> <i32 0, i32 2>
+; CHECK-NEXT:    [[SHUFFLE1:%.*]] = shufflevector <2 x i32> [[TMP2]], <2 x i32> undef, <4 x i32> <i32 0, i32 0, i32 1, i32 1>
+; CHECK-NEXT:    [[TMP3:%.*]] = add <4 x i32> [[SHUFFLE]], [[SHUFFLE1]]
+; CHECK-NEXT:    [[TMP4:%.*]] = sub <4 x i32> [[SHUFFLE]], [[SHUFFLE1]]
+; CHECK-NEXT:    [[TMP5:%.*]] = shufflevector <4 x i32> [[TMP3]], <4 x i32> [[TMP4]], <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+; CHECK-NEXT:    [[TMP6:%.*]] = shufflevector <4 x i32> %v0, <4 x i32> undef, <2 x i32> <i32 1, i32 3>
+; CHECK-NEXT:    [[SHUFFLE2:%.*]] = shufflevector <2 x i32> [[TMP6]], <2 x i32> undef, <4 x i32> <i32 0, i32 0, i32 1, i32 1>
+; CHECK-NEXT:    [[TMP7:%.*]] = shufflevector <4 x i32> %v1, <4 x i32> undef, <2 x i32> <i32 1, i32 3>
+; CHECK-NEXT:    [[SHUFFLE3:%.*]] = shufflevector <2 x i32> [[TMP7]], <2 x i32> undef, <4 x i32> <i32 0, i32 0, i32 1, i32 1>
+; CHECK-NEXT:    [[TMP8:%.*]] = add <4 x i32> [[SHUFFLE2]], [[SHUFFLE3]]
+; CHECK-NEXT:    [[TMP9:%.*]] = sub <4 x i32> [[SHUFFLE2]], [[SHUFFLE3]]
+; CHECK-NEXT:    [[TMP10:%.*]] = shufflevector <4 x i32> [[TMP8]], <4 x i32> [[TMP9]], <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+; CHECK-NEXT:    [[TMP11:%.*]] = add <4 x i32> [[TMP5]], [[TMP10]]
+; CHECK-NEXT:    ret <4 x i32> [[TMP11]]
 ;
   %v0.0 = extractelement <4 x i32> %v0, i32 0
   %v0.1 = extractelement <4 x i32> %v0, i32 1
   %v0.2 = extractelement <4 x i32> %v0, i32 2
   %v0.3 = extractelement <4 x i32> %v0, i32 3
   %v1.0 = extractelement <4 x i32> %v1, i32 0
   %v1.1 = extractelement <4 x i32> %v1, i32 1
   %v1.2 = extractelement <4 x i32> %v1, i32 2
;
   %tmp3.1 = insertelement <4 x i32> %tmp3.0, i32 %tmp2.1, i32 1
   %tmp3.2 = insertelement <4 x i32> %tmp3.1, i32 %tmp2.2, i32 2
   %tmp3.3 = insertelement <4 x i32> %tmp3.2, i32 %tmp2.3, i32 3
   ret <4 x i32> %tmp3.3
 }
 
 define <4 x i32> @build_vec_v4i32_reuse_0(<2 x i32> %v0, <2 x i32> %v1) {
 ; CHECK-LABEL: @build_vec_v4i32_reuse_0(
-; CHECK-NEXT:    [[V0_0:%.*]] = extractelement <2 x i32> %v0, i32 0
-; CHECK-NEXT:    [[V0_1:%.*]] = extractelement <2 x i32> %v0, i32 1
-; CHECK-NEXT:    [[V1_0:%.*]] = extractelement <2 x i32> %v1, i32 0
-; CHECK-NEXT:    [[V1_1:%.*]] = extractelement <2 x i32> %v1, i32 1
-; CHECK-NEXT:    [[TMP0_0:%.*]] = add i32 [[V0_0]], [[V1_0]]
-; CHECK-NEXT:    [[TMP0_1:%.*]] = add i32 [[V0_1]], [[V1_1]]
-; CHECK-NEXT:    [[TMP1_0:%.*]] = sub i32 [[V0_0]], [[V1_0]]
-; CHECK-NEXT:    [[TMP1_1:%.*]] = sub i32 [[V0_1]], [[V1_1]]
-; CHECK-NEXT:    [[TMP2_0:%.*]] = add i32 [[TMP0_0]], [[TMP0_1]]
-; CHECK-NEXT:    [[TMP2_1:%.*]] = add i32 [[TMP1_0]], [[TMP1_1]]
-; CHECK-NEXT:    [[TMP3_0:%.*]] = insertelement <4 x i32> undef, i32 [[TMP2_0]], i32 0
-; CHECK-NEXT:    [[TMP3_1:%.*]] = insertelement <4 x i32> [[TMP3_0]], i32 [[TMP2_1]], i32 1
-; CHECK-NEXT:    [[TMP3_2:%.*]] = insertelement <4 x i32> [[TMP3_1]], i32 [[TMP2_0]], i32 2
-; CHECK-NEXT:    [[TMP3_3:%.*]] = insertelement <4 x i32> [[TMP3_2]], i32 [[TMP2_1]], i32 3
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <2 x i32> %v0, <2 x i32> undef, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <2 x i32> %v1, <2 x i32> undef, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP3:%.*]] = add <2 x i32> [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[TMP4:%.*]] = sub <2 x i32> [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[TMP5:%.*]] = shufflevector <2 x i32> [[TMP3]], <2 x i32> [[TMP4]], <2 x i32> <i32 0, i32 3>
+; CHECK-NEXT:    [[TMP6:%.*]] = shufflevector <2 x i32> %v0, <2 x i32> undef, <2 x i32> <i32 1, i32 1>
+; CHECK-NEXT:    [[TMP7:%.*]] = shufflevector <2 x i32> %v1, <2 x i32> undef, <2 x i32> <i32 1, i32 1>
+; CHECK-NEXT:    [[TMP8:%.*]] = add <2 x i32> [[TMP6]], [[TMP7]]
+; CHECK-NEXT:    [[TMP9:%.*]] = sub <2 x i32> [[TMP6]], [[TMP7]]
+; CHECK-NEXT:    [[TMP10:%.*]] = shufflevector <2 x i32> [[TMP8]], <2 x i32> [[TMP9]], <2 x i32> <i32 0, i32 3>
+; CHECK-NEXT:    [[TMP11:%.*]] = add <2 x i32> [[TMP5]], [[TMP10]]
+; CHECK-NEXT:    [[TMP3_3:%.*]] = shufflevector <2 x i32> [[TMP11]], <2 x i32> undef, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
 ; CHECK-NEXT:    ret <4 x i32> [[TMP3_3]]
 ;
   %v0.0 = extractelement <2 x i32> %v0, i32 0
   %v0.1 = extractelement <2 x i32> %v0, i32 1
   %v1.0 = extractelement <2 x i32> %v1, i32 0
   %v1.1 = extractelement <2 x i32> %v1, i32 1
   %tmp0.0 = add i32 %v0.0, %v1.0
   %tmp0.1 = add i32 %v0.1, %v1.1