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/tracking_allocator.h"
17
18 #include <unordered_map>
19
20 #include "tensorflow/core/framework/allocator.h"
21 #include "tensorflow/core/platform/logging.h"
22 #include "tensorflow/core/platform/mem.h"
23 #include "tensorflow/core/platform/test.h"
24
25 namespace tensorflow {
26
27 class TestableSizeTrackingAllocator : public Allocator {
28 public:
Name()29 string Name() override { return "test"; }
AllocateRaw(size_t,size_t num_bytes)30 void* AllocateRaw(size_t /*alignment*/, size_t num_bytes) override {
31 void* ptr = port::Malloc(num_bytes);
32 size_map_[ptr] = num_bytes;
33 return ptr;
34 }
DeallocateRaw(void * ptr)35 void DeallocateRaw(void* ptr) override {
36 const auto& iter = size_map_.find(ptr);
37 EXPECT_NE(size_map_.end(), iter);
38 size_map_.erase(iter);
39 port::Free(ptr);
40 }
TracksAllocationSizes()41 bool TracksAllocationSizes() override { return true; }
RequestedSize(const void * ptr)42 size_t RequestedSize(const void* ptr) override {
43 const auto& iter = size_map_.find(ptr);
44 EXPECT_NE(size_map_.end(), iter);
45 return iter->second;
46 }
GetStats()47 absl::optional<AllocatorStats> GetStats() override { return absl::nullopt; }
48
49 private:
50 std::unordered_map<const void*, size_t> size_map_;
51 };
52
53 class NoMemoryAllocator : public Allocator {
54 public:
Name()55 string Name() override { return "test"; }
AllocateRaw(size_t,size_t num_bytes)56 void* AllocateRaw(size_t /*alignment*/, size_t num_bytes) override {
57 return nullptr;
58 }
DeallocateRaw(void * ptr)59 void DeallocateRaw(void* ptr) override {}
TracksAllocationSizes()60 bool TracksAllocationSizes() override { return true; }
GetStats()61 absl::optional<AllocatorStats> GetStats() override { return absl::nullopt; }
62 };
63
TEST(TrackingAllocatorTest,SimpleNoTracking)64 TEST(TrackingAllocatorTest, SimpleNoTracking) {
65 Allocator* a = cpu_allocator();
66
67 EXPECT_FALSE(a->TracksAllocationSizes());
68
69 // Don't enable the tracking inside the tracking allocator. Since
70 // the cpu_allocator doesn't track allocations itself the tracking
71 // will be partial
72 TrackingAllocator* ta = new TrackingAllocator(a, false);
73
74 void* p1 = ta->AllocateRaw(4, 4);
75 ta->DeallocateRaw(p1);
76 void* p2 = ta->AllocateRaw(4, 12);
77
78 std::tuple<size_t, size_t, size_t> sizes = ta->GetSizes();
79
80 EXPECT_EQ(16, std::get<0>(sizes));
81 EXPECT_EQ(0, std::get<1>(sizes));
82 EXPECT_EQ(0, std::get<2>(sizes));
83
84 ta->DeallocateRaw(p2);
85 auto records = ta->GetRecordsAndUnRef();
86 EXPECT_EQ(4, records[0].alloc_bytes);
87 EXPECT_EQ(12, records[1].alloc_bytes);
88
89 // This time enable the tracking inside the tracking allocator
90 ta = new TrackingAllocator(a, true);
91 p1 = ta->AllocateRaw(4, 4);
92 EXPECT_EQ(4, ta->RequestedSize(p1));
93 EXPECT_LE(4, ta->AllocatedSize(p1));
94 EXPECT_EQ(1, ta->AllocationId(p1));
95
96 ta->DeallocateRaw(p1);
97 p2 = ta->AllocateRaw(4, 12);
98 EXPECT_EQ(12, ta->RequestedSize(p2));
99 EXPECT_LE(12, ta->AllocatedSize(p2));
100 EXPECT_EQ(2, ta->AllocationId(p2));
101
102 sizes = ta->GetSizes();
103
104 EXPECT_LE(16, std::get<0>(sizes));
105 EXPECT_LE(12, std::get<1>(sizes));
106 EXPECT_LE(12, std::get<2>(sizes));
107
108 ta->DeallocateRaw(p2);
109 records = ta->GetRecordsAndUnRef();
110 EXPECT_LE(4, records[0].alloc_bytes);
111 EXPECT_GE(-4, records[1].alloc_bytes);
112 EXPECT_LE(12, records[2].alloc_bytes);
113 EXPECT_GE(-12, records[3].alloc_bytes);
114 }
115
TEST(TrackingAllocatorTest,SimpleTracking)116 TEST(TrackingAllocatorTest, SimpleTracking) {
117 TestableSizeTrackingAllocator a = TestableSizeTrackingAllocator();
118
119 EXPECT_TRUE(a.TracksAllocationSizes());
120
121 TrackingAllocator* ta = new TrackingAllocator(&a, false);
122
123 void* p1 = ta->AllocateRaw(4, 12);
124 ta->DeallocateRaw(p1);
125 void* p2 = ta->AllocateRaw(4, 4);
126
127 std::tuple<size_t, size_t, size_t> sizes = ta->GetSizes();
128
129 EXPECT_EQ(16, std::get<0>(sizes));
130 EXPECT_EQ(12, std::get<1>(sizes));
131 EXPECT_EQ(4, std::get<2>(sizes));
132
133 ta->DeallocateRaw(p2);
134
135 auto records = ta->GetRecordsAndUnRef();
136 EXPECT_EQ(12, records[0].alloc_bytes);
137 EXPECT_EQ(-12, records[1].alloc_bytes);
138 EXPECT_EQ(4, records[2].alloc_bytes);
139 EXPECT_EQ(-4, records[3].alloc_bytes);
140 }
141
TEST(TrackingAllocatorTest,OutOfMemory)142 TEST(TrackingAllocatorTest, OutOfMemory) {
143 NoMemoryAllocator a;
144
145 EXPECT_TRUE(a.TracksAllocationSizes());
146
147 TrackingAllocator* ta = new TrackingAllocator(&a, false);
148
149 void* p1 = ta->AllocateRaw(4, 12);
150 EXPECT_EQ(nullptr, p1);
151
152 std::tuple<size_t, size_t, size_t> sizes = ta->GetSizes();
153
154 EXPECT_EQ(0, std::get<0>(sizes));
155 EXPECT_EQ(0, std::get<1>(sizes));
156 EXPECT_EQ(0, std::get<2>(sizes));
157
158 EXPECT_EQ(0, ta->GetRecordsAndUnRef().size());
159 }
160
TEST(TrackingAllocatorTest,FreeNullPtr)161 TEST(TrackingAllocatorTest, FreeNullPtr) {
162 NoMemoryAllocator a;
163
164 EXPECT_TRUE(a.TracksAllocationSizes());
165
166 TrackingAllocator* ta = new TrackingAllocator(&a, false);
167
168 ta->DeallocateRaw(nullptr);
169
170 std::tuple<size_t, size_t, size_t> sizes = ta->GetSizes();
171
172 EXPECT_EQ(0, std::get<0>(sizes));
173 EXPECT_EQ(0, std::get<1>(sizes));
174 EXPECT_EQ(0, std::get<2>(sizes));
175
176 EXPECT_EQ(0, ta->GetRecordsAndUnRef().size());
177 }
178
179 } // namespace tensorflow
180