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

Add a comment somewhere around here briefly documenting the expected format of each line

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