1#include "snode_tree_buffer_manager.h"
2#include "taichi/runtime/llvm/llvm_runtime_executor.h"
3
4namespace taichi::lang {
5
6SNodeTreeBufferManager::SNodeTreeBufferManager(
7 LlvmRuntimeExecutor *runtime_exec)
8 : runtime_exec_(runtime_exec) {
9 TI_TRACE("SNode tree buffer manager created.");
10}
11
12void SNodeTreeBufferManager::merge_and_insert(Ptr ptr, std::size_t size) {
13 // merge with right block
14 if (ptr_map_[ptr + size]) {
15 std::size_t tmp = ptr_map_[ptr + size];
16 size_set_.erase(std::make_pair(tmp, ptr + size));
17 ptr_map_.erase(ptr + size);
18 size += tmp;
19 }
20 // merge with left block
21 auto map_it = ptr_map_.lower_bound(ptr);
22 if (map_it != ptr_map_.begin()) {
23 auto x = *--map_it;
24 if (x.first + x.second == ptr) {
25 size_set_.erase(std::make_pair(x.second, x.first));
26 ptr_map_.erase(x.first);
27 ptr = x.first;
28 size += x.second;
29 }
30 }
31 size_set_.insert(std::make_pair(size, ptr));
32 ptr_map_[ptr] = size;
33}
34
35Ptr SNodeTreeBufferManager::allocate(JITModule *runtime_jit,
36 void *runtime,
37 std::size_t size,
38 std::size_t alignment,
39 const int snode_tree_id,
40 uint64 *result_buffer) {
41 TI_TRACE("allocating memory for SNode Tree {}", snode_tree_id);
42 TI_ASSERT_INFO(snode_tree_id < kMaxNumSnodeTreesLlvm,
43 "LLVM backend supports up to {} snode trees",
44 kMaxNumSnodeTreesLlvm);
45 auto set_it = size_set_.lower_bound(std::make_pair(size, nullptr));
46 if (set_it == size_set_.end()) {
47 runtime_jit->call<void *, std::size_t, std::size_t>(
48 "runtime_memory_allocate_aligned", runtime, size, alignment);
49 auto ptr = runtime_exec_->fetch_result<Ptr>(
50 taichi_result_buffer_runtime_query_id, result_buffer);
51 roots_[snode_tree_id] = ptr;
52 sizes_[snode_tree_id] = size;
53 return ptr;
54 } else {
55 auto x = *set_it;
56 size_set_.erase(x);
57 ptr_map_.erase(x.second);
58 if (x.first - size > 0) {
59 size_set_.insert(std::make_pair(x.first - size, x.second + size));
60 ptr_map_[x.second + size] = x.first - size;
61 }
62 TI_ASSERT(x.second);
63 roots_[snode_tree_id] = x.second;
64 sizes_[snode_tree_id] = size;
65 return x.second;
66 }
67}
68
69void SNodeTreeBufferManager::destroy(SNodeTree *snode_tree) {
70 int snode_tree_id = snode_tree->id();
71 TI_TRACE("Destroying SNode tree {}.", snode_tree_id);
72 std::size_t size = sizes_[snode_tree_id];
73 if (size == 0) {
74 TI_DEBUG("SNode tree {} destroy failed.", snode_tree_id);
75 return;
76 }
77 Ptr ptr = roots_[snode_tree_id];
78 merge_and_insert(ptr, size);
79 TI_DEBUG("SNode tree {} destroyed.", snode_tree_id);
80}
81
82} // namespace taichi::lang
83