[llvm] 5186c2f - [Attributor][NFC] Introduce IRP helpers for Attr and Arg handling

[Johannes Doerfert via llvm-commits]

Author: Johannes Doerfert
Date: 2023-06-29T13:32:06-07:00
New Revision: 5186c2f9f83131f87ee3dc308ac5f18c019998c5

URL: https://github.com/llvm/llvm-project/commit/5186c2f9f83131f87ee3dc308ac5f18c019998c5
DIFF: https://github.com/llvm/llvm-project/commit/5186c2f9f83131f87ee3dc308ac5f18c019998c5.diff

LOG: [Attributor][NFC] Introduce IRP helpers for Attr and Arg handling

Added: 
    

Modified: 
    llvm/include/llvm/Transforms/IPO/Attributor.h
    llvm/lib/Transforms/IPO/Attributor.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Transforms/IPO/Attributor.h b/llvm/include/llvm/Transforms/IPO/Attributor.h
index 6509c0549d228..538b2a1cf31b9 100644
--- a/llvm/include/llvm/Transforms/IPO/Attributor.h
+++ b/llvm/include/llvm/Transforms/IPO/Attributor.h
@@ -821,6 +821,42 @@ struct IRPosition {
         "There is no attribute index for a floating or invalid position!");
   }
 
+  /// Return the attributes associated with this function or call site scope.
+  AttributeList getAttrList() const {
+    if (auto *CB = dyn_cast<CallBase>(&getAnchorValue()))
+      return CB->getAttributes();
+    return getAssociatedFunction()->getAttributes();
+  }
+
+  /// Update the attributes associated with this function or call site scope.
+  void setAttrList(const AttributeList &AttrList) const {
+    if (auto *CB = dyn_cast<CallBase>(&getAnchorValue()))
+      return CB->setAttributes(AttrList);
+    return getAssociatedFunction()->setAttributes(AttrList);
+  }
+
+  /// Return the number of arguments associated with this function or call site
+  /// scope.
+  unsigned getNumArgs() const {
+    assert((getPositionKind() == IRP_CALL_SITE ||
+            getPositionKind() == IRP_FUNCTION) &&
+           "Only valid for function/call site positions!");
+    if (auto *CB = dyn_cast<CallBase>(&getAnchorValue()))
+      return CB->arg_size();
+    return getAssociatedFunction()->arg_size();
+  }
+
+  /// Return theargument \p ArgNo associated with this function or call site
+  /// scope.
+  Value *getArg(unsigned ArgNo) const {
+    assert((getPositionKind() == IRP_CALL_SITE ||
+            getPositionKind() == IRP_FUNCTION) &&
+           "Only valid for function/call site positions!");
+    if (auto *CB = dyn_cast<CallBase>(&getAnchorValue()))
+      return CB->getArgOperand(ArgNo);
+    return getAssociatedFunction()->getArg(ArgNo);
+  }
+
   /// Return the associated position kind.
   Kind getPositionKind() const {
     char EncodingBits = getEncodingBits();
@@ -872,21 +908,19 @@ struct IRPosition {
     if (getPositionKind() == IRP_INVALID || getPositionKind() == IRP_FLOAT)
       return;
 
-    AttributeList AttrList;
-    auto *CB = dyn_cast<CallBase>(&getAnchorValue());
-    if (CB)
-      AttrList = CB->getAttributes();
-    else
-      AttrList = getAssociatedFunction()->getAttributes();
+    AttributeList AttrList = getAttrList();
 
+    bool Changed = false;
+    unsigned Idx = getAttrIdx();
     LLVMContext &Ctx = getAnchorValue().getContext();
-    for (Attribute::AttrKind AK : AKs)
-      AttrList = AttrList.removeAttributeAtIndex(Ctx, getAttrIdx(), AK);
-
-    if (CB)
-      CB->setAttributes(AttrList);
-    else
-      getAssociatedFunction()->setAttributes(AttrList);
+    for (Attribute::AttrKind AK : AKs) {
+      if (!AttrList.hasAttributeAtIndex(Idx, AK))
+        continue;
+      Changed = true;
+      AttrList = AttrList.removeAttributeAtIndex(Ctx, Idx, AK);
+    }
+    if (Changed)
+      setAttrList(AttrList);
   }
 
   bool isAnyCallSitePosition() const {

diff  --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index d1a96f07175e3..6b39efb25ba09 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -1015,29 +1015,18 @@ ChangeStatus
 IRAttributeManifest::manifestAttrs(Attributor &A, const IRPosition &IRP,
                                    const ArrayRef<Attribute> &DeducedAttrs,
                                    bool ForceReplace) {
-  Function *ScopeFn = IRP.getAnchorScope();
-  IRPosition::Kind PK = IRP.getPositionKind();
+  switch (IRP.getPositionKind()) {
+  case IRPosition::IRP_FLOAT:
+  case IRPosition::IRP_INVALID:
+    return ChangeStatus::UNCHANGED;
+  default:
+    break;
+  };
 
   // In the following some generic code that will manifest attributes in
   // DeducedAttrs if they improve the current IR. Due to the 
diff erent
   // annotation positions we use the underlying AttributeList interface.
-
-  AttributeList Attrs;
-  switch (PK) {
-  case IRPosition::IRP_INVALID:
-  case IRPosition::IRP_FLOAT:
-    return ChangeStatus::UNCHANGED;
-  case IRPosition::IRP_ARGUMENT:
-  case IRPosition::IRP_FUNCTION:
-  case IRPosition::IRP_RETURNED:
-    Attrs = ScopeFn->getAttributes();
-    break;
-  case IRPosition::IRP_CALL_SITE:
-  case IRPosition::IRP_CALL_SITE_RETURNED:
-  case IRPosition::IRP_CALL_SITE_ARGUMENT:
-    Attrs = cast<CallBase>(IRP.getAnchorValue()).getAttributes();
-    break;
-  }
+  AttributeList Attrs = IRP.getAttrList();
 
   ChangeStatus HasChanged = ChangeStatus::UNCHANGED;
   LLVMContext &Ctx = IRP.getAnchorValue().getContext();
@@ -1051,22 +1040,7 @@ IRAttributeManifest::manifestAttrs(Attributor &A, const IRPosition &IRP,
   if (HasChanged == ChangeStatus::UNCHANGED)
     return HasChanged;
 
-  switch (PK) {
-  case IRPosition::IRP_ARGUMENT:
-  case IRPosition::IRP_FUNCTION:
-  case IRPosition::IRP_RETURNED:
-    ScopeFn->setAttributes(Attrs);
-    break;
-  case IRPosition::IRP_CALL_SITE:
-  case IRPosition::IRP_CALL_SITE_RETURNED:
-  case IRPosition::IRP_CALL_SITE_ARGUMENT:
-    cast<CallBase>(IRP.getAnchorValue()).setAttributes(Attrs);
-    break;
-  case IRPosition::IRP_INVALID:
-  case IRPosition::IRP_FLOAT:
-    break;
-  }
-
+  IRP.setAttrList(Attrs);
   return HasChanged;
 }
 
@@ -1183,11 +1157,7 @@ bool IRPosition::getAttrsFromIRAttr(Attribute::AttrKind AK,
   if (getPositionKind() == IRP_INVALID || getPositionKind() == IRP_FLOAT)
     return false;
 
-  AttributeList AttrList;
-  if (const auto *CB = dyn_cast<CallBase>(&getAnchorValue()))
-    AttrList = CB->getAttributes();
-  else
-    AttrList = getAssociatedFunction()->getAttributes();
+  AttributeList AttrList = getAttrList();
 
   bool HasAttr = AttrList.hasAttributeAtIndex(getAttrIdx(), AK);
   if (HasAttr)