[AMDGPU][CostModel] Improve cost estimation for fused {fadd|fsub}(a,fmul(b,c))

[dfukalov]

Summary

If result of fmul(b,c) has one use, in almost all cases (except denormals are
IEEE) the pair of operations will be fused in one fma/mad/mac/etc.

[rampitec]

LGTM

[arsenm]

llvm/test/Analysis/CostModel/AMDGPU/fused_costs.ll

50

Extra newlines

[dfukalov]

removed extra new lines

llvm/include/llvm/CodeGen/BasicTTIImpl.h

unsigned getArithmeticInstrCost(
     }
 
     // Else, assume that we need to scalarize this op.
     // TODO: If one of the types get legalized by splitting, handle this
     // similarly to what getCastInstrCost() does.
     if (auto *VTy = dyn_cast<VectorType>(Ty)) {
       unsigned Num = cast<FixedVectorType>(VTy)->getNumElements();
       unsigned Cost = thisT()->getArithmeticInstrCost(
-          Opcode, VTy->getScalarType(), CostKind);
+          Opcode, VTy->getScalarType(), CostKind, Opd1Info, Opd2Info,
+          Opd1PropInfo, Opd2PropInfo, Args, CxtI);
       // Return the cost of multiple scalar invocation plus the cost of
       // inserting and extracting the values.
       return getScalarizationOverhead(VTy, Args) + Num * Cost;
     }
 
     // We don't know anything about this scalar instruction.
     return OpCost;
   }

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

int GCNTTIImpl::getArithmeticInstrCost(unsigned Opcode, Type *Ty,
                                        ArrayRef<const Value *> Args,
                                        const Instruction *CxtI) {
   EVT OrigTy = TLI->getValueType(DL, Ty);
   if (!OrigTy.isSimple()) {
     // FIXME: We're having to query the throughput cost so that the basic
     // implementation tries to generate legalize and scalarization costs. Maybe
     // we could hoist the scalarization code here?
     return BaseT::getArithmeticInstrCost(Opcode, Ty, TTI::TCK_RecipThroughput,
-                                         Opd1Info, Opd2Info,
-                                         Opd1PropInfo, Opd2PropInfo);
+                                         Opd1Info, Opd2Info, Opd1PropInfo,
+                                         Opd2PropInfo, Args, CxtI);
   }
 
   // Legalize the type.
   std::pair<int, MVT> LT = TLI->getTypeLegalizationCost(DL, Ty);
   int ISD = TLI->InstructionOpcodeToISD(Opcode);
 
   // Because we don't have any legal vector operations, but the legal types, we
   // need to account for split vectors.
case ISD::MUL: {
     }
 
     if (ST->has16BitInsts() && SLT == MVT::i16)
       NElts = (NElts + 1) / 2;
 
     // i32
     return QuarterRateCost * NElts * LT.first;
   }
+  case ISD::FMUL:
+    // Check possible fuse {fadd|fsub}(a,fmul(b,c)) and return zero cost for
+    // fmul(b,c) supposing the fadd|fsub will get estimated cost for the whole
+    // fused operation.
+    if (!HasFP32Denormals && SLT == MVT::f32 && CxtI && CxtI->hasOneUse())
+      if (const auto *FAdd = dyn_cast<BinaryOperator>(*CxtI->user_begin())) {
+        const int OPC = TLI->InstructionOpcodeToISD(FAdd->getOpcode());
+        if (OPC == ISD::FADD || OPC == ISD::FSUB) {
+          return TargetTransformInfo::TCC_Free;
+        }
+      }
+    LLVM_FALLTHROUGH;
   case ISD::FADD:
   case ISD::FSUB:
-  case ISD::FMUL:
     if (SLT == MVT::f64)
       return LT.first * NElts * get64BitInstrCost();
 
     if (ST->has16BitInsts() && SLT == MVT::f16)
       NElts = (NElts + 1) / 2;
 
     if (SLT == MVT::f32 || SLT == MVT::f16)
       return LT.first * NElts * getFullRateInstrCost();
int GCNTTIImpl::getArithmeticInstrCost(unsigned Opcode, Type *Ty,
   case ISD::FNEG:
     // Use the backend' estimation. If fneg is not free each element will cost
     // one additional instruction.
     return TLI->isFNegFree(SLT) ? 0 : NElts;
   default:
     break;
   }
 
-  return BaseT::getArithmeticInstrCost(Opcode, Ty, CostKind, Opd1Info,
-                                       Opd2Info,
-                                       Opd1PropInfo, Opd2PropInfo);
+  return BaseT::getArithmeticInstrCost(Opcode, Ty, CostKind, Opd1Info, Opd2Info,
+                                       Opd1PropInfo, Opd2PropInfo, Args, CxtI);
 }
 
 // Return true if there's a potential benefit from using v2f16 instructions for
 // an intrinsic, even if it requires nontrivial legalization.
 static bool intrinsicHasPackedVectorBenefit(Intrinsic::ID ID) {
   switch (ID) {
   case Intrinsic::fma: // TODO: fmuladd
   // There's a small benefit to using vector ops in the legalized code.

llvm/test/Analysis/CostModel/AMDGPU/fused_costs.ll

+; RUN: opt -cost-model -analyze -mtriple=amdgcn-unknown-amdhsa -mcpu=gfx900 -denormal-fp-math-f32=preserve-sign < %s | FileCheck -check-prefixes=FUSED,ALL %s
+; RUN: opt -cost-model -analyze -mtriple=amdgcn-unknown-amdhsa -mcpu=gfx900 -denormal-fp-math-f32=ieee < %s | FileCheck -check-prefixes=SLOW,ALL %s
+; RUN: opt -cost-model -cost-kind=code-size -analyze -mtriple=amdgcn-unknown-amdhsa -mcpu=gfx900 -denormal-fp-math-f32=preserve-sign < %s | FileCheck -check-prefixes=FUSED,ALL %s
+; RUN: opt -cost-model -cost-kind=code-size -analyze -mtriple=amdgcn-unknown-amdhsa -mcpu=gfx900 -denormal-fp-math-f32=ieee < %s | FileCheck -check-prefixes=SLOW,ALL %s
+
+target triple = "amdgcn--"
+
+; ALL-LABEL: 'fmul_fadd_f32':
+; FUSED: estimated cost of 0 for instruction:   %mul = fmul float
+; SLOW: estimated cost of 1 for instruction:   %mul = fmul float
+; ALL: estimated cost of 1 for instruction:   %add = fadd float
+define float @fmul_fadd_f32(float %r0, float %r1, float %r2) #0 {
+  %mul = fmul float %r0, %r1
+  %add = fadd float %mul, %r2
+  ret float %add
+}
+
+; ALL-LABEL: 'fmul_fadd_v2f32':
+; FUSED: estimated cost of 0 for instruction:   %mul = fmul <2 x float>
+; SLOW: estimated cost of 2 for instruction:   %mul = fmul <2 x float>
+; ALL: estimated cost of 2 for instruction:   %add = fadd <2 x float>
+define <2 x float> @fmul_fadd_v2f32(<2 x float> %r0, <2 x float> %r1, <2 x float> %r2) #0 {
+  %mul = fmul <2 x float> %r0, %r1
+  %add = fadd <2 x float> %mul, %r2
+  ret <2 x float> %add
+}
+
+; ALL-LABEL: 'fmul_fsub_f32':
+; FUSED: estimated cost of 0 for instruction:   %mul = fmul float
+; SLOW: estimated cost of 1 for instruction:   %mul = fmul float
+; ALL: estimated cost of 1 for instruction:   %sub = fsub float
+define float @fmul_fsub_f32(float %r0, float %r1, float %r2) #0 {
+  %mul = fmul float %r0, %r1
+  %sub = fsub float %mul, %r2
+  ret float %sub
+}
+
+; ALL-LABEL: 'fmul_fsub_v2f32':
+; FUSED: estimated cost of 0 for instruction:   %mul = fmul <2 x float>
+; SLOW: estimated cost of 2 for instruction:   %mul = fmul <2 x float>
+; ALL: estimated cost of 2 for instruction:   %sub = fsub <2 x float>
+define <2 x float> @fmul_fsub_v2f32(<2 x float> %r0, <2 x float> %r1, <2 x float> %r2) #0 {
+  %mul = fmul <2 x float> %r0, %r1
+  %sub = fsub <2 x float> %mul, %r2
+  ret <2 x float> %sub
+}
+
+attributes #0 = { nounwind }