1 | #include "taichi/runtime/llvm/llvm_aot_module_builder.h" |
2 | |
3 | #include <algorithm> |
4 | #include "taichi/runtime/llvm/launch_arg_info.h" |
5 | #include "taichi/runtime/program_impls/llvm/llvm_program.h" |
6 | #include "taichi/runtime/llvm/aot_graph_data.h" |
7 | |
8 | namespace taichi::lang { |
9 | |
10 | void LlvmAotModuleBuilder::dump(const std::string &output_dir, |
11 | const std::string &filename) const { |
12 | LlvmOfflineCacheFileWriter writer; |
13 | writer.set_data(std::move(cache_)); |
14 | writer.dump(output_dir); |
15 | |
16 | dump_graph(output_dir); |
17 | } |
18 | |
19 | void LlvmAotModuleBuilder::add_per_backend(const std::string &identifier, |
20 | Kernel *kernel) { |
21 | auto compiled = compile_kernel(kernel); |
22 | LlvmOfflineCache::KernelCacheData kcache; |
23 | kcache.kernel_key = identifier; |
24 | kcache.compiled_data = std::move(compiled); |
25 | kcache.args = infer_launch_args(kernel); |
26 | kcache.last_used_at = std::time(nullptr); |
27 | kcache.created_at = std::time(nullptr); |
28 | cache_.kernels[identifier] = std::move(kcache); |
29 | } |
30 | |
31 | void LlvmAotModuleBuilder::add_field_per_backend(const std::string &identifier, |
32 | const SNode *rep_snode, |
33 | bool is_scalar, |
34 | DataType dt, |
35 | std::vector<int> shape, |
36 | int row_num, |
37 | int column_num) { |
38 | // Field refers to a leaf node(Place SNode) in a SNodeTree. |
39 | // It makes no sense to just serialize the leaf node or its corresponding |
40 | // branch. Instead, the minimal unit we have to serialize is the entire |
41 | // SNodeTree. Note that SNodeTree's uses snode_tree_id as its identifier, |
42 | // rather than the field's name. (multiple fields may end up referring to the |
43 | // same SNodeTree) |
44 | |
45 | // 1. Find snode_tree_id |
46 | int snode_tree_id = rep_snode->get_snode_tree_id(); |
47 | |
48 | // 2. Fetch Cache from the Program |
49 | // Kernel compilation is not allowed until all the Fields are finalized, |
50 | // so we finished SNodeTree compilation during AOTModuleBuilder construction. |
51 | // |
52 | // By the time "add_field_per_backend()" is called, |
53 | // SNodeTrees should have already been finalized, |
54 | // with compiled info stored in LlvmProgramImpl::cache_data_. |
55 | TI_ASSERT(prog_ != nullptr); |
56 | LlvmOfflineCache::FieldCacheData field_cache = |
57 | prog_->get_cached_field(snode_tree_id); |
58 | |
59 | // 3. Update AOT Cache |
60 | cache_.fields[snode_tree_id] = std::move(field_cache); |
61 | } |
62 | |
63 | } // namespace taichi::lang |
64 | |