1 /* Copyright 2020 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/lite/micro/simple_memory_allocator.h"
17
18 #include <cstdint>
19
20 #include "tensorflow/lite/micro/micro_error_reporter.h"
21 #include "tensorflow/lite/micro/test_helpers.h"
22 #include "tensorflow/lite/micro/testing/micro_test.h"
23
24 TF_LITE_MICRO_TESTS_BEGIN
25
TF_LITE_MICRO_TEST(TestEnsureHeadSizeSimpleAlignment)26 TF_LITE_MICRO_TEST(TestEnsureHeadSizeSimpleAlignment) {
27 constexpr size_t arena_size = 1024;
28 uint8_t arena[arena_size];
29 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
30 arena, arena_size);
31
32 TF_LITE_MICRO_EXPECT_EQ(
33 kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/100, /*alignment=*/1));
34 TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(100),
35 allocator.GetHeadUsedBytes());
36
37 TF_LITE_MICRO_EXPECT_EQ(
38 kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/10, /*alignment=*/1));
39 TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(10),
40 allocator.GetHeadUsedBytes());
41
42 TF_LITE_MICRO_EXPECT_EQ(
43 kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/1000, /*alignment=*/1));
44 TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(1000),
45 allocator.GetHeadUsedBytes());
46 }
47
TF_LITE_MICRO_TEST(TestAdjustHeadSizeMisalignment)48 TF_LITE_MICRO_TEST(TestAdjustHeadSizeMisalignment) {
49 constexpr size_t arena_size = 1024;
50 uint8_t arena[arena_size];
51 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
52 arena, arena_size);
53
54 // First head adjustment of 100 bytes (aligned 12):
55 TF_LITE_MICRO_EXPECT_EQ(
56 kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/100, /*alignment=*/12));
57
58 // Offset alignment of 12 can lead to allocation within 8 byte range of
59 // requested bytes based to arena alignment at runtime:
60 TF_LITE_MICRO_EXPECT_GE(allocator.GetHeadUsedBytes(), 100);
61 TF_LITE_MICRO_EXPECT_LE(allocator.GetHeadUsedBytes(), 100 + 11);
62
63 TF_LITE_MICRO_EXPECT_EQ(
64 kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/10, /*alignment=*/12));
65 TF_LITE_MICRO_EXPECT_GE(allocator.GetHeadUsedBytes(), 10);
66 TF_LITE_MICRO_EXPECT_LE(allocator.GetHeadUsedBytes(), 100 + 11);
67
68 TF_LITE_MICRO_EXPECT_EQ(
69 kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/1000, /*alignment=*/12));
70 TF_LITE_MICRO_EXPECT_GE(allocator.GetHeadUsedBytes(), 1000);
71 TF_LITE_MICRO_EXPECT_LE(allocator.GetHeadUsedBytes(), 1000 + 11);
72 }
73
TF_LITE_MICRO_TEST(TestAdjustHeadSizeMisalignedHandlesCorrectBytesAvailable)74 TF_LITE_MICRO_TEST(TestAdjustHeadSizeMisalignedHandlesCorrectBytesAvailable) {
75 constexpr size_t arena_size = 1024;
76 uint8_t arena[arena_size];
77 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
78 arena, arena_size);
79
80 // First head adjustment of 100 bytes (aligned 12):
81 TF_LITE_MICRO_EXPECT_EQ(
82 kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/100, /*alignment=*/12));
83
84 // allocator.GetAvailableMemory() should also report the actual amount of
85 // memory available based on a requested offset (12):
86 size_t aligned_available_bytes =
87 allocator.GetAvailableMemory(/*alignment=*/12);
88 TF_LITE_MICRO_EXPECT_LE(aligned_available_bytes, arena_size - 100);
89 TF_LITE_MICRO_EXPECT_GE(aligned_available_bytes, arena_size - 100 - 24);
90
91 TF_LITE_MICRO_EXPECT_EQ(
92 kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/10, /*alignment=*/12));
93 aligned_available_bytes = allocator.GetAvailableMemory(/*alignment=*/12);
94
95 TF_LITE_MICRO_EXPECT_LE(aligned_available_bytes, arena_size - 10);
96 TF_LITE_MICRO_EXPECT_GE(aligned_available_bytes, arena_size - 10 - 24);
97
98 TF_LITE_MICRO_EXPECT_EQ(
99 kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/1000, /*alignment=*/12));
100 aligned_available_bytes = allocator.GetAvailableMemory(/*alignment=*/12);
101 TF_LITE_MICRO_EXPECT_LE(aligned_available_bytes, arena_size - 1000);
102 TF_LITE_MICRO_EXPECT_GE(aligned_available_bytes, arena_size - 1000 - 24);
103 }
104
TF_LITE_MICRO_TEST(TestGetAvailableMemory)105 TF_LITE_MICRO_TEST(TestGetAvailableMemory) {
106 constexpr size_t arena_size = 1024;
107 uint8_t arena[arena_size];
108 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
109 arena, arena_size);
110
111 constexpr size_t allocation_size = 100;
112 allocator.SetHeadBufferSize(/*size=*/allocation_size,
113 /*alignment=*/1);
114 allocator.AllocateFromTail(/*size=*/allocation_size,
115 /*alignment=*/1);
116
117 TF_LITE_MICRO_EXPECT_EQ(allocator.GetAvailableMemory(/*alignment=*/1),
118 arena_size - allocation_size * 2);
119 }
120
TF_LITE_MICRO_TEST(TestGetAvailableMemoryWithTempAllocations)121 TF_LITE_MICRO_TEST(TestGetAvailableMemoryWithTempAllocations) {
122 constexpr size_t arena_size = 1024;
123 uint8_t arena[arena_size];
124 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
125 arena, arena_size);
126
127 constexpr size_t allocation_size = 100;
128 allocator.AllocateTemp(/*size=*/allocation_size,
129 /*alignment=*/1);
130
131 TF_LITE_MICRO_EXPECT_EQ(allocator.GetAvailableMemory(/*alignment=*/1),
132 arena_size - allocation_size);
133
134 // Reset temp allocations and ensure GetAvailableMemory() is back to the
135 // starting size:
136 allocator.ResetTempAllocations();
137
138 TF_LITE_MICRO_EXPECT_EQ(allocator.GetAvailableMemory(/*alignment=*/1),
139 arena_size);
140 }
141
TF_LITE_MICRO_TEST(TestGetUsedBytes)142 TF_LITE_MICRO_TEST(TestGetUsedBytes) {
143 constexpr size_t arena_size = 1024;
144 uint8_t arena[arena_size];
145 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
146 arena, arena_size);
147 TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), static_cast<size_t>(0));
148
149 constexpr size_t allocation_size = 100;
150 allocator.SetHeadBufferSize(/*size=*/allocation_size,
151 /*alignment=*/1);
152 allocator.AllocateFromTail(/*size=*/allocation_size,
153 /*alignment=*/1);
154
155 TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), allocation_size * 2);
156 }
157
TF_LITE_MICRO_TEST(TestGetUsedBytesTempAllocations)158 TF_LITE_MICRO_TEST(TestGetUsedBytesTempAllocations) {
159 constexpr size_t arena_size = 1024;
160 uint8_t arena[arena_size];
161 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
162 arena, arena_size);
163
164 constexpr size_t allocation_size = 100;
165 allocator.AllocateTemp(/*size=*/allocation_size,
166 /*alignment=*/1);
167
168 TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), allocation_size);
169
170 // Reset temp allocations and ensure GetUsedBytes() is back to the starting
171 // size:
172 allocator.ResetTempAllocations();
173
174 TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), static_cast<size_t>(0));
175 }
176
TF_LITE_MICRO_TEST(TestJustFits)177 TF_LITE_MICRO_TEST(TestJustFits) {
178 constexpr size_t arena_size = 1024;
179 uint8_t arena[arena_size];
180 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
181 arena, arena_size);
182
183 uint8_t* result = allocator.AllocateFromTail(arena_size, 1);
184 TF_LITE_MICRO_EXPECT(nullptr != result);
185 }
186
TF_LITE_MICRO_TEST(TestAligned)187 TF_LITE_MICRO_TEST(TestAligned) {
188 constexpr size_t arena_size = 1024;
189 uint8_t arena[arena_size];
190 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
191 arena, arena_size);
192
193 uint8_t* result = allocator.AllocateFromTail(1, 1);
194 TF_LITE_MICRO_EXPECT(nullptr != result);
195
196 result = allocator.AllocateFromTail(16, 4);
197 TF_LITE_MICRO_EXPECT(nullptr != result);
198 TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(0),
199 reinterpret_cast<std::uintptr_t>(result) & 3);
200 }
201
TF_LITE_MICRO_TEST(TestMultipleTooLarge)202 TF_LITE_MICRO_TEST(TestMultipleTooLarge) {
203 constexpr size_t arena_size = 1024;
204 uint8_t arena[arena_size];
205 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
206 arena, arena_size);
207
208 uint8_t* result = allocator.AllocateFromTail(768, 1);
209 TF_LITE_MICRO_EXPECT(nullptr != result);
210
211 result = allocator.AllocateFromTail(768, 1);
212 TF_LITE_MICRO_EXPECT(nullptr == result);
213 }
214
TF_LITE_MICRO_TEST(TestTempAllocations)215 TF_LITE_MICRO_TEST(TestTempAllocations) {
216 constexpr size_t arena_size = 1024;
217 uint8_t arena[arena_size];
218 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
219 arena, arena_size);
220
221 uint8_t* temp1 = allocator.AllocateTemp(100, 1);
222 TF_LITE_MICRO_EXPECT(nullptr != temp1);
223
224 uint8_t* temp2 = allocator.AllocateTemp(100, 1);
225 TF_LITE_MICRO_EXPECT(nullptr != temp2);
226
227 // Expect that the next micro allocation is 100 bytes away from each other.
228 TF_LITE_MICRO_EXPECT_EQ(temp2 - temp1, 100);
229 }
230
TF_LITE_MICRO_TEST(TestResetTempAllocations)231 TF_LITE_MICRO_TEST(TestResetTempAllocations) {
232 constexpr size_t arena_size = 1024;
233 uint8_t arena[arena_size];
234 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
235 arena, arena_size);
236
237 uint8_t* temp1 = allocator.AllocateTemp(100, 1);
238 TF_LITE_MICRO_EXPECT(nullptr != temp1);
239
240 allocator.ResetTempAllocations();
241
242 uint8_t* temp2 = allocator.AllocateTemp(100, 1);
243 TF_LITE_MICRO_EXPECT(nullptr != temp2);
244
245 // Reset temp allocations should have the same start address:
246 TF_LITE_MICRO_EXPECT_EQ(temp2 - temp1, 0);
247 }
248
TF_LITE_MICRO_TEST(TestEnsureHeadSizeWithoutResettingTemp)249 TF_LITE_MICRO_TEST(TestEnsureHeadSizeWithoutResettingTemp) {
250 constexpr size_t arena_size = 1024;
251 uint8_t arena[arena_size];
252 tflite::SimpleMemoryAllocator allocator(tflite::GetMicroErrorReporter(),
253 arena, arena_size);
254
255 uint8_t* temp = allocator.AllocateTemp(100, 1);
256 TF_LITE_MICRO_EXPECT(nullptr != temp);
257
258 // Adjustment to head should fail since temp allocation was not followed by a
259 // call to ResetTempAllocations().
260 TF_LITE_MICRO_EXPECT_EQ(kTfLiteError, allocator.SetHeadBufferSize(100, 1));
261
262 allocator.ResetTempAllocations();
263
264 // Reduce head size back to zero.
265 TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.SetHeadBufferSize(0, 1));
266
267 // The most recent head allocation should be in the same location as the
268 // original temp allocation pointer.
269 TF_LITE_MICRO_EXPECT(temp == allocator.GetHeadBuffer());
270 }
271
272 TF_LITE_MICRO_TESTS_END
273