[SVE] Only combine (fneg (fma)) => FNMLA with nsz

[peterwaller-arm]

Summary

-(Za + Zm * Zn) != (-Za + Zm * (-Zn))
when the FMA produces a zero output (e.g. all zero inputs can produce -0
output)

Add a PatFrag to check presence of nsz on the fneg, add tests which
ensure the combine does not fire in the absense of nsz.

See https://reviews.llvm.org/D90901 for a similar discussion on X86.

[dmgreen]

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

405–408

These are all fast...?

[craig.topper]

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

405–408

That's true, but the patterns weren't checking that. Would you prefer I add a PatFrag to check the fast math flags and keep the patterns?

[dmgreen]

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

405–408

It would seem prudent to not delete obviously useful patterns. Having non-fast tests would be useful too, if we don't have them already.

Same goes for D109523 but the tests might be reversed there, where we have non-fast tests, but no nsz tests.

[peterwaller-arm]

Hi @craig.topper, are you actively working on this or would you like us to pick it up?

[craig.topper]

In D109525#3077765, @peterwaller-arm wrote:

Hi @craig.topper, are you actively working on this or would you like us to pick it up?

I'm not actively working on it. Would you mind picking it up?

[peterwaller-arm]

Sure, we'll pick it up (apols for the delayed response). Thanks for raising it.

[peterwaller-arm]

• Update commit message / title
• Introduce PatFrag to check for nsz on the fneg.
• fmad tests no longer change, because -(fma) only arises as a consequence of nsz.

[peterwaller-arm]

Couple of additional notes:

• On main I landed a change to substitute fast for contract, which left the test output unmodified.
• I had to introduce generic code to propagate fast math flags for unary op building, which as far as I can tell only affects the neg.
• Likewise in LowerToPredicatedOp, it was necessary to propagate the fneg node's flags down to instructions.

[dmgreen]

LGTM, with a small suggestion. Thanks for the updates.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19120–19121

Can this just be return DAG.getNode(NewOp, DL, VT, Operands, Op->getFlags());?

[peterwaller-arm]

Thanks for the improvement suggestion.

• Pass Op->getFlags() via getNode.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

if (DAG.getTarget().Options.NoTrapAfterNoreturn) {
     }
   }
 
   DAG.setRoot(DAG.getNode(ISD::TRAP, getCurSDLoc(), MVT::Other, DAG.getRoot()));
 }
 
 void SelectionDAGBuilder::visitUnary(const User &I, unsigned Opcode) {
   SDNodeFlags Flags;
+  if (auto *FPOp = dyn_cast<FPMathOperator>(&I))
+    Flags.copyFMF(*FPOp);
 
   SDValue Op = getValue(I.getOperand(0));
   SDValue UnNodeValue = DAG.getNode(Opcode, getCurSDLoc(), Op.getValueType(),
                                     Op, Flags);
   setValue(&I, UnNodeValue);
 }
 
 void SelectionDAGBuilder::visitBinary(const User &I, unsigned Opcode) {

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

assert((!V.getValueType().isVector() ||
             V.getValueType().isScalableVector()) &&
            "Only scalable vectors are supported!");
     Operands.push_back(V);
   }
 
   if (isMergePassthruOpcode(NewOp))
     Operands.push_back(DAG.getUNDEF(VT));
 
-  return DAG.getNode(NewOp, DL, VT, Operands);
+  return DAG.getNode(NewOp, DL, VT, Operands, Op->getFlags());
 }
 
 // If a fixed length vector operation has no side effects when applied to
 // undefined elements, we can safely use scalable vectors to perform the same
 // operation without needing to worry about predication.
 SDValue AArch64TargetLowering::LowerToScalableOp(SDValue Op,
                                                  SelectionDAG &DAG) const {
   EVT VT = Op.getValueType();
   assert(useSVEForFixedLengthVectorVT(VT) &&

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

 
 def reinterpret_cast : SDNode<"AArch64ISD::REINTERPRET_CAST", SDTUnaryOp>;
 
 def AArch64mul_p_oneuse : PatFrag<(ops node:$pred, node:$src1, node:$src2),
                                   (AArch64mul_p node:$pred, node:$src1, node:$src2), [{
   return N->hasOneUse();
 }]>;
 
+def AArch64fneg_mt_nsz : PatFrag<(ops node:$pred, node:$op, node:$pt),
+                                  (AArch64fneg_mt node:$pred, node:$op, node:$pt), [{
+  return N->getFlags().hasNoSignedZeros();
+}]>;
+
 def SDT_AArch64Arith_Unpred : SDTypeProfile<1, 2, [
   SDTCisVec<0>, SDTCisVec<1>, SDTCisVec<2>,
   SDTCisSameAs<0,1>, SDTCisSameAs<1,2>
 ]>;
 
 def AArch64bic : SDNode<"AArch64ISD::BIC",  SDT_AArch64Arith_Unpred>;
 
 let Predicates = [HasSVE] in {
multiclass fma<ValueType Ty, ValueType PredTy, string Suffix> {
     def : Pat<(Ty (AArch64fma_p PredTy:$P, Ty:$Zn, Ty:$Zm, (AArch64fneg_mt PredTy:$P, Ty:$Za, (Ty (undef))))),
               (!cast<Instruction>("FNMLS_ZPZZZ_UNDEF_"#Suffix) $P, ZPR:$Za, ZPR:$Zn, ZPR:$Zm)>;
 
     // Zd = -Za + -Zn * Zm
     def : Pat<(Ty (AArch64fma_p PredTy:$P, (AArch64fneg_mt PredTy:$P, Ty:$Zn, (Ty (undef))), Ty:$Zm, (AArch64fneg_mt PredTy:$P, Ty:$Za, (Ty (undef))))),
               (!cast<Instruction>("FNMLA_ZPZZZ_UNDEF_"#Suffix) $P, ZPR:$Za, ZPR:$Zn, ZPR:$Zm)>;
 
     // Zd = -(Za + Zn * Zm)
-    def : Pat<(AArch64fneg_mt PredTy:$P, (AArch64fma_p PredTy:$P, Ty:$Zn, Ty:$Zm, Ty:$Za), (Ty (undef))),
+    // (with nsz neg.)
+    def : Pat<(AArch64fneg_mt_nsz PredTy:$P, (AArch64fma_p PredTy:$P, Ty:$Zn, Ty:$Zm, Ty:$Za), (Ty (undef))),
               (!cast<Instruction>("FNMLA_ZPZZZ_UNDEF_"#Suffix) $P, ZPR:$Za, ZPR:$Zn, ZPR:$Zm)>;
 
     // Zda = Zda + Zn * Zm
     def : Pat<(vselect (PredTy PPR:$Pg), (Ty (AArch64fma_p (PredTy (AArch64ptrue 31)), ZPR:$Zn, ZPR:$Zm, ZPR:$Za)), ZPR:$Za),
               (!cast<Instruction>("FMLA_ZPmZZ_"#Suffix) PPR:$Pg, ZPR:$Za, ZPR:$Zn, ZPR:$Zm)>;
 
     // Zda = Zda + -Zn * Zm
     def : Pat<(vselect (PredTy PPR:$Pg), (Ty (AArch64fma_p (PredTy (AArch64ptrue 31)), (AArch64fneg_mt (PredTy (AArch64ptrue 31)), Ty:$Zn, (Ty (undef))), ZPR:$Zm, ZPR:$Za)), ZPR:$Za),

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

; CHECK-NEXT:    ret
   ret <vscale x 2 x double> %res
 }
 
 define <vscale x 8 x half> @fnmad_h(<vscale x 8 x half> %m1, <vscale x 8 x half> %m2, <vscale x 8 x half> %acc) {
 ; CHECK-LABEL: fnmad_h:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.h
 ; CHECK-NEXT:    fnmad z0.h, p0/m, z1.h, z2.h
 ; CHECK-NEXT:    ret
   %neg_m1 = fneg contract <vscale x 8 x half> %m1
   %mul = fmul contract <vscale x 8 x half> %neg_m1, %m2
   %res = fsub contract <vscale x 8 x half> %mul, %acc
   ret <vscale x 8 x half> %res
 }
 
 define <vscale x 4 x half> @fnmad_hx4(<vscale x 4 x half> %m1, <vscale x 4 x half> %m2, <vscale x 4 x half> %acc) {
 ; CHECK-LABEL: fnmad_hx4:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    fnmad z0.h, p0/m, z1.h, z2.h

 define <vscale x 8 x half> @fnmla_h_reversed(<vscale x 8 x half> %acc, <vscale x 8 x half> %m1, <vscale x 8 x half> %m2) {
 ; CHECK-LABEL: fnmla_h_reversed:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.h
 ; CHECK-NEXT:    fnmla z0.h, p0/m, z1.h, z2.h
 ; CHECK-NEXT:    ret
   %mul = fmul contract <vscale x 8 x half> %m1, %m2
   %add = fadd contract <vscale x 8 x half> %mul, %acc
-  %res = fneg contract <vscale x 8 x half> %add
+  %res = fneg contract nsz <vscale x 8 x half> %add
   ret <vscale x 8 x half> %res
 }
 
 define <vscale x 4 x half> @fnmla_hx4_reversed(<vscale x 4 x half> %acc, <vscale x 4 x half> %m1, <vscale x 4 x half> %m2) {
 ; CHECK-LABEL: fnmla_hx4_reversed:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    fnmla z0.h, p0/m, z1.h, z2.h
 ; CHECK-NEXT:    ret
   %mul = fmul contract <vscale x 4 x half> %m1, %m2
   %add = fadd contract <vscale x 4 x half> %mul, %acc
-  %res = fneg contract <vscale x 4 x half> %add
+  %res = fneg contract nsz <vscale x 4 x half> %add
   ret <vscale x 4 x half> %res
 }
 
 define <vscale x 2 x half> @fnmla_hx2_reversed(<vscale x 2 x half> %acc, <vscale x 2 x half> %m1, <vscale x 2 x half> %m2) {
 ; CHECK-LABEL: fnmla_hx2_reversed:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    fnmla z0.h, p0/m, z1.h, z2.h
 ; CHECK-NEXT:    ret
   %mul = fmul contract <vscale x 2 x half> %m1, %m2
   %add = fadd contract <vscale x 2 x half> %mul, %acc
-  %res = fneg contract <vscale x 2 x half> %add
+  %res = fneg contract nsz <vscale x 2 x half> %add
   ret <vscale x 2 x half> %res
 }
 
 define <vscale x 4 x float> @fnmla_s_reversed(<vscale x 4 x float> %acc, <vscale x 4 x float> %m1, <vscale x 4 x float> %m2) {
 ; CHECK-LABEL: fnmla_s_reversed:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    fnmla z0.s, p0/m, z1.s, z2.s
 ; CHECK-NEXT:    ret
   %mul = fmul contract <vscale x 4 x float> %m1, %m2
   %add = fadd contract <vscale x 4 x float> %mul, %acc
-  %res = fneg contract <vscale x 4 x float> %add
+  %res = fneg contract nsz <vscale x 4 x float> %add
   ret <vscale x 4 x float> %res
 }
 
 define <vscale x 2 x float> @fnmla_sx2_reversed(<vscale x 2 x float> %acc, <vscale x 2 x float> %m1, <vscale x 2 x float> %m2) {
 ; CHECK-LABEL: fnmla_sx2_reversed:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    fnmla z0.s, p0/m, z1.s, z2.s
 ; CHECK-NEXT:    ret
   %mul = fmul contract <vscale x 2 x float> %m1, %m2
   %add = fadd contract <vscale x 2 x float> %mul, %acc
-  %res = fneg contract <vscale x 2 x float> %add
+  %res = fneg contract nsz <vscale x 2 x float> %add
   ret <vscale x 2 x float> %res
 }
 
 define <vscale x 2 x double> @fnmla_d_reversed(<vscale x 2 x double> %acc, <vscale x 2 x double> %m1, <vscale x 2 x double> %m2) {
 ; CHECK-LABEL: fnmla_d_reversed:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    fnmla z0.d, p0/m, z1.d, z2.d
 ; CHECK-NEXT:    ret
   %mul = fmul contract <vscale x 2 x double> %m1, %m2
   %add = fadd contract <vscale x 2 x double> %mul, %acc
+  %res = fneg contract nsz <vscale x 2 x double> %add
+  ret <vscale x 2 x double> %res
+}
+
+define <vscale x 8 x half> @signed_zeros_negtest_fnmla_h_reversed(<vscale x 8 x half> %acc, <vscale x 8 x half> %m1, <vscale x 8 x half> %m2) {
+; CHECK-LABEL: signed_zeros_negtest_fnmla_h_reversed:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.h
+; CHECK-NEXT:    fmla z0.h, p0/m, z1.h, z2.h
+; CHECK-NEXT:    fneg z0.h, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %mul = fmul contract <vscale x 8 x half> %m1, %m2
+  %add = fadd contract <vscale x 8 x half> %mul, %acc
+  %res = fneg contract <vscale x 8 x half> %add
+  ret <vscale x 8 x half> %res
+}
+
+define <vscale x 4 x half> @signed_zeros_negtest_fnmla_hx4_reversed(<vscale x 4 x half> %acc, <vscale x 4 x half> %m1, <vscale x 4 x half> %m2) {
+; CHECK-LABEL: signed_zeros_negtest_fnmla_hx4_reversed:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    fmla z0.h, p0/m, z1.h, z2.h
+; CHECK-NEXT:    fneg z0.h, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %mul = fmul contract <vscale x 4 x half> %m1, %m2
+  %add = fadd contract <vscale x 4 x half> %mul, %acc
+  %res = fneg contract <vscale x 4 x half> %add
+  ret <vscale x 4 x half> %res
+}
+
+define <vscale x 2 x half> @signed_zeros_negtest_fnmla_hx2_reversed(<vscale x 2 x half> %acc, <vscale x 2 x half> %m1, <vscale x 2 x half> %m2) {
+; CHECK-LABEL: signed_zeros_negtest_fnmla_hx2_reversed:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    fmla z0.h, p0/m, z1.h, z2.h
+; CHECK-NEXT:    fneg z0.h, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %mul = fmul contract <vscale x 2 x half> %m1, %m2
+  %add = fadd contract <vscale x 2 x half> %mul, %acc
+  %res = fneg contract <vscale x 2 x half> %add
+  ret <vscale x 2 x half> %res
+}
+
+define <vscale x 4 x float> @signed_zeros_negtest_fnmla_s_reversed(<vscale x 4 x float> %acc, <vscale x 4 x float> %m1, <vscale x 4 x float> %m2) {
+; CHECK-LABEL: signed_zeros_negtest_fnmla_s_reversed:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    fmla z0.s, p0/m, z1.s, z2.s
+; CHECK-NEXT:    fneg z0.s, p0/m, z0.s
+; CHECK-NEXT:    ret
+  %mul = fmul contract <vscale x 4 x float> %m1, %m2
+  %add = fadd contract <vscale x 4 x float> %mul, %acc
+  %res = fneg contract <vscale x 4 x float> %add
+  ret <vscale x 4 x float> %res
+}
+
+define <vscale x 2 x float> @signed_zeros_negtest_fnmla_sx2_reversed(<vscale x 2 x float> %acc, <vscale x 2 x float> %m1, <vscale x 2 x float> %m2) {
+; CHECK-LABEL: signed_zeros_negtest_fnmla_sx2_reversed:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    fmla z0.s, p0/m, z1.s, z2.s
+; CHECK-NEXT:    fneg z0.s, p0/m, z0.s
+; CHECK-NEXT:    ret
+  %mul = fmul contract <vscale x 2 x float> %m1, %m2
+  %add = fadd contract <vscale x 2 x float> %mul, %acc
+  %res = fneg contract <vscale x 2 x float> %add
+  ret <vscale x 2 x float> %res
+}
+
+define <vscale x 2 x double> @signed_zeros_negtest_fnmla_d_reversed(<vscale x 2 x double> %acc, <vscale x 2 x double> %m1, <vscale x 2 x double> %m2) {
+; CHECK-LABEL: signed_zeros_negtest_fnmla_d_reversed:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    fmla z0.d, p0/m, z1.d, z2.d
+; CHECK-NEXT:    fneg z0.d, p0/m, z0.d
+; CHECK-NEXT:    ret
+  %mul = fmul contract <vscale x 2 x double> %m1, %m2
+  %add = fadd contract <vscale x 2 x double> %mul, %acc
   %res = fneg contract <vscale x 2 x double> %add
   ret <vscale x 2 x double> %res
 }
 
 define <vscale x 8 x half> @fnmls_h(<vscale x 8 x half> %acc, <vscale x 8 x half> %m1, <vscale x 8 x half> %m2) {
 ; CHECK-LABEL: fnmls_h:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.h