1 | #pragma once |
2 | |
3 | #include <cstring> |
4 | #include <mutex> |
5 | #include <unordered_map> |
6 | |
7 | #include <c10/core/Allocator.h> |
8 | #include <c10/util/Flags.h> |
9 | |
10 | // TODO: rename to c10 |
11 | C10_DECLARE_bool(caffe2_report_cpu_memory_usage); |
12 | |
13 | namespace c10 { |
14 | |
15 | using MemoryDeleter = void (*)(void*); |
16 | |
17 | // A helper function that is basically doing nothing. |
18 | C10_API void NoDelete(void*); |
19 | |
20 | // A simple struct that is used to report C10's memory allocation, |
21 | // deallocation status and out-of-memory events to the profiler |
22 | class C10_API ProfiledCPUMemoryReporter { |
23 | public: |
24 | ProfiledCPUMemoryReporter() = default; |
25 | void New(void* ptr, size_t nbytes); |
26 | void OutOfMemory(size_t nbytes); |
27 | void Delete(void* ptr); |
28 | |
29 | private: |
30 | std::mutex mutex_; |
31 | std::unordered_map<void*, size_t> size_table_; |
32 | size_t allocated_ = 0; |
33 | size_t log_cnt_ = 0; |
34 | }; |
35 | |
36 | C10_API ProfiledCPUMemoryReporter& profiledCPUMemoryReporter(); |
37 | |
38 | // Get the CPU Allocator. |
39 | C10_API at::Allocator* GetCPUAllocator(); |
40 | // Sets the CPU allocator to the given allocator: the caller gives away the |
41 | // ownership of the pointer. |
42 | C10_API void SetCPUAllocator(at::Allocator* alloc, uint8_t priority = 0); |
43 | |
44 | // Get the Default CPU Allocator |
45 | C10_API at::Allocator* GetDefaultCPUAllocator(); |
46 | |
47 | // Get the Default Mobile CPU Allocator |
48 | C10_API at::Allocator* GetDefaultMobileCPUAllocator(); |
49 | |
50 | // The CPUCachingAllocator is experimental and might disappear in the future. |
51 | // The only place that uses it is in StaticRuntime. |
52 | // Set the CPU Caching Allocator |
53 | C10_API void SetCPUCachingAllocator(Allocator* alloc, uint8_t priority = 0); |
54 | // Get the CPU Caching Allocator |
55 | C10_API Allocator* GetCPUCachingAllocator(); |
56 | |
57 | } // namespace c10 |
58 | |