[GlobalISel] Add constant folding support for G_FMA/G_FMAD in the combiner. (#65659)

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

@@ -657,6 +657,9 @@ public:
   /// Do constant FP folding when opportunities are exposed after MIR building.
   bool matchConstantFoldFPBinOp(MachineInstr &MI, ConstantFP* &MatchInfo);
 
+  /// Constant fold G_FMA/G_FMAD.
+  bool matchConstantFoldFMA(MachineInstr &MI, ConstantFP *&MatchInfo);
+
   /// \returns true if it is possible to narrow the width of a scalar binop
   /// feeding a G_AND instruction \p MI.
   bool matchNarrowBinopFeedingAnd(MachineInstr &MI, BuildFnTy &MatchInfo);

llvm/include/llvm/Target/GlobalISel/Combine.td

@@ -985,6 +985,13 @@ def constant_fold_fp_binop : GICombineRule<
    [{ return Helper.matchConstantFoldFPBinOp(*${d}, ${matchinfo}); }]),
   (apply [{ Helper.replaceInstWithFConstant(*${d}, ${matchinfo}); }])>;
 
+
+def constant_fold_fma : GICombineRule<
+  (defs root:$d, constantfp_matchinfo:$matchinfo),
+  (match (wip_match_opcode G_FMAD, G_FMA):$d,
+   [{ return Helper.matchConstantFoldFMA(*${d}, ${matchinfo}); }]),
+  (apply [{ Helper.replaceInstWithFConstant(*${d}, ${matchinfo}); }])>;
+
 def constant_fold_cast_op : GICombineRule<
   (defs root:$d, apint_matchinfo:$matchinfo),
   (match (wip_match_opcode G_ZEXT, G_SEXT, G_ANYEXT):$d,
@@ -1253,7 +1260,8 @@ def all_combines : GICombineGroup<[trivial_combines, insert_vec_elt_combines,
     const_combines, xor_of_and_with_same_reg, ptr_add_with_zero,
     shift_immed_chain, shift_of_shifted_logic_chain, load_or_combine,
     div_rem_to_divrem, funnel_shift_combines, commute_shift,
-    form_bitfield_extract, constant_fold_binops, constant_fold_cast_op, fabs_fneg_fold,
+    form_bitfield_extract, constant_fold_binops, constant_fold_fma,
+    constant_fold_cast_op, fabs_fneg_fold,
     intdiv_combines, mulh_combines, redundant_neg_operands,
     and_or_disjoint_mask, fma_combines, fold_binop_into_select,
     sub_add_reg, select_to_minmax, redundant_binop_in_equality,

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 #include "llvm/CodeGen/GlobalISel/CombinerHelper.h"
+#include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/CodeGen/GlobalISel/GISelChangeObserver.h"
@@ -4621,6 +4622,31 @@ bool CombinerHelper::matchConstantFoldFPBinOp(MachineInstr &MI, ConstantFP* &Mat
   return true;
 }
 
+bool CombinerHelper::matchConstantFoldFMA(MachineInstr &MI,
+                                          ConstantFP *&MatchInfo) {
+  assert(MI.getOpcode() == TargetOpcode::G_FMA ||
+         MI.getOpcode() == TargetOpcode::G_FMAD);
+  auto [_, Op1, Op2, Op3] = MI.getFirst4Regs();
+
+  const ConstantFP *Op3Cst = getConstantFPVRegVal(Op3, MRI);
+  if (!Op3Cst)
+    return false;
+
+  const ConstantFP *Op2Cst = getConstantFPVRegVal(Op2, MRI);
+  if (!Op2Cst)
+    return false;
+
+  const ConstantFP *Op1Cst = getConstantFPVRegVal(Op1, MRI);
+  if (!Op1Cst)
+    return false;
+
+  APFloat Op1F = Op1Cst->getValueAPF();
+  Op1F.fusedMultiplyAdd(Op2Cst->getValueAPF(), Op3Cst->getValueAPF(),
+                        APFloat::rmNearestTiesToEven);
+  MatchInfo = ConstantFP::get(MI.getMF()->getFunction().getContext(), Op1F);
+  return true;
+}
+
 bool CombinerHelper::matchNarrowBinopFeedingAnd(
     MachineInstr &MI, std::function<void(MachineIRBuilder &)> &MatchInfo) {
   // Look for a binop feeding into an AND with a mask:

llvm/test/CodeGen/AArch64/GlobalISel/combine-constant-fold-fma.mir

@@ -0,0 +1,48 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple aarch64 -run-pass=aarch64-prelegalizer-combiner -verify-machineinstrs %s -o - | FileCheck %s
+
+---
+name:            fma
+liveins:
+  - { reg: '$d0' }
+body:             |
+  bb.1.entry:
+    liveins: $d0
+
+    ; CHECK-LABEL: name: fma
+    ; CHECK: liveins: $d0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %res:_(s64) = G_FCONSTANT double 8.100000e+01
+    ; CHECK-NEXT: $d0 = COPY %res(s64)
+    ; CHECK-NEXT: RET_ReallyLR implicit $d0
+    %a:_(s64) = G_FCONSTANT double 40.0
+    %b:_(s64) = G_FCONSTANT double 2.0
+    %c:_(s64) = G_FCONSTANT double 1.0
+    %res:_(s64) = G_FMA %a, %b, %c
+    $d0 = COPY %res(s64)
+    RET_ReallyLR implicit $d0
+
+...
+
+---
+name:            fmad
+liveins:
+  - { reg: '$d0' }
+body:             |
+  bb.1.entry:
+    liveins: $d0
+
+    ; CHECK-LABEL: name: fmad
+    ; CHECK: liveins: $d0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %res:_(s64) = G_FCONSTANT double 8.100000e+01
+    ; CHECK-NEXT: $d0 = COPY %res(s64)
+    ; CHECK-NEXT: RET_ReallyLR implicit $d0
+    %a:_(s64) = G_FCONSTANT double 40.0
+    %b:_(s64) = G_FCONSTANT double 2.0
+    %c:_(s64) = G_FCONSTANT double 1.0
+    %res:_(s64) = G_FMAD %a, %b, %c
+    $d0 = COPY %res(s64)
+    RET_ReallyLR implicit $d0
+
+...