AMDGPU: Reduce number of calls to computeKnownFPClass and pass all arguments

[arsenm]

Summary

Makes assumes work for this case.

[foad]

LGTM with comment fixed.

llvm/include/llvm/Analysis/ValueTracking.h

378

"never contain"

[arsenm]

1bce1beac4b7980ed01c07d6a6d789f8ec55adac

llvm/include/llvm/Analysis/ValueTracking.h

inline bool isKnownNeverInfinity(const Value *V, const DataLayout &DL,
                                  const DominatorTree *DT = nullptr,
                                  OptimizationRemarkEmitter *ORE = nullptr,
                                  bool UseInstrInfo = true) {
   KnownFPClass Known = computeKnownFPClass(V, DL, fcInf, Depth, TLI, AC, CtxI,
                                            DT, ORE, UseInstrInfo);
   return Known.isKnownNeverInfinity();
 }
 
+/// Return true if the floating-point value can neve rcontain a NaN or infinity.
+inline bool isKnownNeverInfOrNaN(
+    const Value *V, const DataLayout &DL, const TargetLibraryInfo *TLI,
+    unsigned Depth = 0, AssumptionCache *AC = nullptr,
+    const Instruction *CtxI = nullptr, const DominatorTree *DT = nullptr,
+    OptimizationRemarkEmitter *ORE = nullptr, bool UseInstrInfo = true) {
+  KnownFPClass Known = computeKnownFPClass(V, DL, fcInf | fcNan, Depth, TLI, AC,
+                                           CtxI, DT, ORE, UseInstrInfo);
+  return Known.isKnownNeverNaN() && Known.isKnownNeverInfinity();
+}
+
 /// Return true if the floating-point scalar value is not a NaN or if the
 /// floating-point vector value has no NaN elements. Return false if a value
 /// could ever be NaN.
 inline bool isKnownNeverNaN(const Value *V, const DataLayout &DL,
                             const TargetLibraryInfo *TLI, unsigned Depth = 0,
                             AssumptionCache *AC = nullptr,
                             const Instruction *CtxI = nullptr,
                             const DominatorTree *DT = nullptr,

llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp

return modifyIntrinsicCall(
         // Convert the bias
         if (!OnlyDerivatives && ImageDimIntr->NumBiasArgs != 0) {
           Value *Bias = II.getOperand(ImageDimIntr->BiasIndex);
           Args[ImageDimIntr->BiasIndex] = convertTo16Bit(*Bias, IC.Builder);
         }
       });
 }
 
-bool GCNTTIImpl::canSimplifyLegacyMulToMul(const Value *Op0, const Value *Op1,
+bool GCNTTIImpl::canSimplifyLegacyMulToMul(const Instruction &I,
+                                           const Value *Op0, const Value *Op1,
                                            InstCombiner &IC) const {
   // The legacy behaviour is that multiplying +/-0.0 by anything, even NaN or
   // infinity, gives +0.0. If we can prove we don't have one of the special
   // cases then we can use a normal multiply instead.
   // TODO: Create and use isKnownFiniteNonZero instead of just matching
   // constants here.
   if (match(Op0, PatternMatch::m_FiniteNonZero()) ||
       match(Op1, PatternMatch::m_FiniteNonZero())) {
     // One operand is not zero or infinity or NaN.
     return true;
   }
 
-  // TODO: Use computeKnownFPClass
   auto *TLI = &IC.getTargetLibraryInfo();
-  if (isKnownNeverInfinity(Op0, IC.getDataLayout(), TLI) &&
-      isKnownNeverNaN(Op0, IC.getDataLayout(), TLI) &&
-      isKnownNeverInfinity(Op1, IC.getDataLayout(), TLI) &&
-      isKnownNeverNaN(Op1, IC.getDataLayout(), TLI)) {
+  if (isKnownNeverInfOrNaN(Op0, IC.getDataLayout(), TLI, 0,
+                           &IC.getAssumptionCache(), &I, &IC.getDominatorTree(),
+                           &IC.getOptimizationRemarkEmitter()) &&
+      isKnownNeverInfOrNaN(Op1, IC.getDataLayout(), TLI, 0,
+                           &IC.getAssumptionCache(), &I, &IC.getDominatorTree(),
+                           &IC.getOptimizationRemarkEmitter())) {
     // Neither operand is infinity or NaN.
     return true;
   }
   return false;
 }
 
 std::optional<Instruction *>
 GCNTTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
case Intrinsic::amdgcn_fmul_legacy: {
     // infinity, gives +0.0.
     // TODO: Move to InstSimplify?
     if (match(Op0, PatternMatch::m_AnyZeroFP()) ||
         match(Op1, PatternMatch::m_AnyZeroFP()))
       return IC.replaceInstUsesWith(II, ConstantFP::getZero(II.getType()));
 
     // If we can prove we don't have one of the special cases then we can use a
     // normal fmul instruction instead.
-    if (canSimplifyLegacyMulToMul(Op0, Op1, IC)) {
+    if (canSimplifyLegacyMulToMul(II, Op0, Op1, IC)) {
       auto *FMul = IC.Builder.CreateFMulFMF(Op0, Op1, &II);
       FMul->takeName(&II);
       return IC.replaceInstUsesWith(II, FMul);
     }
     break;
   }
   case Intrinsic::amdgcn_fma_legacy: {
     Value *Op0 = II.getArgOperand(0);
if (match(Op0, PatternMatch::m_AnyZeroFP()) ||
       auto *Zero = ConstantFP::getZero(II.getType());
       auto *FAdd = IC.Builder.CreateFAddFMF(Zero, Op2, &II);
       FAdd->takeName(&II);
       return IC.replaceInstUsesWith(II, FAdd);
     }
 
     // If we can prove we don't have one of the special cases then we can use a
     // normal fma instead.
-    if (canSimplifyLegacyMulToMul(Op0, Op1, IC)) {
+    if (canSimplifyLegacyMulToMul(II, Op0, Op1, IC)) {
       II.setCalledOperand(Intrinsic::getDeclaration(
           II.getModule(), Intrinsic::fma, II.getType()));
       return &II;
     }
     break;
   }
   case Intrinsic::amdgcn_is_shared:
   case Intrinsic::amdgcn_is_private: {

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.h

public:
   bool canHaveNonUndefGlobalInitializerInAddressSpace(unsigned AS) const {
     return AS != AMDGPUAS::LOCAL_ADDRESS && AS != AMDGPUAS::REGION_ADDRESS &&
            AS != AMDGPUAS::PRIVATE_ADDRESS;
   }
 
   Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II, Value *OldV,
                                           Value *NewV) const;
 
-  bool canSimplifyLegacyMulToMul(const Value *Op0, const Value *Op1,
-                                 InstCombiner &IC) const;
+  bool canSimplifyLegacyMulToMul(const Instruction &I, const Value *Op0,
+                                 const Value *Op1, InstCombiner &IC) const;
   std::optional<Instruction *> instCombineIntrinsic(InstCombiner &IC,
                                                     IntrinsicInst &II) const;
   std::optional<Value *> simplifyDemandedVectorEltsIntrinsic(
       InstCombiner &IC, IntrinsicInst &II, APInt DemandedElts, APInt &UndefElts,
       APInt &UndefElts2, APInt &UndefElts3,
       std::function<void(Instruction *, unsigned, APInt, APInt &)>
           SimplifyAndSetOp) const;
 

llvm/test/Transforms/InstCombine/AMDGPU/fma_legacy.ll

 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %xf = sitofp i32 %x to float
   %yf = sitofp i32 %y to float
   %call = call float @llvm.amdgcn.fma.legacy(float %xf, float %yf, float %z)
   ret float %call
 }
 
+; Combine to fma because neither argument can be infinity or NaN based on assumptions
+define float @test_finite_assumed(float %x, float %y, float %z) {
+; CHECK-LABEL: @test_finite_assumed(
+; CHECK-NEXT:    [[FABS_X:%.*]] = call float @llvm.fabs.f32(float [[X:%.*]])
+; CHECK-NEXT:    [[IS_FINITE_X:%.*]] = fcmp one float [[FABS_X]], 0x7FF0000000000000
+; CHECK-NEXT:    [[FABS_Y:%.*]] = call float @llvm.fabs.f32(float [[Y:%.*]])
+; CHECK-NEXT:    [[IS_FINITE_Y:%.*]] = fcmp one float [[FABS_Y]], 0x7FF0000000000000
+; CHECK-NEXT:    call void @llvm.assume(i1 [[IS_FINITE_X]])
+; CHECK-NEXT:    call void @llvm.assume(i1 [[IS_FINITE_Y]])
+; CHECK-NEXT:    [[CALL:%.*]] = call float @llvm.fma.f32(float [[X]], float [[Y]], float [[Z:%.*]])
+; CHECK-NEXT:    ret float [[CALL]]
+;
+  %fabs.x = call float @llvm.fabs.f32(float %x)
+  %is.finite.x = fcmp one float %fabs.x, 0x7FF0000000000000
+  %fabs.y = call float @llvm.fabs.f32(float %y)
+  %is.finite.y = fcmp one float %fabs.y, 0x7FF0000000000000
+  call void @llvm.assume(i1 %is.finite.x)
+  call void @llvm.assume(i1 %is.finite.y)
+  %call = call float @llvm.amdgcn.fma.legacy(float %x, float %y, float %z)
+  ret float %call
+}
+
 declare float @llvm.amdgcn.fma.legacy(float, float, float)
+declare float @llvm.fabs.f32(float)
+declare void @llvm.assume(i1 noundef)

llvm/test/Transforms/InstCombine/AMDGPU/fmul_legacy.ll

 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %xf = sitofp i32 %x to float
   %yf = sitofp i32 %y to float
   %call = call float @llvm.amdgcn.fmul.legacy(float %xf, float %yf)
   ret float %call
 }
 
+; Combine to fmul because neither argument can be infinity or NaN based on assumptions
+define float @test_finite_assumed(float %x, float %y) {
+; CHECK-LABEL: @test_finite_assumed(
+; CHECK-NEXT:    [[FABS_X:%.*]] = call float @llvm.fabs.f32(float [[X:%.*]])
+; CHECK-NEXT:    [[IS_FINITE_X:%.*]] = fcmp one float [[FABS_X]], 0x7FF0000000000000
+; CHECK-NEXT:    [[FABS_Y:%.*]] = call float @llvm.fabs.f32(float [[Y:%.*]])
+; CHECK-NEXT:    [[IS_FINITE_Y:%.*]] = fcmp one float [[FABS_Y]], 0x7FF0000000000000
+; CHECK-NEXT:    call void @llvm.assume(i1 [[IS_FINITE_X]])
+; CHECK-NEXT:    call void @llvm.assume(i1 [[IS_FINITE_Y]])
+; CHECK-NEXT:    [[CALL:%.*]] = fmul float [[X]], [[Y]]
+; CHECK-NEXT:    ret float [[CALL]]
+;
+  %fabs.x = call float @llvm.fabs.f32(float %x)
+  %is.finite.x = fcmp one float %fabs.x, 0x7FF0000000000000
+  %fabs.y = call float @llvm.fabs.f32(float %y)
+  %is.finite.y = fcmp one float %fabs.y, 0x7FF0000000000000
+  call void @llvm.assume(i1 %is.finite.x)
+  call void @llvm.assume(i1 %is.finite.y)
+  %call = call float @llvm.amdgcn.fmul.legacy(float %x, float %y)
+  ret float %call
+}
+
 declare float @llvm.amdgcn.fmul.legacy(float, float)
+declare float @llvm.fabs.f32(float)
+declare void @llvm.assume(i1 noundef)