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
22namespace taichi::lang {
23namespace {
24
25using Format = LlvmOfflineCache::Format;
26constexpr char kMetadataFilename[] = "metadata";
27constexpr char kMetadataFileLockName[] = "metadata.lock";
28
29static std::string get_llvm_cache_metadata_file_path(const std::string &dir) {
30 return taichi::join_path(dir, std::string(kMetadataFilename) + ".tcb");
31}
32
33static 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
38static 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
48namespace offline_cache {
49
50template <>
51struct 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
100std::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
111bool 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
146LlvmOfflineCacheFileReader::LlvmOfflineCacheFileReader(
147 const std::string &path,
148 LlvmOfflineCache &&data,
149 LlvmOfflineCache::Format format)
150 : path_(path), data_(std::move(data)), format_(format) {
151}
152
153size_t LlvmOfflineCacheFileReader::get_num_snode_trees() {
154 return data_.fields.size();
155}
156
157bool 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
171bool 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
219std::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
247void 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
357void 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
381void 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
396void 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
412LlvmOfflineCache::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