Caffe2 - C++ API
A deep learning, cross platform ML framework
allocator.h
1 
17 #ifndef CAFFE2_CORE_ALLOCATOR_H_
18 #define CAFFE2_CORE_ALLOCATOR_H_
19 
20 #include <unordered_map>
21 
22 #include "caffe2/core/logging.h"
23 
24 CAFFE2_DECLARE_bool(caffe2_report_cpu_memory_usage);
25 CAFFE2_DECLARE_bool(caffe2_cpu_allocator_do_zero_fill);
26 
27 namespace caffe2 {
28 
29 // Use 32-byte alignment should be enough for computation up to AVX512.
30 constexpr size_t gCaffe2Alignment = 32;
31 
32 using MemoryDeleter = void (*)(void*);
33 
34 // A helper function that is basically doing nothing.
35 void NoDelete(void*);
36 
37 // A virtual allocator class to do memory allocation and deallocation.
38 struct CPUAllocator {
39  CPUAllocator() {}
40  virtual ~CPUAllocator() noexcept {}
41  virtual std::pair<void*, MemoryDeleter> New(size_t nbytes) = 0;
42  virtual MemoryDeleter GetDeleter() = 0;
43 };
44 
45 // A virtual struct that is used to report Caffe2's memory allocation and
46 // deallocation status
48  public:
49  MemoryAllocationReporter() : allocated_(0) {}
50  void New(void* ptr, size_t nbytes);
51  void Delete(void* ptr);
52 
53  private:
54  std::mutex mutex_;
55  std::unordered_map<void*, size_t> size_table_;
56  size_t allocated_;
57 };
58 
61  ~DefaultCPUAllocator() override {}
62  std::pair<void*, MemoryDeleter> New(size_t nbytes) override {
63  void* data = nullptr;
64 #ifdef __ANDROID__
65  data = memalign(gCaffe2Alignment, nbytes);
66 #elif defined(_MSC_VER)
67  data = _aligned_malloc(nbytes, gCaffe2Alignment);
68 #else
69  CAFFE_ENFORCE_EQ(posix_memalign(&data, gCaffe2Alignment, nbytes), 0);
70 #endif
71  CAFFE_ENFORCE(data);
72  if (FLAGS_caffe2_cpu_allocator_do_zero_fill) {
73  memset(data, 0, nbytes);
74  }
75  return {data, Delete};
76  }
77 
78 #ifdef _MSC_VER
79  static void Delete(void* data) {
80  _aligned_free(data);
81  }
82 #else
83  static void Delete(void* data) {
84  free(data);
85  }
86 #endif
87 
88  MemoryDeleter GetDeleter() override {
89  return Delete;
90  }
91 };
92 
93 // Get the CPU Alloctor.
94 CPUAllocator* GetCPUAllocator();
95 // Sets the CPU allocator to the given allocator: the caller gives away the
96 // ownership of the pointer.
97 void SetCPUAllocator(CPUAllocator* alloc);
98 
99 } // namespace caffe2
100 
101 #endif // CAFFE2_CORE_ALLOCATOR_H_
Copyright (c) 2016-present, Facebook, Inc.