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 | |
18 | namespace taichi::lang { |
19 | namespace offline_cache { |
20 | |
21 | constexpr char kLlvmCacheFilenameLLExt[] = "ll" ; |
22 | constexpr char kLlvmCacheFilenameBCExt[] = "bc" ; |
23 | constexpr char kSpirvCacheFilenameExt[] = "spv" ; |
24 | constexpr char kMetalCacheFilenameExt[] = "metal" ; |
25 | constexpr char kLlvmCachSubPath[] = "llvm" ; |
26 | constexpr char kSpirvCacheSubPath[] = "gfx" ; |
27 | constexpr char kMetalCacheSubPath[] = "metal" ; |
28 | |
29 | using Version = std::uint16_t[3]; // {MAJOR, MINOR, PATCH} |
30 | |
31 | enum CleanCacheFlags { |
32 | NotClean = 0b000, |
33 | CleanOldVersion = 0b001, |
34 | CleanOldUsed = 0b010, |
35 | CleanOldCreated = 0b100 |
36 | }; |
37 | |
38 | enum CleanCachePolicy { |
39 | Never = NotClean, |
40 | OnlyOldVersion = CleanOldVersion, |
41 | LRU = CleanOldVersion | CleanOldUsed, |
42 | FIFO = CleanOldVersion | CleanOldCreated |
43 | }; |
44 | |
45 | inline 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 | |
58 | template <typename KernelMetadataType> |
59 | struct 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 | |
70 | enum class LoadMetadataError { |
71 | kNoError, |
72 | kCorrupted, |
73 | kFileNotFound, |
74 | kVersionNotMatched, |
75 | }; |
76 | |
77 | template <typename MetadataType> |
78 | inline 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 | |
106 | struct 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 | |
116 | template <typename MetadataType> |
117 | struct 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 | |
150 | template <typename MetadataType> |
151 | class 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 | |
296 | void disable_offline_cache_if_needed(CompileConfig *config); |
297 | std::string get_cache_path_by_arch(const std::string &base_path, Arch arch); |
298 | std::string mangle_name(const std::string &primal_name, const std::string &key); |
299 | bool 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 |
304 | std::size_t clean_offline_cache_files(const std::string &path); |
305 | |
306 | } // namespace offline_cache |
307 | } // namespace taichi::lang |
308 | |