[CLI] Add "ti cache clean" command to clean the offline cache files manually

python/taichi/_main.py

@@ -18,6 +18,7 @@
 from rich.syntax import Syntax
 from taichi._lib import core as _ti_core
 from taichi._lib import utils
+from taichi.lang import impl
 from taichi.tools import cc_compose, diagnose, video
 
 import taichi as ti
@@ -826,6 +827,28 @@ def lint(arguments: list = sys.argv[2:]):
         import pylint  # pylint: disable=C0415
         pylint.lint.Run(options)
 
+    @staticmethod
+    @register
+    def ticache(arguments: list = sys.argv[2:]):
+        """Manage the ticache files manually"""
+        if len(arguments) < 1:
+            return
+
+        subcmd = arguments[0]
+        arguments = arguments[1:]
+        if subcmd == 'clean':
+            parser = argparse.ArgumentParser(
+                prog='ti ticache',
+                description='Clean all ticache files in given path')
+            parser.add_argument(
+                '-p',
+                '--offline-cache-file-path',
+                dest='offline_cache_file_path',
+                default=impl.default_cfg().offline_cache_file_path)
+            args = parser.parse_args(arguments)
+            path = args.offline_cache_file_path
+            _ti_core.clean_offline_cache_files(path)
+
 
 def main():
     cli = TaichiMain()

taichi/cache/gfx/cache_manager.cpp

@@ -79,7 +79,8 @@ struct CacheCleanerUtils<gfx::CacheManager::Metadata> {
       const KernelMetaData &kernel_meta) {
     std::vector<std::string> result;
     for (std::size_t i = 0; i < kernel_meta.num_files; ++i) {
-      result.push_back(kernel_meta.kernel_key + std::to_string(i) + ".spv");
+      result.push_back(kernel_meta.kernel_key + std::to_string(i) + "." +
+                       kSpirvCacheFilenameExt);
     }
     return result;
   }
@@ -95,7 +96,7 @@ struct CacheCleanerUtils<gfx::CacheManager::Metadata> {
   // To check if a file is cache file
   static bool is_valid_cache_file(const CacheCleanerConfig &config,
                                   const std::string &name) {
-    return filename_extension(name) == "spv";
+    return filename_extension(name) == kSpirvCacheFilenameExt;
   }
 };
 
@@ -121,20 +122,27 @@ CacheManager::CacheManager(Params &&init_params)
       auto exists =
           taichi::path_exists(taichi::join_path(path_, kAotMetadataFilename)) &&
           taichi::path_exists(taichi::join_path(path_, kGraphMetadataFilename));
-      if (exists && lock_with_file(lock_path)) {
-        auto _ = make_cleanup([&lock_path]() {
-          if (!unlock_with_file(lock_path)) {
-            TI_WARN(
-                "Unlock {} failed. You can remove this .lock file manually and "
-                "try again.",
-                lock_path);
-          }
-        });
-        gfx::AotModuleParams params;
-        params.module_path = path_;
-        params.runtime = runtime_;
-        params.enable_lazy_loading = true;
-        cached_module_ = gfx::make_aot_module(params, init_params.arch);
+      if (exists) {
+        if (lock_with_file(lock_path)) {
+          auto _ = make_cleanup([&lock_path]() {
+            if (!unlock_with_file(lock_path)) {
+              TI_WARN(
+                  "Unlock {} failed. You can remove this .lock file manually "
+                  "and try again.",
+                  lock_path);
+            }
+          });
+          gfx::AotModuleParams params;
+          params.module_path = path_;
+          params.runtime = runtime_;
+          params.enable_lazy_loading = true;
+          cached_module_ = gfx::make_aot_module(params, init_params.arch);
+        } else {
+          TI_WARN(
+              "Lock {} failed. You can run 'ti ticache clean -p {}' and try "
+              "again.",
+              lock_path, path_);
+        }
       }
     }
   }

taichi/cache/metal/cache_manager.cpp

@@ -33,7 +33,7 @@ struct CacheCleanerUtils<metal::CacheManager::Metadata> {
   static std::vector<std::string> get_cache_files(
       const CacheCleanerConfig &config,
       const KernelMetaData &kernel_meta) {
-    std::string fn = kernel_meta.kernel_key + ".metal";
+    std::string fn = kernel_meta.kernel_key + "." + kMetalCacheFilenameExt;
     return {fn};
   }
 
@@ -46,7 +46,7 @@ struct CacheCleanerUtils<metal::CacheManager::Metadata> {
   // To check if a file is cache file
   static bool is_valid_cache_file(const CacheCleanerConfig &config,
                                   const std::string &name) {
-    return filename_extension(name) == "metal";
+    return filename_extension(name) == kMetalCacheFilenameExt;
   }
 };
 
@@ -62,6 +62,10 @@ CacheManager::CacheManager(Params &&init_params)
     if (lock_with_file(lock_path)) {
       auto _ = make_unlocker(lock_path);
       offline_cache::load_metadata_with_checking(cached_data_, filepath);
+    } else {
+      TI_WARN(
+          "Lock {} failed. You can run 'ti ticache clean -p {}' and try again.",
+          lock_path, config_.cache_path);
     }
   }
 }

taichi/python/export_misc.cpp

@@ -19,6 +19,7 @@
 #include "taichi/system/dynamic_loader.h"
 #include "taichi/system/hacked_signal_handler.h"
 #include "taichi/system/profiler.h"
+#include "taichi/util/offline_cache.h"
 #if defined(TI_WITH_CUDA)
 #include "taichi/rhi/cuda/cuda_driver.h"
 #endif
@@ -178,6 +179,9 @@ void export_misc(py::module &m) {
   m.def("with_cc", []() { return false; });
 #endif
 
+  m.def("clean_offline_cache_files",
+        lang::offline_cache::clean_offline_cache_files);
+
   py::class_<HackedSignalRegister>(m, "HackedSignalRegister").def(py::init<>());
 }
 

taichi/runtime/llvm/llvm_offline_cache.cpp

@@ -38,8 +38,8 @@ static std::string get_llvm_cache_metadata_json_file_path(
 static std::vector<std::string> get_possible_llvm_cache_filename_by_key(
     const std::string &key) {
   return {
-      key + ".ll",
-      key + ".bc",
+      key + "." + offline_cache::kLlvmCacheFilenameLLExt,
+      key + "." + offline_cache::kLlvmCacheFilenameBCExt,
   };
 }
 
@@ -90,7 +90,7 @@ struct CacheCleanerUtils<LlvmOfflineCache> {
   static bool is_valid_cache_file(const CacheCleanerConfig &config,
                                   const std::string &name) {
     std::string ext = filename_extension(name);
-    return ext == "ll" || ext == "bc";
+    return ext == kLlvmCacheFilenameLLExt || ext == kLlvmCacheFilenameBCExt;
   }
 };
 
@@ -138,9 +138,8 @@ bool LlvmOfflineCacheFileReader::load_meta_data(
     });
     return Error::kNoError == load_metadata_with_checking(data, tcb_path);
   }
-  TI_WARN(
-      "Lock {} failed. You can remove this .lock file manually and try again.",
-      lock_path);
+  TI_WARN("Lock {} failed. You can run 'ti ticache clean -p {}' and try again.",
+          lock_path, cache_file_path);
   return false;
 }
 
@@ -224,12 +223,15 @@ std::unique_ptr<llvm::Module> LlvmOfflineCacheFileReader::load_module(
   TI_AUTO_PROF;
   if (format_ & Format::BC) {
     LlvmModuleBitcodeLoader loader;
-    return loader.set_bitcode_path(path_prefix + ".bc")
+    return loader
+        .set_bitcode_path(path_prefix + "." +
+                          offline_cache::kLlvmCacheFilenameBCExt)
         .set_buffer_id(key)
         .set_inline_funcs(false)
         .load(&llvm_ctx);
   } else if (format_ & Format::LL) {
-    const std::string filename = path_prefix + ".ll";
+    const std::string filename =
+        path_prefix + "." + offline_cache::kLlvmCacheFilenameLLExt;
     llvm::SMDiagnostic err;
     auto ret = llvm::parseAssemblyFile(filename, err, llvm_ctx);
     if (!ret) {  // File not found or Parse failed
@@ -270,7 +272,8 @@ void LlvmOfflineCacheFileWriter::dump(const std::string &path,
       auto *mod = data.module.get();
       TI_ASSERT(mod != nullptr);
       if (format & Format::LL) {
-        std::string filename = filename_prefix + ".ll";
+        std::string filename =
+            filename_prefix + "." + offline_cache::kLlvmCacheFilenameLLExt;
         if (!merge_with_old || try_lock_with_file(filename)) {
           size += write_llvm_module(filename, [mod](llvm::raw_os_ostream &os) {
             mod->print(os, /*AAW=*/nullptr);
@@ -280,7 +283,8 @@ void LlvmOfflineCacheFileWriter::dump(const std::string &path,
         }
       }
       if (format & Format::BC) {
-        std::string filename = filename_prefix + ".bc";
+        std::string filename =
+            filename_prefix + "." + offline_cache::kLlvmCacheFilenameBCExt;
         if (!merge_with_old || try_lock_with_file(filename)) {
           size += write_llvm_module(filename, [mod](llvm::raw_os_ostream &os) {
             llvm::WriteBitcodeToFile(*mod, os);
@@ -319,9 +323,8 @@ void LlvmOfflineCacheFileWriter::dump(const std::string &path,
     std::string lock_path = taichi::join_path(path, kMetadataFileLockName);
     if (!lock_with_file(lock_path)) {
       TI_WARN(
-          "Lock {} failed. You can remove this .lock file manually and try "
-          "again.",
-          lock_path);
+          "Lock {} failed. You can run 'ti ticache clean -p {}' and try again.",
+          lock_path, path);
       return;
     }
     auto _ = make_cleanup([&lock_path]() {

taichi/util/offline_cache.cpp

@@ -21,11 +21,11 @@ void disable_offline_cache_if_needed(CompileConfig *config) {
 std::string get_cache_path_by_arch(const std::string &base_path, Arch arch) {
   std::string subdir;
   if (arch_uses_llvm(arch)) {
-    subdir = "llvm";
+    subdir = kLlvmCachSubPath;
   } else if (arch == Arch::vulkan || arch == Arch::opengl) {
-    subdir = "gfx";
+    subdir = kSpirvCacheSubPath;
   } else if (arch == Arch::metal) {
-    subdir = "metal";
+    subdir = kMetalCacheSubPath;
   } else if (arch == Arch::dx12) {
     subdir = "dx12";
   } else {
@@ -87,4 +87,39 @@ bool try_demangle_name(const std::string &mangled_name,
   return true;
 }
 
+void clean_offline_cache_files(const std::string &path) {
+  std::vector<const char *> sub_dirs = {kLlvmCachSubPath, kSpirvCacheSubPath,
+                                        kMetalCacheSubPath};
+  auto is_cache_filename = [](const std::string &name) {
+    const auto ext = taichi::filename_extension(name);
+    return ext == kLlvmCacheFilenameBCExt || ext == kLlvmCacheFilenameLLExt ||
+           ext == kSpirvCacheFilenameExt || ext == kMetalCacheFilenameExt ||
+           ext == "lock" || ext == "tcb";
+  };
+  // Temp implementation. We will refactor the offline cache
+  taichi::traverse_directory(path, [&sub_dirs, &is_cache_filename, &path](
+                                       const std::string &name, bool is_dir) {
+    if (is_dir) {  // ~/.cache/taichi/ticache/llvm ...
+      for (auto subdir : sub_dirs) {
+        auto subpath = taichi::join_path(path, subdir);
+        if (taichi::path_exists(subpath)) {
+          taichi::traverse_directory(
+              subpath, [&is_cache_filename, &subpath](const std::string &name,
+                                                      bool is_dir) {
+                if (is_cache_filename(name) && !is_dir) {
+                  const auto fpath = taichi::join_path(subpath, name);
+                  TI_TRACE("Remove {}", fpath);
+                  taichi::remove(fpath);
+                }
+              });
+        }
+      }
+    } else if (is_cache_filename(name)) {
+      const auto fpath = taichi::join_path(path, name);
+      TI_TRACE("Remove {}", fpath);
+      taichi::remove(fpath);
+    }
+  });
+}
+
 }  // namespace taichi::lang::offline_cache

taichi/util/offline_cache.h

@@ -18,6 +18,14 @@
 namespace taichi::lang {
 namespace offline_cache {
 
+constexpr char kLlvmCacheFilenameLLExt[] = "ll";
+constexpr char kLlvmCacheFilenameBCExt[] = "bc";
+constexpr char kSpirvCacheFilenameExt[] = "spv";
+constexpr char kMetalCacheFilenameExt[] = "metal";
+constexpr char kLlvmCachSubPath[] = "llvm";
+constexpr char kSpirvCacheSubPath[] = "gfx";
+constexpr char kMetalCacheSubPath[] = "metal";
+
 using Version = std::uint16_t[3];  // {MAJOR, MINOR, PATCH}
 
 enum CleanCacheFlags {
@@ -174,9 +182,9 @@ class CacheCleaner {
           taichi::join_path(path, config.metadata_lock_name);
       if (!lock_with_file(lock_path)) {
         TI_WARN(
-            "Lock {} failed. You can remove this .lock file manually and try "
+            "Lock {} failed. You can run 'ti ticache clean -p {}' and try "
             "again.",
-            lock_path);
+            lock_path, path);
         return;
       }
       auto _ = make_cleanup([&lock_path]() {
@@ -293,5 +301,8 @@ bool try_demangle_name(const std::string &mangled_name,
                        std::string &primal_name,
                        std::string &key);
 
+// utils to manage ticache files
+void clean_offline_cache_files(const std::string &path);
+
 }  // namespace offline_cache
 }  // namespace taichi::lang

tests/python/test_cli.py

@@ -1,6 +1,9 @@
 import argparse
 import copy
+import os
+import shutil
 import sys
+import tempfile
 from contextlib import contextmanager
 from pathlib import Path
 from unittest.mock import patch
@@ -9,6 +12,7 @@
 from taichi._main import TaichiMain
 
 import taichi as ti
+from tests import test_utils
 
 
 @contextmanager
@@ -206,3 +210,45 @@ def test_cli_run():
         cli = TaichiMain(test_mode=True)
         args = cli()
         assert args.filename == "a.py"
+
+
+def test_cli_ticache():
+    archs = {ti.cpu, ti.cuda, ti.opengl, ti.vulkan, ti.metal}
+    archs = {v for v in archs if v in test_utils.expected_archs()}
+    exts = ('ll', 'bc', 'spv', 'metal', 'tcb', 'lock')
+    tmp_path = tempfile.mkdtemp()
+
+    @ti.kernel
+    def simple_kernel(a: ti.i32) -> ti.i32:
+        return -a
+
+    def launch_kernel(arch):
+        ti.init(arch=arch,
+                offline_cache=True,
+                offline_cache_file_path=tmp_path)
+        simple_kernel(128)
+        ti.reset()
+
+    for arch in archs:
+        launch_kernel(arch)
+
+    found = False
+    for root, dirs, files in os.walk(tmp_path):
+        for file in files:
+            if file.endswith(exts):
+                found = True
+                break
+        if found:
+            break
+    assert found
+
+    with patch_sys_argv_helper(["ti", "ticache", "clean", "-p",
+                                tmp_path]) as custom_argv:
+        cli = TaichiMain()
+        cli()
+
+    for root, dirs, files in os.walk(tmp_path):
+        for file in files:
+            assert not file.endswith(exts)
+
+    shutil.rmtree(tmp_path)