IR: Add nofpclass parameter attribute

This carries a bitmask indicating forbidden floating-point value kinds
in the argument or return value. This will enable interprocedural
-ffinite-math-only optimizations. This is primarily to cover the
no-nans and no-infinities cases, but also covers the other floating
point classes for free. Textually, this provides a number of names
corresponding to bits in FPClassTest, e.g.

  call nofpclass(nan inf) @must_be_finite()
  call nofpclass(snan) @cannot_be_snan()

This is more expressive than the existing nnan and ninf fast math
flags. As an added bonus, you can represent fun things like nanf:

  declare nofpclass(inf zero sub norm) float @only_nans()

Compared to nnan/ninf:
  - Can be applied to individual call operands as well as the return value
  - Can distinguish signaling and quiet nans
  - Distinguishes the sign of infinities
  - Can be safely propagated since it doesn't imply anything about
    other operands.
  - Does not apply to FP instructions; it's not a flag

This is one step closer to being able to retire "no-nans-fp-math" and
"no-infs-fp-math". The one remaining situation where we have no way to
represent no-nans/infs is for loads (if we wanted to solve this we
could introduce !nofpclass metadata, following along with
noundef/!noundef).

This is to help simplify the GPU builtin math library
distribution. Currently the library code has explicit finite math only
checks, read from global constants the compiler driver needs to set
based on the compiler flags during linking. We end up having to
internalize the library into each translation unit in case different
linked modules have different math flags. By propagating known-not-nan
and known-not-infinity information, we can automatically prune the
edge case handling in most functions if the function is only reached
from fast math uses.

1 files changed, 101 insertions, 0 deletions

diff --git a/llvm/lib/IR/Attributes.cpp b/llvm/lib/IR/Attributes.cpp
index 8c989c4..21f119d 100644
--- a/llvm/lib/IR/Attributes.cpp
+++ b/llvm/lib/IR/Attributes.cpp
@@ -216,6 +216,11 @@ Attribute Attribute::getWithMemoryEffects(LLVMContext &Context,
   return get(Context, Memory, ME.toIntValue());
 }
 
+Attribute Attribute::getWithNoFPClass(LLVMContext &Context,
+                                      FPClassTest ClassMask) {
+  return get(Context, NoFPClass, ClassMask);
+}
+
 Attribute
 Attribute::getWithAllocSizeArgs(LLVMContext &Context, unsigned ElemSizeArg,
                                 const std::optional<unsigned> &NumElemsArg) {
@@ -261,6 +266,16 @@ bool Attribute::isExistingAttribute(StringRef Name) {
       .Default(false);
 }
 
+/// Returns true if this is a type legal for the 'nofpclass' attribute. This
+/// follows the same type rules as FPMathOperator.
+///
+/// TODO: Consider relaxing to any FP type struct fields.
+static bool isNoFPClassCompatibleType(Type *Ty) {
+  while (ArrayType *ArrTy = dyn_cast<ArrayType>(Ty))
+    Ty = ArrTy->getElementType();
+  return Ty->isFPOrFPVectorTy();
+}
+
 //===----------------------------------------------------------------------===//
 // Attribute Accessor Methods
 //===----------------------------------------------------------------------===//
@@ -396,6 +411,12 @@ MemoryEffects Attribute::getMemoryEffects() const {
   return MemoryEffects::createFromIntValue(pImpl->getValueAsInt());
 }
 
+FPClassTest Attribute::getNoFPClass() const {
+  assert(hasAttribute(Attribute::NoFPClass) &&
+         "Can only call getNoFPClass() on nofpclass attribute");
+  return static_cast<FPClassTest>(pImpl->getValueAsInt());
+}
+
 static const char *getModRefStr(ModRefInfo MR) {
   switch (MR) {
   case ModRefInfo::NoModRef:
@@ -410,6 +431,56 @@ static const char *getModRefStr(ModRefInfo MR) {
   llvm_unreachable("Invalid ModRefInfo");
 }
 
+// Every bitfield has a unique name and one or more aliasing names that cover
+// multiple bits. Names should be listed in order of preference, with higher
+// popcounts listed first.
+//
+// Bits are consumed as printed. Each field should only be represented in one
+// printed field.
+static constexpr std::pair<unsigned, StringLiteral> NoFPClassName[] = {
+  {fcAllFlags, "all"},
+  {fcNan, "nan"},
+  {fcSNan, "snan"},
+  {fcQNan, "qnan"},
+  {fcInf, "inf"},
+  {fcNegInf, "ninf"},
+  {fcPosInf, "pinf"},
+  {fcZero, "zero"},
+  {fcNegZero, "nzero"},
+  {fcPosZero, "pzero"},
+  {fcSubnormal, "sub"},
+  {fcNegSubnormal, "nsub"},
+  {fcPosSubnormal, "psub"},
+  {fcNormal, "norm"},
+  {fcNegNormal, "nnorm"},
+  {fcPosNormal, "pnorm"}
+};
+
+static std::string getNoFPClassAttrAsString(unsigned Mask) {
+  std::string Result("nofpclass(");
+  raw_string_ostream OS(Result);
+
+  if (Mask == 0) {
+    OS << "none)";
+    return Result;
+  }
+
+  ListSeparator LS(" ");
+  for (auto [BitTest, Name] : NoFPClassName) {
+    if ((Mask & BitTest) == BitTest) {
+      OS << LS << Name;
+
+      // Clear the bits so we don't print any aliased names later.
+      Mask &= ~BitTest;
+    }
+  }
+
+  assert(Mask == 0 && "didn't print some mask bits");
+
+  OS << ')';
+  return Result;
+}
+
 std::string Attribute::getAsString(bool InAttrGrp) const {
   if (!pImpl) return {};
 
@@ -543,6 +614,9 @@ std::string Attribute::getAsString(bool InAttrGrp) const {
     return Result;
   }
 
+  if (hasAttribute(Attribute::NoFPClass))
+    return getNoFPClassAttrAsString(getValueAsInt());
+
   // Convert target-dependent attributes to strings of the form:
   //
   //   "kind"
@@ -840,6 +914,10 @@ MemoryEffects AttributeSet::getMemoryEffects() const {
   return SetNode ? SetNode->getMemoryEffects() : MemoryEffects::unknown();
 }
 
+FPClassTest AttributeSet::getNoFPClass() const {
+  return SetNode ? SetNode->getNoFPClass() : fcNone;
+}
+
 std::string AttributeSet::getAsString(bool InAttrGrp) const {
   return SetNode ? SetNode->getAsString(InAttrGrp) : "";
 }
@@ -1024,6 +1102,12 @@ MemoryEffects AttributeSetNode::getMemoryEffects() const {
   return MemoryEffects::unknown();
 }
 
+FPClassTest AttributeSetNode::getNoFPClass() const {
+  if (auto A = findEnumAttribute(Attribute::NoFPClass))
+    return A->getNoFPClass();
+  return fcNone;
+}
+
 std::string AttributeSetNode::getAsString(bool InAttrGrp) const {
   std::string Str;
   for (iterator I = begin(), E = end(); I != E; ++I) {
@@ -1560,6 +1644,14 @@ AttributeList::getParamDereferenceableOrNullBytes(unsigned Index) const {
   return getParamAttrs(Index).getDereferenceableOrNullBytes();
 }
 
+FPClassTest AttributeList::getRetNoFPClass() const {
+  return getRetAttrs().getNoFPClass();
+}
+
+FPClassTest AttributeList::getParamNoFPClass(unsigned Index) const {
+  return getParamAttrs(Index).getNoFPClass();
+}
+
 UWTableKind AttributeList::getUWTableKind() const {
   return getFnAttrs().getUWTableKind();
 }
@@ -1803,6 +1895,10 @@ AttrBuilder &AttrBuilder::addMemoryAttr(MemoryEffects ME) {
   return addRawIntAttr(Attribute::Memory, ME.toIntValue());
 }
 
+AttrBuilder &AttrBuilder::addNoFPClassAttr(FPClassTest Mask) {
+  return addRawIntAttr(Attribute::NoFPClass, Mask);
+}
+
 AttrBuilder &AttrBuilder::addAllocKindAttr(AllocFnKind Kind) {
   return addRawIntAttr(Attribute::AllocKind, static_cast<uint64_t>(Kind));
 }
@@ -1926,6 +2022,11 @@ AttributeMask AttributeFuncs::typeIncompatible(Type *Ty,
       Incompatible.addAttribute(Attribute::Alignment);
   }
 
+  if (ASK & ASK_SAFE_TO_DROP) {
+    if (!isNoFPClassCompatibleType(Ty))
+      Incompatible.addAttribute(Attribute::NoFPClass);
+  }
+
   // Some attributes can apply to all "values" but there are no `void` values.
   if (Ty->isVoidTy()) {
     if (ASK & ASK_SAFE_TO_DROP)