1#include "taichi/util/offline_cache.h"
2
3namespace taichi::lang::offline_cache {
4
5constexpr std::size_t offline_cache_key_length = 65;
6constexpr std::size_t min_mangled_name_length = offline_cache_key_length + 2;
7
8void disable_offline_cache_if_needed(CompileConfig *config) {
9 TI_ASSERT(config);
10 if (config->offline_cache) {
11 if (config->print_preprocessed_ir || config->print_ir ||
12 config->print_accessor_ir) {
13 config->offline_cache = false;
14 TI_WARN(
15 "Disable offline_cache because print_preprocessed_ir or print_ir or "
16 "print_accessor_ir is enabled");
17 }
18 }
19}
20
21std::string get_cache_path_by_arch(const std::string &base_path, Arch arch) {
22 std::string subdir;
23 if (arch_uses_llvm(arch)) {
24 subdir = kLlvmCachSubPath;
25 } else if (arch == Arch::vulkan || arch == Arch::opengl ||
26 arch == Arch::gles || arch == Arch::metal) {
27 subdir = kSpirvCacheSubPath;
28 } else if (arch == Arch::dx12) {
29 subdir = "dx12";
30 } else {
31 return base_path;
32 }
33 return taichi::join_path(base_path, subdir);
34}
35
36std::string mangle_name(const std::string &primal_name,
37 const std::string &key) {
38 // Result: {primal_name}{key: char[65]}_{(checksum(primal_name)) ^
39 // checksum(key)}
40 if (key.size() != offline_cache_key_length) {
41 return primal_name;
42 }
43 std::size_t checksum1{0}, checksum2{0};
44 for (auto &e : primal_name) {
45 checksum1 += std::size_t(e);
46 }
47 for (auto &e : key) {
48 checksum2 += std::size_t(e);
49 }
50 return fmt::format("{}{}_{}", primal_name, key, checksum1 ^ checksum2);
51}
52
53bool try_demangle_name(const std::string &mangled_name,
54 std::string &primal_name,
55 std::string &key) {
56 if (mangled_name.size() < min_mangled_name_length) {
57 return false;
58 }
59
60 std::size_t checksum{0}, checksum1{0}, checksum2{0};
61 auto pos = mangled_name.find_last_of('_');
62 if (pos == std::string::npos) {
63 return false;
64 }
65 try {
66 checksum = std::stoull(mangled_name.substr(pos + 1));
67 } catch (const std::exception &) {
68 return false;
69 }
70
71 std::size_t i = 0, primal_len = pos - offline_cache_key_length;
72 for (i = 0; i < primal_len; ++i) {
73 checksum1 += (int)mangled_name[i];
74 }
75 for (; i < pos; ++i) {
76 checksum2 += (int)mangled_name[i];
77 }
78 if ((checksum1 ^ checksum2) != checksum) {
79 return false;
80 }
81
82 primal_name = mangled_name.substr(0, primal_len);
83 key = mangled_name.substr(primal_len, offline_cache_key_length);
84 TI_ASSERT(key.size() == offline_cache_key_length);
85 TI_ASSERT(primal_name.size() + key.size() == pos);
86 return true;
87}
88
89std::size_t clean_offline_cache_files(const std::string &path) {
90 std::vector<const char *> sub_dirs = {kLlvmCachSubPath, kSpirvCacheSubPath,
91 kMetalCacheSubPath};
92
93 auto is_cache_filename = [](const std::string &name) {
94 const auto ext = taichi::filename_extension(name);
95 return ext == kLlvmCacheFilenameBCExt || ext == kLlvmCacheFilenameLLExt ||
96 ext == kSpirvCacheFilenameExt || ext == kMetalCacheFilenameExt ||
97 ext == "lock" || ext == "tcb";
98 };
99
100 std::size_t count = 0;
101
102 // Temp implementation. We will refactor the offline cache
103 taichi::traverse_directory(
104 path, [&count, &sub_dirs, &is_cache_filename, &path](
105 const std::string &name, bool is_dir) {
106 if (is_dir) { // ~/.cache/taichi/ticache/llvm ...
107 for (auto subdir : sub_dirs) {
108 auto subpath = taichi::join_path(path, subdir);
109
110 if (taichi::path_exists(subpath)) {
111 taichi::traverse_directory(
112 subpath, [&count, &is_cache_filename, &subpath](
113 const std::string &name, bool is_dir) {
114 if (is_cache_filename(name) && !is_dir) {
115 const auto fpath = taichi::join_path(subpath, name);
116 TI_TRACE("Removing {}", fpath);
117 bool ok = taichi::remove(fpath);
118 count += ok ? 1 : 0;
119 TI_WARN_IF(!ok, "Remove {} failed", fpath);
120 }
121 });
122 }
123 }
124 } else if (is_cache_filename(name)) {
125 const auto fpath = taichi::join_path(path, name);
126 TI_TRACE("Removing {}", fpath);
127 bool ok = taichi::remove(fpath);
128 count += ok ? 1 : 0;
129 TI_WARN_IF(!ok, "Remove {} failed", fpath);
130 }
131 });
132
133 return count;
134}
135
136} // namespace taichi::lang::offline_cache
137