1 | #include "snode_tree_buffer_manager.h" |
2 | #include "taichi/runtime/llvm/llvm_runtime_executor.h" |
3 | |
4 | namespace taichi::lang { |
5 | |
6 | SNodeTreeBufferManager::SNodeTreeBufferManager( |
7 | LlvmRuntimeExecutor *runtime_exec) |
8 | : runtime_exec_(runtime_exec) { |
9 | TI_TRACE("SNode tree buffer manager created." ); |
10 | } |
11 | |
12 | void 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 | |
35 | Ptr 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 | |
69 | void 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 | |