[misc] Implement KernelCompialtionManager::clean_offline_cache (taich…

…i-dev#7515)

Issue: taichi-dev#7002

taichi/compilation_manager/kernel_compilation_manager.cpp

@@ -2,9 +2,54 @@
 
 #include "taichi/analysis/offline_cache_util.h"
 #include "taichi/codegen/compiled_kernel_data.h"
+#include "taichi/util/offline_cache.h"
 
 namespace taichi::lang {
 
+namespace offline_cache {
+
+template <>
+struct CacheCleanerUtils<CacheData> {
+  using MetadataType = CacheData;
+  using KernelMetaData = typename MetadataType::KernelMetadata;
+
+  // To save metadata as file
+  static bool save_metadata(const CacheCleanerConfig &config,
+                            const MetadataType &data) {
+    write_to_binary_file(
+        data, taichi::join_path(config.path, config.metadata_filename));
+    return true;
+  }
+
+  static bool save_debugging_metadata(const CacheCleanerConfig &config,
+                                      const MetadataType &data) {
+    return true;
+  }
+
+  // To get cache files name
+  static std::vector<std::string> get_cache_files(
+      const CacheCleanerConfig &config,
+      const KernelMetaData &kernel_meta) {
+    auto fn = fmt::format(KernelCompilationManager::kCacheFilenameFormat,
+                          kernel_meta.kernel_key);
+    return {fn};
+  }
+
+  // To remove other files except cache files and offline cache metadta files
+  static void remove_other_files(const CacheCleanerConfig &config) {
+    // Do nothing
+  }
+
+  // To check if a file is cache file
+  static bool is_valid_cache_file(const CacheCleanerConfig &config,
+                                  const std::string &name) {
+    std::string ext = filename_extension(name);
+    return ext == kTiCacheFilenameExt;
+  }
+};
+
+}  // namespace offline_cache
+
 KernelCompilationManager::KernelCompilationManager(Config config)
     : config_(std::move(config)) {
   TI_DEBUG("Create KernelCompilationManager with offline_cache_file_path = {}",
@@ -95,15 +140,23 @@ void KernelCompilationManager::clean_offline_cache(
     offline_cache::CleanCachePolicy policy,
     int max_bytes,
     double cleaning_factor) const {
-  // TODO(PGZXB): Impl
+  using CacheCleaner = offline_cache::CacheCleaner<CacheData>;
+  offline_cache::CacheCleanerConfig config;
+  config.path = config_.offline_cache_path;
+  config.policy = policy;
+  config.cleaning_factor = cleaning_factor;
+  config.max_size = max_bytes;
+  config.metadata_filename = kMetadataFilename;
+  config.debugging_metadata_filename = "";
+  config.metadata_lock_name = kMetadataLockName;
+  CacheCleaner::run(config);
 }
 
 std::string KernelCompilationManager::make_filename(
     const std::string &kernel_key,
     Arch arch) const {
-  return join_path(
-      config_.offline_cache_path,
-      fmt::format(kCacheFilenameFormat, kernel_key, arch_name(arch)));
+  return join_path(config_.offline_cache_path,
+                   fmt::format(kCacheFilenameFormat, kernel_key));
 }
 
 std::unique_ptr<CompiledKernelData> KernelCompilationManager::compile_kernel(

taichi/compilation_manager/kernel_compilation_manager.h

@@ -32,6 +32,8 @@ struct CacheData {
     TI_IO_DEF(kernel_key, size, created_at, last_used_at);
   };
 
+  using KernelMetadata = KernelData;  // Required by CacheCleaner
+
   Version version{};
   std::size_t size{0};
   std::unordered_map<std::string, KernelData> kernels;
@@ -41,14 +43,14 @@ struct CacheData {
 };
 
 class KernelCompilationManager final {
+ public:
   static constexpr char kMetadataFilename[] = "ticache.tcb";
-  static constexpr char kCacheFilenameFormat[] = "{}-{}.tic";
+  static constexpr char kCacheFilenameFormat[] = "{}.tic";
   static constexpr char kMetadataLockName[] = "ticache.lock";
 
   using KernelCacheData = CacheData::KernelData;
   using CachingKernels = std::unordered_map<std::string, KernelCacheData>;
 
- public:
   struct Config {
     std::string offline_cache_path;
     std::unique_ptr<KernelCompiler> kernel_compiler;

tests/python/test_offline_cache.py

@@ -492,8 +492,7 @@ def helper():
         curr_arch, [2, 2])
 
 
-@pytest.mark.parametrize('curr_arch',
-                         supported_archs_offline_cache - supported_gfx_archs)
+@pytest.mark.parametrize('curr_arch', supported_archs_offline_cache)
 @pytest.mark.parametrize('factor', [0.0, 0.25, 0.85, 1.0])
 @pytest.mark.parametrize('policy', ['never', 'version', 'lru', 'fifo'])
 @_test_offline_cache_dec
@@ -511,8 +510,6 @@ def run_simple_kernels(max_size):
         only_init(max_size)
         for kernel, args, get_res, num_offloads in simple_kernels_to_test:
             assert kernel(*args) == test_utils.approx(get_res(*args))
-            # The timestamp used by cache cleaning is at second precision, so we should make sure the kernels are not used in the same second
-            sleep(1)
 
     kernel_count = len(simple_kernels_to_test)
     count_of_cache_file = cache_files_cnt(curr_arch)
@@ -522,7 +519,7 @@ def added_files(arch):
 
     assert added_files(curr_arch) == expected_num_cache_files(curr_arch)
 
-    run_simple_kernels(1024**3)  # 1GB
+    run_simple_kernels(1024**3)  # 1GB (>> size_of_cache_files)
     ti.reset()  # Dumping cache data
     size_of_cache_files = cache_files_size(
         backend_specified_cache_path(curr_arch))
@@ -534,7 +531,7 @@ def added_files(arch):
     assert added_files(curr_arch) == expected_num_cache_files(
         curr_arch, [kern[3] for kern in simple_kernels_to_test])
 
-    only_init(size_of_cache_files)
+    only_init(1)  # 1B (<< size_of_cache_files)
     ti.reset()
     rem = []
     if policy in ['never', 'version']: