1#include "taichi/jit/jit_session.h"
2
3#ifdef TI_WITH_LLVM
4#include "llvm/IR/DataLayout.h"
5#endif
6
7namespace taichi::lang {
8
9#ifdef TI_WITH_LLVM
10std::unique_ptr<JITSession> create_llvm_jit_session_cpu(
11 TaichiLLVMContext *tlctx,
12 const CompileConfig &config,
13 Arch arch);
14
15std::unique_ptr<JITSession> create_llvm_jit_session_cuda(
16 TaichiLLVMContext *tlctx,
17 const CompileConfig &config,
18 Arch arch);
19
20std::unique_ptr<JITSession> create_llvm_jit_session_amdgpu(
21 TaichiLLVMContext *tlctx,
22 const CompileConfig &config,
23 Arch arch);
24#endif
25
26JITSession::JITSession(TaichiLLVMContext *tlctx, const CompileConfig &config)
27 : tlctx_(tlctx), config_(config) {
28}
29
30std::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