1 | #pragma once |
---|---|
2 | |
3 | #include "taichi/aot/module_loader.h" |
4 | #include "taichi/runtime/llvm/llvm_runtime_executor.h" |
5 | |
6 | namespace taichi::lang { |
7 | |
8 | /* TODO(zhanlue) refactor this interface once SNodeTreeType is available |
9 | The "aot::Field" created by "make_new_field()" is a SNodeTree in essense. |
10 | Therefore we're actually initializing the entire SNodeTree. |
11 | */ |
12 | TI_DLL_EXPORT void allocate_aot_snode_tree_type(aot::Module *aot_module, |
13 | aot::Field *aot_field, |
14 | uint64 *result_buffer); |
15 | |
16 | class LlvmAotModule : public aot::Module { |
17 | public: |
18 | explicit LlvmAotModule(const std::string &module_path, |
19 | LlvmRuntimeExecutor *executor) |
20 | : executor_(executor), |
21 | cache_reader_(LlvmOfflineCacheFileReader::make(module_path)) { |
22 | TI_ASSERT(executor_ != nullptr); |
23 | |
24 | const std::string graph_path = fmt::format("{}/graphs.tcb", module_path); |
25 | read_from_binary_file(graphs_, graph_path); |
26 | } |
27 | |
28 | Arch arch() const override { |
29 | return executor_->get_config().arch; |
30 | } |
31 | |
32 | uint64_t version() const override { |
33 | return 0; |
34 | } |
35 | |
36 | size_t get_root_size() const override { |
37 | return 0; |
38 | } |
39 | |
40 | LlvmRuntimeExecutor *const get_runtime_executor() { |
41 | return executor_; |
42 | } |
43 | |
44 | size_t get_num_snode_trees() { |
45 | return cache_reader_->get_num_snode_trees(); |
46 | } |
47 | |
48 | void set_initialized_snode_tree(int snode_tree_id) { |
49 | initialized_snode_tree_ids.insert(snode_tree_id); |
50 | } |
51 | |
52 | bool is_snode_tree_initialized(int snode_tree_id) { |
53 | return initialized_snode_tree_ids.count(snode_tree_id); |
54 | } |
55 | |
56 | std::unique_ptr<aot::CompiledGraph> get_graph( |
57 | const std::string &name) override; |
58 | |
59 | protected: |
60 | virtual FunctionType convert_module_to_function( |
61 | const std::string &name, |
62 | LlvmOfflineCache::KernelCacheData &&loaded) = 0; |
63 | |
64 | LlvmOfflineCache::KernelCacheData load_kernel_from_cache( |
65 | const std::string &name); |
66 | |
67 | std::unique_ptr<aot::Kernel> make_new_kernel( |
68 | const std::string &name) override; |
69 | |
70 | /* TODO(zhanlue): replace "make_new_field()" with "make_snode_tree()" once |
71 | SNodeTreeType is available Field is not a standalone data structure - it is |
72 | essentially part of a SNodeTree object. User should always operate on a |
73 | "SNodeTree" instead of a "Field". |
74 | */ |
75 | std::unique_ptr<aot::Field> make_new_field(const std::string &name) override; |
76 | |
77 | LlvmRuntimeExecutor *const executor_{nullptr}; |
78 | std::unique_ptr<LlvmOfflineCacheFileReader> cache_reader_{nullptr}; |
79 | |
80 | // To prevent repeated SNodeTree initialization |
81 | std::unordered_set<int> initialized_snode_tree_ids; |
82 | }; |
83 | |
84 | } // namespace taichi::lang |
85 |