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