Update ThinLTO Indexing logic

[rdhindsa]

Summary

1. Instead of writing empty index for file, this patch tracks the state of files in ObjectToIndexFileState. If the files are not indexed , only then we emit the empty files.
2. For ThinLTO Index only, skip cache pruning and iterating over tasks, since there is at most one output object file

[tejohnson]

lld/ELF/LTO.cpp

300

There won't always be a regular LTO file created via splitting for ThinLTO builds, so Buff[0] should be empty some of the time. Do you have a test for this case?

The other thing being done by the loop below is a call to createObjectFile, does that need to be done for Buff[0] if non-empty too?

[pcc]

This seems fine for now, but instead of copying this behaviour into LLD, it really ought to be a feature of the LTO implementation.

lld/ELF/LTO.cpp

262–269

I think this needs to be moved after the call to LTO::run.

300

Buff[0] should always be non-empty, see D41267.

[rdhindsa]

lld/ELF/LTO.cpp

300

I have added the test as well to ensure that one file is generated in this case. We just need to save this file and, we don't need to call createObjectFile, since we exit from LLD just after this step.

[tejohnson]

This LGTM, but wait until pcc or ruiu accept it.

[pcc]

I think this will need a rebase.

[rdhindsa]

Rebased diff.

[pcc]

LGTM

lld/ELF/LTO.h

public:
   std::vector<InputFile *> compile();
   void addLazyObjFile(LazyObjFile *File);
 
 private:
   std::unique_ptr<llvm::lto::LTO> LTOObj;
   std::vector<SmallString<0>> Buff;
   std::vector<std::unique_ptr<MemoryBuffer>> Files;
   llvm::DenseSet<StringRef> UsedStartStop;
+  llvm::StringMap<bool> ObjectToIndexFileState;
   std::unique_ptr<llvm::raw_fd_ostream> LinkedObjects;
 };
 } // namespace elf
 } // namespace lld
 
 #endif

lld/ELF/LTO.cpp

auto LinkedObjectsFile = llvm::make_unique<raw_fd_ostream>(
       Config->ThinLTOIndexOnlyObjectsFile, EC, sys::fs::OpenFlags::F_None);
   if (EC)
     error("cannot create " + Config->ThinLTOIndexOnlyObjectsFile + ": " +
           EC.message());
   return LinkedObjectsFile;
 }
 
 // Creates instance of LTO.
+// OnIndexWrite is callback to let caller know when LTO writes index files.
 // LinkedObjectsFile is an output stream to write the list of object files for
 // the final ThinLTO linking. Can be nullptr.
-static std::unique_ptr<lto::LTO> createLTO(raw_fd_ostream *LinkedObjectsFile) {
+static std::unique_ptr<lto::LTO> createLTO(lto::IndexWriteCallback OnIndexWrite,
+                                           raw_fd_ostream *LinkedObjectsFile) {
   lto::Config Conf;
 
   // LLD supports the new relocations.
   Conf.Options = InitTargetOptionsFromCodeGenFlags();
   Conf.Options.RelaxELFRelocations = true;
 
   // Always emit a section per function/datum with LTO.
   Conf.Options.FunctionSections = true;
static std::unique_ptr<lto::LTO> createLTO(lto::IndexWriteCallback OnIndexWrite,
 
   lto::ThinBackend Backend;
   if (Config->ThinLTOJobs != -1u)
     Backend = lto::createInProcessThinBackend(Config->ThinLTOJobs);
 
   if (Config->ThinLTOIndexOnly) {
     std::string OldPrefix, NewPrefix;
     std::tie(OldPrefix, NewPrefix) = Config->ThinLTOPrefixReplace.split(';');
-    Backend = lto::createWriteIndexesThinBackend(OldPrefix, NewPrefix, true,
-                                                 LinkedObjectsFile, nullptr);
+    Backend = lto::createWriteIndexesThinBackend(
+        OldPrefix, NewPrefix, true, LinkedObjectsFile, OnIndexWrite);
   }
 
   Conf.SampleProfile = Config->LTOSampleProfile;
   Conf.UseNewPM = Config->LTONewPassManager;
   Conf.DebugPassManager = Config->LTODebugPassManager;
 
   return llvm::make_unique<lto::LTO>(std::move(Conf), Backend,
                                      Config->LTOPartitions);
 }
 
 BitcodeCompiler::BitcodeCompiler() {
   LinkedObjects = createLinkedObjectsFile();
-  LTOObj = createLTO(LinkedObjects.get());
+  LTOObj = createLTO(
+      [&](const std::string &Identifier) {
+        ObjectToIndexFileState[Identifier] = true;
+      },
+      LinkedObjects.get());
   for (Symbol *Sym : Symtab->getSymbols()) {
     StringRef Name = Sym->getName();
     for (StringRef Prefix : {"__start_", "__stop_"})
       if (Name.startswith(Prefix))
         UsedStartStop.insert(Name.substr(Prefix.size()));
   }
 }
 
 BitcodeCompiler::~BitcodeCompiler() = default;
 
 static void undefine(Symbol *S) {
   replaceSymbol<Undefined>(S, nullptr, S->getName(), STB_GLOBAL, STV_DEFAULT,
                            S->Type);
 }
 
 void BitcodeCompiler::add(BitcodeFile &F) {
   lto::InputFile &Obj = *F.Obj;
 
-  std::string OldPrefix, NewPrefix;
-  std::tie(OldPrefix, NewPrefix) = Config->ThinLTOPrefixReplace.split(';');
-
-  // Create the empty files which, if indexed, will be overwritten later.
   if (Config->ThinLTOIndexOnly)
-    writeEmptyDistributedBuildOutputs(Obj.getName(), OldPrefix, NewPrefix,
-                                      false);
+    ObjectToIndexFileState.insert({Obj.getName(), false});
 
   unsigned SymNum = 0;
   std::vector<Symbol *> Syms = F.getSymbols();
   std::vector<lto::SymbolResolution> Resols(Syms.size());
 
   bool IsExecutable = !Config->Shared && !Config->Relocatable;
 
   // Provide a resolution to the LTO API for each symbol.

 // Merge all the bitcode files we have seen, codegen the result
 // and return the resulting ObjectFile(s).
 std::vector<InputFile *> BitcodeCompiler::compile() {
   std::vector<InputFile *> Ret;
   unsigned MaxTasks = LTOObj->getMaxTasks();
   Buff.resize(MaxTasks);
   Files.resize(MaxTasks);
 
+  // Emit empty index files for non-indexed files
+  if (Config->ThinLTOIndexOnly) {
+    std::string OldPrefix, NewPrefix;
+    std::tie(OldPrefix, NewPrefix) = Config->ThinLTOPrefixReplace.split(';');
+    for (auto &Identifier : ObjectToIndexFileState)
+      if (!Identifier.getValue())
+        writeEmptyDistributedBuildOutputs(Identifier.getKey(), OldPrefix,
+                                          NewPrefix, /* SkipModule */ false);
+  }
+
   // If LazyObjFile has not been added to link, emit empty index files
   if (Config->ThinLTOIndexOnly)
     for (LazyObjFile *F : LazyObjFiles)
       if (!F->AddedToLink && isBitcode(F->MB))
         addLazyObjFile(F);
 
   // The --thinlto-cache-dir option specifies the path to a directory in which
   // to cache native object files for ThinLTO incremental builds. If a path was
   // specified, configure LTO to use it as the cache directory.
   lto::NativeObjectCache Cache;
   if (!Config->ThinLTOCacheDir.empty())
     Cache = check(
         lto::localCache(Config->ThinLTOCacheDir,
                         [&](size_t Task, std::unique_ptr<MemoryBuffer> MB) {
                           Files[Task] = std::move(MB);
                         }));
 
   checkError(LTOObj->run(
       [&](size_t Task) {
         return llvm::make_unique<lto::NativeObjectStream>(
             llvm::make_unique<raw_svector_ostream>(Buff[Task]));
       },
       Cache));
 
+  // ThinLTO with index only option is required to generate only the index
+  // files. After that, we exit from linker and ThinLTO backend runs in a
+  // distributed environment.
+  if (Config->ThinLTOIndexOnly) {
+    if (Config->SaveTemps)
+      saveBuffer(Buff[0], Config->OutputFile + ".lto.o");
+    exit(0);
+  }
+
   if (!Config->ThinLTOCacheDir.empty())
     pruneCache(Config->ThinLTOCacheDir, Config->ThinLTOCachePolicy);
 
   for (unsigned I = 0; I != MaxTasks; ++I) {
     if (Buff[I].empty())
       continue;
     if (Config->SaveTemps) {
       if (I == 0)
         saveBuffer(Buff[I], Config->OutputFile + ".lto.o");
       else
         saveBuffer(Buff[I], Config->OutputFile + Twine(I) + ".lto.o");
     }
     InputFile *Obj = createObjectFile(MemoryBufferRef(Buff[I], "lto.tmp"));
     Ret.push_back(Obj);
   }
 
-  // ThinLTO with index only option is required to generate only the index
-  // files. After that, we exit from linker and ThinLTO backend runs in a
-  // distributed environment.
-  if (Config->ThinLTOIndexOnly)
-    exit(0);
-
   for (std::unique_ptr<MemoryBuffer> &File : Files)
     if (File)
       Ret.push_back(createObjectFile(*File));
 
   return Ret;
 }
 
 // For lazy object files not added to link, adds empty index files
 void BitcodeCompiler::addLazyObjFile(LazyObjFile *File) {
   StringRef Identifier = File->getBuffer().getBufferIdentifier();
   std::string OldPrefix, NewPrefix;
   std::tie(OldPrefix, NewPrefix) = Config->ThinLTOPrefixReplace.split(';');
   writeEmptyDistributedBuildOutputs(Identifier, OldPrefix, NewPrefix,
                                     /* SkipModule */ true);
 }