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
19namespace taichi {
20class DynamicLoader;
21} // namespace taichi
22
23namespace taichi::lang {
24
25using namespace taichi::lang::cccp;
26using CCFuncEntryType = void(cccp::CCContext *);
27
28class 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