[AArch64][SVE] Fold fadda(ptrue, x, select(mask, y, -0.0)) into fadda(mask, x, y)

[RosieSumpter]

Summary

This patch adds an SVE pattern to recognize the use of a select with an
fadda in the form fadda(ptrue, x, select(mask, y, -0.0)). In this case
the select can be folded away, with the select mask used as the
predicate for fadda. This improves the codegen when vectorizing loops
with ordered fp reductions.

Tests have been added for the f16 cases, although the folding doesn't
currently work for these as -0.0 isn't treated as a legal f16 immediate.

[paulwalker-arm]

llvm/lib/Target/AArch64/SVEInstrFormats.td

5129–5134

It is possible to change AArch64fadda_p into a PatFrags to remove the need for these? See AArch64bic_node and AArch64bic.

[RosieSumpter]

• Used PatFrags

llvm/lib/Target/AArch64/SVEInstrFormats.td

5129–5134

Hey Paul, I've had a go at using PatFrags instead - is this what you had in mind?

[paulwalker-arm]

llvm/lib/Target/AArch64/SVEInstrFormats.td

5129–5134

Not quite, which make me think I've asked for something that's not possible. I'll investigate further.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

288–291 ↗

(On Diff #445008)

You shouldn't need duplicate patterns within the PatFrag based on types as this kind of defeats the point of a PatFrags. The main exception to this is the scalar operand of splat_vector due to the way it's defined.

The problem here is that you cannot infer the types which is down to the loose definition of AArch64fadda_p_node or rather SDT_AArch64ReduceWithInit. You'll have to double check it's correct but when I tighten it up to say

SDTCisVec<1>, SDTCVecEltisVT<1,i1>, SDTCisVec<3>, SDTCisSameNumEltsAs<1,3>

i.e. The first operand is a predicate vector and has the same number of elements as the third operand.

I only needed a single f32 and f64 pattern without any other types mentioned.

[RosieSumpter]

• Change definition of SDT_AArch64ReduceWithInit so that fadda patterns don't need to be replicated for different predicate vector types.

[paulwalker-arm]

Given AArch64ISD::FADDA_PRED exists I guess there's an argument this could be a DAG combine. However, this isel route looks clean enough and you're adding things that will be beneficial regardless.

llvm/lib/Target/AArch64/AArch64InstrFormats.td

def fpimm8 : Operand<i32> {
   let ParserMatchClass = FPImmOperand;
   let PrintMethod = "printFPImmOperand";
 }
 
 def fpimm0 : FPImmLeaf<fAny, [{
   return Imm.isExactlyValue(+0.0);
 }]>;
 
+def fpimm_minus0 : FPImmLeaf<fAny, [{
+  return Imm.isExactlyValue(-0.0);
+}]>;
+
 def fpimm_half : FPImmLeaf<fAny, [{
   return Imm.isExactlyValue(+0.5);
 }]>;
 
 def fpimm_one : FPImmLeaf<fAny, [{
   return Imm.isExactlyValue(+1.0);
 }]>;
 

llvm/lib/Target/AArch64/SVEInstrFormats.td

 
 class SVE_3_Op_Pat_Shift_Imm_SelZero<ValueType vtd, SDPatternOperator op,
                                      ValueType vt1, ValueType vt2,
                                      Operand vt3, Instruction inst>
 : Pat<(vtd (op vt1:$Op1, (vselect vt1:$Op1, vt2:$Op2, (SVEDup0)), (i32 (vt3:$Op3)))),
       (inst $Op1, $Op2, vt3:$Op3)>;
 }
 
+class SVE_3_Op_Pat_Pred_FPSelMinusZero<ValueType vtd, SDPatternOperator op,
+                   ValueType vt1, ValueType vt2, ValueType vt3,
+                   ValueType scalarty, Instruction inst>
+: Pat<(vtd (op (vt1 (SVEAllActive)), vt2:$Op2, (vselect vt1:$Pg, vt3:$Op3,
+               (vt3 (splat_vector (scalarty fpimm_minus0)))))),
+      (inst $Pg, $Op2, $Op3)>;
+
 //
 // Common but less generic patterns.
 //
 
 class SVE_1_Op_AllActive_Pat<ValueType vtd, SDPatternOperator op, ValueType vt1,
                              Instruction inst, Instruction ptrue>
 : Pat<(vtd (op vt1:$Op1)),
       (inst (IMPLICIT_DEF), (ptrue 31), $Op1)>;
multiclass sve_fp_2op_p_vd<bits<3> opc, string asm, SDPatternOperator op> {
   def _D : sve_fp_2op_p_vd<0b11, opc, asm, ZPR64, FPR64asZPR>;
 
   def : SVE_3_Op_Pat<nxv2f16, op, nxv2i1, nxv2f16, nxv2f16, !cast<Instruction>(NAME # _H)>;
   def : SVE_3_Op_Pat<nxv4f16, op, nxv4i1, nxv4f16, nxv4f16, !cast<Instruction>(NAME # _H)>;
   def : SVE_3_Op_Pat<nxv8f16, op, nxv8i1, nxv8f16, nxv8f16, !cast<Instruction>(NAME # _H)>;
   def : SVE_3_Op_Pat<nxv2f32, op, nxv2i1, nxv2f32, nxv2f32, !cast<Instruction>(NAME # _S)>;
   def : SVE_3_Op_Pat<nxv4f32, op, nxv4i1, nxv4f32, nxv4f32, !cast<Instruction>(NAME # _S)>;
   def : SVE_3_Op_Pat<nxv2f64, op, nxv2i1, nxv2f64, nxv2f64, !cast<Instruction>(NAME # _D)>;
+
+  def : SVE_3_Op_Pat_Pred_FPSelMinusZero<nxv2f32, op, nxv2i1, nxv2f32, nxv2f32, f32,
+                                                            !cast<Instruction>(NAME # _S)>;
+  def : SVE_3_Op_Pat_Pred_FPSelMinusZero<nxv4f32, op, nxv4i1, nxv4f32, nxv4f32, f32,
+                                                            !cast<Instruction>(NAME # _S)>;
+  def : SVE_3_Op_Pat_Pred_FPSelMinusZero<nxv2f64, op, nxv2i1, nxv2f64, nxv2f64, f64,
+                                                            !cast<Instruction>(NAME # _D)>;
+
 }
 
 //===----------------------------------------------------------------------===//
 // SVE Floating Point Compare - Vectors Group
 //===----------------------------------------------------------------------===//
 
 class sve_fp_3op_p_pd<bits<2> sz, bits<3> opc, string asm, PPRRegOp pprty,
                       ZPRRegOp zprty>

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

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
+
+; Fold fadda(ptrue, x, select(mask, y, -0.0)) -> fadda(mask, x, y)
+
+define float @pred_fadda_nxv2f32(float %x, <vscale x 2 x float> %y, <vscale x 2 x i1> %mask) {
+; CHECK-LABEL: pred_fadda_nxv2f32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    // kill: def $s0 killed $s0 def $z0
+; CHECK-NEXT:    fadda s0, p0, s0, z1.s
+; CHECK-NEXT:    // kill: def $s0 killed $s0 killed $z0
+; CHECK-NEXT:    ret
+  %i = insertelement <vscale x 2 x float> poison, float -0.000000e+00, i32 0
+  %minus0 = shufflevector <vscale x 2 x float> %i, <vscale x 2 x float> poison, <vscale x 2 x i32> zeroinitializer
+  %sel = select <vscale x 2 x i1> %mask, <vscale x 2 x float> %y, <vscale x 2 x float> %minus0
+  %fadda = call float @llvm.vector.reduce.fadd.nxv2f32(float %x, <vscale x 2 x float> %sel)
+  ret float %fadda
+}
+
+define float @pred_fadda_nxv4f32(float %x, <vscale x 4 x float> %y, <vscale x 4 x i1> %mask) {
+; CHECK-LABEL: pred_fadda_nxv4f32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    // kill: def $s0 killed $s0 def $z0
+; CHECK-NEXT:    fadda s0, p0, s0, z1.s
+; CHECK-NEXT:    // kill: def $s0 killed $s0 killed $z0
+; CHECK-NEXT:    ret
+  %i = insertelement <vscale x 4 x float> poison, float -0.000000e+00, i32 0
+  %minus0 = shufflevector <vscale x 4 x float> %i, <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer
+  %sel = select <vscale x 4 x i1> %mask, <vscale x 4 x float> %y, <vscale x 4 x float> %minus0
+  %fadda = call float @llvm.vector.reduce.fadd.nxv4f32(float %x, <vscale x 4 x float> %sel)
+  ret float %fadda
+}
+
+define double @pred_fadda_nxv2f64(double %x, <vscale x 2 x double> %y, <vscale x 2 x i1> %mask) {
+; CHECK-LABEL: pred_fadda_nxv2f64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    // kill: def $d0 killed $d0 def $z0
+; CHECK-NEXT:    fadda d0, p0, d0, z1.d
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $z0
+; CHECK-NEXT:    ret
+  %i = insertelement <vscale x 2 x double> poison, double -0.000000e+00, i32 0
+  %minus0 = shufflevector <vscale x 2 x double> %i, <vscale x 2 x double> poison, <vscale x 2 x i32> zeroinitializer
+  %sel = select <vscale x 2 x i1> %mask, <vscale x 2 x double> %y, <vscale x 2 x double> %minus0
+  %fadda = call double @llvm.vector.reduce.fadd.nxv2f64(double %x, <vscale x 2 x double> %sel)
+  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:    // 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:    // 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:    // 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>)