1 | #include "taichi/ir/transforms.h" |
2 | #include "taichi/program/extension.h" |
3 | #include "taichi/program/kernel.h" |
4 | |
5 | namespace taichi::lang { |
6 | namespace irpass { |
7 | |
8 | static bool supports_lowering(Arch arch) { |
9 | return arch_is_cpu(arch) || (arch == Arch::cuda) || (arch == Arch::dx12) || |
10 | (arch == Arch::metal) || (arch == Arch::amdgpu); |
11 | } |
12 | |
13 | void ast_to_ir(const CompileConfig &config, |
14 | Kernel &kernel, |
15 | bool to_executable) { |
16 | TI_ASSERT(supports_lowering(config.arch)); |
17 | if (kernel.lowered()) { |
18 | return; |
19 | } |
20 | |
21 | bool verbose = config.print_ir; |
22 | if ((kernel.is_accessor && !config.print_accessor_ir) || |
23 | (kernel.is_evaluator && !config.print_evaluator_ir)) |
24 | verbose = false; |
25 | |
26 | if (to_executable) { |
27 | irpass::compile_to_executable( |
28 | kernel.ir.get(), config, &kernel, |
29 | /*autodiff_mode=*/kernel.autodiff_mode, |
30 | /*ad_use_stack=*/true, |
31 | /*verbose*/ verbose, |
32 | /*lower_global_access=*/to_executable, |
33 | /*make_thread_local=*/config.make_thread_local, |
34 | /*make_block_local=*/ |
35 | is_extension_supported(config.arch, Extension::bls) && |
36 | config.make_block_local, |
37 | /*start_from_ast=*/kernel.ir_is_ast()); |
38 | } else { |
39 | irpass::compile_to_offloads(kernel.ir.get(), config, &kernel, verbose, |
40 | /*autodiff_mode=*/kernel.autodiff_mode, |
41 | /*ad_use_stack=*/true, |
42 | /*start_from_ast=*/kernel.ir_is_ast()); |
43 | } |
44 | |
45 | kernel.set_lowered(true); |
46 | } |
47 | |
48 | } // namespace irpass |
49 | } // namespace taichi::lang |
50 | |