1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
27namespace 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.
46struct 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
55class 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