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
8namespace taichi::lang {
9
10class 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 &params) 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