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
11namespace taichi {
12
13// Cross-platform virtual memory allocator
14class 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
56float64 get_memory_usage_gb(int pid = -1);
57uint64 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