1 #ifndef CAFFE2_CORE_ALLOCATOR_H_ 2 #define CAFFE2_CORE_ALLOCATOR_H_ 4 #include <unordered_map> 6 #include "caffe2/core/logging.h" 7 #include "caffe2/core/numa.h" 9 CAFFE2_DECLARE_bool(caffe2_report_cpu_memory_usage);
10 CAFFE2_DECLARE_bool(caffe2_cpu_allocator_do_zero_fill);
15 constexpr
size_t gCaffe2Alignment = 32;
17 using MemoryDeleter = void (*)(
void*);
26 virtual std::pair<void*, MemoryDeleter> New(
size_t nbytes) = 0;
27 virtual MemoryDeleter GetDeleter() = 0;
35 void New(
void* ptr,
size_t nbytes);
36 void Delete(
void* ptr);
40 std::unordered_map<void*, size_t> size_table_;
47 std::pair<void*, MemoryDeleter> New(
size_t nbytes)
override {
50 data = memalign(gCaffe2Alignment, nbytes);
51 #elif defined(_MSC_VER) 52 data = _aligned_malloc(nbytes, gCaffe2Alignment);
54 CAFFE_ENFORCE_EQ(posix_memalign(&data, gCaffe2Alignment, nbytes), 0);
58 NUMAMove(data, nbytes, GetCurrentNUMANode());
59 if (FLAGS_caffe2_cpu_allocator_do_zero_fill) {
60 memset(data, 0, nbytes);
62 return {data, Delete};
66 static void Delete(
void* data) {
70 static void Delete(
void* data) {
75 MemoryDeleter GetDeleter()
override {
88 #endif // CAFFE2_CORE_ALLOCATOR_H_
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...