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
7namespace taichi::lang {
8
9class 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