Simple name-based object cache

Author: weliveindetail

CMakeLists.txt

@@ -40,6 +40,7 @@ add_executable(JitFromScratch
   main.cpp
   JitFromScratch.cpp
   SimpleOptimizer.cpp
+  SimpleObjectCache.cpp
 )
 
 target_include_directories(JitFromScratch PRIVATE
@@ -125,4 +126,9 @@ add_custom_target(run
   COMMENT "Running JitFromScratch"
 )
 
+add_custom_target(run-cached
+  COMMAND $<TARGET_FILE:JitFromScratch> ${debug_args} -cache-dir=cache/
+  COMMENT "Running JitFromScratch (with object cache)"
+)
+
 add_subdirectory(test)

JitFromScratch.cpp

@@ -15,12 +15,14 @@
 using namespace llvm;
 using namespace llvm::orc;
 
-JitFromScratch::JitFromScratch(ExitOnError ExitOnErr)
+JitFromScratch::JitFromScratch(ExitOnError ExitOnErr,
+                               const std::string &CacheDir)
     : ES(std::make_unique<ExecutionSession>()),
       TM(createTargetMachine(ExitOnErr)),
+      ObjCache(std::make_unique<SimpleObjectCache>(CacheDir)),
       GDBListener(JITEventListener::createGDBRegistrationListener()),
       ObjLinkingLayer(*ES, createMemoryManagerFtor()),
-      CompileLayer(*ES, ObjLinkingLayer, SimpleCompiler(*TM)),
+      CompileLayer(*ES, ObjLinkingLayer, SimpleCompiler(*TM, ObjCache.get())),
       OptimizeLayer(*ES, CompileLayer) {
   ObjLinkingLayer.setNotifyLoaded(createNotifyLoadedFtor());
   if (auto R = createHostProcessResolver())
@@ -94,7 +96,22 @@ Error JitFromScratch::applyDataLayout(Module &M) {
 
 Error JitFromScratch::submitModule(std::unique_ptr<Module> M,
                                    std::unique_ptr<LLVMContext> C,
-                                   unsigned OptLevel) {
+                                   unsigned OptLevel, bool AddToCache) {
+  if (AddToCache)
+    ObjCache->setCacheModuleName(*M);
+
+  auto Obj = ObjCache->getCachedObject(*M);
+  if (!Obj) {
+    M.~unique_ptr();
+    return Obj.takeError();
+  }
+
+  if (Obj->hasValue()) {
+    M.~unique_ptr();
+    return ObjLinkingLayer.add(ES->getMainJITDylib(),
+                               std::move(Obj->getValue()));
+  }
+
   LLVM_DEBUG(dbgs() << "Submit IR module:\n\n" << *M << "\n\n");
 
   if (auto Err = applyDataLayout(*M))

JitFromScratch.h

@@ -17,9 +17,12 @@
 #include <memory>
 #include <string>
 
+#include "SimpleObjectCache.h"
+
 class JitFromScratch {
 public:
-  JitFromScratch(llvm::ExitOnError ExitOnErr);
+  JitFromScratch(llvm::ExitOnError ExitOnErr,
+                 const std::string &CacheDir = "");
 
   // Not a value type.
   JitFromScratch(const JitFromScratch &) = delete;
@@ -37,7 +40,7 @@ class JitFromScratch {
 
   llvm::Error submitModule(std::unique_ptr<llvm::Module> M,
                            std::unique_ptr<llvm::LLVMContext> C,
-                           unsigned OptLevel);
+                           unsigned OptLevel, bool AddToCache);
 
   template <class Signature_t>
   llvm::Expected<std::function<Signature_t>> getFunction(llvm::StringRef Name) {
@@ -51,6 +54,8 @@ class JitFromScratch {
 private:
   std::unique_ptr<llvm::orc::ExecutionSession> ES;
   std::unique_ptr<llvm::TargetMachine> TM;
+
+  std::unique_ptr<SimpleObjectCache> ObjCache;
   llvm::JITEventListener *GDBListener;
 
   llvm::orc::RTDyldObjectLinkingLayer ObjLinkingLayer;

SimpleObjectCache.cpp

@@ -0,0 +1,126 @@
+#include "SimpleObjectCache.h"
+
+#include <llvm/ADT/None.h>
+#include <llvm/ADT/SmallString.h>
+#include <llvm/ADT/Twine.h>
+#include <llvm/IR/Module.h>
+#include <llvm/Support/Debug.h>
+#include <llvm/Support/FileSystem.h>
+#include <llvm/Support/MemoryBuffer.h>
+#include <llvm/Support/Path.h>
+#include <llvm/Support/raw_ostream.h>
+
+#define DEBUG_TYPE "jitfromscratch"
+
+using namespace llvm;
+
+SimpleObjectCache::SimpleObjectCache(std::string Dir)
+    : Enabled(!Dir.empty()), CacheDir(endWithSeparator(std::move(Dir))) {
+  if (!sys::fs::exists(CacheDir)) {
+    LLVM_DEBUG(dbgs() << format("Create new cache directory '%s'\n\n",
+                                CacheDir.c_str()));
+    std::error_code EC = sys::fs::create_directories(CacheDir);
+    if (EC) {
+      LLVM_DEBUG(dbgs() << format(
+              "Creating new cache directory '%s' failed with "
+              "error code %d; Caching disabled\n\n",
+              CacheDir.c_str(), EC.value()));
+      Enabled = false;
+    }
+  }
+}
+
+std::string SimpleObjectCache::endWithSeparator(std::string Path) {
+  return Path.back() == '/' ? Path : Path + "/";
+}
+
+// Implements llvm::ObjectCache::notifyObjectCompiled, called from CompileLayer
+void SimpleObjectCache::notifyObjectCompiled(const Module *M,
+                                             MemoryBufferRef Obj) {
+  assert(M && "Caching requires module");
+
+  auto R = getCacheFileName(M->getModuleIdentifier());
+  if (!R.hasValue())
+    return;
+
+  std::string F = std::move(R.getValue());
+  if (auto EC = sys::fs::create_directories(sys::path::parent_path(F))) {
+    LLVM_DEBUG(dbgs() << format(
+                   "Writing cached object '%s' failed with error code %d\n\n",
+                   F.c_str(), EC.value()));
+    return;
+  }
+
+  std::error_code EC;
+  raw_fd_ostream OS(F, EC, sys::fs::F_None);
+  if (EC) {
+    LLVM_DEBUG(dbgs() << format(
+                   "Writing cached object '%s' failed with error code %d\n\n",
+                   F.c_str(), EC.value()));
+    return;
+  }
+
+  LLVM_DEBUG(dbgs() << format("Write cached object '%s'\n\n", F.c_str()));
+
+  OS.write(Obj.getBufferStart(), Obj.getBufferSize());
+  OS.close();
+}
+
+// Implements llvm::ObjectCache::getObject, called from CompileLayer
+std::unique_ptr<MemoryBuffer> SimpleObjectCache::getObject(const Module *M) {
+  assert(M && "Lookup requires module");
+
+  auto R = getCachedObject(*M);
+  if (!R) {
+    logAllUnhandledErrors(R.takeError(), dbgs(), "SimpleObjectCache: ");
+    return nullptr; // Error
+  }
+
+  if (!R->hasValue())
+    return nullptr; // No cache entry
+
+  return std::forward<std::unique_ptr<MemoryBuffer>>(R->getValue());
+}
+
+Expected<Optional<std::unique_ptr<MemoryBuffer>>>
+SimpleObjectCache::getCachedObject(const Module &M) const {
+  auto R = getCacheFileName(M.getModuleIdentifier());
+  if (!R.hasValue())
+    return None;
+
+  std::string F = std::move(R.getValue());
+  if (!sys::fs::exists(F))
+    return None;
+
+  auto B = MemoryBuffer::getFile(F, -1, false);
+  if (!B)
+    return createStringError(
+        B.getError(),
+        "Reading cached object '%s' failed with error code %d\n\n", F.c_str(),
+        B.getError().value());
+
+  LLVM_DEBUG(dbgs() << format("Read cached object '%s'\n\n", F.c_str()));
+  return std::forward<std::unique_ptr<MemoryBuffer>>(*B);
+}
+
+void SimpleObjectCache::setCacheModuleName(Module &M) const {
+  if (Enabled && !M.getName().startswith("file:"))
+    M.setModuleIdentifier("file:" + M.getModuleIdentifier() + ".o");
+}
+
+Optional<std::string>
+SimpleObjectCache::getCacheFileName(StringRef ModID) const {
+  if (!Enabled)
+    return None;
+
+  StringRef Prefix = "file:";
+  if (!ModID.startswith(Prefix))
+    return None;
+
+  std::string Name = Twine(CacheDir + ModID.substr(Prefix.size())).str();
+  size_t DotPos = Name.rfind('.');
+  if (DotPos != std::string::npos)
+    Name.replace(DotPos, Name.size() - DotPos, ".o");
+
+  return Name;
+}

SimpleObjectCache.h

@@ -0,0 +1,32 @@
+#pragma once
+
+#include <llvm/ADT/Optional.h>
+#include <llvm/ADT/StringRef.h>
+#include <llvm/ExecutionEngine/ObjectCache.h>
+
+#include <memory>
+#include <string>
+
+// Simple name-based lookup, based on lli's implementation
+class SimpleObjectCache : public llvm::ObjectCache {
+public:
+  SimpleObjectCache(std::string Dir);
+  ~SimpleObjectCache() = default;
+
+  void setCacheModuleName(llvm::Module &M) const;
+
+  llvm::Expected<llvm::Optional<std::unique_ptr<llvm::MemoryBuffer>>>
+  getCachedObject(const llvm::Module &M) const;
+
+protected:
+  std::unique_ptr<llvm::MemoryBuffer> getObject(const llvm::Module *M) override;
+  void notifyObjectCompiled(const llvm::Module *M,
+                            llvm::MemoryBufferRef Obj) override;
+
+private:
+  bool Enabled;
+  std::string CacheDir;
+
+  static std::string endWithSeparator(std::string Path);
+  llvm::Optional<std::string> getCacheFileName(llvm::StringRef ModID) const;
+};

main.cpp

@@ -27,6 +27,11 @@ static cl::opt<char>
                   "(default = '-O2')"),
         cl::Prefix, cl::ZeroOrMore, cl::init(' '));
 
+static cl::opt<std::string>
+    CacheDir("cache-dir",
+        cl::desc("Pass a directory name to enable object file cache"),
+        cl::value_desc("directory"), cl::init(""));
+
 Expected<unsigned> getOptLevel() {
   switch (OptLevel) {
     case '0': return 0;
@@ -141,7 +146,7 @@ int main(int argc, char **argv) {
   int x[]{0, 1, 2};
   int y[]{3, 1, -1};
 
-  JitFromScratch Jit(ExitOnErr);
+  JitFromScratch Jit(ExitOnErr, CacheDir);
 
   LLVM_DEBUG(dbgs() << "JITing for host target: "
                     << Jit.getTargetTriple().normalize() << "\n\n");
@@ -152,8 +157,9 @@ int main(int argc, char **argv) {
 
   std::string JitedFnName = ExitOnErr(codegenIR(*M, arrayElements(x)));
   unsigned OptLevel = ExitOnErr(getOptLevel());
+  bool AddToCache = !CacheDir.empty();
 
-  ExitOnErr(Jit.submitModule(std::move(M), std::move(C), OptLevel));
+  ExitOnErr(Jit.submitModule(std::move(M), std::move(C), OptLevel, AddToCache));
 
   // Request function; this compiles to machine code and links.
   auto integerDistances =

test/objcache.test-debug

@@ -0,0 +1,20 @@
+# RUN: rm -rf cache/
+
+# Object cache is disabled by default
+# RUN: JitFromScratch -debug -debug-only=jitfromscratch 2>&1 | FileCheck -check-prefix=CACHELESS-CHECK %s
+# CACHELESS-CHECK-NOT: Read cached object 'cache/JitFromScratch.o'
+# CACHELESS-CHECK-NOT: Write cached object 'cache/JitFromScratch.o'
+# CACHELESS-CHECK: Integer Distances: 3, 0, 3
+
+# Object files are written to the cache
+# RUN: JitFromScratch -cache-dir=cache/ -debug -debug-only=jitfromscratch 2>&1 | FileCheck -check-prefix=WRITE-CHECK %s
+# WRITE-CHECK-NOT: Read cached object 'cache/JitFromScratch.o'
+# WRITE-CHECK: Create new cache directory 'cache/'
+# WRITE-CHECK: Write cached object 'cache/JitFromScratch.o'
+# WRITE-CHECK: Integer Distances: 3, 0, 3
+
+# Object files are read from the cache
+# RUN: JitFromScratch -cache-dir=cache/ -debug -debug-only=jitfromscratch 2>&1 | FileCheck -check-prefix=READ-CHECK %s
+# READ-CHECK: Read cached object 'cache/JitFromScratch.o'
+# READ-CHECK-NOT: Write cached object 'cache/JitFromScratch.o'
+# READ-CHECK: Integer Distances: 3, 0, 3