1#pragma once
2
3#include <memory>
4#include <unordered_set>
5
6#include "llvm/IR/Module.h"
7
8namespace taichi::lang {
9
10class 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
23struct 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
44struct 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