[llvm] [ThinLTO] Allow importing based on a workload definition (PR #74545)

[Teresa Johnson via llvm-commits]

================
@@ -369,29 +373,238 @@ class GlobalsImporter final {
   }
 };
 
+static const char *getFailureName(FunctionImporter::ImportFailureReason Reason);
+
 /// Determine the list of imports and exports for each module.
-class ModuleImportsManager final {
+class ModuleImportsManager {
+protected:
   function_ref<bool(GlobalValue::GUID, const GlobalValueSummary *)>
       IsPrevailing;
   const ModuleSummaryIndex &Index;
   DenseMap<StringRef, FunctionImporter::ExportSetTy> *const ExportLists;
 
-public:
   ModuleImportsManager(
       function_ref<bool(GlobalValue::GUID, const GlobalValueSummary *)>
           IsPrevailing,
       const ModuleSummaryIndex &Index,
       DenseMap<StringRef, FunctionImporter::ExportSetTy> *ExportLists = nullptr)
       : IsPrevailing(IsPrevailing), Index(Index), ExportLists(ExportLists) {}
 
+public:
+  virtual ~ModuleImportsManager() = default;
+
   /// Given the list of globals defined in a module, compute the list of imports
   /// as well as the list of "exports", i.e. the list of symbols referenced from
   /// another module (that may require promotion).
-  void computeImportForModule(const GVSummaryMapTy &DefinedGVSummaries,
-                              StringRef ModName,
-                              FunctionImporter::ImportMapTy &ImportList);
+  virtual void
+  computeImportForModule(const GVSummaryMapTy &DefinedGVSummaries,
+                         StringRef ModName,
+                         FunctionImporter::ImportMapTy &ImportList);
+
+  static std::unique_ptr<ModuleImportsManager>
+  create(function_ref<bool(GlobalValue::GUID, const GlobalValueSummary *)>
+             IsPrevailing,
+         const ModuleSummaryIndex &Index,
+         DenseMap<StringRef, FunctionImporter::ExportSetTy> *ExportLists =
+             nullptr);
 };
 
+class WorkloadImportsManager : public ModuleImportsManager {
+  // Keep a module name -> defined value infos association. We use it to
+  // determine if a module's import list should be done by the base
+  // ModuleImportsManager or by us.
+  StringMap<DenseSet<ValueInfo>> Workloads;
+
+  void
+  computeImportForModule(const GVSummaryMapTy &DefinedGVSummaries,
+                         StringRef ModName,
+                         FunctionImporter::ImportMapTy &ImportList) override {
+    auto SetIter = Workloads.find(ModName);
+    if (SetIter == Workloads.end()) {
+      LLVM_DEBUG(dbgs() << "[Workload] " << ModName
+                        << " does not contain the root of any context.\n");
+      return ModuleImportsManager::computeImportForModule(DefinedGVSummaries,
+                                                          ModName, ImportList);
+    }
+    LLVM_DEBUG(dbgs() << "[Workload] " << ModName
+                      << " contains the root(s) of context(s).\n");
+
+    GlobalsImporter GVI(Index, DefinedGVSummaries, IsPrevailing, ImportList,
+                        ExportLists);
+    auto &ValueInfos = SetIter->second;
+    SmallVector<EdgeInfo, 128> GlobWorklist;
+    for (auto &VI : llvm::make_early_inc_range(ValueInfos)) {
+      auto Candidates =
+          qualifyCalleeCandidates(Index, VI.getSummaryList(), ModName);
+
+      const GlobalValueSummary *GVS = nullptr;
+      FunctionImporter::ImportFailureReason LastReason =
+          FunctionImporter::ImportFailureReason::None;
+      for (const auto &Candidate : Candidates) {
+        LastReason = Candidate.first;
+        if (Candidate.first == FunctionImporter::ImportFailureReason::None) {
+          const bool Prevailing = IsPrevailing(VI.getGUID(), Candidate.second);
+          if (Prevailing || !GVS) {
+            if (!GVS && !Prevailing)
+              LLVM_DEBUG(dbgs()
+                         << "[Workload] Considering " << VI.name() << " from "
+                         << Candidate.second->modulePath() << " with linkage "
+                         << Candidate.second->linkage()
+                         << " although it's not prevailing, but it's the "
+                            "first available candidate.\n");
+            GVS = Candidate.second;
+            if (Prevailing) {
+              LLVM_DEBUG(dbgs()
+                         << "[Workload] Considering " << VI.name() << " from "
+                         << GVS->modulePath() << " with linkage "
+                         << GVS->linkage() << " because it's prevailing.\n");
+              break;
+            }
+          } else {
+            LLVM_DEBUG(dbgs() << "[Workload] Skipping " << VI.name() << " from "
+                              << Candidate.second->modulePath()
+                              << " with linkage " << Candidate.second->linkage()
+                              << " because it's not prevailing\n");
+          }
+        }
+      }
+      if (!GVS) {
+        LLVM_DEBUG(dbgs() << "[Workload] Not importing " << VI.name()
+                          << " because can't select Callee. Guid is: "
+                          << Function::getGUID(VI.name())
+                          << ". The reason was: " << getFailureName(LastReason)
+                          << "\n");
+        continue;
+      }
+      const auto *CFS = cast<FunctionSummary>(GVS->getBaseObject());
+      auto ExportingModule = CFS->modulePath();
+      if (ExportingModule == ModName) {
+        LLVM_DEBUG(dbgs() << "[Workload] Not importing " << VI.name()
+                          << " because its defining module is the same as the "
+                             "current module\n");
+        continue;
+      }
+      if (!shouldImport(DefinedGVSummaries, VI.getGUID(), CFS)) {
+        LLVM_DEBUG(dbgs() << "[Workload] Not importing " << VI.name()
+                          << " because we have a local copy.\n");
+        continue;
+      }
+
+      LLVM_DEBUG(dbgs() << "[Workload][Including]" << VI.name() << " from "
+                        << ExportingModule << " : "
+                        << Function::getGUID(VI.name()) << "\n");
+      ImportList[ExportingModule].insert(VI.getGUID());
+      GVI.onImportingSummary(*GVS);
+      if (ExportLists)
+        (*ExportLists)[ExportingModule].insert(VI);
+    }
+    LLVM_DEBUG(dbgs() << "[Workload] Done\n");
+  }
+
+  bool shouldImport(const GVSummaryMapTy &DefinedGVSummaries,
+                    Function::GUID Guid, const GlobalValueSummary *Candidate) {
+    auto DefinedSummary = DefinedGVSummaries.find(Guid);
+    if (DefinedSummary == DefinedGVSummaries.end())
+      return true;
+
+    // See shouldImportGlobal for the justificaton of the isInterposableLinkage.
+    if (!IsPrevailing(Guid, DefinedSummary->second) &&
+        GlobalValue::isInterposableLinkage(DefinedSummary->second->linkage()) &&
+        IsPrevailing(Guid, Candidate)) {
+      LLVM_DEBUG(dbgs() << "[Workload] " << Guid
+                        << ": local non-prevailing in module. Importing from "
+                        << Candidate->modulePath() << "\n");
+      return true;
+    }
+    LLVM_DEBUG(dbgs() << "[Workload] " << Guid
+                      << ": ignored! Target already in destination module.\n");
+    return false;
+  }
+
+public:
+  WorkloadImportsManager(
+      function_ref<bool(GlobalValue::GUID, const GlobalValueSummary *)>
+          IsPrevailing,
+      const ModuleSummaryIndex &Index,
+      DenseMap<StringRef, FunctionImporter::ExportSetTy> *ExportLists)
+      : ModuleImportsManager(IsPrevailing, Index, ExportLists) {
+    StringMap<ValueInfo> CtxGuidToValueInfo;
+    for (auto &I : Index) {
+      ValueInfo VI(Index.haveGVs(), &I);
+      CtxGuidToValueInfo[VI.name()] = VI;
+    }
+    std::error_code EC;
+    auto BufferOrErr = MemoryBuffer::getFileOrSTDIN(WorkloadDefinitions);
+    if (std::error_code EC = BufferOrErr.getError()) {
+      report_fatal_error("Failed to open context file");
+      return;
+    }
+    auto Buffer = std::move(BufferOrErr.get());
+    std::map<std::string, std::vector<std::string>> WorkloadDefs;
+    json::Path::Root NullRoot;
+    auto Parsed = json::parse(Buffer->getBuffer());
----------------
teresajohnson wrote:

The comment just says JSON but that is a pretty generic format. Can you add an example (maybe where you document the option, then refer to it here), otherwise someone would have to dig through the code or look at the test case to specifically see what is expected.

https://github.com/llvm/llvm-project/pull/74545