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
8namespace taichi::lang {
9
10void 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
19void 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
31void 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