[AArch64] Consider instruction-level contract FMFs in combiner patterns.

[fhahn]

Summary

Currently, instruction level fast math flags are not considered when
generating patterns for the machine combiner.

This currently leads to some missed opportunities to generate FMAs in
combination with #pragma clang fp contract (fast).

For example, when building the example below with -O3 for AArch64, no
FMADD is generated. If built with -O2 and the DAGCombiner is used
instead of the MachineCombiner for FMAs, an FMADD is generated.

With this patch, the same code is generated in both cases.

float madd_contract(float a, float b, float c) {
#pragma clang fp contract (fast)
  return (a * b) + c;
}

[dmgreen]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

3919

-> We can only fuse ...

4122

Do both instructions need the fast-math flag, or just the add? I guess it's safer to check both, but the code would end up being simpler if it was just the add.

Reading DAGCombiner::visitFADDForFMACombine it's a bit messy, but seems like it only really considers the add:
https://godbolt.org/z/3ejrhE
Reading the code though, I'm not sure that's really what it is intending to do.

llvm/test/CodeGen/AArch64/machine-combiner-instr-fmf.mir

82

-> both

[fhahn]

Only require contract on FADD/FSUB for fusion.

[fhahn]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

4122

Indeed it seems like it should be sufficient on the FADD/FSUB. Given that is what DAGCOmbine does, it should be safe to do so in the machine combiner too (plus it ensures consistency between -O3 and other levels). It also makes the code much simpler.

I also found llvm/test/CodeGen/AArch64/neon-fma-FMF.ll, which tests with different contract flags and it specifically checks for fusion with contract on FADD and no fusion with contract on FMUL. I added -O3, which checks the machine-combiner path as well.

[dmgreen]

Thanks. LGTM, in that this work the same as DAGCombiner.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

static bool isCombineInstrCandidate64(unsigned Opc) {
   case AArch64::SUBv4i32:
     return true;
   default:
     break;
   }
   return false;
 }
 
-// FP Opcodes that can be combined with a FMUL
+// FP Opcodes that can be combined with a FMUL.
 static bool isCombineInstrCandidateFP(const MachineInstr &Inst) {
   switch (Inst.getOpcode()) {
   default:
     break;
   case AArch64::FADDHrr:
   case AArch64::FADDSrr:
   case AArch64::FADDDrr:
   case AArch64::FADDv4f16:
   case AArch64::FADDv8f16:
   case AArch64::FADDv2f32:
   case AArch64::FADDv2f64:
   case AArch64::FADDv4f32:
   case AArch64::FSUBHrr:
   case AArch64::FSUBSrr:
   case AArch64::FSUBDrr:
   case AArch64::FSUBv4f16:
   case AArch64::FSUBv8f16:
   case AArch64::FSUBv2f32:
   case AArch64::FSUBv2f64:
   case AArch64::FSUBv4f32:
     TargetOptions Options = Inst.getParent()->getParent()->getTarget().Options;
-    return (Options.UnsafeFPMath ||
-            Options.AllowFPOpFusion == FPOpFusion::Fast);
+    // We can fuse FADD/FSUB with FMUL, if fusion is either allowed globally by
+    // the target options or if FADD/FSUB has the contract fast-math flag.
+    return Options.UnsafeFPMath ||
+           Options.AllowFPOpFusion == FPOpFusion::Fast ||
+           Inst.getFlag(MachineInstr::FmContract);
+    return true;
   }
   return false;
 }
 
 // Opcodes that can be combined with a MUL
 static bool isCombineInstrCandidate(unsigned Opc) {
   return (isCombineInstrCandidate32(Opc) || isCombineInstrCandidate64(Opc));
 }
if (MO.isReg() && Register::isVirtualRegister(MO.getReg()))
     MI = MRI.getUniqueVRegDef(MO.getReg());
   // And it needs to be in the trace (otherwise, it won't have a depth).
   if (!MI || MI->getParent() != &MBB || (unsigned)MI->getOpcode() != CombineOpc)
     return false;
   // Must only used by the user we combine with.
   if (!MRI.hasOneNonDBGUse(MI->getOperand(0).getReg()))
     return false;
 
   if (CheckZeroReg) {
     assert(MI->getNumOperands() >= 4 && MI->getOperand(0).isReg() &&
            MI->getOperand(1).isReg() && MI->getOperand(2).isReg() &&
            MI->getOperand(3).isReg() && "MAdd/MSub must have a least 4 regs");
     // The third input reg must be zero.
     if (MI->getOperand(3).getReg() != ZeroReg)
       return false;
   }
 

 /// Find instructions that can be turned into madd.
 static bool getFMAPatterns(MachineInstr &Root,
                            SmallVectorImpl<MachineCombinerPattern> &Patterns) {
 
   if (!isCombineInstrCandidateFP(Root))
     return false;
 
   MachineBasicBlock &MBB = *Root.getParent();
   bool Found = false;
 
   auto Match = [&](int Opcode, int Operand,
                    MachineCombinerPattern Pattern) -> bool {
     if (canCombineWithFMUL(MBB, Root.getOperand(Operand), Opcode)) {
       Patterns.push_back(Pattern);
       return true;
     }
     return false;

llvm/test/CodeGen/AArch64/machine-combiner-instr-fmf.mir

 # RUN: llc -run-pass=machine-combiner -o - -simplify-mir -mtriple=arm64-apple-iphoneos %s | FileCheck %s
 
 # Can create FMADD, because both the fmul and fadd have all fast-math flags.
 #
 # CHECK-LABEL: name: scalar_fmadd_fast
 # CHECK:        [[C:%.*]]:fpr32 = COPY $s2
 # CHECK-NEXT:   [[B:%.*]]:fpr32 = COPY $s1
 # CHECK-NEXT:   [[A:%.*]]:fpr32 = COPY $s0
-# CHECK-NEXT:   [[MUL:%.*]]:fpr32 = nnan ninf nsz arcp contract afn reassoc FMULSrr [[B]], [[A]]
-# CHECK-NEXT:   fpr32 = nnan ninf nsz arcp contract afn reassoc FADDSrr killed [[MUL]], [[C]]
+# CHECK-NEXT:   :fpr32 = FMADDSrrr [[B]], [[A]], [[C]]
 ---
 name:            scalar_fmadd_fast
 alignment:       4
 tracksRegLiveness: true
 registers:
   - { id: 0, class: fpr32 }
   - { id: 1, class: fpr32 }
   - { id: 2, class: fpr32 }

 ...
 
 # Can create FMADD, because both the fmul and fadd have the contract fast-math flag.
 #
 # CHECK-LABEL: name: scalar_fmadd_contract
 # CHECK:        [[C:%.*]]:fpr32 = COPY $s2
 # CHECK-NEXT:   [[B:%.*]]:fpr32 = COPY $s1
 # CHECK-NEXT:   [[A:%.*]]:fpr32 = COPY $s0
-# CHECK-NEXT:   [[MUL:%.*]]:fpr32 = contract FMULSrr [[B]], [[A]]
-# CHECK-NEXT:   fpr32 = contract FADDSrr killed [[MUL]], [[C]]
+# CHECK-NEXT:   :fpr32 = FMADDSrrr [[B]], [[A]], [[C]]
 
 ---
 name:            scalar_fmadd_contract
 alignment:       4
 tracksRegLiveness: true
 registers:
   - { id: 0, class: fpr32 }
   - { id: 1, class: fpr32 }
bb.0.entry:
     %0:fpr32 = COPY $s0
     %3:fpr32 = contract FMULSrr %1, %0
     %4:fpr32 = contract FADDSrr killed %3, %2
     $s0 = COPY %4
     RET_ReallyLR implicit $s0
 
 ...
 
-# Do not create FMADD, because we don't have the contract flag on moth instructions.
+# Do not create FMADD, because we don't have the contract flag on the FADD.
 
 # CHECK-LABEL: name: scalar_fmadd_contract_op0
 # CHECK:        [[C:%.*]]:fpr32 = COPY $s2
 # CHECK-NEXT:   [[B:%.*]]:fpr32 = COPY $s1
 # CHECK-NEXT:   [[A:%.*]]:fpr32 = COPY $s0
 # CHECK-NEXT:   [[MUL:%.*]]:fpr32 = contract FMULSrr [[B]], [[A]]
 # CHECK-NEXT:   fpr32 = FADDSrr killed [[MUL]], [[C]]
 ---
bb.0.entry:
     %0:fpr32 = COPY $s0
     %3:fpr32 = contract FMULSrr %1, %0
     %4:fpr32 = FADDSrr killed %3, %2
     $s0 = COPY %4
     RET_ReallyLR implicit $s0
 
 ...
 
-# Do not create FMADD, because we don't have the contract flag on moth instructions.
+# Do create FMADD, because we have the contract flag on the FADD.
 #
 # CHECK-LABEL: name: scalar_fmadd_contract_op1
 # CHECK:        [[C:%.*]]:fpr32 = COPY $s2
 # CHECK-NEXT:   [[B:%.*]]:fpr32 = COPY $s1
 # CHECK-NEXT:   [[A:%.*]]:fpr32 = COPY $s0
-# CHECK-NEXT:   [[MUL:%.*]]:fpr32 = FMULSrr [[B]], [[A]]
-# CHECK-NEXT:   fpr32 = contract FADDSrr killed [[MUL]], [[C]]
+# CHECK-NEXT:   :fpr32 = FMADDSrrr [[B]], [[A]], [[C]]
 
 ---
 name:            scalar_fmadd_contract_op1
 alignment:       4
 tracksRegLiveness: true
 registers:
   - { id: 0, class: fpr32 }
   - { id: 1, class: fpr32 }
bb.0.entry:
     %0:fpr32 = COPY $s0
     %3:fpr32 = nsz FMULSrr %1, %0
     %4:fpr32 = nsz FADDSrr killed %3, %2
     $s0 = COPY %4
     RET_ReallyLR implicit $s0
 
 ...
 
-# Can create FMADD, because both the fmul and fadd have all fast-math flags.
+# Can create FMLA, because both the fmul and fadd have all fast-math flags.
 #
 # CHECK-LABEL: name: vector_fmadd_fast
 # CHECK:       [[C:%.*]]:fpr128 = COPY $q2
 # CHECK-NEXT:  [[B:%.*]]:fpr128 = COPY $q1
 # CHECK-NEXT:  [[A:%.*]]:fpr128 = COPY $q0
-# CHECK-NEXT:  [[MUL:%.*]]:fpr128 = nnan ninf nsz arcp contract afn reassoc FMULv2f64 [[B]], [[A]]
-# CHECK-NEXT:  fpr128 = nnan ninf nsz arcp contract afn reassoc FADDv2f64 killed [[MUL]], [[C]]
+# CHECK-NEXT:  fpr128 = FMLAv2f64 [[C]], [[B]], [[A]]
 ---
 name:            vector_fmadd_fast
 alignment:       4
 tracksRegLiveness: true
 registers:
   - { id: 0, class: fpr128 }
   - { id: 1, class: fpr128 }
   - { id: 2, class: fpr128 }
bb.0.entry:
     %0:fpr128 = COPY $q0
     %3:fpr128 = nnan ninf nsz arcp contract afn reassoc FMULv2f64 %1, %0
     %4:fpr128 = nnan ninf nsz arcp contract afn reassoc FADDv2f64 killed %3, %2
     $q0 = COPY %4
     RET_ReallyLR implicit $q0
 
 ...
 
-# Can create FMADD, because both the fmul and fadd have the contract fast-math flag.
+# Can create FMLA, because both the fmul and fadd have the contract fast-math flag.
 #
 # CHECK-LABEL: name: vector_fmadd_contract
 # CHECK:       [[C:%.*]]:fpr128 = COPY $q2
 # CHECK-NEXT:  [[B:%.*]]:fpr128 = COPY $q1
 # CHECK-NEXT:  [[A:%.*]]:fpr128 = COPY $q0
-# CHECK-NEXT:  [[MUL:%.*]]:fpr128 = contract FMULv2f64 [[B]], [[A]]
-# CHECK-NEXT:  fpr128 = contract FADDv2f64 killed [[MUL]], [[C]]
+# CHECK-NEXT:  fpr128 = FMLAv2f64 [[C]], [[B]], [[A]]
 ---
 name:            vector_fmadd_contract
 alignment:       4
 tracksRegLiveness: true
 registers:
   - { id: 0, class: fpr128 }
   - { id: 1, class: fpr128 }
   - { id: 2, class: fpr128 }
bb.0.entry:
     %0:fpr128 = COPY $q0
     %3:fpr128 = contract FMULv2f64 %1, %0
     %4:fpr128 = contract FADDv2f64 killed %3, %2
     $q0 = COPY %4
     RET_ReallyLR implicit $q0
 
 ...
 
-# Do not create FMADD, because we don't have the contract flag on moth instructions.
+# Do not create FMLA, because we don't have the contract flag on the FADD.
 #
 # CHECK-LABEL: name: vector_fmadd_contract_op0
 # CHECK:       [[C:%.*]]:fpr128 = COPY $q2
 # CHECK-NEXT:  [[B:%.*]]:fpr128 = COPY $q1
 # CHECK-NEXT:  [[A:%.*]]:fpr128 = COPY $q0
 # CHECK-NEXT:  [[MUL:%.*]]:fpr128 = contract FMULv2f64 [[B]], [[A]]
 # CHECK-NEXT:  fpr128 = FADDv2f64 killed [[MUL]], [[C]]
 ---
bb.0.entry:
     %0:fpr128 = COPY $q0
     %3:fpr128 = contract FMULv2f64 %1, %0
     %4:fpr128 = FADDv2f64 killed %3, %2
     $q0 = COPY %4
     RET_ReallyLR implicit $q0
 
 ...
 
-# Do not create FMADD, because we don't have the contract flag on moth instructions.
+# Do create FMLA, because we have the contract flag on the FADD.
 #
 # CHECK-LABEL: name: vector_fmadd_contract_op1
 # CHECK:       [[C:%.*]]:fpr128 = COPY $q2
 # CHECK-NEXT:  [[B:%.*]]:fpr128 = COPY $q1
 # CHECK-NEXT:  [[A:%.*]]:fpr128 = COPY $q0
-# CHECK-NEXT:  [[MUL:%.*]]:fpr128 = FMULv2f64 [[B]], [[A]]
-# CHECK-NEXT:  fpr128 = contract FADDv2f64 killed [[MUL]], [[C]]
+# CHECK-NEXT:  fpr128 = FMLAv2f64 [[C]], [[B]], [[A]]
 
 ---
 name:            vector_fmadd_contract_op1
 alignment:       4
 tracksRegLiveness: true
 registers:
   - { id: 0, class: fpr128 }
   - { id: 1, class: fpr128 }
bb.0.entry:
     %0:fpr128 = COPY $q0
     %3:fpr128 = FMULv2f64 %1, %0
     %4:fpr128 = contract FADDv2f64 killed %3, %2
     $q0 = COPY %4
     RET_ReallyLR implicit $q0
 
 ...
 
-# Do not create FMADD, as nsz flag does not allow it.
+# Do not create FMLA, as nsz flag does not allow it.
 #
 # CHECK-LABEL: name: vector_fmadd_nsz
 # CHECK:       [[C:%.*]]:fpr128 = COPY $q2
 # CHECK-NEXT:  [[B:%.*]]:fpr128 = COPY $q1
 # CHECK-NEXT:  [[A:%.*]]:fpr128 = COPY $q0
 # CHECK-NEXT:  [[MUL:%.*]]:fpr128 = nsz FMULv2f64 [[B]], [[A]]
 # CHECK-NEXT:  fpr128 = nsz FADDv2f64 killed [[MUL]], [[C]]
 ---

llvm/test/CodeGen/AArch64/neon-fma-FMF.ll

 ; RUN: llc < %s -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -mattr=+neon | FileCheck %s
+; RUN: llc < %s -O3 -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -mattr=+neon | FileCheck %s
 
 define <2 x float> @fma_1(<2 x float> %A, <2 x float> %B, <2 x float> %C) {
 ; CHECK-LABEL: fma_1:
 ; CHECK: fmla {{v[0-9]+}}.2s, {{v[0-9]+}}.2s, {{v[0-9]+}}.2s
   %tmp1 = fmul contract <2 x float> %A, %B;
   %tmp2 = fadd contract <2 x float> %C, %tmp1;
   ret <2 x float> %tmp2
 }