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
11C10_DECLARE_bool(caffe2_report_cpu_memory_usage);
12
13namespace c10 {
14
15using MemoryDeleter = void (*)(void*);
16
17// A helper function that is basically doing nothing.
18C10_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
22class 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
36C10_API ProfiledCPUMemoryReporter& profiledCPUMemoryReporter();
37
38// Get the CPU Allocator.
39C10_API at::Allocator* GetCPUAllocator();
40// Sets the CPU allocator to the given allocator: the caller gives away the
41// ownership of the pointer.
42C10_API void SetCPUAllocator(at::Allocator* alloc, uint8_t priority = 0);
43
44// Get the Default CPU Allocator
45C10_API at::Allocator* GetDefaultCPUAllocator();
46
47// Get the Default Mobile CPU Allocator
48C10_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
53C10_API void SetCPUCachingAllocator(Allocator* alloc, uint8_t priority = 0);
54// Get the CPU Caching Allocator
55C10_API Allocator* GetCPUCachingAllocator();
56
57} // namespace c10
58