1 | #pragma once |
---|---|
2 | #include "taichi/inc/constants.h" |
3 | #include "taichi/struct/snode_tree.h" |
4 | #define TI_RUNTIME_HOST |
5 | |
6 | #include <set> |
7 | |
8 | using Ptr = uint8_t *; |
9 | |
10 | namespace taichi::lang { |
11 | |
12 | class JITModule; |
13 | class LlvmRuntimeExecutor; |
14 | |
15 | class SNodeTreeBufferManager { |
16 | public: |
17 | explicit SNodeTreeBufferManager(LlvmRuntimeExecutor *runtime_exec); |
18 | |
19 | void merge_and_insert(Ptr ptr, std::size_t size); |
20 | |
21 | Ptr allocate(JITModule *runtime_jit, |
22 | void *runtime, |
23 | std::size_t size, |
24 | std::size_t alignment, |
25 | const int snode_tree_id, |
26 | uint64 *result_buffer); |
27 | |
28 | void destroy(SNodeTree *snode_tree); |
29 | |
30 | private: |
31 | std::set<std::pair<std::size_t, Ptr>> size_set_; |
32 | std::map<Ptr, std::size_t> ptr_map_; |
33 | LlvmRuntimeExecutor *runtime_exec_; |
34 | Ptr roots_[kMaxNumSnodeTreesLlvm]; |
35 | std::size_t sizes_[kMaxNumSnodeTreesLlvm]; |
36 | }; |
37 | |
38 | } // namespace taichi::lang |
39 |