1 | #pragma once |
---|---|
2 | |
3 | #include <memory> |
4 | #include <functional> |
5 | |
6 | #include "taichi/runtime/llvm/llvm_fwd.h" |
7 | #include "taichi/util/lang_util.h" |
8 | #include "taichi/jit/jit_module.h" |
9 | |
10 | namespace taichi::lang { |
11 | |
12 | // Backend JIT compiler for all archs |
13 | |
14 | class TaichiLLVMContext; |
15 | struct CompileConfig; |
16 | |
17 | class JITSession { |
18 | protected: |
19 | TaichiLLVMContext *tlctx_; |
20 | const CompileConfig &config_; |
21 | |
22 | std::vector<std::unique_ptr<JITModule>> modules; |
23 | |
24 | public: |
25 | JITSession(TaichiLLVMContext *tlctx, const CompileConfig &config); |
26 | |
27 | virtual JITModule *add_module(std::unique_ptr<llvm::Module> M, |
28 | int max_reg = 0) = 0; |
29 | |
30 | // virtual void remove_module(JITModule *module) = 0; |
31 | |
32 | virtual void *lookup(const std::string Name) { |
33 | TI_NOT_IMPLEMENTED |
34 | } |
35 | |
36 | virtual llvm::DataLayout get_data_layout() = 0; |
37 | |
38 | static std::unique_ptr<JITSession> create(TaichiLLVMContext *tlctx, |
39 | const CompileConfig &config, |
40 | Arch arch); |
41 | |
42 | virtual void global_optimize_module(llvm::Module *module) { |
43 | } |
44 | |
45 | virtual ~JITSession() = default; |
46 | }; |
47 | |
48 | } // namespace taichi::lang |
49 |