1 | #include "taichi/program/callable.h" |
2 | #include "taichi/program/program.h" |
3 | |
4 | namespace taichi::lang { |
5 | |
6 | Callable::Callable() = default; |
7 | |
8 | Callable::~Callable() = default; |
9 | |
10 | int Callable::insert_scalar_param(const DataType &dt) { |
11 | parameter_list.emplace_back(dt->get_compute_type(), /*is_array=*/false); |
12 | return (int)parameter_list.size() - 1; |
13 | } |
14 | |
15 | int Callable::insert_ret(const DataType &dt) { |
16 | rets.emplace_back(dt->get_compute_type()); |
17 | return (int)rets.size() - 1; |
18 | } |
19 | |
20 | int Callable::insert_arr_param(const DataType &dt, |
21 | int total_dim, |
22 | std::vector<int> element_shape) { |
23 | parameter_list.emplace_back(dt->get_compute_type(), /*is_array=*/true, |
24 | /*size=*/0, total_dim, element_shape); |
25 | return (int)parameter_list.size() - 1; |
26 | } |
27 | |
28 | int Callable::insert_texture_param(const DataType &dt) { |
29 | // FIXME: we shouldn't abuse is_array for texture parameters |
30 | parameter_list.emplace_back(dt->get_compute_type(), /*is_array=*/true); |
31 | return (int)parameter_list.size() - 1; |
32 | } |
33 | |
34 | void Callable::finalize_rets() { |
35 | if (rets.empty()) { |
36 | return; |
37 | } |
38 | std::vector<StructMember> members; |
39 | members.reserve(rets.size()); |
40 | for (int i = 0; i < rets.size(); i++) { |
41 | members.push_back({rets[i].dt, fmt::format("ret_{}" , i)}); |
42 | } |
43 | auto *type = |
44 | TypeFactory::get_instance().get_struct_type(members)->as<StructType>(); |
45 | std::string layout = program->get_kernel_return_data_layout(); |
46 | ret_type = program->get_struct_type_with_data_layout(type, layout); |
47 | } |
48 | } // namespace taichi::lang |
49 | |