1 | #pragma once |
---|---|
2 | |
3 | #include <vector> |
4 | |
5 | #include "taichi/rhi/device.h" |
6 | #include "taichi/codegen/spirv/snode_struct_compiler.h" |
7 | #include "taichi/struct/snode_tree.h" |
8 | |
9 | namespace taichi::lang { |
10 | namespace gfx { |
11 | |
12 | class GfxRuntime; |
13 | |
14 | /** |
15 | * @brief Manages the SNodeTrees for the underlying backend. |
16 | * |
17 | */ |
18 | class SNodeTreeManager { |
19 | private: |
20 | using CompiledSNodeStructs = taichi::lang::spirv::CompiledSNodeStructs; |
21 | |
22 | public: |
23 | explicit SNodeTreeManager(GfxRuntime *rtm); |
24 | |
25 | const std::vector<CompiledSNodeStructs> &get_compiled_structs() const { |
26 | return compiled_snode_structs_; |
27 | } |
28 | |
29 | void materialize_snode_tree(SNodeTree *tree); |
30 | |
31 | void destroy_snode_tree(SNodeTree *snode_tree); |
32 | |
33 | size_t get_field_in_tree_offset(int tree_id, const SNode *child); |
34 | |
35 | DevicePtr get_snode_tree_device_ptr(int tree_id); |
36 | |
37 | private: |
38 | GfxRuntime *const runtime_; |
39 | std::vector<CompiledSNodeStructs> compiled_snode_structs_; |
40 | }; |
41 | |
42 | } // namespace gfx |
43 | } // namespace taichi::lang |
44 |