[llvm] 9366610 - [Attributor][NFC] Add querying AA to shouldSpecializeCallSiteForCallee

[Johannes Doerfert via llvm-commits]

Author: Johannes Doerfert
Date: 2023-08-29T22:35:16-07:00
New Revision: 936661084c0413d877088f73f2fb7f6d0e1e6643

URL: https://github.com/llvm/llvm-project/commit/936661084c0413d877088f73f2fb7f6d0e1e6643
DIFF: https://github.com/llvm/llvm-project/commit/936661084c0413d877088f73f2fb7f6d0e1e6643.diff

LOG: [Attributor][NFC] Add querying AA to shouldSpecializeCallSiteForCallee

The callback might require an AA, e.g., to ask other AAs for information
in a way that will enfore dependences.

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Transforms/IPO/Attributor.h b/llvm/include/llvm/Transforms/IPO/Attributor.h
index 004570344af2f6..fd6f2e20cd952f 100644
--- a/llvm/include/llvm/Transforms/IPO/Attributor.h
+++ b/llvm/include/llvm/Transforms/IPO/Attributor.h
@@ -1419,7 +1419,8 @@ struct AttributorConfig {
 
   /// Callback function to determine if an indirect call targets should be made
   /// direct call targets (with an if-cascade).
-  std::function<bool(Attributor &A, CallBase &CB, Function &AssummedCallee)>
+  std::function<bool(Attributor &A, const AbstractAttribute &AA, CallBase &CB,
+                     Function &AssummedCallee)>
       IndirectCalleeSpecializationCallback = nullptr;
 
   /// Helper to update an underlying call graph and to delete functions.
@@ -1695,10 +1696,11 @@ struct Attributor {
 
   /// Return true if we should specialize the call site \b CB for the potential
   /// callee \p Fn.
-  bool shouldSpecializeCallSiteForCallee(CallBase &CB, Function &Callee) {
+  bool shouldSpecializeCallSiteForCallee(const AbstractAttribute &AA,
+                                         CallBase &CB, Function &Callee) {
     return Configuration.IndirectCalleeSpecializationCallback
-               ? Configuration.IndirectCalleeSpecializationCallback(*this, CB,
-                                                                    Callee)
+               ? Configuration.IndirectCalleeSpecializationCallback(*this, AA,
+                                                                    CB, Callee)
                : true;
   }
 

diff  --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index 52d2921fd11a95..69f8952e3aaca0 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -3752,18 +3752,19 @@ static bool runAttributorOnFunctions(InformationCache &InfoCache,
   DenseMap<CallBase *, std::unique_ptr<SmallPtrSet<Function *, 8>>>
       IndirectCalleeTrackingMap;
   if (MaxSpecializationPerCB.getNumOccurrences()) {
-    AC.IndirectCalleeSpecializationCallback = [&](Attributor &, CallBase &CB,
-                                                  Function &Callee) {
-      if (MaxSpecializationPerCB == 0)
-        return false;
-      auto &Set = IndirectCalleeTrackingMap[&CB];
-      if (!Set)
-        Set = std::make_unique<SmallPtrSet<Function *, 8>>();
-      if (Set->size() >= MaxSpecializationPerCB)
-        return Set->contains(&Callee);
-      Set->insert(&Callee);
-      return true;
-    };
+    AC.IndirectCalleeSpecializationCallback =
+        [&](Attributor &, const AbstractAttribute &AA, CallBase &CB,
+            Function &Callee) {
+          if (MaxSpecializationPerCB == 0)
+            return false;
+          auto &Set = IndirectCalleeTrackingMap[&CB];
+          if (!Set)
+            Set = std::make_unique<SmallPtrSet<Function *, 8>>();
+          if (Set->size() >= MaxSpecializationPerCB)
+            return Set->contains(&Callee);
+          Set->insert(&Callee);
+          return true;
+        };
   }
 
   Attributor A(Functions, InfoCache, AC);

diff  --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
index ee63ec337c6c05..8bb8c60e2c4b30 100644
--- a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
+++ b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -12204,7 +12204,7 @@ struct AAIndirectCallInfoCallSite : public AAIndirectCallInfo {
     SmallVector<Function *, 8> SkippedAssumedCallees;
     SmallVector<std::pair<CallInst *, Instruction *>> NewCalls;
     for (Function *NewCallee : AssumedCallees) {
-      if (!A.shouldSpecializeCallSiteForCallee(*CB, *NewCallee)) {
+      if (!A.shouldSpecializeCallSiteForCallee(*this, *CB, *NewCallee)) {
         SkippedAssumedCallees.push_back(NewCallee);
         SpecializedForAllCallees = false;
         continue;