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
9namespace taichi::lang {
10namespace gfx {
11
12class GfxRuntime;
13
14/**
15 * @brief Manages the SNodeTrees for the underlying backend.
16 *
17 */
18class 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