[AArch64] Add f16 fpimm patterns

[RosieSumpter]

Summary

Depends on D129623

This patch recognizes f16 immediates as legal and adds the necessary
patterns. This allows the fadda folding introduced in D129623 to be
applied to the f16 cases.

[RosieSumpter]

• Removed comment from sve-fadda-select.ll test

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1129

I think if this is added we need to add extra lowering for FP16 constants of any kind, and it probably shouldn't be guarded by hasSVE but done in general.

Any i16 constant can be materialized with a single mov, so the worst case is that plus a gpr->fpr, which is likely better than the load we currently produce.
https://godbolt.org/z/r59arn5bq

[RosieSumpter]

• Updated title and summary
• Made f16 immediates legal for hasFullFP16 rather than just SVE
• Added f16 fpimm patterns
• Updated tests

[RosieSumpter]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1129

Hi Dave, thanks for having a look. I've had a go at doing this - let me know if it looks alright.

[dmgreen]

Yeah seems OK as far as I can tell.

llvm/lib/Target/AArch64/AArch64InstrInfo.td

1629 ↗

(On Diff #446167)

COPY's are quite loose in their definition, but it is probably best to keep them copying between values of the same size. I think this can use FMOVWHr directly, so long as the instruction has already been defined.
It should probably have a: let Predicates = [HasFullFP16] too.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

1629 ↗

(On Diff #446167)

Is HasFullFP16 relevant? The lowering requires no special instructions so I'd even argue the setOperationAction call requires no protection. What's important is whether the f16 type is legal. Once we get to operation lowering and isel we only need to protect instructions not types.

We made this mistake for bf16 for SVE and it just lead to unnecessary complexity which was eventually removed.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

1629 ↗

(On Diff #446167)

I've checked and FMOVWHr does require HasFullFP16 so the question becomes, should we allow good code generation even for the case when HasFullFP16 is not available? Or is HasFullFP16 so universally available these days for us not to care?

[dmgreen]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

1629 ↗

(On Diff #446167)

FMOVWHr requires FullFP16. It could use FMOVWSr instead, but then the types are a little awkward.

The call to setOperationAction is currently inside a HasFullFP16 block, so It could be moved out and we could always be using generic instructions. If we do go that route then make sure that we have tests for it.

[RosieSumpter]

• Use FMOVWHr instead of COPY_TO_REGCLASS
• Add let Predicates = [HasFullFP16]
• Removed code to copy from GPR32 to FPR16
• Updated a couple of recently added tests

[dmgreen]

LGTM. Thanks.

I've checked and FMOVWHr does require HasFullFP16 so the question becomes, should we allow good code generation even for the case when HasFullFP16 is not available? Or is HasFullFP16 so universally available these days for us not to care?

I think it's fine for this patch. It's a nice general improvement for when fp16 is available, and fixes the original issue of recognizing scalable splat constants.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Subtarget->hasSVE() || Subtarget->hasSME()) {
     for (auto VT :
          {MVT::nxv16i1, MVT::nxv8i1, MVT::nxv4i1, MVT::nxv2i1, MVT::nxv1i1}) {
       setOperationAction(ISD::SPLAT_VECTOR, VT, Custom);
       setOperationAction(ISD::EXTRACT_VECTOR_ELT, VT, Custom);
     }
   }
 
   if (Subtarget->hasSVE()) {
+    setOperationAction(ISD::ConstantFP, MVT::f16, Legal);
+
     for (auto VT : {MVT::nxv16i8, MVT::nxv8i16, MVT::nxv4i32, MVT::nxv2i64}) {
       setOperationAction(ISD::BITREVERSE, VT, Custom);
       setOperationAction(ISD::BSWAP, VT, Custom);
       setOperationAction(ISD::CTLZ, VT, Custom);
       setOperationAction(ISD::CTPOP, VT, Custom);
       setOperationAction(ISD::CTTZ, VT, Custom);
       setOperationAction(ISD::INSERT_SUBVECTOR, VT, Custom);
       setOperationAction(ISD::UINT_TO_FP, VT, Custom);

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

def SDT_AArch64ReduceWithInit : SDTypeProfile<1, 3,
    [SDTCisVec<1>, SDTCVecEltisVT<1,i1>, SDTCisVec<3>, SDTCisSameNumEltsAs<1,3>]>;
 def AArch64clasta_n     : SDNode<"AArch64ISD::CLASTA_N",   SDT_AArch64ReduceWithInit>;
 def AArch64clastb_n     : SDNode<"AArch64ISD::CLASTB_N",   SDT_AArch64ReduceWithInit>;
 def AArch64fadda_p_node : SDNode<"AArch64ISD::FADDA_PRED", SDT_AArch64ReduceWithInit>;
 
 def AArch64fadda_p : PatFrags<(ops node:$op1, node:$op2, node:$op3),
     [(AArch64fadda_p_node node:$op1, node:$op2, node:$op3),
      (AArch64fadda_p_node (SVEAllActive), node:$op2,
+             (vselect node:$op1, node:$op3, (splat_vector (f16 fpimm_minus0)))),
+     (AArch64fadda_p_node (SVEAllActive), node:$op2,
              (vselect node:$op1, node:$op3, (splat_vector (f32 fpimm_minus0)))),
      (AArch64fadda_p_node (SVEAllActive), node:$op2,
              (vselect node:$op1, node:$op3, (splat_vector (f64 fpimm_minus0))))]>;
 
 def SDT_AArch64PTest : SDTypeProfile<0, 2, [SDTCisVec<0>, SDTCisSameAs<0,1>]>;
 def AArch64ptest     : SDNode<"AArch64ISD::PTEST", SDT_AArch64PTest>;
 
 def SDT_AArch64DUP_PRED  : SDTypeProfile<1, 3, [SDTCisVec<0>, SDTCisSameAs<0, 3>, SDTCisVec<1>, SDTCVecEltisVT<1,i1>]>;
let Predicates = [HasSVEorSME] in {
   def : Pat<(nxv8i16 (splat_vector (i32 (SVECpyDupImm16Pat i32:$a, i32:$b)))),
             (DUP_ZI_H $a, $b)>;
   def : Pat<(nxv4i32 (splat_vector (i32 (SVECpyDupImm32Pat i32:$a, i32:$b)))),
             (DUP_ZI_S $a, $b)>;
   def : Pat<(nxv2i64 (splat_vector (i64 (SVECpyDupImm64Pat i32:$a, i32:$b)))),
             (DUP_ZI_D $a, $b)>;
 
   // Duplicate immediate FP into all vector elements.
+  def : Pat<(nxv2f16 (splat_vector (f16 fpimm:$val))),
+            (DUP_ZR_H (MOVi32imm (bitcast_fpimm_to_i32 f16:$val)))>;
+  def : Pat<(nxv4f16 (splat_vector (f16 fpimm:$val))),
+            (DUP_ZR_H (MOVi32imm (bitcast_fpimm_to_i32 f16:$val)))>;
+  def : Pat<(nxv8f16 (splat_vector (f16 fpimm:$val))),
+            (DUP_ZR_H (MOVi32imm (bitcast_fpimm_to_i32 f16:$val)))>;
   def : Pat<(nxv2f32 (splat_vector (f32 fpimm:$val))),
             (DUP_ZR_S (MOVi32imm (bitcast_fpimm_to_i32 f32:$val)))>;
   def : Pat<(nxv4f32 (splat_vector (f32 fpimm:$val))),
             (DUP_ZR_S (MOVi32imm (bitcast_fpimm_to_i32 f32:$val)))>;
   def : Pat<(nxv2f64 (splat_vector (f64 fpimm:$val))),
             (DUP_ZR_D (MOVi64imm (bitcast_fpimm_to_i64 f64:$val)))>;
 
   // Duplicate FP immediate into all vector elements

llvm/test/CodeGen/AArch64/sve-fadda-select.ll

; CHECK-NEXT:    ret
   ret double %fadda
 }
 
 ; Currently the folding doesn't work for f16 element types, since -0.0 is not treated as a legal f16 immediate.
 
 define half @pred_fadda_nxv2f16(half %x, <vscale x 2 x half> %y, <vscale x 2 x i1> %mask) {
 ; CHECK-LABEL: pred_fadda_nxv2f16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    adrp x8, .LCPI3_0
-; CHECK-NEXT:    add x8, x8, :lo12:.LCPI3_0
-; CHECK-NEXT:    ptrue p1.d
 ; CHECK-NEXT:    // kill: def $h0 killed $h0 def $z0
-; CHECK-NEXT:    ld1rh { z2.d }, p1/z, [x8]
-; CHECK-NEXT:    sel z1.d, p0, z1.d, z2.d
-; CHECK-NEXT:    fadda h0, p1, h0, z1.h
+; CHECK-NEXT:    fadda h0, p0, h0, z1.h
 ; CHECK-NEXT:    // kill: def $h0 killed $h0 killed $z0
 ; CHECK-NEXT:    ret
   %i = insertelement <vscale x 2 x half> poison, half -0.000000e+00, i32 0
   %minus0 = shufflevector <vscale x 2 x half> %i, <vscale x 2 x half> poison, <vscale x 2 x i32> zeroinitializer
   %sel = select <vscale x 2 x i1> %mask, <vscale x 2 x half> %y, <vscale x 2 x half> %minus0
   %fadda = call half @llvm.vector.reduce.fadd.nxv2f16(half %x, <vscale x 2 x half> %sel)
   ret half %fadda
 }
 
 define half @pred_fadda_nxv4f16(half %x, <vscale x 4 x half> %y, <vscale x 4 x i1> %mask) {
 ; CHECK-LABEL: pred_fadda_nxv4f16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    adrp x8, .LCPI4_0
-; CHECK-NEXT:    add x8, x8, :lo12:.LCPI4_0
-; CHECK-NEXT:    ptrue p1.s
 ; CHECK-NEXT:    // kill: def $h0 killed $h0 def $z0
-; CHECK-NEXT:    ld1rh { z2.s }, p1/z, [x8]
-; CHECK-NEXT:    sel z1.s, p0, z1.s, z2.s
-; CHECK-NEXT:    fadda h0, p1, h0, z1.h
+; CHECK-NEXT:    fadda h0, p0, h0, z1.h
 ; CHECK-NEXT:    // kill: def $h0 killed $h0 killed $z0
 ; CHECK-NEXT:    ret
   %i = insertelement <vscale x 4 x half> poison, half -0.000000e+00, i32 0
   %minus0 = shufflevector <vscale x 4 x half> %i, <vscale x 4 x half> poison, <vscale x 4 x i32> zeroinitializer
   %sel = select <vscale x 4 x i1> %mask, <vscale x 4 x half> %y, <vscale x 4 x half> %minus0
   %fadda = call half @llvm.vector.reduce.fadd.nxv4f16(half %x, <vscale x 4 x half> %sel)
   ret half %fadda
 }
 
 define half @pred_fadda_nxv8f16(half %x, <vscale x 8 x half> %y, <vscale x 8 x i1> %mask) {
 ; CHECK-LABEL: pred_fadda_nxv8f16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    adrp x8, .LCPI5_0
-; CHECK-NEXT:    add x8, x8, :lo12:.LCPI5_0
-; CHECK-NEXT:    ptrue p1.h
 ; CHECK-NEXT:    // kill: def $h0 killed $h0 def $z0
-; CHECK-NEXT:    ld1rh { z2.h }, p1/z, [x8]
-; CHECK-NEXT:    sel z1.h, p0, z1.h, z2.h
-; CHECK-NEXT:    fadda h0, p1, h0, z1.h
+; CHECK-NEXT:    fadda h0, p0, h0, z1.h
 ; CHECK-NEXT:    // kill: def $h0 killed $h0 killed $z0
 ; CHECK-NEXT:    ret
   %i = insertelement <vscale x 8 x half> poison, half -0.000000e+00, i32 0
   %minus0 = shufflevector <vscale x 8 x half> %i, <vscale x 8 x half> poison, <vscale x 8 x i32> zeroinitializer
   %sel = select <vscale x 8 x i1> %mask, <vscale x 8 x half> %y, <vscale x 8 x half> %minus0
   %fadda = call half @llvm.vector.reduce.fadd.nxv8f16(half %x, <vscale x 8 x half> %sel)
   ret half %fadda
 }
 
 declare float @llvm.vector.reduce.fadd.nxv2f32(float, <vscale x 2 x float>)
 declare float @llvm.vector.reduce.fadd.nxv4f32(float, <vscale x 4 x float>)
 declare double @llvm.vector.reduce.fadd.nxv2f64(double, <vscale x 2 x double>)
 declare half @llvm.vector.reduce.fadd.nxv2f16(half, <vscale x 2 x half>)
 declare half @llvm.vector.reduce.fadd.nxv4f16(half, <vscale x 4 x half>)
 declare half @llvm.vector.reduce.fadd.nxv8f16(half, <vscale x 8 x half>)

llvm/test/CodeGen/AArch64/sve-fp-reduce.ll

 define half @fadda_nxv6f16(<vscale x 6 x half> %v, half %s) {
 ; CHECK-LABEL: fadda_nxv6f16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    .cfi_offset w29, -16
 ; CHECK-NEXT:    addvl sp, sp, #-1
 ; CHECK-NEXT:    .cfi_escape 0x0f, 0x0c, 0x8f, 0x00, 0x11, 0x10, 0x22, 0x11, 0x08, 0x92, 0x2e, 0x00, 0x1e, 0x22 // sp + 16 + 8 * VG
-; CHECK-NEXT:    adrp x8, .LCPI3_0
-; CHECK-NEXT:    add x8, x8, :lo12:.LCPI3_0
+; CHECK-NEXT:    mov w8, #32768
 ; CHECK-NEXT:    ptrue p0.h
 ; CHECK-NEXT:    ptrue p1.d
 ; CHECK-NEXT:    st1h { z0.h }, p0, [sp]
-; CHECK-NEXT:    ld1rh { z0.d }, p1/z, [x8]
-; CHECK-NEXT:    st1h { z0.d }, p1, [sp, #3, mul vl]
 ; CHECK-NEXT:    fmov s0, s1
+; CHECK-NEXT:    mov z2.h, w8
+; CHECK-NEXT:    st1h { z2.d }, p1, [sp, #3, mul vl]
 ; CHECK-NEXT:    ld1h { z2.h }, p0/z, [sp]
 ; CHECK-NEXT:    fadda h0, p0, h0, z2.h
 ; CHECK-NEXT:    // kill: def $h0 killed $h0 killed $z0
 ; CHECK-NEXT:    addvl sp, sp, #1
 ; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
   %res = call half @llvm.vector.reduce.fadd.nxv6f16(half %s, <vscale x 6 x half> %v)
   ret half %res
 }
 
 define half @fadda_nxv10f16(<vscale x 10 x half> %v, half %s) {
 ; CHECK-LABEL: fadda_nxv10f16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    .cfi_offset w29, -16
 ; CHECK-NEXT:    addvl sp, sp, #-3
 ; CHECK-NEXT:    .cfi_escape 0x0f, 0x0c, 0x8f, 0x00, 0x11, 0x10, 0x22, 0x11, 0x18, 0x92, 0x2e, 0x00, 0x1e, 0x22 // sp + 16 + 24 * VG
-; CHECK-NEXT:    adrp x8, .LCPI4_0
-; CHECK-NEXT:    add x8, x8, :lo12:.LCPI4_0
+; CHECK-NEXT:    mov w8, #32768
 ; CHECK-NEXT:    ptrue p0.h
 ; CHECK-NEXT:    ptrue p1.d
 ; CHECK-NEXT:    st1h { z1.h }, p0, [sp]
 ; CHECK-NEXT:    // kill: def $h2 killed $h2 def $z2
-; CHECK-NEXT:    ld1rh { z1.d }, p1/z, [x8]
+; CHECK-NEXT:    mov z3.h, w8
 ; CHECK-NEXT:    addvl x8, sp, #1
+; CHECK-NEXT:    st1h { z3.d }, p1, [sp, #1, mul vl]
 ; CHECK-NEXT:    fadda h2, p0, h2, z0.h
-; CHECK-NEXT:    st1h { z1.d }, p1, [sp, #1, mul vl]
-; CHECK-NEXT:    ld1h { z3.h }, p0/z, [sp]
-; CHECK-NEXT:    st1h { z3.h }, p0, [sp, #1, mul vl]
-; CHECK-NEXT:    st1h { z1.d }, p1, [sp, #6, mul vl]
-; CHECK-NEXT:    ld1h { z3.h }, p0/z, [sp, #1, mul vl]
-; CHECK-NEXT:    st1h { z3.h }, p0, [sp, #2, mul vl]
-; CHECK-NEXT:    st1h { z1.d }, p1, [x8, #7, mul vl]
+; CHECK-NEXT:    ld1h { z1.h }, p0/z, [sp]
+; CHECK-NEXT:    st1h { z1.h }, p0, [sp, #1, mul vl]
+; CHECK-NEXT:    st1h { z3.d }, p1, [sp, #6, mul vl]
+; CHECK-NEXT:    ld1h { z1.h }, p0/z, [sp, #1, mul vl]
+; CHECK-NEXT:    st1h { z1.h }, p0, [sp, #2, mul vl]
+; CHECK-NEXT:    st1h { z3.d }, p1, [x8, #7, mul vl]
 ; CHECK-NEXT:    ld1h { z1.h }, p0/z, [sp, #2, mul vl]
 ; CHECK-NEXT:    fadda h2, p0, h2, z1.h
 ; CHECK-NEXT:    fmov s0, s2
 ; CHECK-NEXT:    addvl sp, sp, #3
 ; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
   %res = call half @llvm.vector.reduce.fadd.nxv10f16(half %s, <vscale x 10 x half> %v)
   ret half %res
 }
 
 define half @fadda_nxv12f16(<vscale x 12 x half> %v, half %s) {
 ; CHECK-LABEL: fadda_nxv12f16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    adrp x8, .LCPI5_0
-; CHECK-NEXT:    add x8, x8, :lo12:.LCPI5_0
-; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    mov w8, #32768
 ; CHECK-NEXT:    // kill: def $h2 killed $h2 def $z2
 ; CHECK-NEXT:    uunpklo z1.s, z1.h
-; CHECK-NEXT:    ld1rh { z3.s }, p0/z, [x8]
 ; CHECK-NEXT:    ptrue p0.h
 ; CHECK-NEXT:    fadda h2, p0, h2, z0.h
+; CHECK-NEXT:    mov z3.h, w8
 ; CHECK-NEXT:    uzp1 z1.h, z1.h, z3.h
 ; CHECK-NEXT:    fadda h2, p0, h2, z1.h
 ; CHECK-NEXT:    fmov s0, s2
 ; CHECK-NEXT:    ret
   %res = call half @llvm.vector.reduce.fadd.nxv12f16(half %s, <vscale x 12 x half> %v)
   ret half %res
 }