[AArch64][SVE] Fold vselect into predicated fmul, fsub and fadd

[MattDevereau]

Summary

[AArch64][SVE] Fold vselect into predicated fmul, fsub and fadd

Fold vselect with an unpredicated fmul/fsub/fadd operand into a predicated fmul/fsub/fadd:

(vselect (p) (op (a) (b)) (a)) => (op -> (p) (a) (b))

[peterwaller-arm]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

501–518

You might be able to condense it to something like this suggestion, by introducing another multiclass in InstrFormats SVE_VSelect_3_Op_Pat, analogous to SVE_3_Op_Pat_SelZero. What you have looks reasonable to me, you might want a second opinion from someone else as to whether my proposal here makes sense.

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

89

Register called %mul for an fsub -- looks like a few register names need updating.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

501–518

If possible I prefer the patterns to be added directly to the class that defines the instruction (i.e. sve_fp_2op_p_zds in this case)?

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

501–518

Just to extend on this. Perhaps you can create a PatFrags that contains both the intrinsic and vselect DAG patterns? That way you won't need to change the sve_fp_2op_p_zds class, you'll just pass in a different op in place of int_aarch64_sve_fadd for example.

[MattDevereau]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

501–518

I don't understand what you mean by the intrinsic. I've not used the int_aarch64_sve_fadd naming so far so I'm not sure which part that refers to, and the only DAG nodes I saw when working on this was AArch64fadd_p. Are you suggesting I use the PatFrags to add the match to defm FADD_ZPmZ on line 454 when it uses sve_fp_2op_p_zds?

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

501–518

The sve_fp_2op_p_zds multiclass is used to define the FADD_ZPmZ instruction. It takes an operator that represents the instruction at the DAG level. Currently this is set to int_aarch64_sve_fadd. With this patch you're effectively saying there's another way to represent this instruction within the DAG (i.e. (vselect (p) (a) (fadd (a) (b)))) and so I'm thinking you can create a PatFrags that contains bother operators which is them passed into sve_fp_2op_p_zds in place of the existing int_aarch64_sve_fadd. This way you shouldn't need to create any extra isel patterns.

[MattDevereau]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

501–518

Me and Peter had a dive into this, and we assumed you were getting at something like that.
We came up with this PatFrags:

def myfirstpatfrag : PatFrags<(ops node:$Pg, node:$Op1, node:$Op2), [
  (int_aarch64_sve_fmul node:$Pg, node:$Op1, node:$Op2),
  (vselect node:$Pg, node:$Op1, 
    (AArch64fmul_p (nxv2i1 (AArch64ptrue 31)), node:$Op1, node:$Op2)
  ),
]>;

but we haven't managed to find a way to generalise the nxv2i1 part so far.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

501–518

You should be able to drop the (nxv2i1 ) part as the predicate type should be inferable. I just gave this a quick try and it works but you also need to tighten up the definition of the predicated binops (i.e. SDT_AArch64Arith, which is used to define AArch64fmul_p)

I extended it to include SDTCVecEltisVT<1,i1>, SDTCisSameNumEltsAs<0,1>,...., which looks to do the trick.

[MattDevereau]

Use patfrags
Update commit message

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

168

Out of interest are the SDTCisVec<1>, SDTCisVec<2>, SDTCisVec<3> bits removable? I'm thinking they're redundant given the next line is now very specific about what the types of those operands are.

243

We already have a defined naming scheme for these operations, so these should be named AArch64fadd_m1, AArch64fmul_m1 & AArch64fsub_m1.

[MattDevereau]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

168

Removing them has no effect on the tests added in this patch or anything in the check-llvm target

[peterwaller-arm]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

246

Can the PatFrags be factored into a class with two parameters to avoid stamping it out repeatedly?

251

s/(AArch64ptrue 31)/(SVEAllActive)/g

[MattDevereau]

Moved repeated patfrag into class

[peterwaller-arm]

llvm/lib/Target/AArch64/SVEInstrFormats.td

8468 ↗

(On Diff #403275)

I don't think this should mimic the form of the other (SVE_...) names, since those are for defining patterns, whereas this is a PatFrags, which is a different kind of thing. So in my view the name should be distinct to try and avoid catching unsuspecting readers into another trap.

How about EitherVSelectOrPassthruPatFrags. I'm thinking along the lines of trying to avoid the error of Or reading as the binary operation by prefacing it with Either.

Also, this doesn't look like the right place for this. I think a better location would be up near this, prehaps right at the end of that block before the getSVEPseudoMap definition. For good measure I might put a comment-block delimiter which separates it from the patterns above, and indicate in the comment that this "matches either an intrinsic, or a predicated operation with an all active predicate".

//===----------------------------------------------------------------------===//
// SVE pattern match helpers.
//===----------------------------------------------------------------------===//

Note also the missing trailing newline phabricator is complaining about.

[peterwaller-arm]

llvm/lib/Target/AArch64/SVEInstrFormats.td

8464 ↗

(On Diff #403275)

I think better names might be:
op => intrinsic
inst => sdnode

Beware copying the names form the other SVE_... classes, which I think would be confusing. Because in those classes, op is the thing to match for and 'inst' is the replacement, which is not the case for your PatFrags where, both cases are actually things to match.

[MattDevereau]

Updated PatFrags class name and location

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

169

I cannot say if it really matters (given SDTCisSameNumEltsAs<0,1> implies 0 must be a vector) but I'd be inclined to keep SDTCisVec<0>.

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

5–13

Are these tests correct? To me it looks like you've got the true and false operands the wrong way round. This IR is effectively saying "when the predicate it true return op1 else return op1 * op2".

The fact the generate code is what we're after suggests the same mistake exists for the new isel patterns.

[MattDevereau]

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

5–13

The tests aren't correct, I missed the fact that fcmeq in the sqrt example is creating a predicate of 1s. Is it possible to reverse the operands for specific patterns in the PatFrags list? If not I'll have to split up parts of the logic introduced in this patch

[paulwalker-arm]

I'd focus on one thing at a time. This patch enables the isel required for the way merged arithmetic is modelled by the DAG. Once that is complete you can see what's needed to improve the reciprocal code. For that follow-on work I'm guessing you'll need a DAG combine for the select that inverts the comparison for the cases where this'll be beneficial.

[MattDevereau]

update commit message, flip operands for vselect pattern

[paulwalker-arm]

The IR within the tests looks like it's been doubly indented and there's a typo in the commit message (should be (vselect (p) (op (a) (b)) (a)) => (op -> (p) (a) (b)) but otherwise looks good.

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

 def AArch64uminv_p   : SDNode<"AArch64ISD::UMINV_PRED",   SDT_AArch64Reduce>;
 def AArch64orv_p     : SDNode<"AArch64ISD::ORV_PRED",     SDT_AArch64Reduce>;
 def AArch64eorv_p    : SDNode<"AArch64ISD::EORV_PRED",    SDT_AArch64Reduce>;
 def AArch64andv_p    : SDNode<"AArch64ISD::ANDV_PRED",    SDT_AArch64Reduce>;
 def AArch64lasta     : SDNode<"AArch64ISD::LASTA",        SDT_AArch64Reduce>;
 def AArch64lastb     : SDNode<"AArch64ISD::LASTB",        SDT_AArch64Reduce>;
 
 def SDT_AArch64Arith : SDTypeProfile<1, 3, [
   SDTCisVec<0>, SDTCisVec<1>, SDTCisVec<2>, SDTCisVec<3>,
   SDTCVecEltisVT<1,i1>, SDTCisSameAs<0,2>, SDTCisSameAs<2,3>
 ]>;
 
 def SDT_AArch64FMA : SDTypeProfile<1, 4, [
   SDTCisVec<0>, SDTCisVec<1>, SDTCisVec<2>, SDTCisVec<3>, SDTCisVec<4>,
   SDTCVecEltisVT<1,i1>, SDTCisSameAs<0,2>, SDTCisSameAs<2,3>, SDTCisSameAs<3,4>
 ]>;
 
 // Predicated operations with the result of inactive lanes being unspecified.

 def AArch64revw_mt   : SDNode<"AArch64ISD::REVW_MERGE_PASSTHRU", SDT_AArch64Arith>;
 
 // These are like the above but we don't yet have need for ISD nodes. They allow
 // a single pattern to match intrinsic and ISD operand layouts.
 def AArch64cls_mt  : PatFrags<(ops node:$pg, node:$op, node:$pt), [(int_aarch64_sve_cls  node:$pt, node:$pg, node:$op)]>;
 def AArch64cnot_mt : PatFrags<(ops node:$pg, node:$op, node:$pt), [(int_aarch64_sve_cnot node:$pt, node:$pg, node:$op)]>;
 def AArch64not_mt  : PatFrags<(ops node:$pg, node:$op, node:$pt), [(int_aarch64_sve_not  node:$pt, node:$pg, node:$op)]>;
 
 def SDT_AArch64FCVT : SDTypeProfile<1, 3, [
   SDTCisVec<0>, SDTCisVec<1>, SDTCisVec<2>, SDTCisVec<3>,
   SDTCVecEltisVT<1,i1>
 ]>;
 
 def SDT_AArch64FCVTR : SDTypeProfile<1, 4, [
   SDTCisVec<0>, SDTCisVec<1>, SDTCisVec<2>, SDTCisInt<3>, SDTCisVec<4>,
   SDTCVecEltisVT<1,i1>
 ]>;
 
 def AArch64fcvtr_mt  : SDNode<"AArch64ISD::FP_ROUND_MERGE_PASSTHRU", SDT_AArch64FCVTR>;
 def AArch64fcvte_mt  : SDNode<"AArch64ISD::FP_EXTEND_MERGE_PASSTHRU", SDT_AArch64FCVT>;
 def AArch64ucvtf_mt  : SDNode<"AArch64ISD::UINT_TO_FP_MERGE_PASSTHRU", SDT_AArch64FCVT>;
 def AArch64scvtf_mt  : SDNode<"AArch64ISD::SINT_TO_FP_MERGE_PASSTHRU", SDT_AArch64FCVT>;
 def AArch64fcvtzu_mt : SDNode<"AArch64ISD::FCVTZU_MERGE_PASSTHRU", SDT_AArch64FCVT>;
 def AArch64fcvtzs_mt : SDNode<"AArch64ISD::FCVTZS_MERGE_PASSTHRU", SDT_AArch64FCVT>;
 

 
 let Predicates = [HasSVEorStreamingSVE] in {
   defm FADD_ZZZ    : sve_fp_3op_u_zd<0b000, "fadd", fadd, AArch64fadd_p>;
   defm FSUB_ZZZ    : sve_fp_3op_u_zd<0b001, "fsub", fsub, AArch64fsub_p>;
   defm FMUL_ZZZ    : sve_fp_3op_u_zd<0b010, "fmul", fmul, AArch64fmul_p>;
 } // End HasSVEorStreamingSVE
 
 let Predicates = [HasSVE] in {
+  multiclass vselect_to_pred_inst<SDPatternOperator Op, string Inst, ValueType Ty, ValueType PredTy> {
+    def : Pat<(vselect (PredTy PPR:$Pg), (Ty ZPR:$Zn), (Ty (Op (PredTy (AArch64ptrue 31)), (Ty ZPR:$Zn), (Ty ZPR:$Zm)))),
+              (!cast<Instruction>(Inst) PPR:$Pg, ZPR:$Zn, ZPR:$Zm)>;
+}
+
+  defm : vselect_to_pred_inst<AArch64fmul_p, "FMUL_ZPmZ_H", nxv8f16, nxv8i1>;
+  defm : vselect_to_pred_inst<AArch64fmul_p, "FMUL_ZPmZ_S", nxv4f32, nxv4i1>;
+  defm : vselect_to_pred_inst<AArch64fmul_p, "FMUL_ZPmZ_D", nxv2f64, nxv2i1>;
+
+  defm : vselect_to_pred_inst<AArch64fadd_p, "FADD_ZPmZ_H", nxv8f16, nxv8i1>;
+  defm : vselect_to_pred_inst<AArch64fadd_p, "FADD_ZPmZ_S", nxv4f32, nxv4i1>;
+  defm : vselect_to_pred_inst<AArch64fadd_p, "FADD_ZPmZ_D", nxv2f64, nxv2i1>;
+
+  defm : vselect_to_pred_inst<AArch64fsub_p, "FSUB_ZPmZ_H", nxv8f16, nxv8i1>;
+  defm : vselect_to_pred_inst<AArch64fsub_p, "FSUB_ZPmZ_S", nxv4f32, nxv4i1>;
+  defm : vselect_to_pred_inst<AArch64fsub_p, "FSUB_ZPmZ_D", nxv2f64, nxv2i1>;
+
+} // End HasSVE
+
+let Predicates = [HasSVE] in {
   defm FTSMUL_ZZZ  : sve_fp_3op_u_zd_ftsmul<0b011, "ftsmul", int_aarch64_sve_ftsmul_x>;
 } // End HasSVE
 
 let Predicates = [HasSVEorStreamingSVE] in {
   defm FRECPS_ZZZ  : sve_fp_3op_u_zd<0b110, "frecps",  AArch64frecps>;
   defm FRSQRTS_ZZZ : sve_fp_3op_u_zd<0b111, "frsqrts", AArch64frsqrts>;
 } // End HasSVEorStreamingSVE
 

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

 ; CHECK-NEXT:    ptrue p0.h
 ; CHECK-NEXT:    fmul z2.h, z1.h, z1.h
 ; CHECK-NEXT:    fcmeq p0.h, p0/z, z0.h, #0.0
 ; CHECK-NEXT:    frsqrts z2.h, z0.h, z2.h
 ; CHECK-NEXT:    fmul z1.h, z1.h, z2.h
 ; CHECK-NEXT:    fmul z2.h, z1.h, z1.h
 ; CHECK-NEXT:    frsqrts z2.h, z0.h, z2.h
 ; CHECK-NEXT:    fmul z1.h, z1.h, z2.h
-; CHECK-NEXT:    fmul z1.h, z0.h, z1.h
-; CHECK-NEXT:    sel z0.h, p0, z0.h, z1.h
+; CHECK-NEXT:    fmul z0.h, p0/m, z0.h, z1.h
 ; CHECK-NEXT:    ret
   %fsqrt = call fast <vscale x 8 x half> @llvm.sqrt.nxv8f16(<vscale x 8 x half> %a)
   ret <vscale x 8 x half> %fsqrt
 }
 
 define <vscale x 4 x float> @fsqrt_4f32(<vscale x 4 x float> %a) {
 ; CHECK-LABEL: fsqrt_4f32:
 ; CHECK:       // %bb.0:

 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    fmul z2.s, z1.s, z1.s
 ; CHECK-NEXT:    fcmeq p0.s, p0/z, z0.s, #0.0
 ; CHECK-NEXT:    frsqrts z2.s, z0.s, z2.s
 ; CHECK-NEXT:    fmul z1.s, z1.s, z2.s
 ; CHECK-NEXT:    fmul z2.s, z1.s, z1.s
 ; CHECK-NEXT:    frsqrts z2.s, z0.s, z2.s
 ; CHECK-NEXT:    fmul z1.s, z1.s, z2.s
-; CHECK-NEXT:    fmul z1.s, z0.s, z1.s
-; CHECK-NEXT:    sel z0.s, p0, z0.s, z1.s
+; CHECK-NEXT:    fmul z0.s, p0/m, z0.s, z1.s
 ; CHECK-NEXT:    ret
   %fsqrt = call fast <vscale x 4 x float> @llvm.sqrt.nxv4f32(<vscale x 4 x float> %a)
   ret <vscale x 4 x float> %fsqrt
 }
 
 define <vscale x 2 x double> @fsqrt_2f64(<vscale x 2 x double> %a) {
 ; CHECK-LABEL: fsqrt_2f64:
 ; CHECK:       // %bb.0:

 ; CHECK-NEXT:    frsqrts z2.d, z0.d, z2.d
 ; CHECK-NEXT:    fmul z1.d, z1.d, z2.d
 ; CHECK-NEXT:    fmul z2.d, z1.d, z1.d
 ; CHECK-NEXT:    frsqrts z2.d, z0.d, z2.d
 ; CHECK-NEXT:    fmul z1.d, z1.d, z2.d
 ; CHECK-NEXT:    fmul z2.d, z1.d, z1.d
 ; CHECK-NEXT:    frsqrts z2.d, z0.d, z2.d
 ; CHECK-NEXT:    fmul z1.d, z1.d, z2.d
-; CHECK-NEXT:    fmul z1.d, z0.d, z1.d
-; CHECK-NEXT:    sel z0.d, p0, z0.d, z1.d
+; CHECK-NEXT:    fmul z0.d, p0/m, z0.d, z1.d
 ; CHECK-NEXT:    ret
   %fsqrt = call fast <vscale x 2 x double> @llvm.sqrt.nxv2f64(<vscale x 2 x double> %a)
   ret <vscale x 2 x double> %fsqrt
 }
 
 declare <vscale x 2 x half> @llvm.sqrt.nxv2f16(<vscale x 2 x half>)
 declare <vscale x 4 x half> @llvm.sqrt.nxv4f16(<vscale x 4 x half>)
 declare <vscale x 8 x half> @llvm.sqrt.nxv8f16(<vscale x 8 x half>)
 declare <vscale x 2 x float> @llvm.sqrt.nxv2f32(<vscale x 2 x float>)
 declare <vscale x 4 x float> @llvm.sqrt.nxv4f32(<vscale x 4 x float>)
 declare <vscale x 2 x double> @llvm.sqrt.nxv2f64(<vscale x 2 x double>)
 
 attributes #0 = { "reciprocal-estimates"="all" }

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

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
+
+define <vscale x 8 x half> @vselect_fmul_f16(<vscale x 8 x i1> %p, <vscale x 8 x half> %a, <vscale x 8 x half> %b) {
+; CHECK-LABEL: vselect_fmul_f16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    fmul z0.h, p0/m, z0.h, z1.h
+; CHECK-NEXT:    ret
+    %mul = fmul <vscale x 8 x half> %a, %b
+    %x = select <vscale x 8 x i1> %p, <vscale x 8 x half> %a, <vscale x 8 x half> %mul
+    ret <vscale x 8 x half> %x
+}
+
+define <vscale x 4 x float> @vselect_fmul_f32(<vscale x 4 x i1> %p, <vscale x 4 x float> %a, <vscale x 4 x float> %b) {
+; CHECK-LABEL: vselect_fmul_f32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    fmul z0.s, p0/m, z0.s, z1.s
+; CHECK-NEXT:    ret
+    %mul = fmul <vscale x 4 x float> %a, %b
+    %x = select <vscale x 4 x i1> %p, <vscale x 4 x float> %a, <vscale x 4 x float> %mul
+    ret <vscale x 4 x float> %x
+}
+
+define <vscale x 2 x double> @vselect_fmul_f64(<vscale x 2 x i1> %p, <vscale x 2 x double> %a, <vscale x 2 x double> %b) {
+; CHECK-LABEL: vselect_fmul_f64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    fmul z0.d, p0/m, z0.d, z1.d
+; CHECK-NEXT:    ret
+    %mul = fmul <vscale x 2 x double> %a, %b
+    %sel = select <vscale x 2 x i1> %p, <vscale x 2 x double> %a, <vscale x 2 x double> %mul
+    ret <vscale x 2 x double> %sel
+}
+
+define <vscale x 8 x half> @vselect_fadd_f16(<vscale x 8 x i1> %p, <vscale x 8 x half> %a, <vscale x 8 x half> %b) {
+; CHECK-LABEL: vselect_fadd_f16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    fadd z0.h, p0/m, z0.h, z1.h
+; CHECK-NEXT:    ret
+    %mul = fadd <vscale x 8 x half> %a, %b
+    %x = select <vscale x 8 x i1> %p, <vscale x 8 x half> %a, <vscale x 8 x half> %mul
+    ret <vscale x 8 x half> %x
+}
+
+define <vscale x 4 x float> @vselect_fadd_f32(<vscale x 4 x i1> %p, <vscale x 4 x float> %a, <vscale x 4 x float> %b) {
+; CHECK-LABEL: vselect_fadd_f32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    fadd z0.s, p0/m, z0.s, z1.s
+; CHECK-NEXT:    ret
+    %mul = fadd <vscale x 4 x float> %a, %b
+    %x = select <vscale x 4 x i1> %p, <vscale x 4 x float> %a, <vscale x 4 x float> %mul
+    ret <vscale x 4 x float> %x
+}
+
+define <vscale x 2 x double> @vselect_fadd_f64(<vscale x 2 x i1> %p, <vscale x 2 x double> %a, <vscale x 2 x double> %b) {
+; CHECK-LABEL: vselect_fadd_f64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    fadd z0.d, p0/m, z0.d, z1.d
+; CHECK-NEXT:    ret
+    %mul = fadd <vscale x 2 x double> %a, %b
+    %sel = select <vscale x 2 x i1> %p, <vscale x 2 x double> %a, <vscale x 2 x double> %mul
+    ret <vscale x 2 x double> %sel
+}
+
+define <vscale x 8 x half> @vselect_fsub_f16(<vscale x 8 x i1> %p, <vscale x 8 x half> %a, <vscale x 8 x half> %b) {
+; CHECK-LABEL: vselect_fsub_f16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    fsub z0.h, p0/m, z0.h, z1.h
+; CHECK-NEXT:    ret
+    %mul = fsub <vscale x 8 x half> %a, %b
+    %x = select <vscale x 8 x i1> %p, <vscale x 8 x half> %a, <vscale x 8 x half> %mul
+    ret <vscale x 8 x half> %x
+}
+
+define <vscale x 4 x float> @vselect_fsub_f32(<vscale x 4 x i1> %p, <vscale x 4 x float> %a, <vscale x 4 x float> %b) {
+; CHECK-LABEL: vselect_fsub_f32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    fsub z0.s, p0/m, z0.s, z1.s
+; CHECK-NEXT:    ret
+    %mul = fsub <vscale x 4 x float> %a, %b
+    %x = select <vscale x 4 x i1> %p, <vscale x 4 x float> %a, <vscale x 4 x float> %mul
+    ret <vscale x 4 x float> %x
+}
+
+define <vscale x 2 x double> @vselect_fsub_f64(<vscale x 2 x i1> %p, <vscale x 2 x double> %a, <vscale x 2 x double> %b) {
+; CHECK-LABEL: vselect_fsub_f64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    fsub z0.d, p0/m, z0.d, z1.d
+; CHECK-NEXT:    ret
+    %mul = fsub <vscale x 2 x double> %a, %b
+    %sel = select <vscale x 2 x i1> %p, <vscale x 2 x double> %a, <vscale x 2 x double> %mul
+    ret <vscale x 2 x double> %sel
+}