1#pragma once
2
3#include <ctime>
4#include <cstdint>
5#include <queue>
6#include <string>
7#include <type_traits>
8#include <unordered_map>
9
10#include "taichi/common/core.h"
11#include "taichi/common/cleanup.h"
12#include "taichi/common/version.h"
13#include "taichi/rhi/arch.h"
14#include "taichi/util/io.h"
15#include "taichi/util/lock.h"
16#include "taichi/program/compile_config.h"
17
18namespace taichi::lang {
19namespace offline_cache {
20
21constexpr char kLlvmCacheFilenameLLExt[] = "ll";
22constexpr char kLlvmCacheFilenameBCExt[] = "bc";
23constexpr char kSpirvCacheFilenameExt[] = "spv";
24constexpr char kMetalCacheFilenameExt[] = "metal";
25constexpr char kLlvmCachSubPath[] = "llvm";
26constexpr char kSpirvCacheSubPath[] = "gfx";
27constexpr char kMetalCacheSubPath[] = "metal";
28
29using Version = std::uint16_t[3]; // {MAJOR, MINOR, PATCH}
30
31enum CleanCacheFlags {
32 NotClean = 0b000,
33 CleanOldVersion = 0b001,
34 CleanOldUsed = 0b010,
35 CleanOldCreated = 0b100
36};
37
38enum CleanCachePolicy {
39 Never = NotClean,
40 OnlyOldVersion = CleanOldVersion,
41 LRU = CleanOldVersion | CleanOldUsed,
42 FIFO = CleanOldVersion | CleanOldCreated
43};
44
45inline CleanCachePolicy string_to_clean_cache_policy(const std::string &str) {
46 if (str == "never") {
47 return Never;
48 } else if (str == "version") {
49 return OnlyOldVersion;
50 } else if (str == "lru") {
51 return LRU;
52 } else if (str == "fifo") {
53 return FIFO;
54 }
55 return Never;
56}
57
58template <typename KernelMetadataType>
59struct Metadata {
60 using KernelMetadata = KernelMetadataType;
61
62 Version version{};
63 std::size_t size{0}; // byte
64 std::unordered_map<std::string, KernelMetadata> kernels;
65
66 // NOTE: The "version" must be the first field to be serialized
67 TI_IO_DEF(version, size, kernels);
68};
69
70enum class LoadMetadataError {
71 kNoError,
72 kCorrupted,
73 kFileNotFound,
74 kVersionNotMatched,
75};
76
77template <typename MetadataType>
78inline LoadMetadataError load_metadata_with_checking(
79 MetadataType &result,
80 const std::string &filepath) {
81 if (!taichi::path_exists(filepath)) {
82 TI_DEBUG("Offline cache metadata file {} not found", filepath);
83 return LoadMetadataError::kFileNotFound;
84 }
85
86 using VerType = std::remove_reference_t<decltype(result.version)>;
87 static_assert(std::is_same_v<VerType, Version>);
88 const std::vector<uint8> bytes = read_data_from_file(filepath);
89
90 VerType ver{};
91 if (!read_from_binary(ver, bytes.data(), bytes.size(), false)) {
92 return LoadMetadataError::kCorrupted;
93 }
94 if (ver[0] != TI_VERSION_MAJOR || ver[1] != TI_VERSION_MINOR ||
95 ver[2] != TI_VERSION_PATCH) {
96 TI_DEBUG("The offline cache metadata file {} is old (version={}.{}.{})",
97 filepath, ver[0], ver[1], ver[2]);
98 return LoadMetadataError::kVersionNotMatched;
99 }
100
101 return !read_from_binary(result, bytes.data(), bytes.size())
102 ? LoadMetadataError::kCorrupted
103 : LoadMetadataError::kNoError;
104}
105
106struct CacheCleanerConfig {
107 std::string path;
108 CleanCachePolicy policy;
109 int max_size{0};
110 double cleaning_factor{0.f};
111 std::string metadata_filename;
112 std::string debugging_metadata_filename;
113 std::string metadata_lock_name;
114};
115
116template <typename MetadataType>
117struct CacheCleanerUtils {
118 using KernelMetaData = typename MetadataType::KernelMetadata;
119
120 // To save metadata as file
121 static bool save_metadata(const CacheCleanerConfig &config,
122 const MetadataType &data) {
123 TI_NOT_IMPLEMENTED;
124 }
125
126 static bool save_debugging_metadata(const CacheCleanerConfig &config,
127 const MetadataType &data) {
128 TI_NOT_IMPLEMENTED;
129 }
130
131 // To get cache files name
132 static std::vector<std::string> get_cache_files(
133 const CacheCleanerConfig &config,
134 const KernelMetaData &kernel_meta) {
135 TI_NOT_IMPLEMENTED;
136 }
137
138 // To remove other files except cache files and offline cache metadta files
139 static void remove_other_files(const CacheCleanerConfig &config) {
140 TI_NOT_IMPLEMENTED;
141 }
142
143 // To check if a file is cache file
144 static bool is_valid_cache_file(const CacheCleanerConfig &config,
145 const std::string &name) {
146 TI_NOT_IMPLEMENTED;
147 }
148};
149
150template <typename MetadataType>
151class CacheCleaner {
152 using Utils = CacheCleanerUtils<MetadataType>;
153 using KernelMetadata = typename MetadataType::KernelMetadata;
154
155 public:
156 static void run(const CacheCleanerConfig &config) {
157 TI_ASSERT(!config.path.empty());
158 TI_ASSERT(config.max_size > 0);
159 TI_ASSERT(!config.metadata_filename.empty());
160 TI_ASSERT(!config.metadata_lock_name.empty());
161 const auto policy = config.policy;
162 const auto &path = config.path;
163 const auto metadata_file =
164 taichi::join_path(path, config.metadata_filename);
165 const auto debugging_metadata_file =
166 taichi::join_path(path, config.debugging_metadata_filename);
167
168 if (policy == (std::size_t)NotClean) {
169 return;
170 }
171 if (!taichi::path_exists(path)) {
172 return;
173 }
174
175 MetadataType cache_data;
176 std::vector<std::string> files_to_rm;
177 bool ok_rm_meta = false;
178
179 // 1. Remove/Update metadata files
180 {
181 std::string lock_path =
182 taichi::join_path(path, config.metadata_lock_name);
183 if (!lock_with_file(lock_path)) {
184 TI_WARN(
185 "Lock {} failed. You can run 'ti cache clean -p {}' and try again.",
186 lock_path, path);
187 return;
188 }
189 auto _ = make_cleanup([&lock_path]() {
190 TI_DEBUG("Stop cleaning cache");
191 if (!unlock_with_file(lock_path)) {
192 TI_WARN(
193 "Unlock {} failed. You can remove this .lock file manually and "
194 "try again.",
195 lock_path);
196 }
197 });
198 TI_DEBUG("Start cleaning cache");
199
200 using Error = LoadMetadataError;
201 Error error = load_metadata_with_checking(cache_data, metadata_file);
202 if (error == Error::kFileNotFound) {
203 return;
204 } else if (error == Error::kCorrupted ||
205 error == Error::kVersionNotMatched) {
206 if (policy &
207 CleanOldVersion) { // Remove cache files and metadata files
208 TI_DEBUG("Removing all cache files");
209 if (taichi::remove(metadata_file)) {
210 taichi::remove(debugging_metadata_file);
211 Utils::remove_other_files(config);
212 bool success = taichi::traverse_directory(
213 config.path, [&config](const std::string &name, bool is_dir) {
214 if (!is_dir && Utils::is_valid_cache_file(config, name)) {
215 taichi::remove(taichi::join_path(config.path, name));
216 }
217 });
218 TI_ASSERT(success);
219 }
220 }
221 return;
222 }
223
224 if (cache_data.size < config.max_size ||
225 static_cast<std::size_t>(config.cleaning_factor *
226 cache_data.kernels.size()) == 0) {
227 return;
228 }
229
230 // LRU or FIFO
231 using KerData = std::pair<const std::string, KernelMetadata>;
232 using Comparator = std::function<bool(const KerData *, const KerData *)>;
233 using PriQueue =
234 std::priority_queue<const KerData *, std::vector<const KerData *>,
235 Comparator>;
236
237 Comparator cmp{nullptr};
238 if (policy & CleanOldUsed) { // LRU
239 cmp = [](const KerData *a, const KerData *b) -> bool {
240 return a->second.last_used_at < b->second.last_used_at;
241 };
242 } else if (policy & CleanOldCreated) { // FIFO
243 cmp = [](const KerData *a, const KerData *b) -> bool {
244 return a->second.created_at < b->second.created_at;
245 };
246 }
247
248 if (cmp) {
249 PriQueue q(cmp);
250 std::size_t cnt = config.cleaning_factor * cache_data.kernels.size();
251 TI_ASSERT(cnt != 0);
252 for (const auto &e : cache_data.kernels) {
253 if (q.size() == cnt && cmp(&e, q.top())) {
254 q.pop();
255 }
256 if (q.size() < cnt) {
257 q.push(&e);
258 }
259 }
260 TI_ASSERT(q.size() <= cnt);
261 while (!q.empty()) {
262 const auto *e = q.top();
263 for (const auto &f : Utils::get_cache_files(config, e->second)) {
264 files_to_rm.push_back(f);
265 }
266 cache_data.size -= e->second.size;
267 cache_data.kernels.erase(e->first);
268 q.pop();
269 }
270
271 if (cache_data.kernels.empty()) { // Remove
272 ok_rm_meta = taichi::remove(metadata_file);
273 taichi::remove(debugging_metadata_file);
274 Utils::remove_other_files(config);
275 } else { // Update
276 Utils::save_metadata(config, cache_data);
277 ok_rm_meta = true;
278 }
279 }
280 }
281
282 // 2. Remove cache files
283 if (ok_rm_meta) {
284 if (!cache_data.kernels.empty()) {
285 // For debugging (Not safe: without locking)
286 Utils::save_debugging_metadata(config, cache_data);
287 }
288 for (const auto &f : files_to_rm) {
289 auto file_path = taichi::join_path(path, f);
290 taichi::remove(file_path);
291 }
292 }
293 }
294};
295
296void disable_offline_cache_if_needed(CompileConfig *config);
297std::string get_cache_path_by_arch(const std::string &base_path, Arch arch);
298std::string mangle_name(const std::string &primal_name, const std::string &key);
299bool try_demangle_name(const std::string &mangled_name,
300 std::string &primal_name,
301 std::string &key);
302
303// utils to manage the offline cache files
304std::size_t clean_offline_cache_files(const std::string &path);
305
306} // namespace offline_cache
307} // namespace taichi::lang
308