1 | #pragma once |
2 | |
3 | #include <memory> |
4 | #include <unordered_set> |
5 | |
6 | #include "llvm/IR/Module.h" |
7 | |
8 | namespace taichi::lang { |
9 | |
10 | class OffloadedTask { |
11 | public: |
12 | std::string name; |
13 | int block_dim{0}; |
14 | int grid_dim{0}; |
15 | |
16 | explicit OffloadedTask(const std::string &name = "" , |
17 | int block_dim = 0, |
18 | int grid_dim = 0) |
19 | : name(name), block_dim(block_dim), grid_dim(grid_dim){}; |
20 | TI_IO_DEF(name, block_dim, grid_dim); |
21 | }; |
22 | |
23 | struct LLVMCompiledTask { |
24 | std::vector<OffloadedTask> tasks; |
25 | std::unique_ptr<llvm::Module> module{nullptr}; |
26 | std::unordered_set<int> used_tree_ids; |
27 | std::unordered_set<int> struct_for_tls_sizes; |
28 | LLVMCompiledTask() = default; |
29 | LLVMCompiledTask(LLVMCompiledTask &&) = default; |
30 | LLVMCompiledTask &operator=(LLVMCompiledTask &&) = default; |
31 | LLVMCompiledTask(std::vector<OffloadedTask> tasks, |
32 | std::unique_ptr<llvm::Module> module, |
33 | std::unordered_set<int> used_tree_ids, |
34 | std::unordered_set<int> struct_for_tls_sizes) |
35 | : tasks(std::move(tasks)), |
36 | module(std::move(module)), |
37 | used_tree_ids(std::move(used_tree_ids)), |
38 | struct_for_tls_sizes(std::move(struct_for_tls_sizes)) { |
39 | } |
40 | LLVMCompiledTask clone() const; |
41 | TI_IO_DEF(tasks); |
42 | }; |
43 | |
44 | struct LLVMCompiledKernel { |
45 | std::vector<OffloadedTask> tasks; |
46 | std::unique_ptr<llvm::Module> module{nullptr}; |
47 | LLVMCompiledKernel() = default; |
48 | LLVMCompiledKernel(LLVMCompiledKernel &&) = default; |
49 | LLVMCompiledKernel &operator=(LLVMCompiledKernel &&) = default; |
50 | LLVMCompiledKernel(std::vector<OffloadedTask> tasks, |
51 | std::unique_ptr<llvm::Module> module) |
52 | : tasks(std::move(tasks)), module(std::move(module)) { |
53 | } |
54 | LLVMCompiledKernel clone() const; |
55 | TI_IO_DEF(tasks); |
56 | }; |
57 | |
58 | } // namespace taichi::lang |
59 | |