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
8using Ptr = uint8_t *;
9
10namespace taichi::lang {
11
12class JITModule;
13class LlvmRuntimeExecutor;
14
15class 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