FunctionImport: add a progressive heuristic to limit importing too deep in the callgraph

[mehdi_amini]

Summary

The current function importer will walk the callgraph, importing
transitively any callee that is below the threshold. This can
lead to import very deep which is costly in compile time and not
necessarily beneficial as most of the inline would happen in
imported function and not necessarilly in user code.

The actual factor has been carefully chosen by flipping a coin ;)
Some tuning need to be done (just at the existing limiting threshold).

[tejohnson]

This is a reasonable heuristic until we have more sophisticated cost/benefit analysis, especially with PGO data. However, there is an issue due to the depth first traversal order, see below.

lib/Transforms/IPO/FunctionImport.cpp

291

I think we can get missed importing opportunities due to the depth-first nature of the current traversal. We push_back newly identified external calls (which have the reduced threshold), then pop_back_val when traversing the Worklist above, resulting in a depth-first traversal of the call graph. What this means is that if we first encounter a call to a given external function foo() deep in a call chain, it will get added with a very low threshold. It may be that another external call to foo() is encountered higher up the call chain in something added to the worklist earlier and popped later, but since we skip any functions that were already seen (in the CalledFunctions set), we will miss the opportunity to add it with the higher threshold and possibly lose out on the import.

To fix this the tracking of what is or was already in the worklist needs to be more sophisticated. Perhaps a simple fix would be to change the CalledFunctions set to a map from function name to the max threshold it was added to the worklist with. Then if we encounter another call to the same function, if it currently has a higher threshold than recorded in the CalledFunctions map, add it to the worklist again and updated the map entry. Then each function popped off the worklist should be compared against the existing ModuleToFunctionsToImportMap entry to see if we have already decided to import it since it may end up in the worklist multiple times.

[mehdi_amini]

lib/Transforms/IPO/FunctionImport.cpp

291

That's a very good point, improving the CalledFunctions data structure seems a good solution to me.

[mehdi_amini]

Fix the issue reported by Teresa: modified the test to expose the DFS issue
(test as it is now is not passing on the previous revision)

[tejohnson]

Couple nits on comments, otherwise LGTM.

lib/Transforms/IPO/FunctionImport.cpp

309

"as long as" doesn't sound right. Did you mean "as well as"?

test/Transforms/FunctionImport/Inputs/adjustable_threshold.ll

21

Size must start out higher than 5 since it looks like the evolution factor had to be reduced to 0.5 to not get this imported.

[mehdi_amini]

lib/Transforms/IPO/FunctionImport.cpp

309

Yes!

test/Transforms/FunctionImport/Inputs/adjustable_threshold.ll

21

Not sure I understand what you mean? I want it to be imported with the default 0.7 factor, but not with the 0.5 one.

[tejohnson]

test/Transforms/FunctionImport/Inputs/adjustable_threshold.ll

21

Nevermind, I see what is going on now.

[mehdi_amini]

Update before committing

[mehdi_amini]

Thanks Teresa!

r260466

lib/Transforms/IPO/FunctionImport.cpp

 
 #define DEBUG_TYPE "function-import"
 
 /// Limit on instruction count of imported functions.
 static cl::opt<unsigned> ImportInstrLimit(
     "import-instr-limit", cl::init(100), cl::Hidden, cl::value_desc("N"),
     cl::desc("Only import functions with less than N instructions"));
 
+static cl::opt<float>
+    ImportInstrFactor("import-instr-evolution-factor", cl::init(0.7),
+                      cl::Hidden, cl::value_desc("x"),
+                      cl::desc("As we import functions, multiply the "
+                               "`import-instr-limit` threshold by this factor "
+                               "before processing newly imported functions"));
+
 // Load lazily a module from \p FileName in \p Context.
 static std::unique_ptr<Module> loadFile(const std::string &FileName,
                                         LLVMContext &Context) {
   SMDiagnostic Err;
   DEBUG(dbgs() << "Loading '" << FileName << "'\n");
   // Metadata isn't loaded until functions are imported, to minimize
   // the memory overhead.
   std::unique_ptr<Module> Result =
       getLazyIRFileModule(FileName, Err, Context,
                           /* ShouldLazyLoadMetadata = */ true);
   if (!Result) {
     Err.print("function-import", errs());
     return nullptr;
   }
 
   return Result;
 }
 
 namespace {
+
+/// Track functions already seen using a map that record the current
+/// Threshold and the importing decision. Since the traversal of the call graph
+/// is DFS, we can revisit a function a second time with a higher threshold. In
+/// this case and if the function was not imported the first time, it is added
+/// back to the worklist with the new threshold
+using VisitedFunctionTrackerTy = StringMap<std::pair<unsigned, bool>>;
+
 /// Helper to load on demand a Module from file and cache it for subsequent
 /// queries. It can be used with the FunctionImporter.
 class ModuleLazyLoaderCache {
   /// Cache of lazily loaded module for import.
   StringMap<std::unique_ptr<Module>> ModuleMap;
 
   /// Retrieve a Module from the cache or lazily load it on demand.
   std::function<std::unique_ptr<Module>(StringRef FileName)> createLazyModule;
Module &ModuleLazyLoaderCache::operator()(StringRef Identifier) {
   auto &Module = ModuleMap[Identifier];
   if (!Module)
     Module = createLazyModule(Identifier);
   return *Module;
 }
 } // anonymous namespace
 
 /// Walk through the instructions in \p F looking for external
-/// calls not already in the \p CalledFunctions set. If any are
+/// calls not already in the \p VisitedFunctions map. If any are
 /// found they are added to the \p Worklist for importing.
-static void findExternalCalls(const Module &DestModule, Function &F,
-                              const FunctionInfoIndex &Index,
-                              StringSet<> &CalledFunctions,
-                              SmallVector<StringRef, 64> &Worklist) {
+static void findExternalCalls(
+    const Module &DestModule, Function &F, const FunctionInfoIndex &Index,
+    VisitedFunctionTrackerTy &VisitedFunctions, unsigned Threshold,
+    SmallVectorImpl<std::pair<StringRef, unsigned>> &Worklist) {
   // We need to suffix internal function calls imported from other modules,
   // prepare the suffix ahead of time.
   std::string Suffix;
   if (F.getParent() != &DestModule)
     Suffix =
         (Twine(".llvm.") +
          Twine(Index.getModuleId(F.getParent()->getModuleIdentifier()))).str();
 
for (auto &I : BB) {
         // Rename internal functions
         if (CalledFunction->hasInternalLinkage()) {
           ImportedName = Renamed;
         }
         // Compute the global identifier used in the function index.
         auto CalledFunctionGlobalID = Function::getGlobalIdentifier(
             CalledFunction->getName(), CalledFunction->getLinkage(),
             CalledFunction->getParent()->getSourceFileName());
-        auto It = CalledFunctions.insert(CalledFunctionGlobalID);
+
+        auto CalledFunctionInfo = std::make_pair(Threshold, false);
+        auto It = VisitedFunctions.insert(
+            std::make_pair(CalledFunctionGlobalID, CalledFunctionInfo));
         if (!It.second) {
-          // This is a call to a function we already considered, skip.
+          // This is a call to a function we already considered, if the function
+          // has been imported the first time, or if the current threshold is
+          // not higher, skip it.
+          auto &FunctionInfo = It.first->second;
+          if (FunctionInfo.second || FunctionInfo.first >= Threshold)
             continue;
+          It.first->second = CalledFunctionInfo;
         }
         // Ignore functions already present in the destination module
         auto *SrcGV = DestModule.getNamedValue(ImportedName);
         if (SrcGV) {
           if (GlobalAlias *SGA = dyn_cast<GlobalAlias>(SrcGV))
             SrcGV = SGA->getBaseObject();
           assert(isa<Function>(SrcGV) && "Name collision during import");
           if (!cast<Function>(SrcGV)->isDeclaration()) {
             DEBUG(dbgs() << DestModule.getModuleIdentifier() << ": Ignoring "
                          << ImportedName << " already in DestinationModule\n");
             continue;
           }
         }
 
-        Worklist.push_back(It.first->getKey());
+        Worklist.push_back(std::make_pair(It.first->getKey(), Threshold));
         DEBUG(dbgs() << DestModule.getModuleIdentifier()
                      << ": Adding callee for : " << ImportedName << " : "
                      << F.getName() << "\n");
       }
     }
   }
 }
 
 // Helper function: given a worklist and an index, will process all the worklist
 // and decide what to import based on the summary information.
 //
 // Nothing is actually imported, functions are materialized in their source
 // module and analyzed there.
 //
 // \p ModuleToFunctionsToImportMap is filled with the set of Function to import
 // per Module.
-static void GetImportList(Module &DestModule,
-                          SmallVector<StringRef, 64> &Worklist,
-                          StringSet<> &CalledFunctions,
-                          std::map<StringRef, DenseSet<const GlobalValue *>>
-                              &ModuleToFunctionsToImportMap,
-                          const FunctionInfoIndex &Index,
-                          ModuleLazyLoaderCache &ModuleLoaderCache) {
+static void
+GetImportList(Module &DestModule,
+              SmallVectorImpl<std::pair<StringRef, unsigned>> &Worklist,
+              VisitedFunctionTrackerTy &VisitedFunctions,
+              std::map<StringRef, DenseSet<const GlobalValue *>> &
+                  ModuleToFunctionsToImportMap,
   while (!Worklist.empty()) {
-    auto CalledFunctionName = Worklist.pop_back_val();
+    StringRef CalledFunctionName;
+    unsigned Threshold;
+    std::tie(CalledFunctionName, Threshold) = Worklist.pop_back_val();
     DEBUG(dbgs() << DestModule.getModuleIdentifier() << ": Process import for "
-                 << CalledFunctionName << "\n");
+                 << CalledFunctionName << " with Threshold " << Threshold
+                 << "\n");
 
     // Try to get a summary for this function call.
     auto InfoList = Index.findFunctionInfoList(CalledFunctionName);
     if (InfoList == Index.end()) {
       DEBUG(dbgs() << DestModule.getModuleIdentifier() << ": No summary for "
                    << CalledFunctionName << " Ignoring.\n");
       continue;
     }
     assert(!InfoList->second.empty() && "No summary, error at import?");
 
     // Comdat can have multiple entries, FIXME: what do we do with them?
     auto &Info = InfoList->second[0];
     assert(Info && "Nullptr in list, error importing summaries?\n");
 
     auto *Summary = Info->functionSummary();
     if (!Summary) {
       // FIXME: in case we are lazyloading summaries, we can do it now.
       DEBUG(dbgs() << DestModule.getModuleIdentifier()
                    << ": Missing summary for  " << CalledFunctionName
                    << ", error at import?\n");
       llvm_unreachable("Missing summary");
     }
 
-    if (Summary->instCount() > ImportInstrLimit) {
+    if (Summary->instCount() > Threshold) {
       DEBUG(dbgs() << DestModule.getModuleIdentifier() << ": Skip import of "
                    << CalledFunctionName << " with " << Summary->instCount()
-                   << " instructions (limit " << ImportInstrLimit << ")\n");
+                   << " instructions (limit " << Threshold << ")\n");
       continue;
     }
 
+    // Mark the function as imported in the VisitedFunctions tracker
+    assert(VisitedFunctions.count(CalledFunctionName));
+    VisitedFunctions[CalledFunctionName].second = true;
+
     // Get the module path from the summary.
     auto ModuleIdentifier = Summary->modulePath();
     DEBUG(dbgs() << DestModule.getModuleIdentifier() << ": Importing "
                  << CalledFunctionName << " from " << ModuleIdentifier << "\n");
 
     auto &SrcModule = ModuleLoaderCache(ModuleIdentifier);
 
     // The function that we will import!
if (SGV->hasWeakAnyLinkage()) {
       continue;
     }
 
     // Add the function to the import list
     auto &Entry = ModuleToFunctionsToImportMap[SrcModule.getModuleIdentifier()];
     Entry.insert(F);
 
     // Process the newly imported functions and add callees to the worklist.
+    // Adjust the threshold
+    Threshold = Threshold * ImportInstrFactor;
     F->materialize();
-    findExternalCalls(DestModule, *F, Index, CalledFunctions, Worklist);
+    findExternalCalls(DestModule, *F, Index, VisitedFunctions, Threshold,
+                      Worklist);
   }
 }
 
 // Automatically import functions in Module \p DestModule based on the summaries
 // index.
 //
 // The current implementation imports every called functions that exists in the
 // summaries index.
 bool FunctionImporter::importFunctions(Module &DestModule) {
   DEBUG(dbgs() << "Starting import for Module "
                << DestModule.getModuleIdentifier() << "\n");
   unsigned ImportedCount = 0;
 
-  /// First step is collecting the called external functions.
-  StringSet<> CalledFunctions;
-  SmallVector<StringRef, 64> Worklist;
+  // First step is collecting the called external functions.
+  // We keep the function name as well as the import threshold for its callees.
+  VisitedFunctionTrackerTy VisitedFunctions;
+  SmallVector<std::pair<StringRef, unsigned>, 64> Worklist;
   for (auto &F : DestModule) {
     if (F.isDeclaration() || F.hasFnAttribute(Attribute::OptimizeNone))
       continue;
-    findExternalCalls(DestModule, F, Index, CalledFunctions, Worklist);
+    findExternalCalls(DestModule, F, Index, VisitedFunctions, ImportInstrLimit,
+                      Worklist);
   }
   if (Worklist.empty())
     return false;
 
   /// Second step: for every call to an external function, try to import it.
 
   // Linker that will be used for importing function
   Linker TheLinker(DestModule);
 
   // Map of Module -> List of Function to import from the Module
   std::map<StringRef, DenseSet<const GlobalValue *>>
       ModuleToFunctionsToImportMap;
 
   // Analyze the summaries and get the list of functions to import by
   // populating ModuleToFunctionsToImportMap
   ModuleLazyLoaderCache ModuleLoaderCache(ModuleLoader);
-  GetImportList(DestModule, Worklist, CalledFunctions,
+  GetImportList(DestModule, Worklist, VisitedFunctions,
                 ModuleToFunctionsToImportMap, Index, ModuleLoaderCache);
   assert(Worklist.empty() && "Worklist hasn't been flushed in GetImportList");
 
   // Do the actual import of functions now, one Module at a time
   for (auto &FunctionsToImportPerModule : ModuleToFunctionsToImportMap) {
     // Get the module for the import
     auto &FunctionsToImport = FunctionsToImportPerModule.second;
     std::unique_ptr<Module> SrcModule =

test/Transforms/FunctionImport/Inputs/adjustable_threshold.ll

+define void @globalfunc1() {
+entry:
+  call void @trampoline()
+  ret void
+}
+; Adds an artificial level in the call graph to reduce the importing threshold
+define void @trampoline() {
+entry:
+  call void @largefunction()
+  ret void
+}
+
+define void @globalfunc2() {
+entry:
+  call void @largefunction()
+  ret void
+}
+
+
+; Size is 5: if two layers below in the call graph the threshold will be 4,
+; but if only one layer below the threshold will be 7.
+define void @largefunction() {
+  entry:
+  call void @staticfunc2()
+  call void @staticfunc2()
+  call void @staticfunc2()
+  call void @staticfunc2()
+  call void @staticfunc2()
+  ret void
+}
+
+define internal void @staticfunc2() {
+entry:
+  ret void
+}
+
+

test/Transforms/FunctionImport/adjustable_threshold.ll

+; Do setup work for all below tests: generate bitcode and combined index
+; RUN: llvm-as -function-summary %s -o %t.bc
+; RUN: llvm-as -function-summary %p/Inputs/adjustable_threshold.ll -o %t2.bc
+; RUN: llvm-lto -thinlto -o %t3 %t.bc %t2.bc
+
+; Test import with default progressive instruction factor
+; RUN: opt -function-import -summary-file %t3.thinlto.bc %s -import-instr-limit=10 -S | FileCheck %s --check-prefix=CHECK --check-prefix=INSTLIM-DEFAULT
+; INSTLIM-DEFAULT: call void @staticfunc2.llvm.2()
+
+; Test import with a reduced progressive instruction factor
+; RUN: opt -function-import -summary-file %t3.thinlto.bc %s -import-instr-limit=10 -import-instr-evolution-factor=0.5 -S | FileCheck %s --check-prefix=CHECK --check-prefix=INSTLIM-PROGRESSIVE
+; INSTLIM-PROGRESSIVE-NOT: call void @staticfunc
+
+
+
+declare void @globalfunc1()
+declare void @globalfunc2()
+
+define void @entry() {
+entry:
+; Call site are processed in reversed order!
+
+; On the direct call, we reconsider @largefunction with a higher threshold and
+; import it
+  call void @globalfunc2()
+; When importing globalfunc1, the threshold was limited and @largefunction was
+; not imported.
+  call void @globalfunc1()
+  ret void
+}
+