1#pragma once
2
3#include "taichi/aot/module_loader.h"
4#include "taichi/runtime/llvm/llvm_runtime_executor.h"
5
6namespace 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*/
12TI_DLL_EXPORT void allocate_aot_snode_tree_type(aot::Module *aot_module,
13 aot::Field *aot_field,
14 uint64 *result_buffer);
15
16class 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