[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

1210

Why did you remove this condition?

1263

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

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

102

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

1210

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.

1263

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

1257

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

1257

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 },
static const CostTblEntry ShuffleTbl[] = {
       // PermuteSingleSrc shuffle kinds.
       // 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.
+      // Broadcast shuffle kinds for scalable vectors
+      { 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::nxv8bf16, 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;
   }
 
   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 void  @broadcast() #0{
+; CHECK-LABEL: 'broadcast':
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %zero = shufflevector <vscale x 16 x i8> undef, <vscale x 16 x i8> undef, <vscale x 16 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction:   %1 = shufflevector <vscale x 32 x i8> undef, <vscale x 32 x i8> undef, <vscale x 32 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %2 = shufflevector <vscale x 8 x i16> undef, <vscale x 8 x i16> undef, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %3 = shufflevector <vscale x 16 x i16> undef, <vscale x 16 x i16> undef, <vscale x 16 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %4 = shufflevector <vscale x 4 x i32> undef, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %5 = shufflevector <vscale x 8 x i32> undef, <vscale x 8 x i32> undef, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %6 = shufflevector <vscale x 2 x i64> undef, <vscale x 2 x i64> undef, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %7 = shufflevector <vscale x 4 x i64> undef, <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %8 = shufflevector <vscale x 8 x half> undef, <vscale x 8 x half> undef, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %9 = shufflevector <vscale x 16 x half> undef, <vscale x 16 x half> undef, <vscale x 16 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %10 = shufflevector <vscale x 4 x float> undef, <vscale x 4 x float> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %11 = shufflevector <vscale x 8 x float> undef, <vscale x 8 x float> undef, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %12 = shufflevector <vscale x 2 x double> undef, <vscale x 2 x double> undef, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %13 = shufflevector <vscale x 4 x double> undef, <vscale x 4 x double> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %14 = shufflevector <vscale x 8 x bfloat> undef, <vscale x 8 x bfloat> undef, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT: Cost Model: Found an estimated cost of 2 for instruction:   %15 = shufflevector <vscale x 16 x bfloat> undef, <vscale x 16 x bfloat> undef, <vscale x 16 x i32> zeroinitializer
+; CHECK-NETX: Cost Model: Found an estimated cost of 0 for instruction:   ret void
+
+  %zero = shufflevector <vscale x 16 x i8> undef, <vscale x 16 x i8> undef, <vscale x 16 x i32> zeroinitializer
+  %1 = shufflevector <vscale x 32 x i8> undef, <vscale x 32 x i8> undef, <vscale x 32 x i32> zeroinitializer
+  %2 = shufflevector <vscale x 8 x i16> undef, <vscale x 8 x i16> undef, <vscale x 8 x i32> zeroinitializer
+  %3 = shufflevector <vscale x 16 x i16> undef, <vscale x 16 x i16> undef, <vscale x 16 x i32> zeroinitializer
+  %4 = shufflevector <vscale x 4 x i32> undef, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+  %5 = shufflevector <vscale x 8 x i32> undef, <vscale x 8 x i32> undef, <vscale x 8 x i32> zeroinitializer
+  %6 = shufflevector <vscale x 2 x i64> undef, <vscale x 2 x i64> undef, <vscale x 2 x i32> zeroinitializer
+  %7 = shufflevector <vscale x 4 x i64> undef, <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer
+  %8 = shufflevector <vscale x 8 x half> undef, <vscale x 8 x half> undef, <vscale x 8 x i32> zeroinitializer
+  %9 = shufflevector <vscale x 16 x half> undef, <vscale x 16 x half> undef, <vscale x 16 x i32> zeroinitializer
+ %10 = shufflevector <vscale x 4 x float> undef, <vscale x 4 x float> undef, <vscale x 4 x i32> zeroinitializer
+  %11 = shufflevector <vscale x 8 x float> undef, <vscale x 8 x float> undef, <vscale x 8 x i32> zeroinitializer
+ %12 = shufflevector <vscale x 2 x double> undef, <vscale x 2 x double> undef, <vscale x 2 x i32> zeroinitializer
+ %13 = shufflevector <vscale x 4 x double> undef, <vscale x 4 x double> undef, <vscale x 4 x i32> zeroinitializer
+ %14 = shufflevector <vscale x 8 x bfloat> undef, <vscale x 8 x bfloat> undef, <vscale x 8 x i32> zeroinitializer
+ %15 = shufflevector <vscale x 16 x bfloat> undef, <vscale x 16 x bfloat> undef, <vscale x 16 x i32> zeroinitializer
+  ret void
+}
+
+attributes #0 = { "target-features"="+sve,+bf16" }