[AArch64][SVE]Add cost model for broadcast shuffle

[CarolineConcatto]

Summary

This patch adds a cost model for SK_Broadcast in
AArch64TTIImpl::getShuffleCost with scalable vector.
Without this patch, the scalable vector type relies on BasicTTIImpl cost
implementation and assert.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

1194

Why did you remove this condition?

1236

Can't these be moved into the table above? (ShuffleTbl)

llvm/test/Analysis/CostModel/AArch64/sve-shuffle-broadcast.ll

101

s/slapt/broadcast/ (same for tests below)

[CarolineConcatto]

-add original if statement
-move the table cost for scalable types to ShuffleTbl

[CarolineConcatto]

-replace test name @slapt_v4f32 by @broadcast_v4f32(

[CarolineConcatto]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

1194

I don't see any reason why we need this if here.
IMO it is a redundancy that I think should be removed if we have more then one cost table.
But as we are merging the cost for scalable vectors with fixed vectors I am leaving this if as it was.

1236

We can, but I thought that having one table cost for fixed vector and one for scalable makes the code clear.

[david-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

1239

I think we also need one for the type MVT::nxv8bf16 here.

1241

nit: Whitespace change

llvm/test/Analysis/CostModel/AArch64/sve-shuffle-broadcast.ll

12

Hi Carol, this is just a suggestion, but a reviewer (David Green) on one of my recent cost model patches suggested that we could do all tests in one function, i.e.

define void @broadcast_legal() {
%0 = shufflevector <vscale x 16 x i8> %v, <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
%1 = shufflevector <vscale x 32 x i8> %v, <vscale x 32 x i8> poison, <vscale x 32 x i32> zeroinitializer
...
}

and so on. I found it really useful as all the CHECK lines can go in the same function and it reduces the test size. He mentioned that in my patch here - https://reviews.llvm.org/D95039. It means we can reduce the number of functions and number of unnecessary checks for the "ret void" instruction.

[CarolineConcatto]

• simplify sve broadcast test

[CarolineConcatto]

• add bfloat test and support

[CarolineConcatto]

• correct style

[CarolineConcatto]

@Thank you all for the suggestion.
I have added the support for MVT::nxv8bf16 and simplify the test.

[david-arm]

LGTM with nit addressed!

llvm/test/Analysis/CostModel/AArch64/sve-shuffle-broadcast.ll

12

nit: It looks like '%v' isn't used anymore here so can delete?

[CarolineConcatto]

-remove non used parameter from the test function broadcast

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

if (const auto *Entry = CostTableLookup(CostTblNoPairwise, ISD, MTy))
     return LT.first * Entry->Cost;
 
   return BaseT::getArithmeticReductionCost(Opcode, ValTy, IsPairwiseForm,
                                            CostKind);
 }
 
 int AArch64TTIImpl::getShuffleCost(TTI::ShuffleKind Kind, VectorType *Tp,
                                    int Index, VectorType *SubTp) {
-  if (Kind == TTI::SK_Broadcast || Kind == TTI::SK_Transpose ||
-      Kind == TTI::SK_Select || Kind == TTI::SK_PermuteSingleSrc) {
     static const CostTblEntry ShuffleTbl[] = {
       // Broadcast shuffle kinds can be performed with 'dup'.
       { TTI::SK_Broadcast, MVT::v8i8,  1 },
       { TTI::SK_Broadcast, MVT::v16i8, 1 },
       { TTI::SK_Broadcast, MVT::v4i16, 1 },
       { TTI::SK_Broadcast, MVT::v8i16, 1 },
       { TTI::SK_Broadcast, MVT::v2i32, 1 },
       { TTI::SK_Broadcast, MVT::v4i32, 1 },
static const CostTblEntry ShuffleTbl[] = {
       // TODO: handle vXi8/vXi16.
       { TTI::SK_PermuteSingleSrc, MVT::v2i32, 1 }, // mov.
       { TTI::SK_PermuteSingleSrc, MVT::v4i32, 3 }, // perfectshuffle worst case.
       { TTI::SK_PermuteSingleSrc, MVT::v2i64, 1 }, // mov.
       { TTI::SK_PermuteSingleSrc, MVT::v2f32, 1 }, // mov.
       { TTI::SK_PermuteSingleSrc, MVT::v4f32, 3 }, // perfectshuffle worst case.
       { TTI::SK_PermuteSingleSrc, MVT::v2f64, 1 }, // mov.
     };
+
+    static const CostTblEntry ScalableShuffleTbl[] = {
+        {TTI::SK_Broadcast, MVT::nxv16i8, 1},
+        {TTI::SK_Broadcast, MVT::nxv8i16, 1},
+        {TTI::SK_Broadcast, MVT::nxv4i32, 1},
+        {TTI::SK_Broadcast, MVT::nxv2i64, 1},
+        {TTI::SK_Broadcast, MVT::nxv8f16, 1},
+        {TTI::SK_Broadcast, MVT::nxv4f32, 1},
+        {TTI::SK_Broadcast, MVT::nxv2f64, 1},
+    };
+
     std::pair<int, MVT> LT = TLI->getTypeLegalizationCost(DL, Tp);
+
     if (const auto *Entry = CostTableLookup(ShuffleTbl, Kind, LT.second))
       return LT.first * Entry->Cost;
-  }
+
+    if (const auto *Entry =
+            CostTableLookup(ScalableShuffleTbl, Kind, LT.second))
+      return LT.first * Entry->Cost;
 
     return BaseT::getShuffleCost(Kind, Tp, Index, SubTp);
 }

llvm/test/Analysis/CostModel/AArch64/sve-shuffle-broadcast.ll

+; Check getShuffleCost for SK_BroadCast with scalable vector
+
+; RUN: opt -cost-model -analyze -mtriple=aarch64--linux-gnu -mattr=sve  < %s 2>%t | FileCheck %s
+
+; RUN: FileCheck --check-prefix=WARN --allow-empty %s <%t
+
+; If this check fails please read test/CodeGen/AArch64/README for instructions on how to resolve it.
+; WARN-NOT: warning
+
+
+define <vscale x 16 x i8> @broadcast_v16i8(<vscale x 16 x i8> %v) {
+; CHECK-LABEL: 'broadcast_v16i8':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %retval = shufflevector <vscale x 16 x i8> %v, <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 16 x i8> %retval
+
+  %retval = shufflevector <vscale x 16 x i8> %v, <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
+  ret <vscale x 16 x i8> %retval
+}
+
+define <vscale x 32 x i8> @broadcast_v32i8(<vscale x 32 x i8> %v) {
+; CHECK-LABEL: 'broadcast_v32i8':
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction:   %retval = shufflevector <vscale x 32 x i8> %v, <vscale x 32 x i8> poison, <vscale x 32 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 32 x i8> %retval
+
+  %retval = shufflevector <vscale x 32 x i8> %v, <vscale x 32 x i8> poison, <vscale x 32 x i32> zeroinitializer
+  ret <vscale x 32 x i8> %retval
+}
+
+define <vscale x 8 x i16> @broadcast_v8i16(<vscale x 8 x i16> %v) {
+; CHECK-LABEL: 'broadcast_v8i16':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %retval = shufflevector <vscale x 8 x i16> %v, <vscale x 8 x i16> poison, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 8 x i16> %retval
+
+  %retval = shufflevector <vscale x 8 x i16> %v, <vscale x 8 x i16> poison, <vscale x 8 x i32> zeroinitializer
+  ret <vscale x 8 x i16> %retval
+}
+
+define <vscale x 16 x i16> @broadcast_v16i16(<vscale x 16 x i16> %v) {
+; CHECK-LABEL: 'broadcast_v16i16':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %retval = shufflevector <vscale x 16 x i16> %v, <vscale x 16 x i16> poison, <vscale x 16 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 16 x i16> %retval
+
+  %retval = shufflevector <vscale x 16 x i16> %v, <vscale x 16 x i16> poison, <vscale x 16 x i32> zeroinitializer
+  ret <vscale x 16 x i16> %retval
+}
+
+define <vscale x 4 x i32> @broadcast_v4i32(<vscale x 4 x i32> %v) {
+; CHECK-LABEL: 'broadcast_v4i32':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %retval = shufflevector <vscale x 4 x i32> %v, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 4 x i32> %retval
+
+  %retval = shufflevector <vscale x 4 x i32> %v, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+  ret <vscale x 4 x i32> %retval
+}
+
+define <vscale x 8 x i32> @broadcast_v8i32(<vscale x 8 x i32> %v) {
+; CHECK-LABEL: 'broadcast_v8i32':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %retval = shufflevector <vscale x 8 x i32> %v, <vscale x 8 x i32> poison, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 8 x i32> %retval
+
+  %retval = shufflevector <vscale x 8 x i32> %v, <vscale x 8 x i32> poison, <vscale x 8 x i32> zeroinitializer
+  ret <vscale x 8 x i32> %retval
+}
+
+define <vscale x 2 x i64> @broadcast_v2i64(<vscale x 2 x i64> %v) {
+; CHECK-LABEL: 'broadcast_v2i64':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %retval = shufflevector <vscale x 2 x i64> %v, <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 2 x i64> %retval
+
+  %retval = shufflevector <vscale x 2 x i64> %v, <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
+  ret <vscale x 2 x i64> %retval
+}
+
+define <vscale x 4 x i64> @broadcast_v4i64(<vscale x 4 x i64> %v) {
+; CHECK-LABEL: 'broadcast_v4i64':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %retval = shufflevector <vscale x 4 x i64> %v, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 4 x i64> %retval
+
+  %retval = shufflevector <vscale x 4 x i64> %v, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  ret <vscale x 4 x i64> %retval
+}
+
+define <vscale x 8 x half> @broadcast_v8f16(<vscale x 8 x half> %v) {
+; CHECK-LABEL: 'broadcast_v8f16':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %retval = shufflevector <vscale x 8 x half> %v, <vscale x 8 x half> poison, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 8 x half> %retval
+
+  %retval = shufflevector <vscale x 8 x half> %v, <vscale x 8 x half> poison, <vscale x 8 x i32> zeroinitializer
+  ret <vscale x 8 x half> %retval
+}
+
+define <vscale x 16 x half> @broadcast_v16f16(<vscale x 16 x half> %v) {
+; CHECK-LABEL: 'broadcast_v16f16':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %retval = shufflevector <vscale x 16 x half> %v, <vscale x 16 x half> poison, <vscale x 16 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 16 x half> %retval
+
+  %retval = shufflevector <vscale x 16 x half> %v, <vscale x 16 x half> poison, <vscale x 16 x i32> zeroinitializer
+  ret <vscale x 16 x half> %retval
+}
+
+define <vscale x 4 x float> @slapt_v4f32(<vscale x 4 x float> %v){
+; CHECK-LABEL: 'slapt_v4f32':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %retval = shufflevector <vscale x 4 x float> %v, <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 4 x float> %retval
+
+   %retval = shufflevector <vscale x 4 x float> %v, <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer
+  ret <vscale x 4 x float> %retval
+}
+
+define <vscale x 8 x float> @slapt_v8f32(<vscale x 8 x float> %v){
+; CHECK-LABEL: 'slapt_v8f32':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %retval = shufflevector <vscale x 8 x float> %v, <vscale x 8 x float> poison, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 8 x float> %retval
+
+   %retval = shufflevector <vscale x 8 x float> %v, <vscale x 8 x float> poison, <vscale x 8 x i32> zeroinitializer
+  ret <vscale x 8 x float> %retval
+}
+define <vscale x 2 x double> @slapt_v2f64(<vscale x 2 x double> %v){
+; CHECK-LABEL: 'slapt_v2f64':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %retval = shufflevector <vscale x 2 x double> %v, <vscale x 2 x double> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 2 x double> %retval
+
+   %retval = shufflevector <vscale x 2 x double> %v, <vscale x 2 x double> poison, <vscale x 2 x i32> zeroinitializer
+  ret <vscale x 2 x double> %retval
+}
+
+define <vscale x 4 x double> @slapt_v4f64(<vscale x 4 x double> %v){
+; CHECK-LABEL: 'slapt_v4f64':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %retval = shufflevector <vscale x 4 x double> %v, <vscale x 4 x double> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 4 x double> %retval
+
+   %retval = shufflevector <vscale x 4 x double> %v, <vscale x 4 x double> poison, <vscale x 4 x i32> zeroinitializer
+  ret <vscale x 4 x double> %retval
+}