[AArch64] Add FP16 broadcast and transpose costs

[dmgreen]

Summary

The FP16 broadcast and transpose can always use the same instructions as are used for i16 vectors, with or without +fullfp16. This fills in some extra costs to make sure we get it right.

[SjoerdMeijer]

Looks like a good fix to me.

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

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},
         {TTI::SK_Broadcast, MVT::v2i64, 1},
+        {TTI::SK_Broadcast, MVT::v4f16, 1},
+        {TTI::SK_Broadcast, MVT::v8f16, 1},
         {TTI::SK_Broadcast, MVT::v2f32, 1},
         {TTI::SK_Broadcast, MVT::v4f32, 1},
         {TTI::SK_Broadcast, MVT::v2f64, 1},
         // Transpose shuffle kinds can be performed with 'trn1/trn2' and
         // 'zip1/zip2' instructions.
         {TTI::SK_Transpose, MVT::v8i8, 1},
         {TTI::SK_Transpose, MVT::v16i8, 1},
         {TTI::SK_Transpose, MVT::v4i16, 1},
         {TTI::SK_Transpose, MVT::v8i16, 1},
         {TTI::SK_Transpose, MVT::v2i32, 1},
         {TTI::SK_Transpose, MVT::v4i32, 1},
         {TTI::SK_Transpose, MVT::v2i64, 1},
+        {TTI::SK_Transpose, MVT::v4f16, 1},
+        {TTI::SK_Transpose, MVT::v8f16, 1},
         {TTI::SK_Transpose, MVT::v2f32, 1},
         {TTI::SK_Transpose, MVT::v4f32, 1},
         {TTI::SK_Transpose, MVT::v2f64, 1},
         // Select shuffle kinds.
         // TODO: handle vXi8/vXi16.
         {TTI::SK_Select, MVT::v2i32, 1}, // mov.
         {TTI::SK_Select, MVT::v4i32, 2}, // rev+trn (or similar).
         {TTI::SK_Select, MVT::v2i64, 1}, // mov.

llvm/test/Analysis/CostModel/AArch64/shuffle-load.ll

 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %sv4i32 = shufflevector <4 x i32> %lv4i32, <4 x i32> undef, <4 x i32> zeroinitializer
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction: %lv8i32 = load <8 x i32>, ptr undef, align 32
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction: %sv8i32 = shufflevector <8 x i32> %lv8i32, <8 x i32> undef, <8 x i32> zeroinitializer
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %lv2i64 = load <2 x i64>, ptr undef, align 16
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %sv2i64 = shufflevector <2 x i64> %lv2i64, <2 x i64> undef, <2 x i32> zeroinitializer
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction: %lv4i64 = load <4 x i64>, ptr undef, align 32
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction: %sv4i64 = shufflevector <4 x i64> %lv4i64, <4 x i64> undef, <4 x i32> zeroinitializer
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %lv2f16 = load <2 x half>, ptr undef, align 4
-; CHECK-NEXT:  Cost Model: Found an estimated cost of 3 for instruction: %sv2f16 = shufflevector <2 x half> %lv2f16, <2 x half> undef, <2 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %sv2f16 = shufflevector <2 x half> %lv2f16, <2 x half> undef, <2 x i32> zeroinitializer
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %lv4f16 = load <4 x half>, ptr undef, align 8
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %sv4f16 = shufflevector <4 x half> %lv4f16, <4 x half> undef, <4 x i32> zeroinitializer
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %lv8f16 = load <8 x half>, ptr undef, align 16
-; CHECK-NEXT:  Cost Model: Found an estimated cost of 21 for instruction: %sv8f16 = shufflevector <8 x half> %lv8f16, <8 x half> undef, <8 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %sv8f16 = shufflevector <8 x half> %lv8f16, <8 x half> undef, <8 x i32> zeroinitializer
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction: %lv16f16 = load <16 x half>, ptr undef, align 32
-; CHECK-NEXT:  Cost Model: Found an estimated cost of 42 for instruction: %sv16f16 = shufflevector <16 x half> %lv16f16, <16 x half> undef, <16 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction: %sv16f16 = shufflevector <16 x half> %lv16f16, <16 x half> undef, <16 x i32> zeroinitializer
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %lv2f32 = load <2 x float>, ptr undef, align 8
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %sv2f32 = shufflevector <2 x float> %lv2f32, <2 x float> undef, <2 x i32> zeroinitializer
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %lv4f32 = load <4 x float>, ptr undef, align 16
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %sv4f32 = shufflevector <4 x float> %lv4f32, <4 x float> undef, <4 x i32> zeroinitializer
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction: %lv8f32 = load <8 x float>, ptr undef, align 32
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction: %sv8f32 = shufflevector <8 x float> %lv8f32, <8 x float> undef, <8 x i32> zeroinitializer
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %lv2f64 = load <2 x double>, ptr undef, align 16
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %sv2f64 = shufflevector <2 x double> %lv2f64, <2 x double> undef, <2 x i32> zeroinitializer
entry:
   %lane = shufflevector <4 x half> %tmp1, <4 x half> undef, <4 x i32> zeroinitializer
   ret <4 x half> %lane
 }
 
 define <8 x half> @ld1r_8h_float_shuff(ptr nocapture %x) {
 ; CHECK-LABEL: 'ld1r_8h_float_shuff'
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %tmp = load half, ptr %x, align 2
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %tmp1 = insertelement <8 x half> undef, half %tmp, i32 0
-; CHECK-NEXT:  Cost Model: Found an estimated cost of 21 for instruction: %lane = shufflevector <8 x half> %tmp1, <8 x half> undef, <8 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %lane = shufflevector <8 x half> %tmp1, <8 x half> undef, <8 x i32> zeroinitializer
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: ret <8 x half> %lane
 ;
 ; CODESIZE-LABEL: 'ld1r_8h_float_shuff'
 ; CODESIZE-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %tmp = load half, ptr %x, align 2
 ; CODESIZE-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %tmp1 = insertelement <8 x half> undef, half %tmp, i32 0
-; CODESIZE-NEXT:  Cost Model: Found an estimated cost of 21 for instruction: %lane = shufflevector <8 x half> %tmp1, <8 x half> undef, <8 x i32> zeroinitializer
+; CODESIZE-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %lane = shufflevector <8 x half> %tmp1, <8 x half> undef, <8 x i32> zeroinitializer
 ; CODESIZE-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: ret <8 x half> %lane
 ;
 entry:
   %tmp = load half, ptr %x, align 2
   %tmp1 = insertelement <8 x half> undef, half %tmp, i32 0
   %lane = shufflevector <8 x half> %tmp1, <8 x half> undef, <8 x i32> zeroinitializer
   ret <8 x half> %lane
 }

llvm/test/Analysis/CostModel/AArch64/shuffle-transpose.ll

 ; COST-LABEL: trn2.v2f64
 ; COST:       Found an estimated cost of 1 for instruction: %tmp0 = shufflevector <2 x double> %v0, <2 x double> %v1, <2 x i32> <i32 1, i32 3>
 ; CODE-LABEL: trn2.v2f64
 ; CODE:       zip2 v0.2d, v0.2d, v1.2d
 define <2 x double> @trn2.v2f64(<2 x double> %v0, <2 x double> %v1) {
   %tmp0 = shufflevector <2 x double> %v0, <2 x double> %v1, <2 x i32> <i32 1, i32 3>
   ret <2 x double> %tmp0
 }
+
+; COST-LABEL: trn1.v4f16
+; COST:       Found an estimated cost of 1 for instruction: %tmp0 = shufflevector <4 x half> %v0, <4 x half> %v1, <4 x i32> <i32 0, i32 4, i32 2, i32 6>
+; CODE-LABEL: trn1.v4f16
+; CODE:       trn1 v0.4h, v0.4h, v1.4h
+define <4 x half> @trn1.v4f16(<4 x half> %v0, <4 x half> %v1) {
+  %tmp0 = shufflevector <4 x half> %v0, <4 x half> %v1, <4 x i32> <i32 0, i32 4, i32 2, i32 6>
+  ret <4 x half> %tmp0
+}
+
+; COST-LABEL: trn2.v4f16
+; COST:       Found an estimated cost of 1 for instruction: %tmp0 = shufflevector <4 x half> %v0, <4 x half> %v1, <4 x i32> <i32 1, i32 5, i32 3, i32 7>
+; CODE-LABEL: trn2.v4f16
+; CODE:       trn2 v0.4h, v0.4h, v1.4h
+define <4 x half> @trn2.v4f16(<4 x half> %v0, <4 x half> %v1) {
+  %tmp0 = shufflevector <4 x half> %v0, <4 x half> %v1, <4 x i32> <i32 1, i32 5, i32 3, i32 7>
+  ret <4 x half> %tmp0
+}
+
+; COST-LABEL: trn1.v8f16
+; COST:       Found an estimated cost of 1 for instruction: %tmp0 = shufflevector <8 x half> %v0, <8 x half> %v1, <8 x i32> <i32 0, i32 8, i32 2, i32 10, i32 4, i32 12, i32 6, i32 14>
+; CODE-LABEL: trn1.v8f16
+; CODE:       trn1 v0.8h, v0.8h, v1.8h
+define <8 x half> @trn1.v8f16(<8 x half> %v0, <8 x half> %v1) {
+  %tmp0 = shufflevector <8 x half> %v0, <8 x half> %v1, <8 x i32> <i32 0, i32 8, i32 2, i32 10, i32 4, i32 12, i32 6, i32 14>
+  ret <8 x half> %tmp0
+}
+
+; COST-LABEL: trn2.v8f16
+; COST:       Found an estimated cost of 1 for instruction: %tmp0 = shufflevector <8 x half> %v0, <8 x half> %v1, <8 x i32> <i32 1, i32 9, i32 3, i32 11, i32 5, i32 13, i32 7, i32 15>
+; CODE-LABEL: trn2.v8f16
+; CODE:       trn2 v0.8h, v0.8h, v1.8h
+define <8 x half> @trn2.v8f16(<8 x half> %v0, <8 x half> %v1) {
+  %tmp0 = shufflevector <8 x half> %v0, <8 x half> %v1, <8 x i32> <i32 1, i32 9, i32 3, i32 11, i32 5, i32 13, i32 7, i32 15>
+  ret <8 x half> %tmp0
+}