1 | #include "taichi/jit/jit_session.h" |
2 | |
3 | #ifdef TI_WITH_LLVM |
4 | #include "llvm/IR/DataLayout.h" |
5 | #endif |
6 | |
7 | namespace taichi::lang { |
8 | |
9 | #ifdef TI_WITH_LLVM |
10 | std::unique_ptr<JITSession> create_llvm_jit_session_cpu( |
11 | TaichiLLVMContext *tlctx, |
12 | const CompileConfig &config, |
13 | Arch arch); |
14 | |
15 | std::unique_ptr<JITSession> create_llvm_jit_session_cuda( |
16 | TaichiLLVMContext *tlctx, |
17 | const CompileConfig &config, |
18 | Arch arch); |
19 | |
20 | std::unique_ptr<JITSession> create_llvm_jit_session_amdgpu( |
21 | TaichiLLVMContext *tlctx, |
22 | const CompileConfig &config, |
23 | Arch arch); |
24 | #endif |
25 | |
26 | JITSession::JITSession(TaichiLLVMContext *tlctx, const CompileConfig &config) |
27 | : tlctx_(tlctx), config_(config) { |
28 | } |
29 | |
30 | std::unique_ptr<JITSession> JITSession::create(TaichiLLVMContext *tlctx, |
31 | const CompileConfig &config, |
32 | Arch arch) { |
33 | #ifdef TI_WITH_LLVM |
34 | if (arch_is_cpu(arch)) { |
35 | return create_llvm_jit_session_cpu(tlctx, config, arch); |
36 | } else if (arch == Arch::cuda) { |
37 | #if defined(TI_WITH_CUDA) |
38 | return create_llvm_jit_session_cuda(tlctx, config, arch); |
39 | #else |
40 | TI_NOT_IMPLEMENTED |
41 | #endif |
42 | } else if (arch == Arch::dx12) { |
43 | #ifdef TI_WITH_DX12 |
44 | // NOTE: there's no jit for dx12. Create cpu session to avoid crash. |
45 | return create_llvm_jit_session_cpu(tlctx, config, Arch::x64); |
46 | #else |
47 | TI_NOT_IMPLEMENTED |
48 | #endif |
49 | } else if (arch == Arch::amdgpu) { |
50 | #ifdef TI_WITH_AMDGPU |
51 | return create_llvm_jit_session_amdgpu(tlctx, config, arch); |
52 | #else |
53 | TI_NOT_IMPLEMENTED |
54 | #endif |
55 | } |
56 | #else |
57 | TI_ERROR("Llvm disabled" ); |
58 | #endif |
59 | return nullptr; |
60 | } |
61 | |
62 | } // namespace taichi::lang |
63 | |