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

Issue facebookresearch#465.
ChrisCummins · Dec 9, 2021 · b9acc4c · b9acc4c
1 parent 98b82bd
commit b9acc4c
Showing 1 changed file with 22 additions and 12 deletions.
diff --git a/compiler_gym/envs/llvm/service/Benchmark.cc b/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;