Merge pull request #510 from ChrisCummins/shm-exe

[llvm] Use on-disk cache as benchmark scratch directory

compiler_gym/envs/llvm/datasets/cbench.py

@@ -288,7 +288,7 @@ def _make_cBench_validator(
     def validator_cb(env: "LlvmEnv") -> Optional[ValidationError]:  # noqa: F821
         """The validation callback."""
         with _CBENCH_DOWNLOAD_THREAD_LOCK:
-            with fasteners.InterProcessLock(cache_path("cbench-v1-runtime-data.LOCK")):
+            with fasteners.InterProcessLock(cache_path(".cbench-v1-runtime-data.LOCK")):
                 download_cBench_runtime_data()
 
         cbench_data = site_data_path("llvm-v0/cbench-v1-runtime-data/runtime_data")
@@ -557,7 +557,7 @@ def __init__(self, site_data_base: Path):
     def install(self):
         super().install()
         with _CBENCH_DOWNLOAD_THREAD_LOCK:
-            with fasteners.InterProcessLock(cache_path("cbench-v1-runtime-data.LOCK")):
+            with fasteners.InterProcessLock(cache_path(".cbench-v1-runtime-data.LOCK")):
                 download_cBench_runtime_data()
 
 

compiler_gym/envs/llvm/service/Benchmark.cc

@@ -67,6 +67,24 @@ RealizedBenchmarkDynamicConfig realizeDynamicConfig(const BenchmarkDynamicConfig
   return RealizedBenchmarkDynamicConfig(cfg);
 }
 
+/**
+ * Create a temporary directory to use as a scratch pad for on-disk storage.
+ * This directory is guaranteed to exist.
+ *
+ * Errors in this function are fatal.
+ *
+ * @return fs::path A path.
+ */
+fs::path createScratchDirectoryOrDie() {
+  const fs::path cacheRoot = util::getCacheRootPath();
+  const fs::path dir = fs::unique_path(cacheRoot / "benchmark-scratch-%%%%-%%%%");
+
+  sys::error_code ec;
+  fs::create_directories(dir, ec);
+  CHECK(!ec) << "Failed to create scratch directory: " << dir;
+  return dir;
+}
+
 }  // anonymous namespace
 
 Status readBitcodeFile(const fs::path& path, Bitcode* bitcode) {
@@ -135,36 +153,28 @@ Benchmark::Benchmark(const std::string& name, const Bitcode& bitcode,
                      const BaselineCosts& baselineCosts)
     : context_(std::make_unique<llvm::LLVMContext>()),
       module_(makeModuleOrDie(*context_, bitcode, name)),
-      scratchDirectory_(fs::path(fs::unique_path(workingDirectory / "scratch-%%%%-%%%%"))),
+      scratchDirectory_(createScratchDirectoryOrDie()),
       dynamicConfigProto_(dynamicConfig),
       dynamicConfig_(realizeDynamicConfig(dynamicConfig, scratchDirectory_)),
       baselineCosts_(baselineCosts),
       name_(name),
       needsRecompile_(true),
       runtimesPerObservationCount_(kDefaultRuntimesPerObservationCount),
       warmupRunsPerRuntimeObservationCount_(kDefaultWarmupRunsPerRuntimeObservationCount),
-      buildtimesPerObservationCount_(kDefaultBuildtimesPerObservationCount) {
-  sys::error_code ec;
-  fs::create_directory(scratchDirectory(), ec);
-  CHECK(!ec) << "Failed to create scratch directory: " << scratchDirectory();
-}
+      buildtimesPerObservationCount_(kDefaultBuildtimesPerObservationCount) {}
 
 Benchmark::Benchmark(const std::string& name, std::unique_ptr<llvm::LLVMContext> context,
                      std::unique_ptr<llvm::Module> module,
                      const BenchmarkDynamicConfig& dynamicConfig, const fs::path& workingDirectory,
                      const BaselineCosts& baselineCosts)
     : context_(std::move(context)),
       module_(std::move(module)),
-      scratchDirectory_(fs::path(fs::unique_path(workingDirectory / "scratch-%%%%-%%%%"))),
+      scratchDirectory_(createScratchDirectoryOrDie()),
       dynamicConfigProto_(dynamicConfig),
       dynamicConfig_(realizeDynamicConfig(dynamicConfig, scratchDirectory_)),
       baselineCosts_(baselineCosts),
       name_(name),
-      needsRecompile_(true) {
-  sys::error_code ec;
-  fs::create_directory(scratchDirectory(), ec);
-  CHECK(!ec) << "Failed to create scratch directory: " << scratchDirectory().string();
-}
+      needsRecompile_(true) {}
 
 void Benchmark::close() {
   sys::error_code ec;

compiler_gym/envs/llvm/service/BenchmarkFactory.cc

@@ -38,6 +38,15 @@ BenchmarkFactory::BenchmarkFactory(const boost::filesystem::path& workingDirecto
   VLOG(2) << "BenchmarkFactory initialized";
 }
 
+BenchmarkFactory::~BenchmarkFactory() { close(); }
+
+void BenchmarkFactory::close() {
+  VLOG(2) << "BenchmarkFactory closing with " << benchmarks_.size() << " entries";
+  for (auto& entry : benchmarks_) {
+    entry.second.close();
+  }
+}
+
 Status BenchmarkFactory::getBenchmark(const BenchmarkProto& benchmarkMessage,
                                       std::unique_ptr<Benchmark>* benchmark) {
   // Check if the benchmark has already been loaded into memory.

compiler_gym/envs/llvm/service/BenchmarkFactory.h

@@ -62,6 +62,10 @@ class BenchmarkFactory {
     return instance;
   }
 
+  ~BenchmarkFactory();
+
+  void close();
+
   /**
    * Get the requested named benchmark.
    *

compiler_gym/util/BUILD

@@ -68,6 +68,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         "@boost//:filesystem",
+        "@fmt",
     ],
 )
 

compiler_gym/util/RunfilesPath.cc

@@ -2,16 +2,24 @@
 //
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
+#include <fmt/format.h>
+
 #include "boost/filesystem.hpp"
 
 namespace fs = boost::filesystem;
 
 namespace compiler_gym::util {
 
-// When running under bazel, the working directory is the root of the
-// CompilerGym workspace. Back up one level so that we can reference other
-// workspaces.
-static const std::string kDefaultBase{"."};
+namespace {
+
+static const char* UNKNOWN_USER_NAME = "unknown";
+
+inline std::string getUser() {
+  const char* base = std::getenv("USER");
+  return base ? base : UNKNOWN_USER_NAME;
+}
+
+}  // namespace
 
 fs::path getRunfilesPath(const std::string& relPath) {
   const char* base = std::getenv("COMPILER_GYM_RUNFILES");
@@ -37,7 +45,26 @@ fs::path getSiteDataPath(const std::string& relPath) {
   } else {
     // $HOME may not be set under testing conditions. In this case, use a
     // throwaway directory.
-    return fs::temp_directory_path() / "compiler_gym" / relPath;
+    return fs::temp_directory_path() / fmt::format("compiler_gym_{}", getUser()) / relPath;
+  }
+}
+
+fs::path getCacheRootPath() {
+  // NOTE(cummins): This function has a related implementation in the Python
+  // sources, compiler_gym.util.runfiles_path.get_cache_path(). Any change to
+  // behavior here must be reflected in the Python version.
+  const char* force = std::getenv("COMPILER_GYM_CACHE");
+  if (force) {
+    return fs::path(force);
+  }
+
+  const char* home = std::getenv("HOME");
+  if (home) {
+    return fs::path(home) / ".local/cache/compiler_gym";
+  } else {
+    // $HOME may not be set under testing conditions. In this case, use a
+    // throwaway directory.
+    return fs::temp_directory_path() / fmt::format("compiler_gym_{}", getUser());
   }
 }
 

compiler_gym/util/RunfilesPath.h

@@ -22,4 +22,11 @@ boost::filesystem::path getRunfilesPath(const std::string& relPath);
  */
 boost::filesystem::path getSiteDataPath(const std::string& relPath);
 
+/**
+ * Resolve the root of the cache path.
+ *
+ * @return boost::filesystem::path A path.
+ */
+boost::filesystem::path getCacheRootPath();
+
 }  // namespace compiler_gym::util