[IR][Attribute] Add support for intersecting AttributeLists; NFC (#109719)

Add support for taking the intersection of two AttributeLists s.t the
result list contains attributes that are valid in the context of both
inputs.

i.e if we have `nonnull align(32) noundef` intersected with `nonnull
align(16) dereferenceable(10)`, the result is `nonnull align(16)`.

Further it handles attributes that are not-droppable. For example
dropping `byval` can change the nature of a callsite/function so its
impossible to correct a correct intersection if its dropped from the
result. i.e `nonnull byval(i64)` intersected with `nonnull` is
invalid.

The motivation for the infrastructure is to enable sinking/hoisting
callsites with differing attributes.

1 files changed, 219 insertions, 14 deletions

diff --git a/llvm/lib/IR/Attributes.cpp b/llvm/lib/IR/Attributes.cpp
index eb61583..692207e 100644
--- a/llvm/lib/IR/Attributes.cpp
+++ b/llvm/lib/IR/Attributes.cpp
@@ -362,8 +362,7 @@ bool Attribute::isConstantRangeListAttribute() const {
 
 Attribute::AttrKind Attribute::getKindAsEnum() const {
   if (!pImpl) return None;
-  assert((isEnumAttribute() || isIntAttribute() || isTypeAttribute() ||
-          isConstantRangeAttribute() || isConstantRangeListAttribute()) &&
+  assert(hasKindAsEnum() &&
          "Invalid attribute type to get the kind as an enum!");
   return pImpl->getKindAsEnum();
 }
@@ -712,6 +711,16 @@ bool Attribute::hasParentContext(LLVMContext &C) const {
   return C.pImpl->AttrsSet.FindNodeOrInsertPos(ID, Unused) == pImpl;
 }
 
+int Attribute::cmpKind(Attribute A) const {
+  if (!pImpl && !A.pImpl)
+    return 0;
+  if (!pImpl)
+    return 1;
+  if (!A.pImpl)
+    return -1;
+  return pImpl->cmp(*A.pImpl, /*KindOnly=*/true);
+}
+
 bool Attribute::operator<(Attribute A) const {
   if (!pImpl && !A.pImpl) return false;
   if (!pImpl) return true;
@@ -727,16 +736,25 @@ enum AttributeProperty {
   FnAttr = (1 << 0),
   ParamAttr = (1 << 1),
   RetAttr = (1 << 2),
+  IntersectPreserve = (0 << 3),
+  IntersectAnd = (1 << 3),
+  IntersectMin = (2 << 3),
+  IntersectCustom = (3 << 3),
+  IntersectPropertyMask = (3 << 3),
 };
 
 #define GET_ATTR_PROP_TABLE
 #include "llvm/IR/Attributes.inc"
 
-static bool hasAttributeProperty(Attribute::AttrKind Kind,
-                                 AttributeProperty Prop) {
+static unsigned getAttributeProperties(Attribute::AttrKind Kind) {
   unsigned Index = Kind - 1;
   assert(Index < std::size(AttrPropTable) && "Invalid attribute kind");
-  return AttrPropTable[Index] & Prop;
+  return AttrPropTable[Index];
+}
+
+static bool hasAttributeProperty(Attribute::AttrKind Kind,
+                                 AttributeProperty Prop) {
+  return getAttributeProperties(Kind) & Prop;
 }
 
 bool Attribute::canUseAsFnAttr(AttrKind Kind) {
@@ -751,6 +769,30 @@ bool Attribute::canUseAsRetAttr(AttrKind Kind) {
   return hasAttributeProperty(Kind, AttributeProperty::RetAttr);
 }
 
+static bool hasIntersectProperty(Attribute::AttrKind Kind,
+                                 AttributeProperty Prop) {
+  assert(Prop == AttributeProperty::IntersectPreserve ||
+         Prop == AttributeProperty::IntersectAnd ||
+         Prop == AttributeProperty::IntersectMin ||
+         Prop == AttributeProperty::IntersectCustom &&
+             "Unknown intersect property");
+  return (getAttributeProperties(Kind) &
+          AttributeProperty::IntersectPropertyMask) == Prop;
+}
+
+bool Attribute::intersectMustPreserve(AttrKind Kind) {
+  return hasIntersectProperty(Kind, AttributeProperty::IntersectPreserve);
+}
+bool Attribute::intersectWithAnd(AttrKind Kind) {
+  return hasIntersectProperty(Kind, AttributeProperty::IntersectAnd);
+}
+bool Attribute::intersectWithMin(AttrKind Kind) {
+  return hasIntersectProperty(Kind, AttributeProperty::IntersectMin);
+}
+bool Attribute::intersectWithCustom(AttrKind Kind) {
+  return hasIntersectProperty(Kind, AttributeProperty::IntersectCustom);
+}
+
 //===----------------------------------------------------------------------===//
 // AttributeImpl Definition
 //===----------------------------------------------------------------------===//
@@ -808,17 +850,21 @@ ArrayRef<ConstantRange> AttributeImpl::getValueAsConstantRangeList() const {
       ->getConstantRangeListValue();
 }
 
-bool AttributeImpl::operator<(const AttributeImpl &AI) const {
+int AttributeImpl::cmp(const AttributeImpl &AI, bool KindOnly) const {
   if (this == &AI)
-    return false;
+    return 0;
 
   // This sorts the attributes with Attribute::AttrKinds coming first (sorted
   // relative to their enum value) and then strings.
   if (!isStringAttribute()) {
     if (AI.isStringAttribute())
-      return true;
+      return -1;
+
     if (getKindAsEnum() != AI.getKindAsEnum())
-      return getKindAsEnum() < AI.getKindAsEnum();
+      return getKindAsEnum() < AI.getKindAsEnum() ? -1 : 1;
+    else if (KindOnly)
+      return 0;
+
     assert(!AI.isEnumAttribute() && "Non-unique attribute");
     assert(!AI.isTypeAttribute() && "Comparison of types would be unstable");
     assert(!AI.isConstantRangeAttribute() && "Unclear how to compare ranges");
@@ -826,14 +872,21 @@ bool AttributeImpl::operator<(const AttributeImpl &AI) const {
            "Unclear how to compare range list");
     // TODO: Is this actually needed?
     assert(AI.isIntAttribute() && "Only possibility left");
-    return getValueAsInt() < AI.getValueAsInt();
+    if (getValueAsInt() < AI.getValueAsInt())
+      return -1;
+    return getValueAsInt() == AI.getValueAsInt() ? 0 : 1;
   }
-
   if (!AI.isStringAttribute())
-    return false;
+    return 1;
+  if (KindOnly)
+    return getKindAsString().compare(AI.getKindAsString());
   if (getKindAsString() == AI.getKindAsString())
-    return getValueAsString() < AI.getValueAsString();
-  return getKindAsString() < AI.getKindAsString();
+    return getValueAsString().compare(AI.getValueAsString());
+  return getKindAsString().compare(AI.getKindAsString());
+}
+
+bool AttributeImpl::operator<(const AttributeImpl &AI) const {
+  return cmp(AI, /*KindOnly=*/false) < 0;
 }
 
 //===----------------------------------------------------------------------===//
@@ -903,6 +956,132 @@ AttributeSet AttributeSet::removeAttributes(LLVMContext &C,
   return get(C, B);
 }
 
+std::optional<AttributeSet>
+AttributeSet::intersectWith(LLVMContext &C, AttributeSet Other) const {
+  if (*this == Other)
+    return *this;
+
+  AttrBuilder Intersected(C);
+  // Iterate over both attr sets at once.
+  auto ItBegin0 = begin();
+  auto ItEnd0 = end();
+  auto ItBegin1 = Other.begin();
+  auto ItEnd1 = Other.end();
+
+  while (ItBegin0 != ItEnd0 || ItBegin1 != ItEnd1) {
+    // Loop through all attributes in both this and Other in sorted order. If
+    // the attribute is only present in one of the sets, it will be set in
+    // Attr0. If it is present in both sets both Attr0 and Attr1 will be set.
+    Attribute Attr0, Attr1;
+    if (ItBegin1 == ItEnd1)
+      Attr0 = *ItBegin0++;
+    else if (ItBegin0 == ItEnd0)
+      Attr0 = *ItBegin1++;
+    else {
+      int Cmp = ItBegin0->cmpKind(*ItBegin1);
+      if (Cmp == 0) {
+        Attr0 = *ItBegin0++;
+        Attr1 = *ItBegin1++;
+      } else if (Cmp < 0)
+        Attr0 = *ItBegin0++;
+      else
+        Attr0 = *ItBegin1++;
+    }
+    assert(Attr0.isValid() && "Iteration should always yield a valid attr");
+
+    auto IntersectEq = [&]() {
+      if (!Attr1.isValid())
+        return false;
+      if (Attr0 != Attr1)
+        return false;
+      Intersected.addAttribute(Attr0);
+      return true;
+    };
+
+    // Non-enum assume we must preserve. Handle early so we can unconditionally
+    // use Kind below.
+    if (!Attr0.hasKindAsEnum()) {
+      if (!IntersectEq())
+        return std::nullopt;
+      continue;
+    }
+
+    Attribute::AttrKind Kind = Attr0.getKindAsEnum();
+    // If we don't have both attributes, then fail if the attribute is
+    // must-preserve or drop it otherwise.
+    if (!Attr1.isValid()) {
+      if (Attribute::intersectMustPreserve(Kind))
+        return std::nullopt;
+      continue;
+    }
+
+    // We have both attributes so apply the intersection rule.
+    assert(Attr1.hasKindAsEnum() && Kind == Attr1.getKindAsEnum() &&
+           "Iterator picked up two different attributes in the same iteration");
+
+    // Attribute we can intersect with "and"
+    if (Attribute::intersectWithAnd(Kind)) {
+      assert(Attribute::isEnumAttrKind(Kind) &&
+             "Invalid attr type of intersectAnd");
+      Intersected.addAttribute(Kind);
+      continue;
+    }
+
+    // Attribute we can intersect with "min"
+    if (Attribute::intersectWithMin(Kind)) {
+      assert(Attribute::isIntAttrKind(Kind) &&
+             "Invalid attr type of intersectMin");
+      uint64_t NewVal = std::min(Attr0.getValueAsInt(), Attr1.getValueAsInt());
+      Intersected.addRawIntAttr(Kind, NewVal);
+      continue;
+    }
+    // Attribute we can intersect but need a custom rule for.
+    if (Attribute::intersectWithCustom(Kind)) {
+      switch (Kind) {
+      case Attribute::Alignment:
+        // If `byval` is present, alignment become must-preserve. This is
+        // handled below if we have `byval`.
+        Intersected.addAlignmentAttr(
+            std::min(Attr0.getAlignment().valueOrOne(),
+                     Attr1.getAlignment().valueOrOne()));
+        break;
+      case Attribute::Memory:
+        Intersected.addMemoryAttr(Attr0.getMemoryEffects() |
+                                  Attr1.getMemoryEffects());
+        break;
+      case Attribute::NoFPClass:
+        Intersected.addNoFPClassAttr(Attr0.getNoFPClass() &
+                                     Attr1.getNoFPClass());
+        break;
+      case Attribute::Range: {
+        ConstantRange Range0 = Attr0.getRange();
+        ConstantRange Range1 = Attr1.getRange();
+        ConstantRange NewRange = Range0.unionWith(Range1);
+        if (!NewRange.isFullSet())
+          Intersected.addRangeAttr(NewRange);
+      } break;
+      default:
+        llvm_unreachable("Unknown attribute with custom intersection rule");
+      }
+      continue;
+    }
+
+    // Attributes with no intersection rule. Only intersect if they are equal.
+    // Otherwise fail.
+    if (!IntersectEq())
+      return std::nullopt;
+
+    // Special handling of `byval`. `byval` essentially turns align attr into
+    // must-preserve
+    if (Kind == Attribute::ByVal &&
+        getAttribute(Attribute::Alignment) !=
+            Other.getAttribute(Attribute::Alignment))
+      return std::nullopt;
+  }
+
+  return get(C, Intersected);
+}
+
 unsigned AttributeSet::getNumAttributes() const {
   return SetNode ? SetNode->getNumAttributes() : 0;
 }
@@ -1614,6 +1793,32 @@ AttributeList AttributeList::addAllocSizeParamAttr(
   return addParamAttributes(C, Index, B);
 }
 
+std::optional<AttributeList>
+AttributeList::intersectWith(LLVMContext &C, AttributeList Other) const {
+  // Trivial case, the two lists are equal.
+  if (*this == Other)
+    return *this;
+
+  // At least for now, only intersect lists with same number of params.
+  if (getNumAttrSets() != Other.getNumAttrSets())
+    return std::nullopt;
+
+  SmallVector<std::pair<unsigned, AttributeSet>> IntersectedAttrs;
+  for (unsigned Idx : indexes()) {
+    auto IntersectedAS =
+        getAttributes(Idx).intersectWith(C, Other.getAttributes(Idx));
+    // If any index fails to intersect, fail.
+    if (!IntersectedAS)
+      return std::nullopt;
+    if (!IntersectedAS->hasAttributes())
+      continue;
+    IntersectedAttrs.push_back(std::make_pair(Idx, *IntersectedAS));
+  }
+
+  llvm::sort(IntersectedAttrs, llvm::less_first());
+  return AttributeList::get(C, IntersectedAttrs);
+}
+
 //===----------------------------------------------------------------------===//
 // AttributeList Accessor Methods
 //===----------------------------------------------------------------------===//