[Attributor] Introduce the closed world flag

The Attributor user can now set the closed world flag
(`AttributorConfig.IsClosedWorldModule` or
`-attributor-assume-closed-world`) in order to specialize call edges
based only on available callees. That means, we assume all functions are
known and hence all potential callees must be declared/defined in the
module. We will use this for GPUs and LTO cases, but for now the user
has to set it via a flag.

Author: jdoerfert

llvm/include/llvm/Transforms/IPO/Attributor.h

@@ -1316,6 +1316,11 @@ struct InformationCache {
     return TargetTriple.isAMDGPU() || TargetTriple.isNVPTX();
   }
 
+  /// Return all functions that might be called indirectly, only valid for
+  /// closed world modules (see isClosedWorldModule).
+  const ArrayRef<Function *>
+  getIndirectlyCallableFunctions(Attributor &A) const;
+
 private:
   struct FunctionInfo {
     ~FunctionInfo();
@@ -1348,6 +1353,10 @@ struct InformationCache {
     return *FI;
   }
 
+  /// Vector of functions that might be callable indirectly, i.a., via a
+  /// function pointer.
+  SmallVector<Function *> IndirectlyCallableFunctions;
+
   /// Initialize the function information cache \p FI for the function \p F.
   ///
   /// This method needs to be called for all function that might be looked at
@@ -1413,6 +1422,10 @@ struct AttributorConfig {
   /// Flag to determine if we should skip all liveness checks early on.
   bool UseLiveness = true;
 
+  /// Flag to indicate if the entire world is contained in this module, that
+  /// is, no outside functions exist.
+  bool IsClosedWorldModule = false;
+
   /// Callback function to be invoked on internal functions marked live.
   std::function<void(Attributor &A, const Function &F)> InitializationCallback =
       nullptr;
@@ -1483,9 +1496,7 @@ struct Attributor {
   /// \param Configuration The Attributor configuration which determines what
   ///                      generic features to use.
   Attributor(SetVector<Function *> &Functions, InformationCache &InfoCache,
-             AttributorConfig Configuration)
-      : Allocator(InfoCache.Allocator), Functions(Functions),
-        InfoCache(InfoCache), Configuration(Configuration) {}
+             AttributorConfig Configuration);
 
   ~Attributor();
 
@@ -1704,6 +1715,10 @@ struct Attributor {
                : true;
   }
 
+  /// Return true if the module contains the whole world, thus, no outside
+  /// functions exist.
+  bool isClosedWorldModule() const;
+
   /// Return true if we derive attributes for \p Fn
   bool isRunOn(Function &Fn) const { return isRunOn(&Fn); }
   bool isRunOn(Function *Fn) const {

llvm/lib/Transforms/IPO/Attributor.cpp

@@ -175,6 +175,10 @@ static cl::opt<bool> SimplifyAllLoads("attributor-simplify-all-loads",
                                       cl::desc("Try to simplify all loads."),
                                       cl::init(true));
 
+static cl::opt<bool> CloseWorldAssumption(
+    "attributor-assume-closed-world", cl::Hidden,
+    cl::desc("Should a closed world be assumed, or not. Default if not set."));
+
 /// Logic operators for the change status enum class.
 ///
 ///{
@@ -1057,6 +1061,21 @@ ChangeStatus AbstractAttribute::update(Attributor &A) {
   return HasChanged;
 }
 
+Attributor::Attributor(SetVector<Function *> &Functions,
+                       InformationCache &InfoCache,
+                       AttributorConfig Configuration)
+    : Allocator(InfoCache.Allocator), Functions(Functions),
+      InfoCache(InfoCache), Configuration(Configuration) {
+  if (!isClosedWorldModule())
+    return;
+  for (Function *Fn : Functions)
+    if (Fn->hasAddressTaken(/*PutOffender=*/nullptr,
+                            /*IgnoreCallbackUses=*/false,
+                            /*IgnoreAssumeLikeCalls=*/true,
+                            /*IgnoreLLVMUsed=*/true))
+      InfoCache.IndirectlyCallableFunctions.push_back(Fn);
+}
+
 bool Attributor::getAttrsFromAssumes(const IRPosition &IRP,
                                      Attribute::AttrKind AK,
                                      SmallVectorImpl<Attribute> &Attrs) {
@@ -3251,6 +3270,12 @@ InformationCache::FunctionInfo::~FunctionInfo() {
     It.getSecond()->~InstructionVectorTy();
 }
 
+const ArrayRef<Function *>
+InformationCache::getIndirectlyCallableFunctions(Attributor &A) const {
+  assert(A.isClosedWorldModule() && "Cannot see all indirect callees!");
+  return IndirectlyCallableFunctions;
+}
+
 void Attributor::recordDependence(const AbstractAttribute &FromAA,
                                   const AbstractAttribute &ToAA,
                                   DepClassTy DepClass) {
@@ -3586,6 +3611,12 @@ void Attributor::identifyDefaultAbstractAttributes(Function &F) {
   assert(Success && "Expected the check call to be successful!");
 }
 
+bool Attributor::isClosedWorldModule() const {
+  if (CloseWorldAssumption.getNumOccurrences())
+    return CloseWorldAssumption;
+  return isModulePass() && Configuration.IsClosedWorldModule;
+}
+
 /// Helpers to ease debugging through output streams and print calls.
 ///
 ///{

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

@@ -12065,11 +12065,22 @@ struct AAIndirectCallInfoCallSite : public AAIndirectCallInfo {
   /// See AbstractAttribute::initialize(...).
   void initialize(Attributor &A) override {
     auto *MD = getCtxI()->getMetadata(LLVMContext::MD_callees);
-    if (!MD)
+    if (!MD && !A.isClosedWorldModule())
       return;
-    for (const auto &Op : MD->operands())
-      if (Function *Callee = mdconst::dyn_extract_or_null<Function>(Op))
-        PotentialCallees.insert(Callee);
+
+    if (MD) {
+      for (const auto &Op : MD->operands())
+        if (Function *Callee = mdconst::dyn_extract_or_null<Function>(Op))
+          PotentialCallees.insert(Callee);
+    } else if (A.isClosedWorldModule()) {
+      ArrayRef<Function *> IndirectlyCallableFunctions =
+          A.getInfoCache().getIndirectlyCallableFunctions(A);
+      PotentialCallees.insert(IndirectlyCallableFunctions.begin(),
+                              IndirectlyCallableFunctions.end());
+    }
+
+    if (PotentialCallees.empty())
+      indicateOptimisticFixpoint();
   }
 
   ChangeStatus updateImpl(Attributor &A) override {