Attributor: Assume handling for nofpclass

[sameerds]

llvm/include/llvm/Analysis/ValueTracking.h

230

Would "LookThroughAbsForLHS" be a more specific name for this argument?

[arsenm]

llvm/include/llvm/Analysis/ValueTracking.h

230

The point is to exactly report the same incoming value, and fabs is more specific about what it's doing. This could perform other tricks by looking at fneg and fcopysign in the future

[sameerds]

I am not familiar enough to assess whether this covers everything. But as far as I can see, the implementation does what is promised. Might want to wait for comments from other reviewers.

[jdoerfert]

LG, two nits.

llvm/lib/Analysis/ValueTracking.cpp

4285

Copy-paste.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

10309

We don't need two places to followUses

[arsenm]

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

10309

This one got dropped since it didn't do anything in the previous patch. Which one would be better?

[arsenm]

d2404ea6ced5fce9442260bde08a02d607fdd50d

llvm/include/llvm/Analysis/ValueTracking.h

 
 /// Map a call instruction to an intrinsic ID.  Libcalls which have equivalent
 /// intrinsics are treated as-if they were intrinsics.
 Intrinsic::ID getIntrinsicForCallSite(const CallBase &CB,
                                       const TargetLibraryInfo *TLI);
 
 /// Returns a pair of values, which if passed to llvm.is.fpclass, returns the
 /// same result as an fcmp with the given operands.
+///
+/// If \p LookThroughSrc is true, consider the input value when computing the
+/// mask.
+///
+/// If \p LookThroughSrc is false, ignore the source value (i.e. the first pair
+/// element will always be LHS.
 std::pair<Value *, FPClassTest> fcmpToClassTest(CmpInst::Predicate Pred,
-                                                Value *LHS, Value *RHS);
+                                                Value *LHS, Value *RHS,
+                                                bool LookThroughSrc = true);
 
 struct KnownFPClass {
   /// Floating-point classes the value could be one of.
   FPClassTest KnownFPClasses = fcAllFlags;
 
   /// std::nullopt if the sign bit is unknown, true if the sign bit is
   /// definitely set or false if the sign bit is definitely unset.
   std::optional<bool> SignBit;

llvm/lib/Analysis/AssumptionCache.cpp

if (CI->getOperandBundleAt(Idx).Inputs.size() > ABA_WasOn &&
         CI->getOperandBundleAt(Idx).getTagName() != IgnoreBundleTag)
       AddAffected(CI->getOperandBundleAt(Idx).Inputs[ABA_WasOn], Idx);
   }
 
   Value *Cond = CI->getArgOperand(0), *A, *B;
   AddAffected(Cond);
 
   CmpInst::Predicate Pred;
-  if (match(Cond, m_ICmp(Pred, m_Value(A), m_Value(B)))) {
+  if (match(Cond, m_Cmp(Pred, m_Value(A), m_Value(B)))) {
     AddAffected(A);
     AddAffected(B);
 
     if (Pred == ICmpInst::ICMP_EQ) {
       // For equality comparisons, we handle the case of bit inversion.
       auto AddAffectedFromEq = [&AddAffected](Value *V) {
         Value *A;
         if (match(V, m_Not(m_Value(A)))) {
if (Pred == ICmpInst::ICMP_EQ) {
       }
     } else if (Pred == ICmpInst::ICMP_ULT) {
       Value *X;
       // Handle (A + C1) u< C2, which is the canonical form of A > C3 && A < C4,
       // and recognized by LVI at least.
       if (match(A, m_Add(m_Value(X), m_ConstantInt())) &&
           match(B, m_ConstantInt()))
         AddAffected(X);
+    } else if (CmpInst::isFPPredicate(Pred)) {
+      // fcmp fneg(x), y
+      // fcmp fabs(x), y
+      // fcmp fneg(fabs(x)), y
+      if (match(A, m_FNeg(m_Value(A))))
+        AddAffected(A);
+      if (match(A, m_FAbs(m_Value(A))))
+        AddAffected(A);
     }
+  } else if (match(Cond, m_Intrinsic<Intrinsic::is_fpclass>(m_Value(A),
+                                                            m_Value(B)))) {
+    AddAffected(A);
   }
 
   if (TTI) {
     const Value *Ptr;
     unsigned AS;
     std::tie(Ptr, AS) = TTI->getPredicatedAddrSpace(Cond);
     if (Ptr)
       AddAffected(const_cast<Value *>(Ptr->stripInBoundsOffsets()));

llvm/lib/Analysis/ValueTracking.cpp

bool llvm::isKnownNeverNaN(const Value *V, const TargetLibraryInfo *TLI,
 
   // Was not able to prove that V never contains NaN
   return false;
 }
 
 /// Returns a pair of values, which if passed to llvm.is.fpclass, returns the
 /// same result as an fcmp with the given operands.
 std::pair<Value *, FPClassTest> llvm::fcmpToClassTest(FCmpInst::Predicate Pred,
-                                                      Value *LHS, Value *RHS) {
+                                                      Value *LHS, Value *RHS,
+                                                      bool LookThroughSrc) {
   const APFloat *ConstRHS;
   if (!match(RHS, m_APFloat(ConstRHS)))
     return {nullptr, fcNone};
 
   if (ConstRHS->isZero()) {
     switch (Pred) {
     // TODO: Compares eq/ne with 0 depends on the denormal handling mode.
     case FCmpInst::FCMP_ORD:
       // Canonical form of ord/uno is with a zero. We could also handle
       // non-canonical other non-NaN constants or LHS == RHS.
       return {LHS, ~fcNan};
     case FCmpInst::FCMP_UNO:
       return {LHS, fcNan};
     default:
       break;
     }
 
     return {nullptr, fcNone};
   }
 
   Value *Src = LHS;
-  const bool IsFabs = match(LHS, m_FAbs(m_Value(Src)));
+  const bool IsFabs = LookThroughSrc && match(LHS, m_FAbs(m_Value(Src)));
 
   // Compute the test mask that would return true for the ordered comparisons.
   FPClassTest Mask;
 
   if (ConstRHS->isInfinity()) {
     switch (Pred) {
     case FCmpInst::FCMP_OEQ:
     case FCmpInst::FCMP_UNE: {
std::pair<Value *, FPClassTest> llvm::fcmpToClassTest(FCmpInst::Predicate Pred,
 
   // Invert the comparison for the unordered cases.
   if (FCmpInst::isUnordered(Pred))
     Mask = ~Mask;
 
   return {Src, Mask};
 }
 
+static FPClassTest computeKnownFPClassFromAssumes(
+  const Value *V, const Query &Q) {
+  FPClassTest KnownFromAssume = fcNone;
+
+  // Try to restrict the range based on information from assumptions.
+  for (auto &AssumeVH : Q.AC->assumptionsFor(V)) {
+    if (!AssumeVH)
+      continue;
+    CallInst *I = cast<CallInst>(AssumeVH);
+    assert(I->getParent()->getParent() == Q.CxtI->getParent()->getParent() &&
+           "Got assumption for the wrong function!");
+    assert(I->getCalledFunction()->getIntrinsicID() == Intrinsic::assume &&
+           "must be an assume intrinsic");
+
+    if (!isValidAssumeForContext(I, Q.CxtI, Q.DT))
+      continue;
+
+    CmpInst::Predicate Pred;
+    Value *LHS, *RHS;
+    uint64_t ClassVal = 0;
+    if (match(I->getArgOperand(0), m_FCmp(Pred, m_Value(LHS), m_Value(RHS)))) {
+      auto [TestedValue, TestedMask] = fcmpToClassTest(Pred, LHS, RHS, true);
+      // First see if we can fold in fabs/fneg into the test.
+      if (TestedValue == V)
+        KnownFromAssume |= TestedMask;
+      else {
+        // Try again without the lookthrough if we found a different source
+        // value.
+        auto [TestedValue, TestedMask] = fcmpToClassTest(Pred, LHS, RHS, false);
+        if (TestedValue == V)
+          KnownFromAssume |= TestedMask;
+      }
+    } else if (match(I->getArgOperand(0),
+                     m_Intrinsic<Intrinsic::is_fpclass>(
+                         m_Value(LHS), m_ConstantInt(ClassVal)))) {
+      KnownFromAssume |= static_cast<FPClassTest>(ClassVal);
+    }
+  }
+
+  return KnownFromAssume;
+}
+
 // TODO: Merge implementations of isKnownNeverNaN, isKnownNeverInfinity,
 // CannotBeNegativeZero, cannotBeOrderedLessThanZero into here.
 void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
                          FPClassTest InterestedClasses, KnownFPClass &Known,
                          unsigned Depth, const Query &Q,
                          const TargetLibraryInfo *TLI) {
   if (!DemandedElts) {
     // No demanded elts, better to assume we don't know anything.
void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
   const Operator *Op = dyn_cast<Operator>(V);
   if (const FPMathOperator *FPOp = dyn_cast_or_null<FPMathOperator>(Op)) {
     if (FPOp->hasNoNaNs())
       KnownNotFromFlags |= fcNan;
     if (FPOp->hasNoInfs())
       KnownNotFromFlags |= fcInf;
   }
 
+  if (Q.AC)
+    KnownNotFromFlags |= computeKnownFPClassFromAssumes(V, Q);
+
   // We no longer need to find out about these bits from inputs if we can
   // assume this from flags/attributes.
   InterestedClasses &= ~KnownNotFromFlags & fcAllFlags;
 
   auto ClearClassesFromFlags = make_scope_exit([=, &Known] {
     Known.knownNot(KnownNotFromFlags);
   });
 

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

if (!Attrs.empty()) {
       return;
     }
 
     const DataLayout &DL = A.getDataLayout();
     if (getPositionKind() != IRPosition::IRP_RETURNED) {
       KnownFPClass KnownFPClass = computeKnownFPClass(&V, DL);
       addKnownBits(~KnownFPClass.KnownFPClasses);
     }
+
+    if (Instruction *CtxI = getCtxI())
+      followUsesInMBEC(*this, A, getState(), *CtxI);
   }
 
   /// See followUsesInMBEC
   bool followUseInMBEC(Attributor &A, const Use *U, const Instruction *I,
                        AANoFPClass::StateType &State) {
     const Value *UseV = U->get();
     const DominatorTree *DT = nullptr;
     AssumptionCache *AC = nullptr;
struct AANoFPClassFloating : public AANoFPClassImpl {
   AANoFPClassFloating(const IRPosition &IRP, Attributor &A)
       : AANoFPClassImpl(IRP, A) {}
 
   /// See AbstractAttribute::initialize(...).
   void initialize(Attributor &A) override {
     AANoFPClassImpl::initialize(A);
     if (!getState().isAtFixpoint())
       if (Instruction *CtxI = getCtxI())
         followUsesInMBEC(*this, A, getState(), *CtxI);
   }
 
   /// See AbstractAttribute::updateImpl(...).
   ChangeStatus updateImpl(Attributor &A) override {
     SmallVector<AA::ValueAndContext> Values;
     bool UsedAssumedInformation = false;
     if (!A.getAssumedSimplifiedValues(getIRPosition(), *this, Values,
                                       AA::AnyScope, UsedAssumedInformation)) {

llvm/test/Transforms/Attributor/nofpclass.ll

bb0:
   br label %bb1
 
 bb1:
   ret void
 }
 
 ; Should be able to infer nofpclass on both %arg uses
 define float @fcmp_ord_assume_callsite_arg_return(float %arg) {
-; CHECK-LABEL: define float @fcmp_ord_assume_callsite_arg_return
-; CHECK-SAME: (float returned [[ARG:%.*]]) {
+; CHECK-LABEL: define nofpclass(inf zero sub norm) float @fcmp_ord_assume_callsite_arg_return
+; CHECK-SAME: (float returned nofpclass(inf zero sub norm) [[ARG:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[IS_NOT_NAN:%.*]] = fcmp ord float [[ARG]], 0.000000e+00
 ; CHECK-NEXT:    call void @llvm.assume(i1 noundef [[IS_NOT_NAN]]) #[[ATTR4:[0-9]+]]
-; CHECK-NEXT:    call void @extern.use(float [[ARG]])
+; CHECK-NEXT:    call void @extern.use(float nofpclass(inf zero sub norm) [[ARG]])
 ; CHECK-NEXT:    ret float [[ARG]]
 ;
 entry:
   %is.not.nan = fcmp ord float %arg, 0.0
   call void @llvm.assume(i1 %is.not.nan)
   call void @extern.use(float %arg)
   ret float %arg
 }

 entry:
   %is.not.zero.or.nan = fcmp une float %arg, 0.0
   call void @llvm.assume(i1 %is.not.zero.or.nan)
   call void @extern.use(float %arg)
   ret float %arg
 }
 
 define half @fcmp_assume_issubnormal_callsite_arg_return(half %arg) {
-; CHECK-LABEL: define half @fcmp_assume_issubnormal_callsite_arg_return
-; CHECK-SAME: (half returned [[ARG:%.*]]) {
+; CHECK-LABEL: define nofpclass(zero sub) half @fcmp_assume_issubnormal_callsite_arg_return
+; CHECK-SAME: (half returned nofpclass(zero sub) [[ARG:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[FABS:%.*]] = call half @llvm.fabs.f16(half [[ARG]]) #[[ATTR4]]
+; CHECK-NEXT:    [[FABS:%.*]] = call half @llvm.fabs.f16(half nofpclass(zero sub) [[ARG]]) #[[ATTR4]]
 ; CHECK-NEXT:    [[IS_SUBNORMAL:%.*]] = fcmp olt half [[FABS]], 0xH0400
 ; CHECK-NEXT:    call void @llvm.assume(i1 noundef [[IS_SUBNORMAL]]) #[[ATTR4]]
-; CHECK-NEXT:    call void @extern.use.f16(half [[ARG]])
+; CHECK-NEXT:    call void @extern.use.f16(half nofpclass(zero sub) [[ARG]])
 ; CHECK-NEXT:    ret half [[ARG]]
 ;
 entry:
   %fabs = call half @llvm.fabs.f16(half %arg)
   %is.subnormal = fcmp olt half %fabs, 0xH0400
   call void @llvm.assume(i1 %is.subnormal)
   call void @extern.use.f16(half %arg)
   ret half %arg
entry:
   call void @extern.use.f16(half %arg)
   %not.inf = fcmp oeq half %arg, 0xH7C00
   call void @llvm.assume(i1 %not.inf)
   ret half %arg
 }
 
 ; Assume not subnormal or zero, and not infinity
 define half @fcmp_assume2_callsite_arg_return(half %arg) {
-; CHECK-LABEL: define half @fcmp_assume2_callsite_arg_return
-; CHECK-SAME: (half returned [[ARG:%.*]]) {
+; CHECK-LABEL: define nofpclass(inf norm) half @fcmp_assume2_callsite_arg_return
+; CHECK-SAME: (half returned nofpclass(inf norm) [[ARG:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[FABS:%.*]] = call half @llvm.fabs.f16(half [[ARG]]) #[[ATTR4]]
+; CHECK-NEXT:    [[FABS:%.*]] = call nofpclass(inf) half @llvm.fabs.f16(half nofpclass(inf norm) [[ARG]]) #[[ATTR4]]
 ; CHECK-NEXT:    [[NOT_SUBNORMAL_OR_ZERO:%.*]] = fcmp oge half [[FABS]], 0xH0400
 ; CHECK-NEXT:    call void @llvm.assume(i1 noundef [[NOT_SUBNORMAL_OR_ZERO]]) #[[ATTR4]]
 ; CHECK-NEXT:    [[NOT_INF:%.*]] = fcmp oeq half [[ARG]], 0xH7C00
 ; CHECK-NEXT:    call void @llvm.assume(i1 noundef [[NOT_INF]]) #[[ATTR4]]
-; CHECK-NEXT:    call void @extern.use.f16(half [[ARG]])
+; CHECK-NEXT:    call void @extern.use.f16(half nofpclass(inf norm) [[ARG]])
 ; CHECK-NEXT:    ret half [[ARG]]
 ;
 entry:
   %fabs = call half @llvm.fabs.f16(half %arg)
   %not.subnormal.or.zero = fcmp oge half %fabs, 0xH0400
   call void @llvm.assume(i1 %not.subnormal.or.zero)
 
   %not.inf = fcmp oeq half %arg, 0xH7C00
   call void @llvm.assume(i1 %not.inf)
 
   call void @extern.use.f16(half %arg)
   ret half %arg
 }
 
 define float @is_fpclass_assume_arg_return(float %arg) {
-; CHECK-LABEL: define float @is_fpclass_assume_arg_return
-; CHECK-SAME: (float returned [[ARG:%.*]]) {
+; CHECK-LABEL: define nofpclass(ninf nzero pnorm) float @is_fpclass_assume_arg_return
+; CHECK-SAME: (float returned nofpclass(ninf nzero pnorm) [[ARG:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CLASS_TEST:%.*]] = call i1 @llvm.is.fpclass.f32(float [[ARG]], i32 noundef 292) #[[ATTR4]]
+; CHECK-NEXT:    [[CLASS_TEST:%.*]] = call i1 @llvm.is.fpclass.f32(float nofpclass(ninf nzero pnorm) [[ARG]], i32 noundef 292) #[[ATTR4]]
 ; CHECK-NEXT:    call void @llvm.assume(i1 noundef [[CLASS_TEST]]) #[[ATTR4]]
-; CHECK-NEXT:    call void @extern.use(float [[ARG]])
+; CHECK-NEXT:    call void @extern.use(float nofpclass(ninf nzero pnorm) [[ARG]])
 ; CHECK-NEXT:    ret float [[ARG]]
 ;
 entry:
   %class.test = call i1 @llvm.is.fpclass.f32(float %arg, i32 292)
   call void @llvm.assume(i1 %class.test)
   call void @extern.use(float %arg)
   ret float %arg
 }
 
 ; Make sure we don't get confused by looking at an unrelated assume
 ; based on the fabs of the value.
 define half @assume_fcmp_fabs_with_other_fabs_assume(half %arg) {
-; CHECK-LABEL: define half @assume_fcmp_fabs_with_other_fabs_assume
-; CHECK-SAME: (half returned [[ARG:%.*]]) {
+; CHECK-LABEL: define nofpclass(zero sub) half @assume_fcmp_fabs_with_other_fabs_assume
+; CHECK-SAME: (half returned nofpclass(zero sub) [[ARG:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[FABS:%.*]] = call half @llvm.fabs.f16(half [[ARG]]) #[[ATTR4]]
+; CHECK-NEXT:    [[FABS:%.*]] = call nofpclass(ninf zero sub nnorm) half @llvm.fabs.f16(half nofpclass(zero sub) [[ARG]]) #[[ATTR4]]
 ; CHECK-NEXT:    [[UNRELATED_FABS:%.*]] = fcmp one half [[FABS]], 0xH0000
 ; CHECK-NEXT:    call void @llvm.assume(i1 noundef [[UNRELATED_FABS]]) #[[ATTR4]]
 ; CHECK-NEXT:    [[IS_SUBNORMAL:%.*]] = fcmp olt half [[FABS]], 0xH0400
 ; CHECK-NEXT:    call void @llvm.assume(i1 noundef [[IS_SUBNORMAL]]) #[[ATTR4]]
-; CHECK-NEXT:    call void @extern.use.f16(half [[ARG]])
-; CHECK-NEXT:    call void @extern.use.f16(half [[FABS]])
+; CHECK-NEXT:    call void @extern.use.f16(half nofpclass(zero sub) [[ARG]])
+; CHECK-NEXT:    call void @extern.use.f16(half nofpclass(ninf zero sub nnorm) [[FABS]])
 ; CHECK-NEXT:    ret half [[ARG]]
 ;
 entry:
 
   %fabs = call half @llvm.fabs.f16(half %arg)
   %unrelated.fabs = fcmp one half %fabs, 0.0
   call void @llvm.assume(i1 %unrelated.fabs)
   %is.subnormal = fcmp olt half %fabs, 0xH0400
   call void @llvm.assume(i1 %is.subnormal)
   call void @extern.use.f16(half %arg)
   call void @extern.use.f16(half %fabs)
   ret half %arg
 }
 
 ; Make sure if looking through the fabs finds a different source
 ; value, we still identify a test mask by ignoring the fabs
 define half @assume_fcmp_fabs_with_other_fabs_assume_fallback(half %arg) {
-; CHECK-LABEL: define half @assume_fcmp_fabs_with_other_fabs_assume_fallback
-; CHECK-SAME: (half returned [[ARG:%.*]]) {
+; CHECK-LABEL: define nofpclass(pinf zero sub) half @assume_fcmp_fabs_with_other_fabs_assume_fallback
+; CHECK-SAME: (half returned nofpclass(pinf zero sub) [[ARG:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[FABS:%.*]] = call half @llvm.fabs.f16(half [[ARG]]) #[[ATTR4]]
+; CHECK-NEXT:    [[FABS:%.*]] = call nofpclass(inf zero sub nnorm) half @llvm.fabs.f16(half nofpclass(pinf zero sub) [[ARG]]) #[[ATTR4]]
 ; CHECK-NEXT:    [[ONE_INF:%.*]] = fcmp oeq half [[ARG]], 0xH7C00
 ; CHECK-NEXT:    call void @llvm.assume(i1 noundef [[ONE_INF]]) #[[ATTR4]]
 ; CHECK-NEXT:    [[UNRELATED_FABS:%.*]] = fcmp oeq half [[FABS]], 0xH0000
 ; CHECK-NEXT:    call void @llvm.assume(i1 noundef [[UNRELATED_FABS]]) #[[ATTR4]]
 ; CHECK-NEXT:    [[IS_SUBNORMAL:%.*]] = fcmp olt half [[FABS]], 0xH0400
 ; CHECK-NEXT:    call void @llvm.assume(i1 noundef [[IS_SUBNORMAL]]) #[[ATTR4]]
-; CHECK-NEXT:    call void @extern.use.f16(half [[ARG]])
-; CHECK-NEXT:    call void @extern.use.f16(half [[FABS]])
+; CHECK-NEXT:    call void @extern.use.f16(half nofpclass(pinf zero sub) [[ARG]])
+; CHECK-NEXT:    call void @extern.use.f16(half nofpclass(inf zero sub nnorm) [[FABS]])
 ; CHECK-NEXT:    ret half [[ARG]]
 ;
 entry:
 
   %fabs = call half @llvm.fabs.f16(half %arg)
 
   %one.inf = fcmp oeq half %arg, 0xH7C00
   call void @llvm.assume(i1 %one.inf)