1 | #pragma once |
---|---|
2 | #include <mutex> |
3 | #include <vector> |
4 | #include <memory> |
5 | |
6 | #include "taichi/rhi/arch.h" |
7 | #include "taichi/rhi/device.h" |
8 | |
9 | namespace taichi { |
10 | class VirtualMemoryAllocator; |
11 | } |
12 | |
13 | namespace taichi::lang { |
14 | |
15 | // This class can only have one instance |
16 | class UnifiedAllocator { |
17 | std::unique_ptr<VirtualMemoryAllocator> cpu_vm_; |
18 | std::size_t size_; |
19 | Arch arch_; |
20 | |
21 | // put these two on the unified memory so that GPU can have access |
22 | public: |
23 | uint8 *data; |
24 | DeviceAllocation alloc{kDeviceNullAllocation}; |
25 | uint8 *head; |
26 | uint8 *tail; |
27 | std::mutex lock; |
28 | |
29 | public: |
30 | UnifiedAllocator(std::size_t size, Arch arch, Device *device); |
31 | |
32 | ~UnifiedAllocator(); |
33 | |
34 | void *allocate(std::size_t size, std::size_t alignment) { |
35 | std::lock_guard<std::mutex> _(lock); |
36 | auto ret = |
37 | head + alignment - 1 - ((std::size_t)head + alignment - 1) % alignment; |
38 | TI_TRACE("UM [data={}] allocate() request={} remain={}", (intptr_t)data, |
39 | size, (tail - head)); |
40 | head = ret + size; |
41 | if (head > tail) { |
42 | // allocation failed |
43 | return nullptr; |
44 | } else { |
45 | // success |
46 | TI_ASSERT((std::size_t)ret % alignment == 0); |
47 | return ret; |
48 | } |
49 | } |
50 | |
51 | void memset(unsigned char val); |
52 | |
53 | bool initialized() const { |
54 | return data != nullptr; |
55 | } |
56 | |
57 | UnifiedAllocator operator=(const UnifiedAllocator &) = delete; |
58 | |
59 | private: |
60 | Device *device_{nullptr}; |
61 | }; |
62 | |
63 | } // namespace taichi::lang |
64 |