AMDGPU: InstCombine amdgcn.rcp(amdgcn.sqrt) -> amdgcn.rsq

[arsenm]

Summary

We currently have some wrong combines in the backend that
approximately do this.

[arsenm]

7c4aa3b37ed1c8ab276d67eb9757e56a2f26790c

llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp

 }
 
 static Value *simplifyAMDGCNMemoryIntrinsicDemanded(InstCombiner &IC,
                                                     IntrinsicInst &II,
                                                     APInt DemandedElts,
                                                     int DMaskIdx = -1,
                                                     bool IsLoad = true);
 
+/// Return true if it's legal to contract llvm.amdgcn.rcp(llvm.sqrt)
+static bool canContractSqrtToRsq(const FPMathOperator *SqrtOp) {
+  return (SqrtOp->getType()->isFloatTy() &&
+          (SqrtOp->hasApproxFunc() || SqrtOp->getFPAccuracy() >= 1.0f)) ||
+         SqrtOp->getType()->isHalfTy();
+}
+
 std::optional<Instruction *>
 GCNTTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
   Intrinsic::ID IID = II.getIntrinsicID();
   switch (IID) {
   case Intrinsic::amdgcn_rcp: {
     Value *Src = II.getArgOperand(0);
 
     // TODO: Move to ConstantFolding/InstSimplify?
if (const ConstantFP *C = dyn_cast<ConstantFP>(Src)) {
 
       // This is more precise than the instruction may give.
       //
       // TODO: The instruction always flushes denormal results (except for f16),
       // should this also?
       return IC.replaceInstUsesWith(II, ConstantFP::get(II.getContext(), Val));
     }
 
+    FastMathFlags FMF = cast<FPMathOperator>(II).getFastMathFlags();
+    if (!FMF.allowContract())
+      break;
+    auto *SrcCI = dyn_cast<IntrinsicInst>(Src);
+    if (!SrcCI)
+      break;
+
+    auto IID = SrcCI->getIntrinsicID();
+    // llvm.amdgcn.rcp(llvm.amdgcn.sqrt(x)) -> llvm.amdgcn.rsq(x) if contractable
+    //
+    // llvm.amdgcn.rcp(llvm.sqrt(x)) -> llvm.amdgcn.rsq(x) if contractable and
+    // relaxed.
+    if (IID == Intrinsic::amdgcn_sqrt || IID == Intrinsic::sqrt) {
+      const FPMathOperator *SqrtOp = cast<FPMathOperator>(SrcCI);
+      FastMathFlags InnerFMF = SqrtOp->getFastMathFlags();
+      if (!InnerFMF.allowContract() || !SrcCI->hasOneUse())
+        break;
+
+      if (IID == Intrinsic::sqrt && !canContractSqrtToRsq(SqrtOp))
+        break;
+
+      Function *NewDecl = Intrinsic::getDeclaration(
+          SrcCI->getModule(), Intrinsic::amdgcn_rsq, {SrcCI->getType()});
+
+      InnerFMF |= FMF;
+      II.setFastMathFlags(InnerFMF);
+
+      II.setCalledFunction(NewDecl);
+      return IC.replaceOperand(II, 0, SrcCI->getArgOperand(0));
+    }
+
     break;
   }
   case Intrinsic::amdgcn_sqrt:
   case Intrinsic::amdgcn_rsq: {
     Value *Src = II.getArgOperand(0);
 
     // TODO: Move to ConstantFolding/InstSimplify?
     if (isa<UndefValue>(Src)) {

llvm/test/Transforms/InstCombine/AMDGPU/rcp-contract-rsq.ll

 declare half @llvm.sqrt.f16(half)
 declare float @llvm.sqrt.f32(float)
 declare double @llvm.sqrt.f64(double)
 
 ; Should allow precision increasing contraction to rsq call
 define float @amdgcn_rcp_amdgcn_sqrt_f32_contract(float %x) {
 ; CHECK-LABEL: define float @amdgcn_rcp_amdgcn_sqrt_f32_contract
 ; CHECK-SAME: (float [[X:%.*]]) #[[ATTR1:[0-9]+]] {
-; CHECK-NEXT:    [[SQRT:%.*]] = call contract float @llvm.amdgcn.sqrt.f32(float [[X]])
-; CHECK-NEXT:    [[RSQ:%.*]] = call contract float @llvm.amdgcn.rcp.f32(float [[SQRT]])
+; CHECK-NEXT:    [[RSQ:%.*]] = call contract float @llvm.amdgcn.rsq.f32(float [[X]])
 ; CHECK-NEXT:    ret float [[RSQ]]
 ;
   %sqrt = call contract float @llvm.amdgcn.sqrt.f32(float %x)
   %rsq = call contract float @llvm.amdgcn.rcp.f32(float %sqrt)
   ret float %rsq
 }
 
 ; contract required on both calls
;
   %rsq = call contract float @llvm.amdgcn.rcp.f32(float %sqrt)
   ret float %rsq
 }
 
 ; Test flags are or'd together
 define float @amdgcn_rcp_nnan_amdgcn_sqrt_ninf_f32_contract(float %x) {
 ; CHECK-LABEL: define float @amdgcn_rcp_nnan_amdgcn_sqrt_ninf_f32_contract
 ; CHECK-SAME: (float [[X:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[SQRT:%.*]] = call nnan contract float @llvm.amdgcn.sqrt.f32(float [[X]])
-; CHECK-NEXT:    [[RSQ:%.*]] = call ninf contract float @llvm.amdgcn.rcp.f32(float [[SQRT]])
+; CHECK-NEXT:    [[RSQ:%.*]] = call nnan ninf contract float @llvm.amdgcn.rsq.f32(float [[X]])
 ; CHECK-NEXT:    ret float [[RSQ]]
 ;
   %sqrt = call nnan contract float @llvm.amdgcn.sqrt.f32(float %x)
   %rsq = call ninf contract float @llvm.amdgcn.rcp.f32(float %sqrt)
   ret float %rsq
 }
 
 ; Should allow precision increasing contraction to rsq call
 define half @amdgcn_rcp_amdgcn_sqrt_f16_contract(half %x) {
 ; CHECK-LABEL: define half @amdgcn_rcp_amdgcn_sqrt_f16_contract
 ; CHECK-SAME: (half [[X:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[SQRT:%.*]] = call contract half @llvm.amdgcn.sqrt.f16(half [[X]])
-; CHECK-NEXT:    [[RSQ:%.*]] = call contract half @llvm.amdgcn.rcp.f16(half [[SQRT]])
+; CHECK-NEXT:    [[RSQ:%.*]] = call contract half @llvm.amdgcn.rsq.f16(half [[X]])
 ; CHECK-NEXT:    ret half [[RSQ]]
 ;
   %sqrt = call contract half @llvm.amdgcn.sqrt.f16(half %x)
   %rsq = call contract half @llvm.amdgcn.rcp.f16(half %sqrt)
   ret half %rsq
 }
 
 ; contract required on both calls
;
   %rsq = call contract half @llvm.amdgcn.rcp.f16(half %sqrt)
   ret half %rsq
 }
 
 ; Test flags are or'd together
 define half @amdgcn_rcp_nnan_amdgcn_sqrt_ninf_f16_contract(half %x) {
 ; CHECK-LABEL: define half @amdgcn_rcp_nnan_amdgcn_sqrt_ninf_f16_contract
 ; CHECK-SAME: (half [[X:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[SQRT:%.*]] = call nnan contract half @llvm.amdgcn.sqrt.f16(half [[X]])
-; CHECK-NEXT:    [[RSQ:%.*]] = call ninf contract half @llvm.amdgcn.rcp.f16(half [[SQRT]])
+; CHECK-NEXT:    [[RSQ:%.*]] = call nnan ninf contract half @llvm.amdgcn.rsq.f16(half [[X]])
 ; CHECK-NEXT:    ret half [[RSQ]]
 ;
   %sqrt = call nnan contract half @llvm.amdgcn.sqrt.f16(half %x)
   %rsq = call ninf contract half @llvm.amdgcn.rcp.f16(half %sqrt)
   ret half %rsq
 }
 
 ; Should allow precision increasing contraction to rsq call
 define double @amdgcn_rcp_amdgcn_sqrt_f64_contract(double %x) {
 ; CHECK-LABEL: define double @amdgcn_rcp_amdgcn_sqrt_f64_contract
 ; CHECK-SAME: (double [[X:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[SQRT:%.*]] = call contract double @llvm.amdgcn.sqrt.f64(double [[X]])
-; CHECK-NEXT:    [[RSQ:%.*]] = call contract double @llvm.amdgcn.rcp.f64(double [[SQRT]])
+; CHECK-NEXT:    [[RSQ:%.*]] = call contract double @llvm.amdgcn.rsq.f64(double [[X]])
 ; CHECK-NEXT:    ret double [[RSQ]]
 ;
   %sqrt = call contract double @llvm.amdgcn.sqrt.f64(double %x)
   %rsq = call contract double @llvm.amdgcn.rcp.f64(double %sqrt)
   ret double %rsq
 }
 
 ; contract required on both calls
;
   %rsq = call contract double @llvm.amdgcn.rcp.f64(double %sqrt)
   ret double %rsq
 }
 
 ; Test flags are or'd together
 define double @amdgcn_rcp_nnan_amdgcn_sqrt_ninf_f64_contract(double %x) {
 ; CHECK-LABEL: define double @amdgcn_rcp_nnan_amdgcn_sqrt_ninf_f64_contract
 ; CHECK-SAME: (double [[X:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[SQRT:%.*]] = call nnan contract double @llvm.amdgcn.sqrt.f64(double [[X]])
-; CHECK-NEXT:    [[RSQ:%.*]] = call ninf contract double @llvm.amdgcn.rcp.f64(double [[SQRT]])
+; CHECK-NEXT:    [[RSQ:%.*]] = call nnan ninf contract double @llvm.amdgcn.rsq.f64(double [[X]])
 ; CHECK-NEXT:    ret double [[RSQ]]
 ;
   %sqrt = call nnan contract double @llvm.amdgcn.sqrt.f64(double %x)
   %rsq = call ninf contract double @llvm.amdgcn.rcp.f64(double %sqrt)
   ret double %rsq
 }
 
 ; Do not contract with regular sqrt
 define float @amdgcn_rcp_sqrt_f32_contract(float %x) {
 ; CHECK-LABEL: define float @amdgcn_rcp_sqrt_f32_contract
 ; CHECK-SAME: (float [[X:%.*]]) #[[ATTR1]] {
 ; CHECK-NEXT:    [[SQRT:%.*]] = call contract float @llvm.sqrt.f32(float [[X]])
 ; CHECK-NEXT:    [[RSQ:%.*]] = call contract float @llvm.amdgcn.rcp.f32(float [[SQRT]])
 ; CHECK-NEXT:    ret float [[RSQ]]
 ;
   %sqrt = call contract float @llvm.sqrt.f32(float %x)
   %rsq = call contract float @llvm.amdgcn.rcp.f32(float %sqrt)
   ret float %rsq
 }
 
 ; Do contract with regular sqrt for f16
 define half @amdgcn_rcp_sqrt_f16_contract(half %x) {
 ; CHECK-LABEL: define half @amdgcn_rcp_sqrt_f16_contract
 ; CHECK-SAME: (half [[X:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[SQRT:%.*]] = call contract half @llvm.sqrt.f16(half [[X]])
-; CHECK-NEXT:    [[RSQ:%.*]] = call contract half @llvm.amdgcn.rcp.f16(half [[SQRT]])
+; CHECK-NEXT:    [[RSQ:%.*]] = call contract half @llvm.amdgcn.rsq.f16(half [[X]])
 ; CHECK-NEXT:    ret half [[RSQ]]
 ;
   %sqrt = call contract half @llvm.sqrt.f16(half %x)
   %rsq = call contract half @llvm.amdgcn.rcp.f16(half %sqrt)
   ret half %rsq
 }
 
 ; Do not contract with regular sqrt
 define double @amdgcn_rcp_sqrt_f64_contract(double %x) {
 ; CHECK-LABEL: define double @amdgcn_rcp_sqrt_f64_contract
 ; CHECK-SAME: (double [[X:%.*]]) #[[ATTR1]] {
 ; CHECK-NEXT:    [[SQRT:%.*]] = call contract double @llvm.sqrt.f64(double [[X]])
 ; CHECK-NEXT:    [[RSQ:%.*]] = call contract double @llvm.amdgcn.rcp.f64(double [[SQRT]])
 ; CHECK-NEXT:    ret double [[RSQ]]
 ;
   %sqrt = call contract double @llvm.sqrt.f64(double %x)
   %rsq = call contract double @llvm.amdgcn.rcp.f64(double %sqrt)
   ret double %rsq
 }
 
 define float @amdgcn_rcp_afn_sqrt_f32_contract(float %x) {
 ; CHECK-LABEL: define float @amdgcn_rcp_afn_sqrt_f32_contract
 ; CHECK-SAME: (float [[X:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[SQRT:%.*]] = call contract afn float @llvm.sqrt.f32(float [[X]])
-; CHECK-NEXT:    [[RSQ:%.*]] = call contract float @llvm.amdgcn.rcp.f32(float [[SQRT]])
+; CHECK-NEXT:    [[RSQ:%.*]] = call contract afn float @llvm.amdgcn.rsq.f32(float [[X]])
 ; CHECK-NEXT:    ret float [[RSQ]]
 ;
   %sqrt = call afn contract float @llvm.sqrt.f32(float %x)
   %rsq = call contract float @llvm.amdgcn.rcp.f32(float %sqrt)
   ret float %rsq
 }
 
 define float @amdgcn_rcp_fpmath3_sqrt_f32_contract(float %x) {
 ; CHECK-LABEL: define float @amdgcn_rcp_fpmath3_sqrt_f32_contract
 ; CHECK-SAME: (float [[X:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[SQRT:%.*]] = call contract float @llvm.sqrt.f32(float [[X]]), !fpmath !0
-; CHECK-NEXT:    [[RSQ:%.*]] = call contract float @llvm.amdgcn.rcp.f32(float [[SQRT]])
+; CHECK-NEXT:    [[RSQ:%.*]] = call contract float @llvm.amdgcn.rsq.f32(float [[X]])
 ; CHECK-NEXT:    ret float [[RSQ]]
 ;
   %sqrt = call contract float @llvm.sqrt.f32(float %x), !fpmath !0
   %rsq = call contract float @llvm.amdgcn.rcp.f32(float %sqrt)
   ret float %rsq
 }
 
 define float @amdgcn_rcp_fpmath1_sqrt_f32_contract(float %x) {
 ; CHECK-LABEL: define float @amdgcn_rcp_fpmath1_sqrt_f32_contract
 ; CHECK-SAME: (float [[X:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[SQRT:%.*]] = call contract float @llvm.sqrt.f32(float [[X]]), !fpmath !1
-; CHECK-NEXT:    [[RSQ:%.*]] = call contract float @llvm.amdgcn.rcp.f32(float [[SQRT]])
+; CHECK-NEXT:    [[RSQ:%.*]] = call contract float @llvm.amdgcn.rsq.f32(float [[X]])
 ; CHECK-NEXT:    ret float [[RSQ]]
 ;
   %sqrt = call contract float @llvm.sqrt.f32(float %x), !fpmath !1
   %rsq = call contract float @llvm.amdgcn.rcp.f32(float %sqrt)
   ret float %rsq
 }
 
 ; Ignore f64 sqrt