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
10namespace taichi::lang {
11
12// Backend JIT compiler for all archs
13
14class TaichiLLVMContext;
15struct CompileConfig;
16
17class 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