[misc] Implement cache file cleaning

taichi/program/compile_config.h

@@ -110,6 +110,10 @@ struct CompileConfig {
   // Offline cache options
   bool offline_cache{false};
   std::string offline_cache_file_path{get_repo_dir() + "ticache"};
+  std::string offline_cache_cleaning_policy{
+      "never"};  // "never"|"version"|"lru"|"fifo"
+  int offline_cache_max_size_of_files{1024 * 1024};  // bytes
+  double offline_cache_cleaning_factor{0.25};        // [0.f, 1.f]
 
   CompileConfig();
 };

taichi/python/export_lang.cpp

@@ -238,7 +238,13 @@ void export_lang(py::module &m) {
                      &CompileConfig::auto_mesh_local_default_occupacy)
       .def_readwrite("offline_cache", &CompileConfig::offline_cache)
       .def_readwrite("offline_cache_file_path",
-                     &CompileConfig::offline_cache_file_path);
+                     &CompileConfig::offline_cache_file_path)
+      .def_readwrite("offline_cache_cleaning_policy",
+                     &CompileConfig::offline_cache_cleaning_policy)
+      .def_readwrite("offline_cache_max_size_of_files",
+                     &CompileConfig::offline_cache_max_size_of_files)
+      .def_readwrite("offline_cache_cleaning_factor",
+                     &CompileConfig::offline_cache_cleaning_factor);
 
   m.def("reset_default_compile_config",
         [&]() { default_compile_config = CompileConfig(); });

taichi/runtime/llvm/llvm_offline_cache.cpp

@@ -1,7 +1,16 @@
 #include "llvm_offline_cache.h"
 
-#include <fstream>
-#include <sstream>
+#include <queue>
+
+#if __has_include(<filesystem>)
+#include <filesystem>
+namespace fs = std::filesystem;
+#elif __has_include(<experimental/filesystem>)
+#include <experimental/filesystem>
+namespace fs = std::experimental::filesystem;
+#else
+error "Missing the <filesystem> header."
+#endif  //  __has_include(<filesystem>)
 
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/Bitcode/BitcodeReader.h"
@@ -10,7 +19,9 @@
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_os_ostream.h"
 #include "llvm/Transforms/Utils/Cloning.h"
+#include "taichi/common/version.h"
 #include "taichi/ir/transforms.h"
+#include "taichi/program/kernel.h"
 #include "taichi/runtime/llvm/llvm_context.h"
 
 namespace taichi {
@@ -20,6 +31,34 @@ namespace {
 using Format = LlvmOfflineCache::Format;
 constexpr char kMetadataFilename[] = "metadata";
 
+static bool is_current_llvm_cache_version(
+    const LlvmOfflineCache::Version &ver) {
+  // TODO(PGZXB): Do more detailed checking
+  return ver[0] == TI_VERSION_MAJOR && ver[1] == TI_VERSION_MINOR &&
+         ver[2] == TI_VERSION_PATCH;
+}
+
+static std::string get_llvm_cache_metadata_file_path(const std::string &dir) {
+  std::stringstream tcb_ss;
+  tcb_ss << dir << "/" << kMetadataFilename << ".tcb";
+  return tcb_ss.str();
+}
+
+static std::string get_llvm_cache_metadata_json_file_path(
+    const std::string &dir) {
+  std::stringstream tcb_ss;
+  tcb_ss << dir << "/" << kMetadataFilename << ".json";
+  return tcb_ss.str();
+}
+
+static std::vector<std::string> get_possible_llvm_cache_filename_by_key(
+    const std::string &key) {
+  return {
+      key + ".ll",
+      key + ".bc",
+  };
+}
+
 }  // namespace
 
 // static
@@ -37,9 +76,7 @@ std::unique_ptr<LlvmOfflineCacheFileReader> LlvmOfflineCacheFileReader::make(
 bool LlvmOfflineCacheFileReader::load_meta_data(
     LlvmOfflineCache &data,
     const std::string &cache_file_path) {
-  std::stringstream tcb_ss;
-  tcb_ss << cache_file_path << "/" << kMetadataFilename << ".tcb";
-  const auto tcb_path = tcb_ss.str();
+  const auto tcb_path = get_llvm_cache_metadata_file_path(cache_file_path);
   {
     // No the best way to check for filepath existence, but whatever... See
     // https://stackoverflow.com/questions/12774207/fastest-way-to-check-if-a-file-exists-using-standard-c-c11-14-17-c
@@ -123,7 +160,11 @@ void LlvmOfflineCacheFileWriter::dump(const std::string &path,
                                       LlvmOfflineCache::Format format,
                                       bool merge_with_old) {
   taichi::create_directories(path);
+  std::time_t now = std::time(nullptr);
+  std::size_t new_kernels_size = 0;  // bytes
+
   for (auto &[k, v] : data_.kernels) {
+    std::size_t size = 0;  // bytes
     std::stringstream filename_ss;
     filename_ss << path << "/" << k;
     std::string filename_prefix = filename_ss.str();
@@ -137,6 +178,7 @@ void LlvmOfflineCacheFileWriter::dump(const std::string &path,
           TI_ERROR_IF(!os.is_open(), "File {} open failed", filename);
           llvm::raw_os_ostream llvm_os{os};
           writer(llvm_os);
+          return llvm_os.tell();
         };
     {
       auto *mod = v.module;
@@ -147,38 +189,45 @@ void LlvmOfflineCacheFileWriter::dump(const std::string &path,
 
       mangle_offloaded_task_name(k, mod, v.offloaded_task_list);
       if (format & Format::LL) {
-        write_llvm_module(".ll", [mod](llvm::raw_os_ostream &os) {
+        size += write_llvm_module(".ll", [mod](llvm::raw_os_ostream &os) {
           mod->print(os, /*AAW=*/nullptr);
         });
       }
       if (format & Format::BC) {
-        write_llvm_module(".bc", [mod](llvm::raw_os_ostream &os) {
+        size += write_llvm_module(".bc", [mod](llvm::raw_os_ostream &os) {
           llvm::WriteBitcodeToFile(*mod, os);
         });
       }
     }
+
+    // Set meta info
+    v.created_at = now;
+    v.last_used_at = now;
+    v.size = size;
+    TI_ASSERT(v.size > 0);
+    new_kernels_size += v.size;
   }
-  {
-    // Merge with old metadata
-    if (merge_with_old) {
-      LlvmOfflineCache old_data;
-      if (LlvmOfflineCacheFileReader::load_meta_data(old_data, path)) {
-        add_data(std::move(old_data));
-      }
+
+  data_.version[0] = TI_VERSION_MAJOR;
+  data_.version[1] = TI_VERSION_MINOR;
+  data_.version[2] = TI_VERSION_PATCH;
+  data_.size = new_kernels_size;
+  // Merge with old metadata
+  if (merge_with_old) {
+    LlvmOfflineCache old_data;
+    if (LlvmOfflineCacheFileReader::load_meta_data(old_data, path)) {
+      merge_with(std::move(old_data));
     }
-    // Dump metadata
-    std::stringstream prefix_ss;
-    prefix_ss << path << "/" << kMetadataFilename;
-    const std::string file_prefix = prefix_ss.str();
-    write_to_binary_file(data_, file_prefix + ".tcb");
-    // For debugging
-    TextSerializer ts;
-    ts.serialize_to_json("cache", data_);
-    ts.write_to_file(file_prefix + ".json");
   }
+  // Dump metadata
+  write_to_binary_file(data_, get_llvm_cache_metadata_file_path(path));
+  // For debugging
+  TextSerializer ts;
+  ts.serialize_to_json("cache", data_);
+  ts.write_to_file(get_llvm_cache_metadata_json_file_path(path));
 }
 
-void LlvmOfflineCacheFileWriter::add_data(LlvmOfflineCache &&data) {
+void LlvmOfflineCacheFileWriter::merge_with(LlvmOfflineCache &&data) {
   // Note: merge this->data_ with data, new cover old
   auto &new_kernels = data_.kernels;
   auto &new_fields = data_.fields;
@@ -189,7 +238,14 @@ void LlvmOfflineCacheFileWriter::add_data(LlvmOfflineCache &&data) {
     old_fields[k] = std::move(v);
   }
   for (auto &[k, v] : new_kernels) {
-    old_kernels[k] = std::move(v);
+    auto iter = old_kernels.find(k);
+    if (iter == old_kernels.end()) {
+      data.size += v.size;
+      old_kernels[k] = std::move(v);
+    } else {
+      data.size += v.size - iter->second.size;
+      iter->second = std::move(v);
+    }
   }
 
   data_ = std::move(data);
@@ -212,5 +268,96 @@ void LlvmOfflineCacheFileWriter::mangle_offloaded_task_name(
   }
 }
 
+void LlvmOfflineCacheFileWriter::clean_cache(const std::string &path,
+                                             CleanCachePolicy policy,
+                                             int max_bytes,
+                                             double cleaning_factor) {
+  if (policy == (std::size_t)NotClean) {
+    return;
+  }
+
+  LlvmOfflineCache cache_data;
+  LlvmOfflineCacheFileReader::load_meta_data(cache_data, path);
+
+  if ((policy & CleanOldVersion) &&
+      !is_current_llvm_cache_version(cache_data.version)) {
+    if (bool ok = fs::remove(get_llvm_cache_metadata_file_path(path)) &&
+                  fs::remove(get_llvm_cache_metadata_json_file_path(path));
+        ok) {
+      auto root_path = fs::path(path);
+      for (const auto &[k, v] : cache_data.kernels) {
+        const auto files = get_possible_llvm_cache_filename_by_key(k);
+        for (const auto &f : files) {
+          fs::remove(root_path / f);
+        }
+      }
+    }
+    return;
+  }
+
+  if (cache_data.size < max_bytes ||
+      static_cast<std::size_t>(cleaning_factor * cache_data.kernels.size()) ==
+          0) {
+    return;
+  }
+
+  // LRU or FIFO
+  using KerData = LlvmOfflineCache::KernelCacheData;
+  using Comparator = std::function<bool(const KerData &, const KerData &)>;
+  using PriQueue =
+      std::priority_queue<KerData, std::vector<KerData>, Comparator>;
+
+  Comparator cmp{nullptr};
+  if (policy & CleanOldUsed) {  // LRU
+    cmp = [](const KerData &a, const KerData &b) -> bool {
+      return a.last_used_at < b.last_used_at;
+    };
+  } else if (policy & CleanOldCreated) {  // FIFO
+    cmp = [](const KerData &a, const KerData &b) -> bool {
+      return a.created_at < b.created_at;
+    };
+  }
+  if (cmp) {
+    PriQueue q(cmp);
+    std::size_t cnt = cleaning_factor * cache_data.kernels.size();
+    TI_ASSERT(cnt != 0);
+    for (auto &[k, v] : cache_data.kernels) {
+      if (q.size() == cnt && cmp(v, q.top()))
+        q.pop();
+      if (q.size() < cnt)
+        q.push(std::move(v));
+    }
+    TI_ASSERT(q.size() <= cnt);
+    auto root_path = fs::path(path);
+    while (!q.empty()) {
+      for (const auto &f :
+           get_possible_llvm_cache_filename_by_key(q.top().kernel_key)) {
+        fs::remove(root_path / f);
+      }
+      q.pop();
+    }
+    if (cnt == cache_data.kernels.size()) {  // Removed all
+      fs::remove(get_llvm_cache_metadata_file_path(path));
+      fs::remove(get_llvm_cache_metadata_json_file_path(path));
+    }
+  }
+}
+
+LlvmOfflineCacheFileWriter::CleanCachePolicy
+LlvmOfflineCacheFileWriter::string_to_clean_cache_policy(
+    const std::string &str) {
+  if (str == "never")
+    return Never;
+  if (str == "version")
+    return OnlyOldVersion;
+  if (str == "lru")
+    return LRU;
+  if (str == "fifo")
+    return FIFO;
+
+  TI_WARN("Invalid CleanCachePolicy");
+  return Never;
+}
+
 }  // namespace lang
 }  // namespace taichi

taichi/runtime/llvm/llvm_offline_cache.h

@@ -13,6 +13,8 @@ namespace taichi {
 namespace lang {
 
 struct LlvmOfflineCache {
+  using Version = uint16[3];  // {MAJOR, MINOR, PATCH}
+
   enum Format {
     LL = 0x01,
     BC = 0x10,
@@ -34,12 +36,22 @@ struct LlvmOfflineCache {
     std::unique_ptr<llvm::Module> owned_module{nullptr};
     llvm::Module *module{nullptr};
 
+    // For cache cleaning
+    std::size_t size{0};          // byte
+    std::time_t created_at{0};    // millsec
+    std::time_t last_used_at{0};  // millsec
+
     KernelCacheData() = default;
     KernelCacheData(KernelCacheData &&) = default;
     KernelCacheData &operator=(KernelCacheData &&) = default;
     ~KernelCacheData() = default;
 
-    TI_IO_DEF(kernel_key, args, offloaded_task_list);
+    TI_IO_DEF(kernel_key,
+              args,
+              offloaded_task_list,
+              size,
+              created_at,
+              last_used_at);
   };
 
   struct FieldCacheData {
@@ -86,6 +98,9 @@ struct LlvmOfflineCache {
     // other
   };
 
+  Version version{};
+  std::size_t size{0};  // byte
+
   // TODO(zhanlue): we need a better identifier for each FieldCacheData
   // (SNodeTree) Given that snode_tree_id is not continuous, it is ridiculous to
   // ask the users to remember each of the snode_tree_ids
@@ -95,7 +110,7 @@ struct LlvmOfflineCache {
   std::unordered_map<std::string, KernelCacheData>
       kernels;  // key = kernel_name
 
-  TI_IO_DEF(fields, kernels);
+  TI_IO_DEF(version, size, fields, kernels);
 };
 
 class LlvmOfflineCacheFileReader {
@@ -129,7 +144,21 @@ class LlvmOfflineCacheFileReader {
 };
 
 class LlvmOfflineCacheFileWriter {
+  enum CleanCacheFlags {
+    NotClean = 0b000,
+    CleanOldVersion = 0b001,
+    CleanOldUsed = 0b010,
+    CleanOldCreated = 0b100
+  };
+
  public:
+  enum CleanCachePolicy {
+    Never = NotClean,
+    OnlyOldVersion = CleanOldVersion,
+    LRU = CleanOldVersion | CleanOldUsed,
+    FIFO = CleanOldVersion | CleanOldCreated
+  };
+
   void set_data(LlvmOfflineCache &&data) {
     this->mangled_ = false;
     this->data_ = std::move(data);
@@ -148,8 +177,15 @@ class LlvmOfflineCacheFileWriter {
     mangled_ = true;
   }
 
+  static void clean_cache(const std::string &path,
+                          CleanCachePolicy policy,
+                          int max_bytes,
+                          double cleaning_factor);
+
+  static CleanCachePolicy string_to_clean_cache_policy(const std::string &str);
+
  private:
-  void add_data(LlvmOfflineCache &&data);
+  void merge_with(LlvmOfflineCache &&data);
 
   void mangle_offloaded_task_name(
       const std::string &kernel_key,