[ThinLTO] Allow importing based on a workload definition

llvm/lib/Transforms/IPO/FunctionImport.cpp

@@ -451,76 +451,79 @@ class WorkloadImportsManager : public ModuleImportsManager {
     auto &ValueInfos = SetIter->second;
     SmallVector<EdgeInfo, 128> GlobWorklist;
     for (auto &VI : llvm::make_early_inc_range(ValueInfos)) {
+      auto It = DefinedGVSummaries.find(VI.getGUID());
+      if (It != DefinedGVSummaries.end() &&
+          IsPrevailing(VI.getGUID(), It->second)) {
+        LLVM_DEBUG(
+            dbgs() << "[Workload] " << VI.name()
+                   << " has the prevailing variant already in the module "
+                   << ModName << ". No need to import\n");
+        continue;
+      }
       auto Candidates =
           qualifyCalleeCandidates(Index, VI.getSummaryList(), ModName);
 
       const GlobalValueSummary *GVS = nullptr;
-      FunctionImporter::ImportFailureReason LastReason =
-          FunctionImporter::ImportFailureReason::None;
-      for (const auto &Candidate : Candidates) {
-        /// We will prefer importing the prevailing candidate, if not, we'll
-        /// still pick the first available candidate. The reason we want to make
-        /// sure we do import the prevailing candidate is because the goal of
-        /// workload-awareness is to enable optimizations specializing the call
-        /// graph of that workload. Suppose a function is already defined in the
-        /// module, but it's not the prevailing variant. Suppose also we do not
-        /// inline it, but we could specialize it to the workload in other ways.
-        /// However, the linker would drop it in the favor of the prevailing
-        /// copy. Instead, by importing the prevailing variant (assuming also
-        /// the use of `-avail-extern-to-local`), we keep the specialization. We
-        /// could alteranatively make the non-prevailing variant local, but the
-        /// prevailing one is also the one for which we would have previously
-        /// collected profiles, making it preferrable.
-        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) {
+      auto PotentialCandidates = llvm::map_range(
+          llvm::make_filter_range(
+              Candidates,
+              [&](const auto &Candidate) {
+                LLVM_DEBUG(dbgs() << "[Workflow] Candidate for " << VI.name()
+                                  << " from " << Candidate.second->modulePath()
+                                  << " ImportFailureReason: "
+                                  << getFailureName(Candidate.first) << "\n");
+                return Candidate.first ==
+                       FunctionImporter::ImportFailureReason::None;
+              }),
+          [](const auto &Candidate) { return Candidate.second; });
+      if (PotentialCandidates.empty()) {
         LLVM_DEBUG(dbgs() << "[Workload] Not importing " << VI.name()
                           << " because can't find eligible Callee. Guid is: "
-                          << Function::getGUID(VI.name())
-                          << ". The reason was: " << getFailureName(LastReason)
-                          << "\n");
+                          << Function::getGUID(VI.name()) << "\n");
         continue;
       }
-      const auto *CFS = cast<FunctionSummary>(GVS->getBaseObject());
-      auto ExportingModule = CFS->modulePath();
+      /// We will prefer importing the prevailing candidate, if not, we'll
+      /// still pick the first available candidate. The reason we want to make
+      /// sure we do import the prevailing candidate is because the goal of
+      /// workload-awareness is to enable optimizations specializing the call
+      /// graph of that workload. Suppose a function is already defined in the
+      /// module, but it's not the prevailing variant. Suppose also we do not
+      /// inline it (in fact, if it were interposable, we can't inline it),
+      /// but we could specialize it to the workload in other ways. However,
+      /// the linker would drop it in the favor of the prevailing copy.
+      /// Instead, by importing the prevailing variant (assuming also the use
+      /// of `-avail-extern-to-local`), we keep the specialization. We could
+      /// alteranatively make the non-prevailing variant local, but the
+      /// prevailing one is also the one for which we would have previously
+      /// collected profiles, making it preferrable.
+      auto PrevailingCandidates = llvm::make_filter_range(
+          PotentialCandidates, [&](const auto *Candidate) {
+            return IsPrevailing(VI.getGUID(), Candidate);
+          });
+      if (PrevailingCandidates.empty()) {
+        if (!llvm::hasSingleElement(PotentialCandidates))
+          LLVM_DEBUG(
+              dbgs()
+              << "[Workload] Found multiple non-prevailing candidates for "
+              << VI.name()
+              << ". This is unexpected. Are module paths passed to the "
+                 "compiler unique for the modules passed to the linker?");
+        GVS = *PotentialCandidates.begin();
+      } else {
+        assert(llvm::hasSingleElement(PrevailingCandidates));
+        GVS = *PrevailingCandidates.begin();
+      }
+
+      auto ExportingModule = GVS->modulePath();
+      // We checked that for the prevailing case, but if we happen to have for
+      // example an internal that's defined in this module, it'd have no
+      // PrevailingCandidates.
       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 (moduleAlreadyHasPreferredDef(DefinedGVSummaries, VI.getGUID(), CFS)) {
-        LLVM_DEBUG(
-            dbgs() << "[Workload] Not importing " << VI.name()
-                   << " because we have a copy already in this module.\n");
-        continue;
-      }
-
       LLVM_DEBUG(dbgs() << "[Workload][Including]" << VI.name() << " from "
                         << ExportingModule << " : "
                         << Function::getGUID(VI.name()) << "\n");
@@ -532,28 +535,6 @@ class WorkloadImportsManager : public ModuleImportsManager {
     LLVM_DEBUG(dbgs() << "[Workload] Done\n");
   }
 
-  bool moduleAlreadyHasPreferredDef(const GVSummaryMapTy &DefinedGVSummaries,
-                                    Function::GUID Guid,
-                                    const FunctionSummary *Candidate) {
-    auto DefinedSummary = DefinedGVSummaries.find(Guid);
-    if (DefinedSummary == DefinedGVSummaries.end())
-      return false;
-
-    // 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
-                        << ": non-prevailing preexisting definition in module. "
-                           "Importing from "
-                        << Candidate->modulePath() << "\n");
-      return false;
-    }
-    LLVM_DEBUG(dbgs() << "[Workload] " << Guid
-                      << ": ignored! Target already in destination module.\n");
-    return true;
-  }
-
 public:
   WorkloadImportsManager(
       function_ref<bool(GlobalValue::GUID, const GlobalValueSummary *)>
@@ -568,15 +549,14 @@ class WorkloadImportsManager : public ModuleImportsManager {
     for (auto &I : Index) {
       ValueInfo VI = Index.getValueInfo(I);
       if (!NameToValueInfo.insert(std::make_pair(VI.name(), VI)).second)
-        AmbiguousNames.insert(VI.name());
+        LLVM_DEBUG(AmbiguousNames.insert(VI.name()));
     }
     auto DbgReportIfAmbiguous = [&](StringRef Name) {
-      if (AmbiguousNames.count(Name) > 0)
-        LLVM_DEBUG(
-            dbgs()
-            << "[Workload] Function name " << Name
-            << " present in the workload definition is ambiguous. Consider "
-               "compiling with -funique-internal-linkage-names.");
+      LLVM_DEBUG(if (AmbiguousNames.count(Name) > 0) {
+        dbgs() << "[Workload] Function name " << Name
+               << " present in the workload definition is ambiguous. Consider "
+                  "compiling with -funique-internal-linkage-names.";
+      });
     };
     std::error_code EC;
     auto BufferOrErr = MemoryBuffer::getFileOrSTDIN(WorkloadDefinitions);
@@ -588,7 +568,11 @@ class WorkloadImportsManager : public ModuleImportsManager {
     std::map<std::string, std::vector<std::string>> WorkloadDefs;
     json::Path::Root NullRoot;
     // The JSON is supposed to contain a dictionary matching the type of
-    // WorkloadDefs.
+    // WorkloadDefs. For example:
+    // {
+    //   "rootFunction_1": ["function_to_import_1", "function_to_import_2"],
+    //   "rootFunction_2": ["function_to_import_3", "function_to_import_4"]
+    // }
     auto Parsed = json::parse(Buffer->getBuffer());
     if (!Parsed)
       report_fatal_error(Parsed.takeError());
@@ -627,14 +611,14 @@ class WorkloadImportsManager : public ModuleImportsManager {
         }
         Set.insert(ElemIt->second);
       }
-      LLVM_DEBUG( //
-          dbgs() << "[Workload] Root: " << Root << " we have " << Set.size()
-                 << " distinct callees.\n";
-          for (const auto &VI
-               : Set) {
-            dbgs() << "[Workload] Root: " << Root
-                   << " Would include: " << VI.getGUID() << "\n";
-          });
+      LLVM_DEBUG({
+        dbgs() << "[Workload] Root: " << Root << " we have " << Set.size()
+               << " distinct callees.\n";
+        for (const auto &VI : Set) {
+          dbgs() << "[Workload] Root: " << Root
+                 << " Would include: " << VI.getGUID() << "\n";
+        }
+      });
     }
   }
 };