[DAGCombiner] More opportunities to fuse fmul and fadd to fma aggressively

[slydiman]

Summary

This patch is adding a new pattern matching in DAGCombiner to look for opportunities to fuse fmul and fadd to fma aggressively.

The matching is very specific to the following pattern:
  FADD X, (FMA A, B, (FMA C, D, (MUL Y, Z)))
can be converted to
  FMA Y, Z (FMA A, B, (FMA C, D, X))

This transformation is also changing the orders of operations, so it will only be enabled for FADD, when it's fast-math attribute has reassoc or contract flag enabled.

[slydiman]

Ping

[tra]

I've been burned by FP numerics in the past, so I'm pretty sure I'm not qualified to tell whether the change will break something in an unexpected way.

FADD X, (FMA A, B, (FMA C, D, (MUL Y, Z))) -> FMA Y, Z (FMA A, B, (FMA C, D, X))

I'm curious, is there a reason for this particular arrangement? Are there any benefits/disadvantages vs, let's say FMA A, B (FMA C, D, (FMA Y, Z, X)) ?

[slydiman]

It seems already implemented here https://reviews.llvm.org/D132837

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (isContractableFMUL(N10) &&
                          DAG.getNode(ISD::FP_EXTEND, SL, VT, N10.getOperand(0)),
                          DAG.getNode(ISD::FP_EXTEND, SL, VT, N10.getOperand(1)),
                          N0);
     }
   }
 
   // More folding opportunities when target permits.
   if (Aggressive) {
+    // fold (fadd x, (fma a, b, (fma c, d, (fmul y, z))))
+    //   -> (fma y, z (fma a, b, (fma c, d, x)))
+    auto FoldFAddFMAFMAFMul = [&](SDValue A, SDValue B, SDValue C, SDValue D,
+                                  SDValue X, SDValue Y, SDValue Z) {
+      return DAG.getNode(
+          PreferredFusedOpcode, SL, VT, Y, Z,
+          DAG.getNode(PreferredFusedOpcode, SL, VT, A, B,
+                      DAG.getNode(PreferredFusedOpcode, SL, VT, C, D, X)));
+    };
+    if (isFusedOp(N1)) {
+      SDValue FMA1_N2 = N1.getOperand(2);
+      if (isFusedOp(FMA1_N2)) {
+        SDValue FMA2_N2 = FMA1_N2.getOperand(2);
+        if (FMA2_N2.getOpcode() == ISD::FMUL) {
+          return FoldFAddFMAFMAFMul(
+              N1.getOperand(0), N1.getOperand(1), FMA1_N2.getOperand(0),
+              FMA1_N2.getOperand(1), N0, FMA2_N2.getOperand(0),
+              FMA2_N2.getOperand(1));
+        }
+      }
+    }
+
     // fold (fadd (fma x, y, (fpext (fmul u, v))), z)
     //   -> (fma x, y, (fma (fpext u), (fpext v), z))
     auto FoldFAddFMAFPExtFMul = [&](SDValue X, SDValue Y, SDValue U, SDValue V,
                                     SDValue Z) {
       return DAG.getNode(PreferredFusedOpcode, SL, VT, X, Y,
                          DAG.getNode(PreferredFusedOpcode, SL, VT,
                                      DAG.getNode(ISD::FP_EXTEND, SL, VT, U),
                                      DAG.getNode(ISD::FP_EXTEND, SL, VT, V),

llvm/test/CodeGen/NVPTX/fma-assoc.ll

 ; CHECK: fma.rn.f64
   %1 = fmul double %val1, %val2
   %2 = fmul double %1, %1
   %3 = fadd double %1, %2
 
   ret double %3
 }
 
+define float @FoldFAddFMAFMAFMul(float %a, float %b, float %c, float %d,
+                                 float %x, float %y, float %z) {
+; CHECK-LABEL: FoldFAddFMAFMAFMul(
+; fold (fadd x, (fma a, b, (fma c, d, (fmul y, z))))
+;   -> (fma y, z (fma a, b, (fma c, d, x)))
+; CHECK-NOT: mul.f32
+; CHECK: fma.rn.f32
+; CHECK: fma.rn.f32
+; CHECK: fma.rn.f32
+; CHECK-NOT: add.f32
+  %mul = fmul float %y, %z
+  %fma1 = call float @llvm.fma.f32(float %c, float %d, float %mul)
+  %fma2 = call float @llvm.fma.f32(float %a, float %b, float %fma1)
+  %res = fadd float %x, %fma2
+  ret float %res
+}
+
+declare float @llvm.fma.f32(float, float, float) #0
+
+attributes #0 = { nounwind readnone }