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 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 36 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 49 Allocator::~Allocator() {} 50 51 void RunResourceCtor(ResourceHandle* p, size_t n) { 52 for (size_t i = 0; i < n; ++p, ++i) new (p) ResourceHandle(); 53 } 54 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 64 void EnableCPUAllocatorStats(bool enable) { 65 cpu_allocator_collect_stats = enable; 66 } 67 void EnableCPUAllocatorFullStats(bool enable) { 68 cpu_allocator_collect_full_stats = enable; 69 } 70 71 class CPUAllocator : public Allocator { 72 public: 73 CPUAllocator() {} 74 75 ~CPUAllocator() override {} 76 77 string Name() override { return "cpu"; } 78 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 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 104 void GetStats(AllocatorStats* stats) override { 105 mutex_lock l(mu_); 106 *stats = stats_; 107 } 108 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 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 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