[ThinLTO/gold] Enable summary-based internalization

[tejohnson]

Summary

Enable existing summary-based importing support in the gold-plugin.

[pcc]

Test cases please (for the changes to FunctionImport and the gold plugin).

[tejohnson]

• Add new tests.

[tejohnson]

In D21080#451305, @pcc wrote:

Test cases please (for the changes to FunctionImport and the gold plugin).

Done - 2 test cases added.

[mehdi_amini]

(Forgot to click submit in phab)

lib/Transforms/IPO/FunctionImport.cpp

534

Note: we plan to change the alias representation to https://llvm.org/bugs/show_bug.cgi?id=27866 ; which I believe will make this issue obsolete. It may even remove the need for the aliases representation in the summary (not sure what are the implications about comdats).

Also, this seems like a bug fix that could be exposed with llvm-lto, and thus be split out of this patch?

[mehdi_amini]

tools/gold/gold-plugin.cpp

1326

Could these contain directly GUIDs?

1367

Isn't there an API to subtract a set from another one?

[tejohnson]

lib/Transforms/IPO/FunctionImport.cpp

534

Ok, I can add a FIXME to reevaluate this after that bug is fixed.

Let me see if I can provoke this with llvm-lto and split it out.

tools/gold/gold-plugin.cpp

1326

Yes, will switch.

1367

Not that I could find in either StringSet or DenseSet.

[tejohnson]

lib/Transforms/IPO/FunctionImport.cpp

534

Can't reproduce with llvm-lto. I think it is specific to the way the gold plugin handles setting up the list of symbols to link in (it doesn't link in the preempted weak symbol).

[tejohnson]

• Address Mehdi's comments.

[mehdi_amini]

LGTM.

lib/Transforms/IPO/FunctionImport.cpp

if (GS == DefinedGlobals.end()) {
       // FIXME: Eventually we should control promotion instead of promoting
       // and internalizing again.
       StringRef OrigName =
           ModuleSummaryIndex::getOriginalNameBeforePromote(GV.getName());
       std::string OrigId = GlobalValue::getGlobalIdentifier(
           OrigName, GlobalValue::InternalLinkage,
           TheModule.getSourceFileName());
       const auto &GS = DefinedGlobals.find(GlobalValue::getGUID(OrigId));
+      if (GS == DefinedGlobals.end()) {
+        // Also check the original non-promoted non-globalized name. In some
+        // cases a preempted weak value is linked in as a local copy because
+        // it is referenced by an alias (IRLinker::linkGlobalValueProto).
+        // In that case, since it was originally not a local value, it was
+        // recorded in the index using the original name.
+        // FIXME: This may not be needed once PR27866 is fixed.
+        const auto &GS = DefinedGlobals.find(GlobalValue::getGUID(OrigName));
         assert(GS != DefinedGlobals.end());
         Linkage = GS->second->linkage();
+      } else {
+        Linkage = GS->second->linkage();
+      }
     } else
       Linkage = GS->second->linkage();
     return !GlobalValue::isLocalLinkage(Linkage);
   };
 
   // FIXME: See if we can just internalize directly here via linkage changes
   // based on the index, rather than invoking internalizeModule.
   llvm::internalizeModule(TheModule, MustPreserveGV);

test/tools/gold/X86/Inputs/thinlto_alias.ll

+define weak void @weakfunc() {
+entry:
+  ret void
+}

test/tools/gold/X86/Inputs/thinlto_internalize.ll

+target triple = "x86_64-unknown-linux-gnu"
+declare i32 @g()
+define i32 @main() {
+  call i32 @g()
+  ret i32 0
+}

test/tools/gold/X86/thinlto_alias.ll

+; RUN: opt -module-summary %s -o %t.o
+; RUN: opt -module-summary %p/Inputs/thinlto_alias.ll -o %t2.o
+
+; Ensure that a preempted weak symbol that is linked in as a local
+; copy is handled properly. Specifically, the local copy will be promoted,
+; and internalization should be able to use the original non-promoted
+; name to locate the summary (otherwise internalization will abort because
+; it expects to locate summaries for all definitions).
+; Note that gold picks the first copy of weakfunc() as the prevailing one,
+; so listing %t2.o first is sufficient to ensure that this copy is
+; preempted.
+; RUN: %gold -m elf_x86_64 -plugin %llvmshlibdir/LLVMgold.so \
+; RUN:     --plugin-opt=thinlto \
+; RUN:     --plugin-opt=save-temps \
+; RUN:     -o %t3.o %t2.o %t.o
+; RUN: llvm-nm %t3.o | FileCheck %s
+; RUN: llvm-dis %t.o.opt.bc -o - | FileCheck --check-prefix=OPT %s
+; RUN: llvm-dis %t2.o.opt.bc -o - | FileCheck --check-prefix=OPT2 %s
+
+; CHECK-NOT: U f
+; OPT: define hidden void @weakfunc.llvm.0()
+; OPT2: define weak void @weakfunc()
+
+target triple = "x86_64-unknown-linux-gnu"
+
+@weakfuncAlias = alias void (...), bitcast (void ()* @weakfunc to void (...)*)
+define weak void @weakfunc() {
+entry:
+  ret void
+}

test/tools/gold/X86/thinlto_internalize.ll

+; RUN: opt -module-summary %s -o %t.o
+; RUN: opt -module-summary %p/Inputs/thinlto_internalize.ll -o %t2.o
+
+; RUN: %gold -m elf_x86_64 -plugin %llvmshlibdir/LLVMgold.so \
+; RUN:     --plugin-opt=thinlto \
+; RUN:     --plugin-opt=-import-instr-limit=0 \
+; RUN:     --plugin-opt=save-temps \
+; RUN:     -o %t3.o %t2.o %t.o
+; RUN: llvm-dis %t.o.opt.bc -o - | FileCheck %s
+
+; f() should be internalized and eliminated after inlining
+; CHECK-NOT: @f()
+
+target triple = "x86_64-unknown-linux-gnu"
+define i32 @g() {
+  call void @f()
+  ret i32 0
+}
+define void @f() {
+  ret void
+}

test/tools/gold/X86/thinlto_linkonceresolution.ll

 ; RUN: opt -module-summary %s -o %t.o
 ; RUN: opt -module-summary %p/Inputs/thinlto_linkonceresolution.ll -o %t2.o
 
 ; Ensure the plugin ensures that for ThinLTO the prevailing copy of a
 ; linkonce symbol is changed to weak to ensure it is not eliminated.
 ; Note that gold picks the first copy of f() as the prevailing one,
 ; so listing %t2.o first is sufficient to ensure that this copy is
 ; preempted. Also, set the import-instr-limit to 0 to prevent f() from
 ; being imported from %t2.o which hides the problem.
 ; RUN: %gold -m elf_x86_64 -plugin %llvmshlibdir/LLVMgold.so \
 ; RUN:     --plugin-opt=thinlto \
 ; RUN:     --plugin-opt=-import-instr-limit=0 \
 ; RUN:     --plugin-opt=save-temps \
+; RUN:     -shared \
 ; RUN:     -o %t3.o %t2.o %t.o
 ; RUN: llvm-nm %t3.o | FileCheck %s
 ; RUN: llvm-dis %t.o.opt.bc -o - | FileCheck --check-prefix=OPT %s
 ; RUN: llvm-dis %t2.o.opt.bc -o - | FileCheck --check-prefix=OPT2 %s
 
 ; Ensure that f() is defined in resulting object file, and also
 ; confirm the weak linkage directly in the saved opt bitcode files.
 ; CHECK-NOT: U f

tools/gold/gold-plugin.cpp

class CodeGen {
 
   /// Map from a module name to the corresponding buffer holding a view of the
   /// bitcode provided via the get_view gold callback.
   StringMap<MemoryBufferRef> *ModuleMap;
 
   // Functions to import into this module.
   FunctionImporter::ImportMapTy *ImportList;
 
+  // Map of globals defined in this module to their summary.
+  std::map<GlobalValue::GUID, GlobalValueSummary *> *DefinedGlobals;
+
 public:
   /// Constructor used by full LTO.
   CodeGen(std::unique_ptr<llvm::Module> M)
       : M(std::move(M)), OS(nullptr), TaskID(-1), CombinedIndex(nullptr),
         ModuleMap(nullptr) {
     initTargetMachine();
   }
   /// Constructor used by ThinLTO.
   CodeGen(std::unique_ptr<llvm::Module> M, raw_fd_ostream *OS, int TaskID,
           const ModuleSummaryIndex *CombinedIndex, std::string Filename,
           StringMap<MemoryBufferRef> *ModuleMap,
-          FunctionImporter::ImportMapTy *ImportList)
+          FunctionImporter::ImportMapTy *ImportList,
+          std::map<GlobalValue::GUID, GlobalValueSummary *> *DefinedGlobals)
       : M(std::move(M)), OS(OS), TaskID(TaskID), CombinedIndex(CombinedIndex),
         SaveTempsFilename(std::move(Filename)), ModuleMap(ModuleMap),
-        ImportList(ImportList) {
+        ImportList(ImportList), DefinedGlobals(DefinedGlobals) {
     assert(options::thinlto == !!CombinedIndex &&
            "Expected module summary index iff performing ThinLTO");
     initTargetMachine();
   }
 
   /// Invoke LTO passes and the code generator for the module.
   void runAll();
 
return std::unique_ptr<TargetMachine>(TheTarget->createTargetMachine(
       TripleStr, options::mcpu, FeaturesString, Options, RelocationModel,
       CodeModel::Default, CGOptLevel));
 }
 
 void CodeGen::runLTOPasses() {
   M->setDataLayout(TM->createDataLayout());
 
   if (CombinedIndex) {
+    // Apply summary-based internalization decisions. Skip if there are no
+    // defined globals from the summary since not only is it unnecessary, but
+    // if this module did not have a summary section the internalizer will
+    // assert if it finds any definitions in this module that aren't in the
+    // DefinedGlobals set.
+    if (!DefinedGlobals->empty())
+      thinLTOInternalizeModule(*M, *DefinedGlobals);
+
     // Create a loader that will parse the bitcode from the buffers
     // in the ModuleMap.
     ModuleLoader Loader(M->getContext(), *ModuleMap);
 
     // Perform function importing.
     FunctionImporter Importer(*CombinedIndex, Loader);
     Importer.importFunctions(*M, *ImportList);
   }
else
     runSplitCodeGen(OptFilename);
 }
 
 /// Links the module in \p View from file \p F into the combined module
 /// saved in the IRMover \p L.
 static void linkInModule(LLVMContext &Context, IRMover &L, claimed_file &F,
                          const void *View, StringRef Name,
                          raw_fd_ostream *ApiFile, StringSet<> &Internalize,
-                         StringSet<> &Maybe) {
+                         StringSet<> &Maybe, bool SetName = false) {
   std::vector<GlobalValue *> Keep;
   StringMap<unsigned> Realign;
   std::unique_ptr<Module> M = getModuleForFile(
       Context, F, View, Name, ApiFile, Internalize, Maybe, Keep, Realign);
   if (!M.get())
     return;
   if (!options::triple.empty())
     M->setTargetTriple(options::triple.c_str());
   else if (M->getTargetTriple().empty()) {
     M->setTargetTriple(DefaultTriple);
   }
 
+  // For ThinLTO we want to propagate the source file name to ensure
+  // we can create the correct global identifiers matching those in the
+  // original module.
+  if (SetName)
+    L.getModule().setSourceFileName(M->getSourceFileName());
+
   if (Error E = L.move(std::move(M), Keep,
                        [](GlobalValue &, IRMover::ValueAdder) {})) {
     handleAllErrors(std::move(E), [&](const llvm::ErrorInfoBase &EIB) {
       message(LDPL_FATAL, "Failed to link module %s: %s", Name.str().c_str(),
               EIB.message().c_str());
     });
   }
 
   for (const auto &I : Realign) {
     GlobalValue *Dst = L.getModule().getNamedValue(I.first());
     if (!Dst)
       continue;
     cast<GlobalVariable>(Dst)->setAlignment(I.second);
   }
 }
 
 /// Perform the ThinLTO backend on a single module, invoking the LTO and codegen
 /// pipelines.
 static void thinLTOBackendTask(claimed_file &F, const void *View,
                                StringRef Name, raw_fd_ostream *ApiFile,
                                const ModuleSummaryIndex &CombinedIndex,
                                raw_fd_ostream *OS, unsigned TaskID,
                                StringMap<MemoryBufferRef> &ModuleMap,
-                               FunctionImporter::ImportMapTy &ImportList) {
+                               FunctionImporter::ImportMapTy &ImportList,
+                               std::map<GlobalValue::GUID, GlobalValueSummary *> &DefinedGlobals) {
   // Need to use a separate context for each task
   LLVMContext Context;
   Context.setDiscardValueNames(options::TheOutputType !=
                                options::OT_SAVE_TEMPS);
   Context.enableDebugTypeODRUniquing(); // Merge debug info types.
   Context.setDiagnosticHandler(diagnosticHandlerForContext, nullptr, true);
 
   std::unique_ptr<llvm::Module> NewModule(new llvm::Module(Name, Context));
   IRMover L(*NewModule.get());
 
   StringSet<> Dummy;
-  linkInModule(Context, L, F, View, Name, ApiFile, Dummy, Dummy);
+  linkInModule(Context, L, F, View, Name, ApiFile, Dummy, Dummy, true);
   if (renameModuleForThinLTO(*NewModule, CombinedIndex))
     message(LDPL_FATAL, "Failed to rename module for ThinLTO");
 
   CodeGen codeGen(std::move(NewModule), OS, TaskID, &CombinedIndex, Name,
-                  &ModuleMap, &ImportList);
+                  &ModuleMap, &ImportList, &DefinedGlobals);
   codeGen.runAll();
 }
 
 /// Launch each module's backend pipeline in a separate task in a thread pool.
 static void
 thinLTOBackends(raw_fd_ostream *ApiFile,
                 const ModuleSummaryIndex &CombinedIndex,
                 StringMap<MemoryBufferRef> &ModuleMap,
-                StringMap<FunctionImporter::ImportMapTy> &ImportLists) {
+                StringMap<FunctionImporter::ImportMapTy> &ImportLists,
+  StringMap<std::map<GlobalValue::GUID, GlobalValueSummary *>>
+      &ModuleToDefinedGVSummaries) {
   unsigned TaskCount = 0;
   std::vector<ThinLTOTaskInfo> Tasks;
   Tasks.reserve(Modules.size());
   unsigned int MaxThreads = options::Parallelism
                                 ? options::Parallelism
                                 : thread::hardware_concurrency();
 
   // Create ThreadPool in nested scope so that threads will be joined
for (claimed_file &F : Modules) {
       TaskCount++;
       std::unique_ptr<raw_fd_ostream> OS =
           llvm::make_unique<raw_fd_ostream>(FD, true);
 
       // Enqueue the task
       ThinLTOThreadPool.async(thinLTOBackendTask, std::ref(F), View, F.name,
                               ApiFile, std::ref(CombinedIndex), OS.get(),
                               TaskCount, std::ref(ModuleMap),
-                              std::ref(ImportLists[F.name]));
+                              std::ref(ImportLists[F.name]),
+                              std::ref(ModuleToDefinedGVSummaries[F.name]));
 
       // Record the information needed by the task or during its cleanup
       // to a ThinLTOTaskInfo instance. For information needed by the task
       // the unique_ptr ownership is transferred to the ThinLTOTaskInfo.
       Tasks.emplace_back(std::move(OS), NewFilename.c_str(), TempOutFile);
     }
   }
 

 static ld_plugin_status thinLTOLink(raw_fd_ostream *ApiFile) {
   // Map from a module name to the corresponding buffer holding a view of the
   // bitcode provided via the get_view gold callback.
   StringMap<MemoryBufferRef> ModuleMap;
   // Map to own RAII objects that manage the file opening and releasing
   // interfaces with gold.
   DenseMap<void *, std::unique_ptr<PluginInputFile>> HandleToInputFile;
 
+  // Keep track of internalization candidates as well as those that may not
+  // be internalized because they are refereneced from other IR modules.
+  DenseSet<GlobalValue::GUID> Internalize;
+  DenseSet<GlobalValue::GUID> CrossReferenced;
+
   ModuleSummaryIndex CombinedIndex;
   uint64_t NextModuleId = 0;
   for (claimed_file &F : Modules) {
     if (!HandleToInputFile.count(F.leader_handle))
       HandleToInputFile.insert(std::make_pair(
           F.leader_handle, llvm::make_unique<PluginInputFile>(F.handle)));
     // Pass this into getModuleSummaryIndexForFile
     const void *View = getSymbolsAndView(F);
     if (!View)
       continue;
 
     MemoryBufferRef ModuleBuffer(StringRef((const char *)View, F.filesize),
                                  F.name);
     assert(ModuleMap.find(ModuleBuffer.getBufferIdentifier()) ==
                ModuleMap.end() &&
            "Expect unique Buffer Identifier");
     ModuleMap[ModuleBuffer.getBufferIdentifier()] = ModuleBuffer;
 
     std::unique_ptr<ModuleSummaryIndex> Index = getModuleSummaryIndexForFile(F);
 
     // Skip files without a module summary.
     if (Index)
       CombinedIndex.mergeFrom(std::move(Index), ++NextModuleId);
+
+    // Look for internalization candidates based on gold's symbol resolution
+    // information. Also track symbols referenced from other IR modules.
+    for (auto &Sym : F.syms) {
+      ld_plugin_symbol_resolution Resolution =
+          (ld_plugin_symbol_resolution)Sym.resolution;
+      if (Resolution == LDPR_PREVAILING_DEF_IRONLY)
+        Internalize.insert(GlobalValue::getGUID(Sym.name));
+      if (Resolution == LDPR_RESOLVED_IR || Resolution == LDPR_PREEMPTED_IR)
+        CrossReferenced.insert(GlobalValue::getGUID(Sym.name));
+    }
   }
 
-  if (options::thinlto_emit_imports_files && !options::thinlto_index_only)
-    message(LDPL_WARNING,
-            "thinlto-emit-imports-files ignored unless thinlto-index-only");
+  // Remove symbols referenced from other IR modules from the internalization
+  // candidate set.
+  for (auto &S : CrossReferenced)
+    Internalize.erase(S);
 
   // Collect for each module the list of function it defines (GUID ->
   // Summary).
   StringMap<std::map<GlobalValue::GUID, GlobalValueSummary *>>
       ModuleToDefinedGVSummaries(NextModuleId);
   CombinedIndex.collectDefinedGVSummariesPerModule(ModuleToDefinedGVSummaries);
 
   StringMap<FunctionImporter::ImportMapTy> ImportLists(NextModuleId);
   StringMap<FunctionImporter::ExportSetTy> ExportLists(NextModuleId);
   ComputeCrossModuleImport(CombinedIndex, ModuleToDefinedGVSummaries,
                            ImportLists, ExportLists);
 
+  // Callback for internalization, to prevent internalization of symbols
+  // that were not candidates initially, and those that are being imported
+  // (which introduces new cross references).
+  auto isExported = [&](StringRef ModuleIdentifier, GlobalValue::GUID GUID) {
+    const auto &ExportList = ExportLists.find(ModuleIdentifier);
+    return (ExportList != ExportLists.end() &&
+            ExportList->second.count(GUID)) ||
+           !Internalize.count(GUID);
+  };
+
+  // Use global summary-based analysis to identify symbols that can be
+  // internalized (because they aren't exported or preserved as per callback).
+  // Changes are made in the index, consumed in the ThinLTO backends.
+  thinLTOInternalizeAndPromoteInIndex(CombinedIndex, isExported);
+
+  if (options::thinlto_emit_imports_files && !options::thinlto_index_only)
+    message(LDPL_WARNING,
+            "thinlto-emit-imports-files ignored unless thinlto-index-only");
+
   if (options::thinlto_index_only) {
     // If the thinlto-prefix-replace option was specified, parse it and
     // extract the old and new prefixes.
     std::string OldPrefix, NewPrefix;
     getThinLTOOldAndNewPrefix(OldPrefix, NewPrefix);
 
     // For each input bitcode file, generate an individual index that
     // contains summaries only for its own global values, and for any that
// Create OS in nested scope so that it will be closed on destruction.
     raw_fd_ostream OS(output_name + ".thinlto.bc", EC,
                       sys::fs::OpenFlags::F_None);
     if (EC)
       message(LDPL_FATAL, "Unable to open %s.thinlto.bc for writing: %s",
               output_name.data(), EC.message().c_str());
     WriteIndexToFile(CombinedIndex, OS);
   }
 
-  thinLTOBackends(ApiFile, CombinedIndex, ModuleMap, ImportLists);
+  thinLTOBackends(ApiFile, CombinedIndex, ModuleMap, ImportLists,
+                  ModuleToDefinedGVSummaries);
   return LDPS_OK;
 }
 
 /// gold informs us that all symbols have been read. At this point, we use
 /// get_symbols to see if any of our definitions have been overridden by a
 /// native object file. Then, perform optimization and codegen.
 static ld_plugin_status allSymbolsReadHook(raw_fd_ostream *ApiFile) {
   if (Modules.empty())