1 | #include "llvm_offline_cache.h" |
2 | |
3 | #include <queue> |
4 | |
5 | #include "llvm/AsmParser/Parser.h" |
6 | #include "llvm/Bitcode/BitcodeReader.h" |
7 | #include "llvm/Bitcode/BitcodeWriter.h" |
8 | #include "llvm/IR/Module.h" |
9 | #include "llvm/Support/SourceMgr.h" |
10 | #include "llvm/Support/raw_os_ostream.h" |
11 | #include "llvm/Transforms/Utils/Cloning.h" |
12 | #include "taichi/analysis/offline_cache_util.h" |
13 | #include "taichi/common/cleanup.h" |
14 | #include "taichi/common/version.h" |
15 | #include "taichi/ir/transforms.h" |
16 | #include "taichi/program/kernel.h" |
17 | #include "taichi/runtime/llvm/llvm_context.h" |
18 | #include "taichi/util/io.h" |
19 | #include "taichi/util/lock.h" |
20 | #include "taichi/util/offline_cache.h" |
21 | |
22 | namespace taichi::lang { |
23 | namespace { |
24 | |
25 | using Format = LlvmOfflineCache::Format; |
26 | constexpr char kMetadataFilename[] = "metadata" ; |
27 | constexpr char kMetadataFileLockName[] = "metadata.lock" ; |
28 | |
29 | static std::string get_llvm_cache_metadata_file_path(const std::string &dir) { |
30 | return taichi::join_path(dir, std::string(kMetadataFilename) + ".tcb" ); |
31 | } |
32 | |
33 | static std::string get_llvm_cache_metadata_json_file_path( |
34 | const std::string &dir) { |
35 | return taichi::join_path(dir, std::string(kMetadataFilename) + ".json" ); |
36 | } |
37 | |
38 | static std::vector<std::string> get_possible_llvm_cache_filename_by_key( |
39 | const std::string &key) { |
40 | return { |
41 | key + "." + offline_cache::kLlvmCacheFilenameLLExt, |
42 | key + "." + offline_cache::kLlvmCacheFilenameBCExt, |
43 | }; |
44 | } |
45 | |
46 | } // namespace |
47 | |
48 | namespace offline_cache { |
49 | |
50 | template <> |
51 | struct CacheCleanerUtils<LlvmOfflineCache> { |
52 | using MetadataType = LlvmOfflineCache; |
53 | using KernelMetaData = typename MetadataType::KernelMetadata; |
54 | |
55 | // To save metadata as file |
56 | static bool save_metadata(const CacheCleanerConfig &config, |
57 | const MetadataType &data) { |
58 | write_to_binary_file( |
59 | data, taichi::join_path(config.path, config.metadata_filename)); |
60 | return true; |
61 | } |
62 | |
63 | static bool save_debugging_metadata(const CacheCleanerConfig &config, |
64 | const MetadataType &data) { |
65 | TextSerializer ts; |
66 | ts.serialize_to_json("cache" , data); |
67 | ts.write_to_file( |
68 | taichi::join_path(config.path, config.debugging_metadata_filename)); |
69 | return true; |
70 | } |
71 | |
72 | // To get cache files name |
73 | static std::vector<std::string> get_cache_files( |
74 | const CacheCleanerConfig &config, |
75 | const KernelMetaData &kernel_meta) { |
76 | std::vector<std::string> result; |
77 | for (const auto &f : |
78 | get_possible_llvm_cache_filename_by_key(kernel_meta.kernel_key)) { |
79 | result.push_back(f); |
80 | } |
81 | return result; |
82 | } |
83 | |
84 | // To remove other files except cache files and offline cache metadta files |
85 | static void remove_other_files(const CacheCleanerConfig &config) { |
86 | // Do nothing |
87 | } |
88 | |
89 | // To check if a file is cache file |
90 | static bool is_valid_cache_file(const CacheCleanerConfig &config, |
91 | const std::string &name) { |
92 | std::string ext = filename_extension(name); |
93 | return ext == kLlvmCacheFilenameLLExt || ext == kLlvmCacheFilenameBCExt; |
94 | } |
95 | }; |
96 | |
97 | } // namespace offline_cache |
98 | |
99 | // static |
100 | std::unique_ptr<LlvmOfflineCacheFileReader> LlvmOfflineCacheFileReader::make( |
101 | const std::string &path, |
102 | LlvmOfflineCache::Format format) { |
103 | LlvmOfflineCache data; |
104 | if (!load_meta_data(data, path)) { |
105 | return nullptr; |
106 | } |
107 | return std::unique_ptr<LlvmOfflineCacheFileReader>( |
108 | new LlvmOfflineCacheFileReader(path, std::move(data), format)); |
109 | } |
110 | |
111 | bool LlvmOfflineCacheFileReader::load_meta_data( |
112 | LlvmOfflineCache &data, |
113 | const std::string &cache_file_path, |
114 | bool with_lock) { |
115 | using offline_cache::load_metadata_with_checking; |
116 | using Error = offline_cache::LoadMetadataError; |
117 | const auto tcb_path = get_llvm_cache_metadata_file_path(cache_file_path); |
118 | |
119 | if (!taichi::path_exists(tcb_path)) { |
120 | TI_DEBUG("File {} not found" , tcb_path); |
121 | return false; |
122 | } |
123 | |
124 | if (!with_lock) { |
125 | return Error::kNoError == load_metadata_with_checking(data, tcb_path); |
126 | } |
127 | |
128 | std::string lock_path = |
129 | taichi::join_path(cache_file_path, kMetadataFileLockName); |
130 | if (lock_with_file(lock_path)) { |
131 | auto _ = make_cleanup([&lock_path]() { |
132 | if (!unlock_with_file(lock_path)) { |
133 | TI_WARN( |
134 | "Unlock {} failed. You can remove this .lock file manually and try " |
135 | "again." , |
136 | lock_path); |
137 | } |
138 | }); |
139 | return Error::kNoError == load_metadata_with_checking(data, tcb_path); |
140 | } |
141 | TI_WARN("Lock {} failed. You can run 'ti cache clean -p {}' and try again." , |
142 | lock_path, cache_file_path); |
143 | return false; |
144 | } |
145 | |
146 | LlvmOfflineCacheFileReader::LlvmOfflineCacheFileReader( |
147 | const std::string &path, |
148 | LlvmOfflineCache &&data, |
149 | LlvmOfflineCache::Format format) |
150 | : path_(path), data_(std::move(data)), format_(format) { |
151 | } |
152 | |
153 | size_t LlvmOfflineCacheFileReader::get_num_snode_trees() { |
154 | return data_.fields.size(); |
155 | } |
156 | |
157 | bool LlvmOfflineCacheFileReader::get_field_cache( |
158 | LlvmOfflineCache::FieldCacheData &res, |
159 | int snode_tree_id) { |
160 | auto itr = data_.fields.find(snode_tree_id); |
161 | if (itr == data_.fields.end()) { |
162 | TI_DEBUG("Cannot find field with snode_tree_id={}" , snode_tree_id); |
163 | return false; |
164 | } |
165 | |
166 | const auto &loaded_field_cache = itr->second; |
167 | res = loaded_field_cache; // copy assign |
168 | return true; |
169 | } |
170 | |
171 | bool LlvmOfflineCacheFileReader::get_kernel_cache( |
172 | LlvmOfflineCache::KernelCacheData &res, |
173 | const std::string &key, |
174 | llvm::LLVMContext &llvm_ctx) { |
175 | TI_AUTO_PROF; |
176 | auto itr = data_.kernels.find(key); |
177 | if (itr == data_.kernels.end()) { |
178 | TI_DEBUG("Cannot find kernel={}" , key); |
179 | return false; |
180 | } |
181 | |
182 | auto &kernel_data = itr->second; |
183 | auto &data = kernel_data.compiled_data; |
184 | if (!data.module) { |
185 | std::string filename_prefix = taichi::join_path(path_, key); |
186 | data.module = load_module(filename_prefix, key, llvm_ctx); |
187 | if (!data.module) { |
188 | data_.kernels.erase(itr); |
189 | return false; // Must return |
190 | } |
191 | } |
192 | res.compiled_data = data.clone(); |
193 | |
194 | kernel_data.last_used_at = std::time(nullptr); |
195 | |
196 | res.created_at = kernel_data.created_at; |
197 | res.last_used_at = kernel_data.last_used_at; |
198 | res.kernel_key = key; |
199 | res.args = kernel_data.args; |
200 | |
201 | // Verify the `res: LlvmOfflineCache::KernelCacheData` |
202 | const auto &compiled_data = res.compiled_data; |
203 | const auto &tasks = compiled_data.tasks; |
204 | bool verified = true; |
205 | for (const auto &t : tasks) { |
206 | if (compiled_data.module->getFunction(t.name) == nullptr) { |
207 | verified = false; |
208 | } |
209 | } |
210 | if (!verified) { |
211 | for (const auto &f : get_possible_llvm_cache_filename_by_key(key)) { |
212 | taichi::remove(taichi::join_path(path_, f)); |
213 | } |
214 | } |
215 | |
216 | return verified; |
217 | } |
218 | |
219 | std::unique_ptr<llvm::Module> LlvmOfflineCacheFileReader::load_module( |
220 | const std::string &path_prefix, |
221 | const std::string &key, |
222 | llvm::LLVMContext &llvm_ctx) const { |
223 | TI_AUTO_PROF; |
224 | if (format_ & Format::BC) { |
225 | LlvmModuleBitcodeLoader loader; |
226 | return loader |
227 | .set_bitcode_path(path_prefix + "." + |
228 | offline_cache::kLlvmCacheFilenameBCExt) |
229 | .set_buffer_id(key) |
230 | .set_inline_funcs(false) |
231 | .load(&llvm_ctx); |
232 | } else if (format_ & Format::LL) { |
233 | const std::string filename = |
234 | path_prefix + "." + offline_cache::kLlvmCacheFilenameLLExt; |
235 | llvm::SMDiagnostic err; |
236 | auto ret = llvm::parseAssemblyFile(filename, err, llvm_ctx); |
237 | if (!ret) { // File not found or Parse failed |
238 | TI_DEBUG("Fail to parse {}: {}" , filename, err.getMessage().str()); |
239 | return nullptr; |
240 | } |
241 | return ret; |
242 | } |
243 | TI_ERROR("Unknown LLVM format={}" , format_); |
244 | return nullptr; |
245 | } |
246 | |
247 | void LlvmOfflineCacheFileWriter::dump(const std::string &path, |
248 | LlvmOfflineCache::Format format, |
249 | bool merge_with_old) { |
250 | auto write_llvm_module = |
251 | [](const std::string &filename, |
252 | std::function<void(llvm::raw_os_ostream & os)> writer) { |
253 | std::ofstream os(filename, std::ios::out | std::ios::binary); |
254 | TI_ERROR_IF(!os.is_open(), "File {} open failed" , filename); |
255 | llvm::raw_os_ostream llvm_os{os}; |
256 | writer(llvm_os); |
257 | return llvm_os.tell(); |
258 | }; |
259 | |
260 | using Iter = typename decltype(data_.kernels)::iterator; |
261 | taichi::create_directories(path); |
262 | std::size_t new_kernels_size = 0; // bytes |
263 | std::vector<Iter> iters_to_erased; // Kernels which have been saved |
264 | |
265 | for (auto iter = data_.kernels.begin(); iter != data_.kernels.end(); ++iter) { |
266 | auto &[k, v] = *iter; |
267 | std::size_t size = 0; // bytes |
268 | std::string filename_prefix = taichi::join_path(path, k); |
269 | { |
270 | mangle_offloaded_task_name(k, v.compiled_data); |
271 | auto &data = v.compiled_data; |
272 | auto *mod = data.module.get(); |
273 | TI_ASSERT(mod != nullptr); |
274 | if (format & Format::LL) { |
275 | std::string filename = |
276 | filename_prefix + "." + offline_cache::kLlvmCacheFilenameLLExt; |
277 | if (!merge_with_old || try_lock_with_file(filename)) { |
278 | size += write_llvm_module(filename, [mod](llvm::raw_os_ostream &os) { |
279 | mod->print(os, /*AAW=*/nullptr); |
280 | }); |
281 | } else { |
282 | TI_DEBUG("Cache file {} exists" , filename); |
283 | } |
284 | } |
285 | if (format & Format::BC) { |
286 | std::string filename = |
287 | filename_prefix + "." + offline_cache::kLlvmCacheFilenameBCExt; |
288 | if (!merge_with_old || try_lock_with_file(filename)) { |
289 | size += write_llvm_module(filename, [mod](llvm::raw_os_ostream &os) { |
290 | llvm::WriteBitcodeToFile(*mod, os); |
291 | }); |
292 | } else { |
293 | TI_DEBUG("Cache file {} exists" , filename); |
294 | } |
295 | } |
296 | } |
297 | |
298 | // Set meta info |
299 | TI_ASSERT(v.created_at); |
300 | TI_ASSERT(v.last_used_at); |
301 | v.size = size; |
302 | new_kernels_size += v.size; |
303 | |
304 | if (v.size == 0) { // The kernel cache has been saved |
305 | iters_to_erased.push_back(iter); |
306 | } |
307 | } |
308 | |
309 | // Erase the kernels which aren't needed to re-saved |
310 | for (auto &iter : iters_to_erased) { |
311 | data_.kernels.erase(iter); |
312 | } |
313 | |
314 | data_.version[0] = TI_VERSION_MAJOR; |
315 | data_.version[1] = TI_VERSION_MINOR; |
316 | data_.version[2] = TI_VERSION_PATCH; |
317 | data_.size = new_kernels_size; |
318 | |
319 | { |
320 | // Lock |
321 | // TODO(PGZXB): High overhead (read -> merge -> write). Redesign the |
322 | // metadata file format to reduce overhead. |
323 | std::string lock_path = taichi::join_path(path, kMetadataFileLockName); |
324 | if (!lock_with_file(lock_path)) { |
325 | TI_WARN( |
326 | "Lock {} failed. You can run 'ti ticache clean -p {}' and try again." , |
327 | lock_path, path); |
328 | return; |
329 | } |
330 | auto _ = make_cleanup([&lock_path]() { |
331 | if (!unlock_with_file(lock_path)) { |
332 | TI_WARN( |
333 | "Unlock {} failed. You can remove this .lock file manually and try " |
334 | "again." , |
335 | lock_path); |
336 | } |
337 | }); |
338 | |
339 | // Merge with old metadata |
340 | if (merge_with_old) { |
341 | LlvmOfflineCache old_data; |
342 | if (LlvmOfflineCacheFileReader::load_meta_data(old_data, path, false)) { |
343 | merge_with(std::move(old_data)); |
344 | } |
345 | } |
346 | |
347 | // Dump metadata |
348 | std::string target_path = get_llvm_cache_metadata_file_path(path); |
349 | write_to_binary_file(data_, target_path); |
350 | } |
351 | // For debugging (Not safe: without locking) |
352 | TextSerializer ts; |
353 | ts.serialize_to_json("cache" , data_); |
354 | ts.write_to_file(get_llvm_cache_metadata_json_file_path(path)); |
355 | } |
356 | |
357 | void LlvmOfflineCacheFileWriter::merge_with(LlvmOfflineCache &&data) { |
358 | // Note: merge this->data_ with data, new cover old |
359 | auto &new_kernels = data_.kernels; |
360 | auto &new_fields = data_.fields; |
361 | auto &old_kernels = data.kernels; |
362 | auto &old_fields = data.fields; |
363 | |
364 | for (auto &[k, v] : new_fields) { |
365 | old_fields[k] = std::move(v); |
366 | } |
367 | for (auto &[k, v] : new_kernels) { |
368 | auto iter = old_kernels.find(k); |
369 | if (iter == old_kernels.end()) { |
370 | data.size += v.size; |
371 | old_kernels[k] = std::move(v); |
372 | } else { |
373 | data.size += v.size - iter->second.size; |
374 | iter->second = std::move(v); |
375 | } |
376 | } |
377 | |
378 | data_ = std::move(data); |
379 | } |
380 | |
381 | void LlvmOfflineCacheFileWriter::mangle_offloaded_task_name( |
382 | const std::string &kernel_key, |
383 | LLVMCompiledKernel &compiled_data) { |
384 | if (!mangled_) { |
385 | for (auto &offload : compiled_data.tasks) { |
386 | std::string mangled_name = |
387 | offline_cache::mangle_name(offload.name, kernel_key); |
388 | auto func = compiled_data.module->getFunction(offload.name); |
389 | TI_ASSERT(func != nullptr); |
390 | func->setName(mangled_name); |
391 | offload.name = mangled_name; |
392 | } |
393 | } |
394 | } |
395 | |
396 | void LlvmOfflineCacheFileWriter::clean_cache(const std::string &path, |
397 | CleanCachePolicy policy, |
398 | int max_bytes, |
399 | double cleaning_factor) { |
400 | using CacheCleaner = offline_cache::CacheCleaner<LlvmOfflineCache>; |
401 | offline_cache::CacheCleanerConfig config; |
402 | config.path = path; |
403 | config.policy = policy; |
404 | config.cleaning_factor = cleaning_factor; |
405 | config.max_size = max_bytes; |
406 | config.metadata_filename = std::string(kMetadataFilename) + ".tcb" ; |
407 | config.debugging_metadata_filename = std::string(kMetadataFilename) + ".json" ; |
408 | config.metadata_lock_name = kMetadataFileLockName; |
409 | CacheCleaner::run(config); |
410 | } |
411 | |
412 | LlvmOfflineCache::KernelCacheData LlvmOfflineCache::KernelCacheData::clone() |
413 | const { |
414 | LlvmOfflineCache::KernelCacheData result; |
415 | result.kernel_key = kernel_key; |
416 | result.args = args; |
417 | result.compiled_data = compiled_data.clone(); |
418 | result.size = size; |
419 | result.created_at = created_at; |
420 | result.last_used_at = last_used_at; |
421 | return result; |
422 | } |
423 | } // namespace taichi::lang |
424 | |