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
9namespace taichi {
10class VirtualMemoryAllocator;
11}
12
13namespace taichi::lang {
14
15// This class can only have one instance
16class 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