Refactor determinePointerAccessAttrs to a dedicated function to colle…

[haopliu]

…ct argument uses.

With this refactoring, getArgumentUses can be used to infer other attributes, like initialized.

https://discourse.llvm.org/t/rfc-llvm-new-initialized-parameter-attribute-for-improved-interprocedural-dse/77337

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Haopeng Liu (haopliu)

Changes

…ct argument uses.

With this refactoring, getArgumentUses can be used to infer other attributes, like initialized.

https://discourse.llvm.org/t/rfc-llvm-new-initialized-parameter-attribute-for-improved-interprocedural-dse/77337

---

Full diff: https://github.com/llvm/llvm-project/pull/84869.diff

1 Files Affected:

• (modified) llvm/lib/Transforms/IPO/FunctionAttrs.cpp (+88-63)

diff --git a/llvm/lib/Transforms/IPO/FunctionAttrs.cpp b/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
index 7ebf265e17ba1f..f28e1774baae6a 100644
--- a/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
@@ -580,56 +580,21 @@ struct ArgumentUsesTracker : public CaptureTracker {
   const SCCNodeSet &SCCNodes;
 };
 
-} // end anonymous namespace
-
-namespace llvm {
-
-template <> struct GraphTraits<ArgumentGraphNode *> {
-  using NodeRef = ArgumentGraphNode *;
-  using ChildIteratorType = SmallVectorImpl<ArgumentGraphNode *>::iterator;
-
-  static NodeRef getEntryNode(NodeRef A) { return A; }
-  static ChildIteratorType child_begin(NodeRef N) { return N->Uses.begin(); }
-  static ChildIteratorType child_end(NodeRef N) { return N->Uses.end(); }
-};
-
-template <>
-struct GraphTraits<ArgumentGraph *> : public GraphTraits<ArgumentGraphNode *> {
-  static NodeRef getEntryNode(ArgumentGraph *AG) { return AG->getEntryNode(); }
-
-  static ChildIteratorType nodes_begin(ArgumentGraph *AG) {
-    return AG->begin();
-  }
-
-  static ChildIteratorType nodes_end(ArgumentGraph *AG) { return AG->end(); }
-};
-
-} // end namespace llvm
-
-/// Returns Attribute::None, Attribute::ReadOnly or Attribute::ReadNone.
-static Attribute::AttrKind
-determinePointerAccessAttrs(Argument *A,
-                            const SmallPtrSet<Argument *, 8> &SCCNodes) {
+/// Get all uses of an argument in the function and store them to different
+/// lists: Reads, Writes, and SpecialUses.
+void getArgumentUses(Argument *A, const SmallPtrSet<Argument *, 8> &SCCNodes,
+                     SmallVector<Instruction *, 16> *Reads,
+                     SmallVector<Instruction *, 16> *Writes,
+                     SmallVector<Instruction *, 16> *SpecialUses) {
   SmallVector<Use *, 32> Worklist;
   SmallPtrSet<Use *, 32> Visited;
 
-  // inalloca arguments are always clobbered by the call.
-  if (A->hasInAllocaAttr() || A->hasPreallocatedAttr())
-    return Attribute::None;
-
-  bool IsRead = false;
-  bool IsWrite = false;
-
   for (Use &U : A->uses()) {
     Visited.insert(&U);
     Worklist.push_back(&U);
   }
 
   while (!Worklist.empty()) {
-    if (IsWrite && IsRead)
-      // No point in searching further..
-      return Attribute::None;
-
     Use *U = Worklist.pop_back_val();
     Instruction *I = cast<Instruction>(U->getUser());
 
@@ -649,7 +614,7 @@ determinePointerAccessAttrs(Argument *A,
     case Instruction::Invoke: {
       CallBase &CB = cast<CallBase>(*I);
       if (CB.isCallee(U)) {
-        IsRead = true;
+        Reads->push_back(I);
         // Note that indirect calls do not capture, see comment in
         // CaptureTracking for context
         continue;
@@ -668,12 +633,15 @@ determinePointerAccessAttrs(Argument *A,
           if (Visited.insert(&UU).second)
             Worklist.push_back(&UU);
       } else if (!CB.doesNotCapture(UseIndex)) {
-        if (!CB.onlyReadsMemory())
+        if (!CB.onlyReadsMemory()) {
           // If the callee can save a copy into other memory, then simply
           // scanning uses of the call is insufficient.  We have no way
           // of tracking copies of the pointer through memory to see
-          // if a reloaded copy is written to, thus we must give up.
-          return Attribute::None;
+          // if a reloaded copy is written to, thus we treat it as special
+          // uses.
+          SpecialUses->push_back(I);
+          continue;
+        }
         // Push users for processing once we finish this one
         if (!I->getType()->isVoidTy())
           for (Use &UU : I->uses())
@@ -698,12 +666,12 @@ determinePointerAccessAttrs(Argument *A,
       if (CB.doesNotAccessMemory(UseIndex)) {
         /* nop */
       } else if (!isModSet(ArgMR) || CB.onlyReadsMemory(UseIndex)) {
-        IsRead = true;
+        Reads->push_back(I);
       } else if (!isRefSet(ArgMR) ||
                  CB.dataOperandHasImpliedAttr(UseIndex, Attribute::WriteOnly)) {
-        IsWrite = true;
+        Writes->push_back(I);
       } else {
-        return Attribute::None;
+        SpecialUses->push_back(I);
       }
       break;
     }
@@ -711,23 +679,29 @@ determinePointerAccessAttrs(Argument *A,
     case Instruction::Load:
       // A volatile load has side effects beyond what readonly can be relied
       // upon.
-      if (cast<LoadInst>(I)->isVolatile())
-        return Attribute::None;
+      if (cast<LoadInst>(I)->isVolatile()) {
+        SpecialUses->push_back(I);
+        continue;
+      }
 
-      IsRead = true;
+      Reads->push_back(I);
       break;
 
     case Instruction::Store:
-      if (cast<StoreInst>(I)->getValueOperand() == *U)
+      if (cast<StoreInst>(I)->getValueOperand() == *U) {
         // untrackable capture
-        return Attribute::None;
+        SpecialUses->push_back(I);
+        continue;
+      }
 
       // A volatile store has side effects beyond what writeonly can be relied
       // upon.
-      if (cast<StoreInst>(I)->isVolatile())
-        return Attribute::None;
+      if (cast<StoreInst>(I)->isVolatile()) {
+        SpecialUses->push_back(I);
+        continue;
+      }
 
-      IsWrite = true;
+      Writes->push_back(I);
       break;
 
     case Instruction::ICmp:
@@ -735,15 +709,54 @@ determinePointerAccessAttrs(Argument *A,
       break;
 
     default:
-      return Attribute::None;
+      SpecialUses->push_back(I);
     }
   }
+}
+
+} // end anonymous namespace
+
+namespace llvm {
+
+template <> struct GraphTraits<ArgumentGraphNode *> {
+  using NodeRef = ArgumentGraphNode *;
+  using ChildIteratorType = SmallVectorImpl<ArgumentGraphNode *>::iterator;
+
+  static NodeRef getEntryNode(NodeRef A) { return A; }
+  static ChildIteratorType child_begin(NodeRef N) { return N->Uses.begin(); }
+  static ChildIteratorType child_end(NodeRef N) { return N->Uses.end(); }
+};
+
+template <>
+struct GraphTraits<ArgumentGraph *> : public GraphTraits<ArgumentGraphNode *> {
+  static NodeRef getEntryNode(ArgumentGraph *AG) { return AG->getEntryNode(); }
 
-  if (IsWrite && IsRead)
+  static ChildIteratorType nodes_begin(ArgumentGraph *AG) {
+    return AG->begin();
+  }
+
+  static ChildIteratorType nodes_end(ArgumentGraph *AG) { return AG->end(); }
+};
+
+} // end namespace llvm
+
+/// Returns Attribute::None, Attribute::ReadOnly or Attribute::ReadNone.
+static Attribute::AttrKind
+determinePointerAccessAttrs(Argument *A, SmallVector<Instruction *, 16> &Reads,
+                            SmallVector<Instruction *, 16> Writes,
+                            SmallVector<Instruction *, 16> SpecialUses) {
+  // inalloca arguments are always clobbered by the call.
+  if (A->hasInAllocaAttr() || A->hasPreallocatedAttr())
+    return Attribute::None;
+
+  if (!SpecialUses.empty())
     return Attribute::None;
-  else if (IsRead)
+
+  if (!Writes.empty() && !Reads.empty())
+    return Attribute::None;
+  else if (!Reads.empty())
     return Attribute::ReadOnly;
-  else if (IsWrite)
+  else if (!Writes.empty())
     return Attribute::WriteOnly;
   else
     return Attribute::ReadNone;
@@ -931,7 +944,11 @@ static void addArgumentAttrs(const SCCNodeSet &SCCNodes,
         // functions in the SCC.
         SmallPtrSet<Argument *, 8> Self;
         Self.insert(&A);
-        Attribute::AttrKind R = determinePointerAccessAttrs(&A, Self);
+        SmallVector<Instruction *, 16> Reads, Writes, SpecialUses;
+        getArgumentUses(&A, Self, &Reads, &Writes, &SpecialUses);
+
+        Attribute::AttrKind R =
+            determinePointerAccessAttrs(&A, Reads, Writes, SpecialUses);
         if (R != Attribute::None)
           if (addAccessAttr(&A, R))
             Changed.insert(F);
@@ -963,7 +980,11 @@ static void addArgumentAttrs(const SCCNodeSet &SCCNodes,
         // Infer the access attributes given the new nocapture one
         SmallPtrSet<Argument *, 8> Self;
         Self.insert(&*A);
-        Attribute::AttrKind R = determinePointerAccessAttrs(&*A, Self);
+        SmallVector<Instruction *, 16> Reads, Writes, SpecialUses;
+        getArgumentUses(A, Self, &Reads, &Writes, &SpecialUses);
+
+        Attribute::AttrKind R =
+            determinePointerAccessAttrs(&*A, Reads, Writes, SpecialUses);
         if (R != Attribute::None)
           addAccessAttr(A, R);
       }
@@ -1032,7 +1053,11 @@ static void addArgumentAttrs(const SCCNodeSet &SCCNodes,
     Attribute::AttrKind AccessAttr = Attribute::ReadNone;
     for (ArgumentGraphNode *N : ArgumentSCC) {
       Argument *A = N->Definition;
-      Attribute::AttrKind K = determinePointerAccessAttrs(A, ArgumentSCCNodes);
+      SmallVector<Instruction *, 16> Reads, Writes, SpecialUses;
+      getArgumentUses(A, ArgumentSCCNodes, &Reads, &Writes, &SpecialUses);
+
+      Attribute::AttrKind K =
+          determinePointerAccessAttrs(A, Reads, Writes, SpecialUses);
       AccessAttr = meetAccessAttr(AccessAttr, K);
       if (AccessAttr == Attribute::None)
         break;

[aeubanks]

since this actually makes the existing code worse without any new functionality, let's keep this change as part of the main FunctionAttrs PR

also, you should be able to return a std::tuple like here and nicely separate it in the caller like here