[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.

Author: jdoerfert

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;
   }
 

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);

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;