1 | #pragma once |
2 | |
3 | #include "taichi/aot/module_builder.h" |
4 | #include "taichi/runtime/llvm/llvm_offline_cache.h" |
5 | #include "taichi/codegen/llvm/codegen_llvm.h" |
6 | |
7 | namespace taichi::lang { |
8 | |
9 | class LlvmAotModuleBuilder : public AotModuleBuilder { |
10 | public: |
11 | explicit LlvmAotModuleBuilder(const CompileConfig &compile_config, |
12 | LlvmProgramImpl *prog, |
13 | TaichiLLVMContext &tlctx) |
14 | : compile_config_(compile_config), prog_(prog), tlctx_(tlctx) { |
15 | } |
16 | |
17 | void dump(const std::string &output_dir, |
18 | const std::string &filename) const override; |
19 | |
20 | protected: |
21 | void add_per_backend(const std::string &identifier, Kernel *kernel) override; |
22 | virtual LLVMCompiledKernel compile_kernel(Kernel *kernel) = 0; |
23 | |
24 | void add_field_per_backend(const std::string &identifier, |
25 | const SNode *rep_snode, |
26 | bool is_scalar, |
27 | DataType dt, |
28 | std::vector<int> shape, |
29 | int row_num, |
30 | int column_num) override; |
31 | |
32 | const LlvmOfflineCache &get_cache() { |
33 | return cache_; |
34 | } |
35 | |
36 | const CompileConfig &get_compile_config() const { |
37 | return compile_config_; |
38 | } |
39 | |
40 | TaichiLLVMContext &get_taichi_llvm_context() { |
41 | return tlctx_; |
42 | } |
43 | |
44 | private: |
45 | mutable LlvmOfflineCache cache_; |
46 | const CompileConfig &compile_config_; |
47 | LlvmProgramImpl *prog_ = nullptr; |
48 | TaichiLLVMContext &tlctx_; |
49 | }; |
50 | |
51 | } // namespace taichi::lang |
52 | |