1 | #pragma once |
2 | |
3 | #include <memory> |
4 | |
5 | #ifdef TI_WITH_LLVM |
6 | #include "llvm/IR/Module.h" |
7 | #include "taichi/common/core.h" |
8 | #include "taichi/common/serialization.h" |
9 | #include "taichi/runtime/llvm/launch_arg_info.h" |
10 | #include "taichi/program/kernel.h" |
11 | #include "taichi/util/offline_cache.h" |
12 | #include "taichi/codegen/llvm/llvm_compiled_data.h" |
13 | |
14 | namespace taichi::lang { |
15 | |
16 | struct LlvmOfflineCache { |
17 | using Version = uint16[3]; // {MAJOR, MINOR, PATCH} |
18 | |
19 | enum Format { |
20 | LL = 0x01, |
21 | BC = 0x10, |
22 | }; |
23 | |
24 | struct KernelCacheData { |
25 | std::string kernel_key; |
26 | std::vector<LlvmLaunchArgInfo> args; |
27 | LLVMCompiledKernel compiled_data; |
28 | |
29 | // For cache cleaning |
30 | std::size_t size{0}; // byte |
31 | std::time_t created_at{0}; // millsec |
32 | std::time_t last_used_at{0}; // millsec |
33 | |
34 | KernelCacheData() = default; |
35 | KernelCacheData(KernelCacheData &&) = default; |
36 | KernelCacheData &operator=(KernelCacheData &&) = default; |
37 | ~KernelCacheData() = default; |
38 | |
39 | KernelCacheData clone() const; |
40 | |
41 | TI_IO_DEF(kernel_key, args, compiled_data, size, created_at, last_used_at); |
42 | }; |
43 | |
44 | struct FieldCacheData { |
45 | struct SNodeCacheData { |
46 | int id{0}; |
47 | SNodeType type = SNodeType::undefined; |
48 | size_t cell_size_bytes{0}; |
49 | size_t chunk_size{0}; |
50 | |
51 | TI_IO_DEF(id, type, cell_size_bytes, chunk_size); |
52 | }; |
53 | |
54 | int tree_id{0}; |
55 | int root_id{0}; |
56 | size_t root_size{0}; |
57 | std::vector<SNodeCacheData> snode_metas; |
58 | |
59 | TI_IO_DEF(tree_id, root_id, root_size, snode_metas); |
60 | |
61 | // TODO(zhanlue): refactor llvm::Modules |
62 | // |
63 | // struct_module will eventually get cloned into each kernel_module, |
64 | // so there's no need to serialize it here. |
65 | // |
66 | // We have three different types of llvm::Module |
67 | // 1. runtime_module: contains runtime functions. |
68 | // 2. struct_module: contains compiled SNodeTree in llvm::Type. |
69 | // 3. kernel_modules: contains compiled kernel codes. |
70 | // |
71 | // The way those modules work rely on a recursive clone mechanism: |
72 | // runtime_module = load("runtime.bc") |
73 | // struct_module = clone(runtime_module) + compiled-SNodeTree |
74 | // kernel_module = clone(struct_module) + compiled-Kernel |
75 | // |
76 | // As a result, every kernel_module contains a copy of struct_module + |
77 | // runtime_module. |
78 | // |
79 | // This recursive clone mechanism is super fragile, |
80 | // which potentially causes inconsistency between modules if not handled |
81 | // properly. |
82 | // |
83 | // Let's turn to use llvm::link to bind the modules, |
84 | // and make runtime_module, struct_module, kernel_module independent of each |
85 | // other |
86 | }; |
87 | |
88 | using KernelMetadata = KernelCacheData; // Required by CacheCleaner |
89 | |
90 | Version version{}; |
91 | std::size_t size{0}; // byte |
92 | |
93 | // TODO(zhanlue): we need a better identifier for each FieldCacheData |
94 | // (SNodeTree) Given that snode_tree_id is not continuous, it is ridiculous to |
95 | // ask the users to remember each of the snode_tree_ids |
96 | // ** Find a way to name each SNodeTree ** |
97 | std::unordered_map<int, FieldCacheData> fields; // key = snode_tree_id |
98 | |
99 | std::unordered_map<std::string, KernelCacheData> |
100 | kernels; // key = kernel_name |
101 | |
102 | // NOTE: The "version" must be the first field to be serialized |
103 | TI_IO_DEF(version, size, fields, kernels); |
104 | }; |
105 | |
106 | class LlvmOfflineCacheFileReader { |
107 | public: |
108 | bool get_kernel_cache(LlvmOfflineCache::KernelCacheData &res, |
109 | const std::string &key, |
110 | llvm::LLVMContext &llvm_ctx); |
111 | |
112 | bool get_field_cache(LlvmOfflineCache::FieldCacheData &res, |
113 | int snode_tree_id); |
114 | |
115 | size_t get_num_snode_trees(); |
116 | |
117 | static std::unique_ptr<LlvmOfflineCacheFileReader> make( |
118 | const std::string &path, |
119 | LlvmOfflineCache::Format format = LlvmOfflineCache::Format::LL); |
120 | |
121 | static bool load_meta_data(LlvmOfflineCache &data, |
122 | const std::string &cache_file_path, |
123 | bool with_lock = true); |
124 | |
125 | private: |
126 | LlvmOfflineCacheFileReader(const std::string &path, |
127 | LlvmOfflineCache &&data, |
128 | LlvmOfflineCache::Format format); |
129 | |
130 | std::unique_ptr<llvm::Module> load_module(const std::string &path_prefix, |
131 | const std::string &key, |
132 | llvm::LLVMContext &llvm_ctx) const; |
133 | |
134 | std::string path_; |
135 | LlvmOfflineCache data_; |
136 | LlvmOfflineCache::Format format_; |
137 | }; |
138 | |
139 | class LlvmOfflineCacheFileWriter { |
140 | public: |
141 | using CleanCachePolicy = offline_cache::CleanCachePolicy; |
142 | |
143 | void set_data(LlvmOfflineCache &&data) { |
144 | this->mangled_ = false; |
145 | this->data_ = std::move(data); |
146 | } |
147 | |
148 | void set_data(std::unique_ptr<LlvmOfflineCache> &&data_ptr) { |
149 | set_data(std::move(*data_ptr.get())); |
150 | } |
151 | |
152 | void add_kernel_cache(const std::string &key, |
153 | LlvmOfflineCache::KernelCacheData &&kernel_cache) { |
154 | data_.kernels[key] = std::move(kernel_cache); |
155 | } |
156 | |
157 | void dump(const std::string &path, |
158 | LlvmOfflineCache::Format format = LlvmOfflineCache::Format::LL, |
159 | bool merge_with_old = false); |
160 | |
161 | void set_no_mangle() { |
162 | mangled_ = true; |
163 | } |
164 | |
165 | static void clean_cache(const std::string &path, |
166 | CleanCachePolicy policy, |
167 | int max_bytes, |
168 | double cleaning_factor); |
169 | |
170 | private: |
171 | void merge_with(LlvmOfflineCache &&data); |
172 | |
173 | void mangle_offloaded_task_name(const std::string &kernel_key, |
174 | LLVMCompiledKernel &compiled_data); |
175 | |
176 | LlvmOfflineCache data_; |
177 | bool mangled_{false}; |
178 | }; |
179 | |
180 | } // namespace taichi::lang |
181 | #endif // TI_WITH_LLVM |
182 | |