1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/core/framework/allocator.h"
17
18 #include "tensorflow/core/framework/allocator_registry.h"
19 #include "tensorflow/core/framework/log_memory.h"
20 #include "tensorflow/core/framework/tracking_allocator.h"
21 #include "tensorflow/core/lib/strings/stringprintf.h"
22 #include "tensorflow/core/platform/mem.h"
23 #include "tensorflow/core/platform/mutex.h"
24 #include "tensorflow/core/platform/types.h"
25
26 namespace tensorflow {
27
Clear()28 void AllocatorStats::Clear() {
29 this->num_allocs = 0;
30 this->bytes_in_use = 0;
31 this->max_bytes_in_use = 0;
32 this->max_alloc_size = 0;
33 this->bytes_limit = 0;
34 }
35
DebugString() const36 string AllocatorStats::DebugString() const {
37 return strings::Printf(
38 "Limit: %20lld\n"
39 "InUse: %20lld\n"
40 "MaxInUse: %20lld\n"
41 "NumAllocs: %20lld\n"
42 "MaxAllocSize: %20lld\n",
43 this->bytes_limit, this->bytes_in_use, this->max_bytes_in_use,
44 this->num_allocs, this->max_alloc_size);
45 }
46
47 constexpr size_t Allocator::kAllocatorAlignment;
48
~Allocator()49 Allocator::~Allocator() {}
50
RunResourceCtor(ResourceHandle * p,size_t n)51 void RunResourceCtor(ResourceHandle* p, size_t n) {
52 for (size_t i = 0; i < n; ++p, ++i) new (p) ResourceHandle();
53 }
54
RunResourceDtor(ResourceHandle * p,size_t n)55 void RunResourceDtor(ResourceHandle* p, size_t n) {
56 for (size_t i = 0; i < n; ++p, ++i) p->~ResourceHandle();
57 }
58
59 // If true, cpu allocator collects more stats.
60 static bool cpu_allocator_collect_stats = false;
61 // If true, cpu allocator collects full stats.
62 static bool cpu_allocator_collect_full_stats = false;
63
EnableCPUAllocatorStats(bool enable)64 void EnableCPUAllocatorStats(bool enable) {
65 cpu_allocator_collect_stats = enable;
66 }
EnableCPUAllocatorFullStats(bool enable)67 void EnableCPUAllocatorFullStats(bool enable) {
68 cpu_allocator_collect_full_stats = enable;
69 }
70
71 class CPUAllocator : public Allocator {
72 public:
CPUAllocator()73 CPUAllocator() {}
74
~CPUAllocator()75 ~CPUAllocator() override {}
76
Name()77 string Name() override { return "cpu"; }
78
AllocateRaw(size_t alignment,size_t num_bytes)79 void* AllocateRaw(size_t alignment, size_t num_bytes) override {
80 void* p = port::AlignedMalloc(num_bytes, alignment);
81 if (cpu_allocator_collect_stats) {
82 const std::size_t alloc_size = port::MallocExtension_GetAllocatedSize(p);
83 mutex_lock l(mu_);
84 ++stats_.num_allocs;
85 stats_.bytes_in_use += alloc_size;
86 stats_.max_bytes_in_use =
87 std::max<int64>(stats_.max_bytes_in_use, stats_.bytes_in_use);
88 stats_.max_alloc_size =
89 std::max<int64>(stats_.max_alloc_size, alloc_size);
90 }
91 return p;
92 }
93
DeallocateRaw(void * ptr)94 void DeallocateRaw(void* ptr) override {
95 if (cpu_allocator_collect_stats) {
96 const std::size_t alloc_size =
97 port::MallocExtension_GetAllocatedSize(ptr);
98 mutex_lock l(mu_);
99 stats_.bytes_in_use -= alloc_size;
100 }
101 port::AlignedFree(ptr);
102 }
103
GetStats(AllocatorStats * stats)104 void GetStats(AllocatorStats* stats) override {
105 mutex_lock l(mu_);
106 *stats = stats_;
107 }
108
ClearStats()109 void ClearStats() override {
110 mutex_lock l(mu_);
111 stats_.num_allocs = 0;
112 stats_.max_bytes_in_use = stats_.bytes_in_use;
113 stats_.max_alloc_size = 0;
114 }
115
AllocatedSizeSlow(const void * ptr)116 size_t AllocatedSizeSlow(const void* ptr) override {
117 return port::MallocExtension_GetAllocatedSize(ptr);
118 }
119
120 private:
121 mutex mu_;
122 AllocatorStats stats_ GUARDED_BY(mu_);
123
124 TF_DISALLOW_COPY_AND_ASSIGN(CPUAllocator);
125 };
126
cpu_allocator()127 Allocator* cpu_allocator() {
128 static Allocator* cpu_alloc = AllocatorRegistry::Global()->GetAllocator();
129 if (cpu_allocator_collect_full_stats && !cpu_alloc->TracksAllocationSizes()) {
130 cpu_alloc = new TrackingAllocator(cpu_alloc, true);
131 }
132 return cpu_alloc;
133 }
134
135 REGISTER_MEM_ALLOCATOR("DefaultCPUAllocator", 100, CPUAllocator);
136
137 } // namespace tensorflow
138