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