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
14namespace taichi::lang {
15
16struct 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
106class 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
139class 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