1 | #pragma once |
---|---|
2 | |
3 | #include "taichi/cache/gfx/cache_manager.h" |
4 | #include "taichi/runtime/gfx/runtime.h" |
5 | #include "taichi/runtime/gfx/snode_tree_manager.h" |
6 | #include "taichi/program/program_impl.h" |
7 | |
8 | namespace taichi::lang { |
9 | |
10 | class OpenglProgramImpl : public ProgramImpl { |
11 | public: |
12 | explicit OpenglProgramImpl(CompileConfig &config); |
13 | ~OpenglProgramImpl() override; |
14 | FunctionType compile(const CompileConfig &compile_config, |
15 | Kernel *kernel) override; |
16 | |
17 | std::size_t get_snode_num_dynamically_allocated( |
18 | SNode *snode, |
19 | uint64 *result_buffer) override { |
20 | return 0; // TODO: support sparse |
21 | } |
22 | |
23 | void compile_snode_tree_types(SNodeTree *tree) override; |
24 | |
25 | void materialize_runtime(MemoryPool *memory_pool, |
26 | KernelProfilerBase *profiler, |
27 | uint64 **result_buffer_ptr) override; |
28 | |
29 | void materialize_snode_tree(SNodeTree *tree, uint64 *result_buffer) override; |
30 | |
31 | void synchronize() override { |
32 | runtime_->synchronize(); |
33 | } |
34 | |
35 | void finalize() override; |
36 | |
37 | StreamSemaphore flush() override { |
38 | return runtime_->flush(); |
39 | } |
40 | |
41 | std::unique_ptr<AotModuleBuilder> make_aot_module_builder( |
42 | const DeviceCapabilityConfig &caps) override; |
43 | |
44 | void destroy_snode_tree(SNodeTree *snode_tree) override { |
45 | TI_ASSERT(snode_tree_mgr_ != nullptr); |
46 | snode_tree_mgr_->destroy_snode_tree(snode_tree); |
47 | } |
48 | |
49 | DeviceAllocation allocate_memory_ndarray(std::size_t alloc_size, |
50 | uint64 *result_buffer) override; |
51 | |
52 | bool used_in_kernel(DeviceAllocationId id) override { |
53 | return runtime_->used_in_kernel(id); |
54 | } |
55 | |
56 | DeviceAllocation allocate_texture(const ImageParams ¶ms) override; |
57 | |
58 | Device *get_compute_device() override { |
59 | return device_.get(); |
60 | } |
61 | |
62 | Device *get_graphics_device() override { |
63 | return device_.get(); |
64 | } |
65 | |
66 | size_t get_field_in_tree_offset(int tree_id, const SNode *child) override { |
67 | return snode_tree_mgr_->get_field_in_tree_offset(tree_id, child); |
68 | } |
69 | |
70 | DevicePtr get_snode_tree_device_ptr(int tree_id) override { |
71 | return snode_tree_mgr_->get_snode_tree_device_ptr(tree_id); |
72 | } |
73 | |
74 | void dump_cache_data_to_disk() override; |
75 | |
76 | const std::unique_ptr<gfx::CacheManager> &get_cache_manager(); |
77 | |
78 | private: |
79 | std::shared_ptr<Device> device_{nullptr}; |
80 | std::unique_ptr<gfx::GfxRuntime> runtime_{nullptr}; |
81 | std::unique_ptr<gfx::SNodeTreeManager> snode_tree_mgr_{nullptr}; |
82 | std::vector<spirv::CompiledSNodeStructs> aot_compiled_snode_structs_; |
83 | std::unique_ptr<gfx::CacheManager> cache_manager_{nullptr}; |
84 | }; |
85 | |
86 | } // namespace taichi::lang |
87 |