LTO: Simplify caching interface.

The NativeObjectOutput class has a design problem: it mixes up the caching
policy with the interface for output streams, which makes the client-side
code hard to follow and would for example make it harder to replace the
cache implementation in an arbitrary client.

This change separates the two aspects by moving the caching policy
to a separate field in Config, replacing NativeObjectOutput with a
NativeObjectStream class which only deals with streams and does not need to
be overridden by most clients and introducing an AddFile callback for adding
files (e.g. from the cache) to the link.

Differential Revision: https://reviews.llvm.org/D24622

llvm-svn: 282299

3 files changed, 101 insertions, 115 deletions

diff --git a/llvm/lib/LTO/Caching.cpp b/llvm/lib/LTO/Caching.cpp
index 728e5164dfa..065239cbb1b 100644
--- a/llvm/lib/LTO/Caching.cpp
+++ b/llvm/lib/LTO/Caching.cpp
@@ -12,13 +12,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/LTO/Caching.h"
-
-#ifdef HAVE_LLVM_REVISION
-#include "LLVMLTORevision.h"
-#endif
-
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -30,6 +26,8 @@ static void commitEntry(StringRef TempFilename, StringRef EntryPath) {
   auto EC = sys::fs::rename(TempFilename, EntryPath);
   if (EC) {
     // Renaming failed, probably not the same filesystem, copy and delete.
+    // FIXME: Avoid needing to do this by creating the temporary file in the
+    // cache directory.
     {
       auto ReloadedBufferOrErr = MemoryBuffer::getFile(TempFilename);
       if (auto EC = ReloadedBufferOrErr.getError())
@@ -48,51 +46,54 @@ static void commitEntry(StringRef TempFilename, StringRef EntryPath) {
   }
 }
 
-CacheObjectOutput::~CacheObjectOutput() {
-  if (EntryPath.empty())
-    // The entry was never used by the client (tryLoadFromCache() wasn't called)
-    return;
-  // TempFilename is only set if getStream() was called, i.e. on cache miss when
-  // tryLoadFromCache() returned false. And EntryPath is valid if a Key was
-  // submitted, otherwise it has been set to CacheDirectoryPath in
-  // tryLoadFromCache.
-  if (!TempFilename.empty()) {
-    if (EntryPath == CacheDirectoryPath)
-      // The Key supplied to tryLoadFromCache was empty, do not commit the temp.
-      EntryPath = TempFilename;
-    else
-      // We commit the tempfile into the cache now, by moving it to EntryPath.
-      commitEntry(TempFilename, EntryPath);
-  }
-  // Supply the cache path to the user.
-  AddBuffer(EntryPath.str());
-}
+NativeObjectCache lto::localCache(std::string CacheDirectoryPath,
+                                  AddFileFn AddFile) {
+  return [=](unsigned Task, StringRef Key) -> AddStreamFn {
+    // First, see if we have a cache hit.
+    SmallString<64> EntryPath;
+    sys::path::append(EntryPath, CacheDirectoryPath, Key);
+    if (sys::fs::exists(EntryPath)) {
+      AddFile(Task, EntryPath);
+      return AddStreamFn();
+    }
 
-// Return an allocated stream for the output, or null in case of failure.
-std::unique_ptr<raw_pwrite_stream> CacheObjectOutput::getStream() {
-  assert(!EntryPath.empty() && "API Violation: client didn't call "
-                               "tryLoadFromCache() before getStream()");
-  // Write to a temporary to avoid race condition
-  int TempFD;
-  std::error_code EC =
-      sys::fs::createTemporaryFile("Thin", "tmp.o", TempFD, TempFilename);
-  if (EC) {
-    errs() << "Error: " << EC.message() << "\n";
-    report_fatal_error("ThinLTO: Can't get a temporary file");
-  }
-  return llvm::make_unique<raw_fd_ostream>(TempFD, /* ShouldClose */ true);
-}
+    // This native object stream is responsible for commiting the resulting
+    // file to the cache and calling AddFile to add it to the link.
+    struct CacheStream : NativeObjectStream {
+      AddFileFn AddFile;
+      std::string TempFilename;
+      std::string EntryPath;
+      unsigned Task;
 
-// Try loading from a possible cache first, return true on cache hit.
-bool CacheObjectOutput::tryLoadFromCache(StringRef Key) {
-  assert(!CacheDirectoryPath.empty() &&
-         "CacheObjectOutput was initialized without a cache path");
-  if (Key.empty()) {
-    // Client didn't compute a valid key. EntryPath has been set to
-    // CacheDirectoryPath.
-    EntryPath = CacheDirectoryPath;
-    return false;
-  }
-  sys::path::append(EntryPath, CacheDirectoryPath, Key);
-  return sys::fs::exists(EntryPath);
+      CacheStream(std::unique_ptr<raw_pwrite_stream> OS, AddFileFn AddFile,
+                  std::string TempFilename, std::string EntryPath,
+                  unsigned Task)
+          : NativeObjectStream(std::move(OS)), AddFile(AddFile),
+            TempFilename(TempFilename), EntryPath(EntryPath), Task(Task) {}
+
+      ~CacheStream() {
+        // Make sure the file is closed before committing it.
+        OS.reset();
+        commitEntry(TempFilename, EntryPath);
+        AddFile(Task, EntryPath);
+      }
+    };
+
+    return [=](size_t Task) -> std::unique_ptr<NativeObjectStream> {
+      // Write to a temporary to avoid race condition
+      int TempFD;
+      SmallString<64> TempFilename;
+      std::error_code EC =
+          sys::fs::createTemporaryFile("Thin", "tmp.o", TempFD, TempFilename);
+      if (EC) {
+        errs() << "Error: " << EC.message() << "\n";
+        report_fatal_error("ThinLTO: Can't get a temporary file");
+      }
+
+      // This CacheStream will move the temporary file into the cache when done.
+      return make_unique<CacheStream>(
+          llvm::make_unique<raw_fd_ostream>(TempFD, /* ShouldClose */ true),
+          AddFile, TempFilename.str(), EntryPath.str(), Task);
+    };
+  };
 }
diff --git a/llvm/lib/LTO/LTO.cpp b/llvm/lib/LTO/LTO.cpp
index b2f0be2b5e8..b2d42f4d2a6 100644
--- a/llvm/lib/LTO/LTO.cpp
+++ b/llvm/lib/LTO/LTO.cpp
@@ -409,19 +409,19 @@ unsigned LTO::getMaxTasks() const {
   return RegularLTO.ParallelCodeGenParallelismLevel + ThinLTO.ModuleMap.size();
 }
 
-Error LTO::run(AddOutputFn AddOutput) {
+Error LTO::run(AddStreamFn AddStream, NativeObjectCache Cache) {
   // Save the status of having a regularLTO combined module, as
   // this is needed for generating the ThinLTO Task ID, and
   // the CombinedModule will be moved at the end of runRegularLTO.
   bool HasRegularLTO = RegularLTO.CombinedModule != nullptr;
   // Invoke regular LTO if there was a regular LTO module to start with.
   if (HasRegularLTO)
-    if (auto E = runRegularLTO(AddOutput))
+    if (auto E = runRegularLTO(AddStream))
       return E;
-  return runThinLTO(AddOutput, HasRegularLTO);
+  return runThinLTO(AddStream, Cache, HasRegularLTO);
 }
 
-Error LTO::runRegularLTO(AddOutputFn AddOutput) {
+Error LTO::runRegularLTO(AddStreamFn AddStream) {
   // Make sure commons have the right size/alignment: we kept the largest from
   // all the prevailing when adding the inputs, and we apply it here.
   const DataLayout &DL = RegularLTO.CombinedModule->getDataLayout();
@@ -478,7 +478,7 @@ Error LTO::runRegularLTO(AddOutputFn AddOutput) {
         !Conf.PostInternalizeModuleHook(0, *RegularLTO.CombinedModule))
       return Error();
   }
-  return backend(Conf, AddOutput, RegularLTO.ParallelCodeGenParallelismLevel,
+  return backend(Conf, AddStream, RegularLTO.ParallelCodeGenParallelismLevel,
                  std::move(RegularLTO.CombinedModule));
 }
 
@@ -507,7 +507,8 @@ public:
 
 class InProcessThinBackend : public ThinBackendProc {
   ThreadPool BackendThreadPool;
-  AddOutputFn AddOutput;
+  AddStreamFn AddStream;
+  NativeObjectCache Cache;
 
   Optional<Error> Err;
   std::mutex ErrMu;
@@ -517,42 +518,40 @@ public:
       Config &Conf, ModuleSummaryIndex &CombinedIndex,
       unsigned ThinLTOParallelismLevel,
       const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
-      AddOutputFn AddOutput)
+      AddStreamFn AddStream, NativeObjectCache Cache)
       : ThinBackendProc(Conf, CombinedIndex, ModuleToDefinedGVSummaries),
         BackendThreadPool(ThinLTOParallelismLevel),
-        AddOutput(std::move(AddOutput)) {}
+        AddStream(std::move(AddStream)), Cache(std::move(Cache)) {}
 
   Error runThinLTOBackendThread(
-      AddOutputFn AddOutput, unsigned Task, MemoryBufferRef MBRef,
-      ModuleSummaryIndex &CombinedIndex,
+      AddStreamFn AddStream, NativeObjectCache Cache, unsigned Task,
+      MemoryBufferRef MBRef, ModuleSummaryIndex &CombinedIndex,
       const FunctionImporter::ImportMapTy &ImportList,
       const FunctionImporter::ExportSetTy &ExportList,
       const std::map<GlobalValue::GUID, GlobalValue::LinkageTypes> &ResolvedODR,
       const GVSummaryMapTy &DefinedGlobals,
       MapVector<StringRef, MemoryBufferRef> &ModuleMap) {
+    auto RunThinBackend = [&](AddStreamFn AddStream) {
+      LTOLLVMContext BackendContext(Conf);
+      ErrorOr<std::unique_ptr<Module>> MOrErr =
+          parseBitcodeFile(MBRef, BackendContext);
+      assert(MOrErr && "Unable to load module in thread?");
+
+      return thinBackend(Conf, Task, AddStream, **MOrErr, CombinedIndex,
+                         ImportList, DefinedGlobals, ModuleMap);
+    };
 
-    auto ModuleIdentifier = MBRef.getBufferIdentifier();
-    auto Output = AddOutput(Task);
-    if (Output->isCachingEnabled()) {
-      SmallString<40> Key;
-      // The module may be cached, this helps handling it.
-      computeCacheKey(Key, CombinedIndex, ModuleIdentifier, ImportList,
-                      ExportList, ResolvedODR, DefinedGlobals);
-      if (Output->tryLoadFromCache(Key))
-        return Error();
-    }
+    if (!Cache)
+      return RunThinBackend(AddStream);
 
-    LTOLLVMContext BackendContext(Conf);
-    ErrorOr<std::unique_ptr<Module>> MOrErr =
-        parseBitcodeFile(MBRef, BackendContext);
-    assert(MOrErr && "Unable to load module in thread?");
+    SmallString<40> Key;
+    // The module may be cached, this helps handling it.
+    computeCacheKey(Key, CombinedIndex, MBRef.getBufferIdentifier(),
+                    ImportList, ExportList, ResolvedODR, DefinedGlobals);
+    if (AddStreamFn CacheAddStream = Cache(Task, Key))
+      return RunThinBackend(CacheAddStream);
 
-    auto AddOutputWrapper = [&](unsigned TaskId) {
-      assert(Task == TaskId && "Unexpexted TaskId mismatch");
-      return std::move(Output);
-    };
-    return thinBackend(Conf, Task, AddOutputWrapper, **MOrErr, CombinedIndex,
-                       ImportList, DefinedGlobals, ModuleMap);
+    return Error();
   }
 
   Error start(
@@ -574,8 +573,8 @@ public:
             const GVSummaryMapTy &DefinedGlobals,
             MapVector<StringRef, MemoryBufferRef> &ModuleMap) {
           Error E = runThinLTOBackendThread(
-              AddOutput, Task, MBRef, CombinedIndex, ImportList, ExportList,
-              ResolvedODR, DefinedGlobals, ModuleMap);
+              AddStream, Cache, Task, MBRef, CombinedIndex, ImportList,
+              ExportList, ResolvedODR, DefinedGlobals, ModuleMap);
           if (E) {
             std::unique_lock<std::mutex> L(ErrMu);
             if (Err)
@@ -602,10 +601,10 @@ public:
 ThinBackend lto::createInProcessThinBackend(unsigned ParallelismLevel) {
   return [=](Config &Conf, ModuleSummaryIndex &CombinedIndex,
              const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
-             AddOutputFn AddOutput) {
+             AddStreamFn AddStream, NativeObjectCache Cache) {
     return llvm::make_unique<InProcessThinBackend>(
         Conf, CombinedIndex, ParallelismLevel, ModuleToDefinedGVSummaries,
-        AddOutput);
+        AddStream, Cache);
   };
 }
 
@@ -693,14 +692,15 @@ ThinBackend lto::createWriteIndexesThinBackend(std::string OldPrefix,
                                                std::string LinkedObjectsFile) {
   return [=](Config &Conf, ModuleSummaryIndex &CombinedIndex,
              const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
-             AddOutputFn AddOutput) {
+             AddStreamFn AddStream, NativeObjectCache Cache) {
     return llvm::make_unique<WriteIndexesThinBackend>(
         Conf, CombinedIndex, ModuleToDefinedGVSummaries, OldPrefix, NewPrefix,
         ShouldEmitImportsFiles, LinkedObjectsFile);
   };
 }
 
-Error LTO::runThinLTO(AddOutputFn AddOutput, bool HasRegularLTO) {
+Error LTO::runThinLTO(AddStreamFn AddStream, NativeObjectCache Cache,
+                      bool HasRegularLTO) {
   if (ThinLTO.ModuleMap.empty())
     return Error();
 
@@ -759,8 +759,9 @@ Error LTO::runThinLTO(AddOutputFn AddOutput, bool HasRegularLTO) {
   thinLTOResolveWeakForLinkerInIndex(ThinLTO.CombinedIndex, isPrevailing,
                                      recordNewLinkage);
 
-  std::unique_ptr<ThinBackendProc> BackendProc = ThinLTO.Backend(
-      Conf, ThinLTO.CombinedIndex, ModuleToDefinedGVSummaries, AddOutput);
+  std::unique_ptr<ThinBackendProc> BackendProc =
+      ThinLTO.Backend(Conf, ThinLTO.CombinedIndex, ModuleToDefinedGVSummaries,
+                      AddStream, Cache);
 
   // Partition numbers for ThinLTO jobs start at 1 (see comments for
   // GlobalResolution in LTO.h). Task numbers, however, start at
diff --git a/llvm/lib/LTO/LTOBackend.cpp b/llvm/lib/LTO/LTOBackend.cpp
index d83b65de4cc..96215b78498 100644
--- a/llvm/lib/LTO/LTOBackend.cpp
+++ b/llvm/lib/LTO/LTOBackend.cpp
@@ -199,34 +199,20 @@ bool opt(Config &Conf, TargetMachine *TM, unsigned Task, Module &Mod,
   return !Conf.PostOptModuleHook || Conf.PostOptModuleHook(Task, Mod);
 }
 
-/// Monolithic LTO does not support caching (yet), this is a convenient wrapper
-/// around AddOutput to workaround this.
-static AddOutputFn getUncachedOutputWrapper(AddOutputFn &AddOutput,
-                                            unsigned Task) {
-  return [Task, &AddOutput](unsigned TaskId) {
-    auto Output = AddOutput(Task);
-    if (Output->isCachingEnabled() && Output->tryLoadFromCache(""))
-      report_fatal_error("Cache hit without a valid key?");
-    assert(Task == TaskId && "Unexpexted TaskId mismatch");
-    return Output;
-  };
-}
-
-void codegen(Config &Conf, TargetMachine *TM, AddOutputFn AddOutput,
+void codegen(Config &Conf, TargetMachine *TM, AddStreamFn AddStream,
              unsigned Task, Module &Mod) {
   if (Conf.PreCodeGenModuleHook && !Conf.PreCodeGenModuleHook(Task, Mod))
     return;
 
-  auto Output = AddOutput(Task);
-  std::unique_ptr<raw_pwrite_stream> OS = Output->getStream();
+  auto Stream = AddStream(Task);
   legacy::PassManager CodeGenPasses;
-  if (TM->addPassesToEmitFile(CodeGenPasses, *OS,
+  if (TM->addPassesToEmitFile(CodeGenPasses, *Stream->OS,
                               TargetMachine::CGFT_ObjectFile))
     report_fatal_error("Failed to setup codegen");
   CodeGenPasses.run(Mod);
 }
 
-void splitCodeGen(Config &C, TargetMachine *TM, AddOutputFn AddOutput,
+void splitCodeGen(Config &C, TargetMachine *TM, AddStreamFn AddStream,
                   unsigned ParallelCodeGenParallelismLevel,
                   std::unique_ptr<Module> Mod) {
   ThreadPool CodegenThreadPool(ParallelCodeGenParallelismLevel);
@@ -260,9 +246,7 @@ void splitCodeGen(Config &C, TargetMachine *TM, AddOutputFn AddOutput,
               std::unique_ptr<TargetMachine> TM =
                   createTargetMachine(C, MPartInCtx->getTargetTriple(), T);
 
-              codegen(C, TM.get(),
-                      getUncachedOutputWrapper(AddOutput, ThreadId), ThreadId,
-                      *MPartInCtx);
+              codegen(C, TM.get(), AddStream, ThreadId, *MPartInCtx);
             },
             // Pass BC using std::move to ensure that it get moved rather than
             // copied into the thread's context.
@@ -299,7 +283,7 @@ static void handleAsmUndefinedRefs(Module &Mod, TargetMachine &TM) {
   updateCompilerUsed(Mod, TM, AsmUndefinedRefs);
 }
 
-Error lto::backend(Config &C, AddOutputFn AddOutput,
+Error lto::backend(Config &C, AddStreamFn AddStream,
                    unsigned ParallelCodeGenParallelismLevel,
                    std::unique_ptr<Module> Mod) {
   Expected<const Target *> TOrErr = initAndLookupTarget(C, *Mod);
@@ -316,15 +300,15 @@ Error lto::backend(Config &C, AddOutputFn AddOutput,
       return Error();
 
   if (ParallelCodeGenParallelismLevel == 1) {
-    codegen(C, TM.get(), getUncachedOutputWrapper(AddOutput, 0), 0, *Mod);
+    codegen(C, TM.get(), AddStream, 0, *Mod);
   } else {
-    splitCodeGen(C, TM.get(), AddOutput, ParallelCodeGenParallelismLevel,
+    splitCodeGen(C, TM.get(), AddStream, ParallelCodeGenParallelismLevel,
                  std::move(Mod));
   }
   return Error();
 }
 
-Error lto::thinBackend(Config &Conf, unsigned Task, AddOutputFn AddOutput,
+Error lto::thinBackend(Config &Conf, unsigned Task, AddStreamFn AddStream,
                        Module &Mod, ModuleSummaryIndex &CombinedIndex,
                        const FunctionImporter::ImportMapTy &ImportList,
                        const GVSummaryMapTy &DefinedGlobals,
@@ -339,7 +323,7 @@ Error lto::thinBackend(Config &Conf, unsigned Task, AddOutputFn AddOutput,
   handleAsmUndefinedRefs(Mod, *TM);
 
   if (Conf.CodeGenOnly) {
-    codegen(Conf, TM.get(), AddOutput, Task, Mod);
+    codegen(Conf, TM.get(), AddStream, Task, Mod);
     return Error();
   }
 
@@ -379,6 +363,6 @@ Error lto::thinBackend(Config &Conf, unsigned Task, AddOutputFn AddOutput,
   if (!opt(Conf, TM.get(), Task, Mod, /*IsThinLto=*/true))
     return Error();
 
-  codegen(Conf, TM.get(), AddOutput, Task, Mod);
+  codegen(Conf, TM.get(), AddStream, Task, Mod);
   return Error();
 }