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 #ifndef TENSORFLOW_CORE_FRAMEWORK_TRACKING_ALLOCATOR_H_ 17 #define TENSORFLOW_CORE_FRAMEWORK_TRACKING_ALLOCATOR_H_ 18 19 #include <unordered_map> 20 #include "tensorflow/core/framework/allocator.h" 21 #include "tensorflow/core/lib/gtl/inlined_vector.h" 22 #include "tensorflow/core/platform/mutex.h" 23 #include "tensorflow/core/platform/thread_annotations.h" 24 #include "tensorflow/core/platform/types.h" 25 26 namespace tensorflow { 27 28 // TrackingAllocator is a wrapper for an Allocator. It keeps a running 29 // count of the number of bytes allocated through the wrapper. It is 30 // used by the Executor to "charge" allocations to particular Op 31 // executions. Each Op gets a separate TrackingAllocator wrapper 32 // around the underlying allocator. 33 // 34 // The implementation assumes the invariant that all calls to 35 // AllocateRaw by an Op (or work items spawned by the Op) will occur 36 // before the Op's Compute method returns. Thus the high watermark is 37 // established once Compute returns. 38 // 39 // DeallocateRaw can be called long after the Op has finished, 40 // e.g. when an output tensor is deallocated, and the wrapper cannot 41 // be deleted until the last of these calls has occurred. The 42 // TrackingAllocator keeps track of outstanding calls using a 43 // reference count, and deletes itself once the last call has been 44 // received and the high watermark has been retrieved. 45 struct AllocRecord { AllocRecordAllocRecord46 AllocRecord(int64_t a_btyes, int64_t a_micros) 47 : alloc_bytes(a_btyes), alloc_micros(a_micros) {} AllocRecordAllocRecord48 AllocRecord() : AllocRecord(0, 0) {} 49 50 int64_t alloc_bytes; 51 int64_t alloc_micros; 52 }; 53 54 class TrackingAllocator : public Allocator { 55 public: 56 explicit TrackingAllocator(Allocator* allocator, bool track_ids); Name()57 std::string Name() override { return allocator_->Name(); } AllocateRaw(size_t alignment,size_t num_bytes)58 void* AllocateRaw(size_t alignment, size_t num_bytes) override { 59 return AllocateRaw(alignment, num_bytes, AllocationAttributes()); 60 } 61 void* AllocateRaw(size_t alignment, size_t num_bytes, 62 const AllocationAttributes& allocation_attr) override; 63 void DeallocateRaw(void* ptr) override; 64 bool TracksAllocationSizes() const override; 65 size_t RequestedSize(const void* ptr) const override; 66 size_t AllocatedSize(const void* ptr) const override; 67 int64_t AllocationId(const void* ptr) const override; 68 absl::optional<AllocatorStats> GetStats() override; 69 bool ClearStats() override; 70 GetMemoryType()71 AllocatorMemoryType GetMemoryType() const override { 72 return allocator_->GetMemoryType(); 73 } 74 75 // If the underlying allocator tracks allocation sizes, this returns 76 // a tuple where the first value is the total number of bytes 77 // allocated through this wrapper, the second value is the high 78 // watermark of bytes allocated through this wrapper and the third value is 79 // the allocated bytes through this wrapper that are still alive. If the 80 // underlying allocator does not track allocation sizes the first 81 // value is the total number of bytes requested through this wrapper 82 // and the second and the third are 0. 83 // 84 std::tuple<size_t, size_t, size_t> GetSizes(); 85 // After GetRecordsAndUnRef is called, the only further calls allowed 86 // on this wrapper are calls to DeallocateRaw with pointers that 87 // were allocated by this wrapper and have not yet been 88 // deallocated. After this call completes and all allocated pointers 89 // have been deallocated the wrapper will delete itself. 90 gtl::InlinedVector<AllocRecord, 4> GetRecordsAndUnRef(); 91 // Returns a copy of allocation records collected so far. 92 gtl::InlinedVector<AllocRecord, 4> GetCurrentRecords(); 93 94 protected: ~TrackingAllocator()95 ~TrackingAllocator() override {} 96 97 private: 98 bool UnRef() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_); 99 100 Allocator* allocator_; // not owned. 101 mutable mutex mu_; 102 // the number of calls to AllocateRaw that have not yet been matched 103 // by a corresponding call to DeAllocateRaw, plus 1 if the Executor 104 // has not yet read out the high watermark. 105 int ref_ TF_GUARDED_BY(mu_); 106 // the current number of outstanding bytes that have been allocated 107 // by this wrapper, or 0 if the underlying allocator does not track 108 // allocation sizes. 109 size_t allocated_ TF_GUARDED_BY(mu_); 110 // the maximum number of outstanding bytes that have been allocated 111 // by this wrapper, or 0 if the underlying allocator does not track 112 // allocation sizes. 113 size_t high_watermark_ TF_GUARDED_BY(mu_); 114 // the total number of bytes that have been allocated by this 115 // wrapper if the underlying allocator tracks allocation sizes, 116 // otherwise the total number of bytes that have been requested by 117 // this allocator. 118 size_t total_bytes_ TF_GUARDED_BY(mu_); 119 120 gtl::InlinedVector<AllocRecord, 4> allocations_ TF_GUARDED_BY(mu_); 121 122 // Track allocations locally if requested in the constructor and the 123 // underlying allocator doesn't already do it for us. 124 const bool track_sizes_locally_; 125 struct Chunk { 126 size_t requested_size; 127 size_t allocated_size; 128 int64_t allocation_id; 129 }; 130 std::unordered_map<const void*, Chunk> in_use_ TF_GUARDED_BY(mu_); 131 int64_t next_allocation_id_ TF_GUARDED_BY(mu_); 132 }; 133 134 } // end namespace tensorflow 135 136 #endif // TENSORFLOW_CORE_FRAMEWORK_TRACKING_ALLOCATOR_H_ 137