[MVE] Fold fadd(select(..., +0.0)) into a predicated fadd

[david-arm]

Summary

We already have patterns for matching fadd(select(..., -0.0)),
but an upcoming patch will lead to patterns using +0.0 as the
identity instead of -0.0. I'm adding support for these patterns
now to avoid any regressions for MVE.

[spatel]

If this is inverting the transform from D126774, do we need 'nsz' to avoid miscompiling -0.0?
https://alive2.llvm.org/ce/z/Z6sngi

[david-arm]

In D127275#3566222, @spatel wrote:

If this is inverting the transform from D126774, do we need 'nsz' to avoid miscompiling -0.0?
https://alive2.llvm.org/ce/z/Z6sngi

I'm assuming that for MVE an fadd of +0 for a given lane is equivalent to simply not performing the fadd for that lane (i.e. by using predication). This is the same assumption that we were making for -0 case, i.e. that adding either -0 or +0 leaves the lane unchanged. Perhaps @dmgreen can confirm this?

[david-arm]

But if necessary I can certainly add an additional check for nsz.

[dmgreen]

I'm assuming that for MVE an fadd of +0 for a given lane is equivalent to simply not performing the fadd for that lane (i.e. by using predication).

Yes, but I'm not sure how that matters. fadd -0.0, 0.0 should produce 0.0. The same should be true if the first argument is in a variable holding -0.0, and the second is from a select that could be 0.0 - it should give the output 0.0, not of the variable.

llvm/lib/Target/ARM/ARMISelLowering.cpp

16708–16709

Can you update this comment. The Lambda name is no longer correct too.

16716

ImmVal == 0 would show the intent a little better.

llvm/test/CodeGen/Thumb2/mve-pred-selectop3.ll

375

Can you add nsz variants of the tests and make sure it is only those that get folded.

[david-arm]

llvm/test/CodeGen/Thumb2/mve-pred-selectop3.ll

375

I was a bit confused by your comment starting Yes, but I'm not sure how that matters. Do you agree with @spatel that we should only apply the fold for '+0.0' if the nsz flag is set on the select? Hopefully the FMF flags on the IR select have been propagated to the VSELECT DAG node.

[dmgreen]

Oh, sorry - I did misread your comment. I should have said No, that's not how floating point fadd works :)
fadd -0.0, 0.0 is 0.0. The identity value for a fadd is -0.0. A vaddt.f32 Qd, Qn, Qm with a false lane predicate will take the original value of Qd unchanged for that lane (so if Qd==Qn, the original value from the input will be used, without any addition).

Luckily though, the transform you are altering is in terms of VSELECT and FADD nodes, which have the standard definitions, and you can think about without considering MVE specifics. The transform just needs to be valid in general.

Old: fadd (select c, y, 0), x
true -> fadd y, x
false -> fadd 0, x

New: select c, (fadd x, y), x
true -> fadd x, y
false -> x

Which is only valid if fadd 0, x is x, which needs nsz.

[david-arm]

• Added checks for nsz flag on the select when the identity is +0.0

[david-arm]

Thanks for the explanation @dmgreen!

[spatel]

llvm/lib/Target/ARM/ARMISelLowering.cpp

16729

This is not strictly correct. If the 'fadd' doesn't have 'nsz', then the transform results in an 'nsz' value where the original sequence did not.

The transform is safe as long as the 'fadd' has 'nsz', so that's the flag we should use to enable the transform.

We can propagate the fadd's 'nsz' to the new 'select' even if the original 'select' was not 'nsz' itself.
It's confusing, but you can run examples with Alive2 to see what's valid:
https://alive2.llvm.org/ce/z/Z-cuhQ

[david-arm]

llvm/lib/Target/ARM/ARMISelLowering.cpp

16729

Doesn't this contradict https://reviews.llvm.org/D126774 then? In my original instcombine patch I think it was suggested that we should add nsz to the select instruction being generated by the vectoriser, whereas from this comment it sounds like we don't need to and can just rely upon the nsz flag on the IR fadd?

[spatel]

llvm/lib/Target/ARM/ARMISelLowering.cpp

16729

No, the transforms are not strictly reversible/bidirectional:
https://alive2.llvm.org/ce/z/G42d8e

In both cases, I think we need 'nsz' on the final value in the pattern to enable the transform safely.
It is possible that we've created some unintended FMF propagation though when the flags are different on each value.

The semantics are fuzzy and having different flags on values within a function seems highly unlikely, but we should probably add more tests to cover those possibilities.

[david-arm]

• Use the flags for the fadd to determine the safety of the transformation and apply the NSZ flag to the select if needed.
• Tweaked the tests to demonstrate the transform succeeds when the nsz flag is only on the fadd.

[david-arm]

llvm/lib/Target/ARM/ARMISelLowering.cpp

16729

OK thanks for explaining and looking into this in detail! I guess differences in flags are much more likely to occur with inlining, especially as part of LTO where different translation units may be compiled differently.

[dmgreen]

Thanks for putting this patch together. That's a great help and the regressions were pretty large otherwise! I am still getting some odd results from the change, but I think that is just folding the select 0 into a masked load, which is OK.

Other than using FaddFlags instead of SelFlags, this LGTM.

llvm/lib/Target/ARM/ARMISelLowering.cpp

16735

I don't think it is valid to transfer the flags to the select, unfortunately. nsz may be valid, but not the existing flags.
https://alive2.llvm.org/ce/z/9u943A

It should be valid to transfer the flags from the fadd, as that is the end result of the old pattern. https://alive2.llvm.org/ce/z/PzGZSw

[david-arm]

llvm/lib/Target/ARM/ARMISelLowering.cpp

16735

Thanks for the LGTM @dmgreen! I addressed this comment before merging.

llvm/lib/Target/ARM/ARMISelLowering.cpp

static SDValue PerformFAddVSelectCombine(SDNode *N, SelectionDAG &DAG,
 
   // Turn (fadd x, (vselect c, y, -0.0)) into (vselect c, (fadd x, y), x)
   // The second form can be more easily turned into a predicated vadd, and
   // possibly combined into a fma to become a predicated vfma.
   SDValue Op0 = N->getOperand(0);
   SDValue Op1 = N->getOperand(1);
   EVT VT = N->getValueType(0);
   SDLoc DL(N);
 
-  // The identity element for a fadd is -0.0, which these VMOV's represent.
-  auto isNegativeZeroSplat = [&](SDValue Op) {
+  // The identity element for a fadd is -0.0 or +0.0 when the nsz flag is set,
+  // which these VMOV's represent.
+  auto isIdentitySplat = [&](SDValue Op, bool NSZ) {
     if (Op.getOpcode() != ISD::BITCAST ||
         Op.getOperand(0).getOpcode() != ARMISD::VMOVIMM)
       return false;
-    if (VT == MVT::v4f32 && Op.getOperand(0).getConstantOperandVal(0) == 1664)
+    uint64_t ImmVal = Op.getOperand(0).getConstantOperandVal(0);
+    if (VT == MVT::v4f32 && (ImmVal == 1664 || (ImmVal == 0 && NSZ)))
       return true;
-    if (VT == MVT::v8f16 && Op.getOperand(0).getConstantOperandVal(0) == 2688)
+    if (VT == MVT::v8f16 && (ImmVal == 2688 || (ImmVal == 0 && NSZ)))
       return true;
     return false;
   };
 
   if (Op0.getOpcode() == ISD::VSELECT && Op1.getOpcode() != ISD::VSELECT)
     std::swap(Op0, Op1);
 
-  if (Op1.getOpcode() != ISD::VSELECT ||
-      !isNegativeZeroSplat(Op1.getOperand(2)))
+  if (Op1.getOpcode() != ISD::VSELECT)
+    return SDValue();
+
+  SDNodeFlags SelFlags = Op1->getFlags();
+  if (!isIdentitySplat(Op1.getOperand(2), SelFlags.hasNoSignedZeros()))
     return SDValue();
   SDValue FAdd =
       DAG.getNode(ISD::FADD, DL, VT, Op0, Op1.getOperand(1), N->getFlags());
-  return DAG.getNode(ISD::VSELECT, DL, VT, Op1.getOperand(0), FAdd, Op0);
+  return DAG.getNode(ISD::VSELECT, DL, VT, Op1.getOperand(0), FAdd, Op0, SelFlags);
 }
 
 /// PerformVDIVCombine - VCVT (fixed-point to floating-point, Advanced SIMD)
 /// can replace combinations of VCVT (integer to floating-point) and VDIV
 /// when the VDIV has a constant operand that is a power of 2.
 ///
 /// Example (assume d17 = <float 8.000000e+00, float 8.000000e+00>):
 ///  vcvt.f32.s32    d16, d16
 ///  vdiv.f32        d16, d17, d16

llvm/test/CodeGen/Thumb2/mve-pred-selectop3.ll

 ; CHECK-NEXT:    bx lr
 entry:
   %c = call <4 x i1> @llvm.arm.mve.vctp32(i32 %n)
   %a = select <4 x i1> %c, <4 x float> %y, <4 x float> <float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00>
   %b = fadd <4 x float> %a, %x
   ret <4 x float> %b
 }
 
+define arm_aapcs_vfpcc <4 x float> @fadd_v4f32_x2(<4 x float> %x, <4 x float> %y, i32 %n) {
+; CHECK-LABEL: fadd_v4f32_x2:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    vmov.i32 q2, #0x0
+; CHECK-NEXT:    vctp.32 r0
+; CHECK-NEXT:    vpst
+; CHECK-NEXT:    vmovt q2, q1
+; CHECK-NEXT:    vadd.f32 q0, q2, q0
+; CHECK-NEXT:    bx lr
+entry:
+  %c = call <4 x i1> @llvm.arm.mve.vctp32(i32 %n)
+  %a = select <4 x i1> %c, <4 x float> %y, <4 x float> <float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00>
+  %b = fadd <4 x float> %a, %x
+  ret <4 x float> %b
+}
+
+define arm_aapcs_vfpcc <4 x float> @fadd_v4f32_x3(<4 x float> %x, <4 x float> %y, i32 %n) {
+; CHECK-LABEL: fadd_v4f32_x3:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    vctp.32 r0
+; CHECK-NEXT:    vpst
+; CHECK-NEXT:    vaddt.f32 q0, q0, q1
+; CHECK-NEXT:    bx lr
+entry:
+  %c = call <4 x i1> @llvm.arm.mve.vctp32(i32 %n)
+  %a = select nsz <4 x i1> %c, <4 x float> %y, <4 x float> <float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00>
+  %b = fadd nsz <4 x float> %a, %x
+  ret <4 x float> %b
+}
+
 define arm_aapcs_vfpcc <8 x half> @fadd_v8f16_x(<8 x half> %x, <8 x half> %y, i32 %n) {
 ; CHECK-LABEL: fadd_v8f16_x:
 ; CHECK:       @ %bb.0: @ %entry
 ; CHECK-NEXT:    vctp.16 r0
 ; CHECK-NEXT:    vpst
 ; CHECK-NEXT:    vaddt.f16 q0, q0, q1
 ; CHECK-NEXT:    bx lr
 entry:
   %c = call <8 x i1> @llvm.arm.mve.vctp16(i32 %n)
   %a = select <8 x i1> %c, <8 x half> %y, <8 x half> <half 0xH8000, half 0xH8000, half 0xH8000, half 0xH8000, half 0xH8000, half 0xH8000, half 0xH8000, half 0xH8000>
   %b = fadd <8 x half> %a, %x
   ret <8 x half> %b
 }
 
+define arm_aapcs_vfpcc <8 x half> @fadd_v8f16_x2(<8 x half> %x, <8 x half> %y, i32 %n) {
+; CHECK-LABEL: fadd_v8f16_x2:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    vmov.i32 q2, #0x0
+; CHECK-NEXT:    vctp.16 r0
+; CHECK-NEXT:    vpst
+; CHECK-NEXT:    vmovt q2, q1
+; CHECK-NEXT:    vadd.f16 q0, q2, q0
+; CHECK-NEXT:    bx lr
+entry:
+  %c = call <8 x i1> @llvm.arm.mve.vctp16(i32 %n)
+  %a = select <8 x i1> %c, <8 x half> %y, <8 x half> <half 0x0000, half 0x00000, half 0x00000, half 0x00000, half 0x00000, half 0x00000, half 0x00000, half 0x00000>
+  %b = fadd <8 x half> %a, %x
+  ret <8 x half> %b
+}
+
+define arm_aapcs_vfpcc <8 x half> @fadd_v8f16_x3(<8 x half> %x, <8 x half> %y, i32 %n) {
+; CHECK-LABEL: fadd_v8f16_x3:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    vctp.16 r0
+; CHECK-NEXT:    vpst
+; CHECK-NEXT:    vaddt.f16 q0, q0, q1
+; CHECK-NEXT:    bx lr
+entry:
+  %c = call <8 x i1> @llvm.arm.mve.vctp16(i32 %n)
+  %a = select nsz <8 x i1> %c, <8 x half> %y, <8 x half> <half 0x0000, half 0x00000, half 0x00000, half 0x00000, half 0x00000, half 0x00000, half 0x00000, half 0x00000>
+  %b = fadd nsz <8 x half> %a, %x
+  ret <8 x half> %b
+}
+
 define arm_aapcs_vfpcc <4 x float> @fsub_v4f32_x(<4 x float> %x, <4 x float> %y, i32 %n) {
 ; CHECK-LABEL: fsub_v4f32_x:
 ; CHECK:       @ %bb.0: @ %entry
 ; CHECK-NEXT:    vctp.32 r0
 ; CHECK-NEXT:    vpst
 ; CHECK-NEXT:    vsubt.f32 q0, q0, q1
 ; CHECK-NEXT:    bx lr
 entry: