1 | #pragma once |
---|---|
2 | |
3 | #include "taichi/common/core.h" |
4 | #include "taichi/program/kernel.h" |
5 | #include "taichi/program/program.h" |
6 | #include "taichi/system/dynamic_loader.h" |
7 | #include "taichi/util/action_recorder.h" |
8 | #include "struct_cc.h" |
9 | #include "cc_runtime.h" |
10 | #include "cc_kernel.h" |
11 | #include "cc_layout.h" |
12 | #include "cc_utils.h" |
13 | #include "codegen_cc.h" |
14 | #include "context.h" |
15 | #include "taichi/util/lang_util.h" |
16 | #include <vector> |
17 | #include <memory> |
18 | |
19 | namespace taichi { |
20 | class DynamicLoader; |
21 | } // namespace taichi |
22 | |
23 | namespace taichi::lang { |
24 | |
25 | using namespace taichi::lang::cccp; |
26 | using CCFuncEntryType = void(cccp::CCContext *); |
27 | |
28 | class CCProgramImpl : public ProgramImpl { |
29 | public: |
30 | explicit CCProgramImpl(CompileConfig &config); |
31 | |
32 | FunctionType compile(const CompileConfig &compile_config, |
33 | Kernel *kernel) override; |
34 | |
35 | std::size_t get_snode_num_dynamically_allocated( |
36 | SNode *snode, |
37 | uint64 *result_buffer) override { |
38 | return 0; // TODO: support sparse in cc. |
39 | } |
40 | |
41 | void materialize_runtime(MemoryPool *memory_pool, |
42 | KernelProfilerBase *, |
43 | uint64 **result_buffer_ptr) override; |
44 | |
45 | void materialize_snode_tree(SNodeTree *tree, uint64 *result_buffer) override; |
46 | |
47 | void synchronize() override { |
48 | // Not implemented yet. |
49 | } |
50 | |
51 | std::unique_ptr<AotModuleBuilder> make_aot_module_builder( |
52 | const DeviceCapabilityConfig &caps) override { |
53 | // Not implemented yet. |
54 | return nullptr; |
55 | } |
56 | |
57 | void destroy_snode_tree(SNodeTree *snode_tree) override { |
58 | // Not implemented yet. |
59 | } |
60 | |
61 | CCLayout *get_layout() { |
62 | return layout_.get(); |
63 | } |
64 | |
65 | CCRuntime *get_runtime() { |
66 | return runtime_.get(); |
67 | } |
68 | |
69 | ~CCProgramImpl() override { |
70 | } |
71 | |
72 | CCFuncEntryType *load_kernel(std::string const &name); |
73 | void relink(); |
74 | |
75 | CCContext *update_context(RuntimeContext *ctx); |
76 | void context_to_result_buffer(); |
77 | |
78 | private: |
79 | void add_kernel(std::unique_ptr<CCKernel> kernel); |
80 | |
81 | std::vector<std::unique_ptr<CCKernel>> kernels_; |
82 | std::unique_ptr<CCContext> context_; |
83 | std::unique_ptr<CCRuntime> runtime_; |
84 | std::unique_ptr<CCLayout> layout_; |
85 | std::unique_ptr<DynamicLoader> dll_; |
86 | std::string dll_path_; |
87 | std::vector<char> args_buf_; |
88 | std::vector<char> root_buf_; |
89 | std::vector<char> gtmp_buf_; |
90 | uint64 *result_buffer_{nullptr}; |
91 | bool need_relink_{true}; |
92 | }; |
93 | } // namespace taichi::lang |
94 |