1 | #pragma once |
2 | |
3 | #include "taichi/common/core.h" |
4 | |
5 | #if defined(TI_PLATFORM_UNIX) |
6 | #include <sys/mman.h> |
7 | #else |
8 | #include "taichi/platform/windows/windows.h" |
9 | #endif |
10 | |
11 | namespace taichi { |
12 | |
13 | // Cross-platform virtual memory allocator |
14 | class VirtualMemoryAllocator { |
15 | public: |
16 | static constexpr size_t page_size = (1 << 12); // 4 KB page size by default |
17 | void *ptr; |
18 | size_t size; |
19 | explicit VirtualMemoryAllocator(size_t size) : size(size) { |
20 | // http://pages.cs.wisc.edu/~sifakis/papers/SPGrid.pdf Sec 3.1 |
21 | #if defined(TI_PLATFORM_UNIX) |
22 | ptr = mmap(nullptr, size, PROT_READ | PROT_WRITE, |
23 | MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); |
24 | TI_ERROR_IF(ptr == MAP_FAILED, "Virtual memory allocation ({} B) failed." , |
25 | size); |
26 | #else |
27 | MEMORYSTATUSEX stat; |
28 | stat.dwLength = sizeof(stat); |
29 | GlobalMemoryStatusEx(&stat); |
30 | if (stat.ullAvailVirtual < size) { |
31 | TI_P(stat.ullAvailVirtual); |
32 | TI_P(size); |
33 | TI_ERROR("Insufficient virtual memory space" ); |
34 | } |
35 | ptr = VirtualAlloc(nullptr, size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE); |
36 | TI_ERROR_IF(ptr == nullptr, "Virtual memory allocation ({} B) failed." , |
37 | size); |
38 | #endif |
39 | TI_ERROR_IF(((uint64_t)ptr) % page_size != 0, |
40 | "Allocated address ({:}) is not aligned by page size {}" , ptr, |
41 | page_size); |
42 | } |
43 | |
44 | ~VirtualMemoryAllocator() { |
45 | #if defined(TI_PLATFORM_UNIX) |
46 | if (munmap(ptr, size) != 0) |
47 | #else |
48 | // https://docs.microsoft.com/en-us/windows/win32/api/memoryapi/nf-memoryapi-virtualfree |
49 | // According to MS Doc: size must be when using MEM_RELEASE |
50 | if (!VirtualFree(ptr, 0, MEM_RELEASE)) |
51 | #endif |
52 | TI_ERROR("Failed to free virtual memory ({} B)" , size); |
53 | } |
54 | }; |
55 | |
56 | float64 get_memory_usage_gb(int pid = -1); |
57 | uint64 get_memory_usage(int pid = -1); |
58 | |
59 | #define TI_MEMORY_USAGE(name) \ |
60 | TI_DEBUG("Memory Usage [{}] = {:.2f} GB", name, get_memory_usage_gb()); |
61 | |
62 | } // namespace taichi |
63 | |