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/recording_micro_allocator.h"
17
18 #include "tensorflow/lite/micro/all_ops_resolver.h"
19 #include "tensorflow/lite/micro/micro_error_reporter.h"
20 #include "tensorflow/lite/micro/test_helpers.h"
21 #include "tensorflow/lite/micro/testing/micro_test.h"
22 #include "tensorflow/lite/micro/testing/test_conv_model.h"
23
24 #define TF_LITE_TENSOR_STRUCT_SIZE sizeof(TfLiteTensor)
25 #define TF_LITE_EVAL_TENSOR_STRUCT_SIZE sizeof(TfLiteEvalTensor)
26 #define TF_LITE_AFFINE_QUANTIZATION_SIZE sizeof(TfLiteAffineQuantization)
27 #define NODE_AND_REGISTRATION_STRUCT_SIZE sizeof(tflite::NodeAndRegistration)
28
29 // TODO(b/158303868): Move tests into anonymous namespace.
30 namespace {
31
32 constexpr int kTestConvArenaSize = 1024 * 12;
33
34 } // namespace
35
36 TF_LITE_MICRO_TESTS_BEGIN
37
TF_LITE_MICRO_TEST(TestRecordsTfLiteEvalTensorArrayData)38 TF_LITE_MICRO_TEST(TestRecordsTfLiteEvalTensorArrayData) {
39 TfLiteEvalTensor* eval_tensors = nullptr;
40 tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
41 tflite::AllOpsResolver all_ops_resolver;
42 tflite::NodeAndRegistration* node_and_registration;
43 const tflite::Model* model = tflite::GetModel(kTestConvModelData);
44 uint8_t arena[kTestConvArenaSize];
45
46 tflite::RecordingMicroAllocator* micro_allocator =
47 tflite::RecordingMicroAllocator::Create(arena, kTestConvArenaSize,
48 tflite::GetMicroErrorReporter());
49 // TODO(b/158102673): ugly workaround for not having fatal assertions. Same
50 // throughout this file.
51 TF_LITE_MICRO_EXPECT_NE(micro_allocator, nullptr);
52 if (micro_allocator == nullptr) return 1;
53
54 TfLiteStatus status;
55 status = micro_allocator->StartModelAllocation(
56 model, all_ops_resolver, &node_and_registration, &eval_tensors);
57 TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
58 if (status != kTfLiteOk) return 1;
59
60 status = micro_allocator->FinishModelAllocation(model, eval_tensors,
61 &scratch_buffer_handles);
62 TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
63 if (status != kTfLiteOk) return 1;
64
65 micro_allocator->PrintAllocations();
66
67 tflite::RecordedAllocation recorded_allocation =
68 micro_allocator->GetRecordedAllocation(
69 tflite::RecordedAllocationType::kTfLiteEvalTensorData);
70
71 micro_allocator->PrintAllocations();
72
73 size_t tensors_count = tflite::testing::GetModelTensorCount(model);
74
75 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count, tensors_count);
76 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
77 tensors_count * TF_LITE_EVAL_TENSOR_STRUCT_SIZE);
78 TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,
79 tensors_count * TF_LITE_EVAL_TENSOR_STRUCT_SIZE);
80 }
81
TF_LITE_MICRO_TEST(TestRecordsNodeAndRegistrationArrayData)82 TF_LITE_MICRO_TEST(TestRecordsNodeAndRegistrationArrayData) {
83 TfLiteEvalTensor* eval_tensors = nullptr;
84 tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
85 tflite::AllOpsResolver all_ops_resolver;
86 tflite::NodeAndRegistration* node_and_registration;
87 const tflite::Model* model = tflite::GetModel(kTestConvModelData);
88 uint8_t arena[kTestConvArenaSize];
89
90 tflite::RecordingMicroAllocator* micro_allocator =
91 tflite::RecordingMicroAllocator::Create(arena, kTestConvArenaSize,
92 tflite::GetMicroErrorReporter());
93 TF_LITE_MICRO_EXPECT_NE(micro_allocator, nullptr);
94 if (micro_allocator == nullptr) return 1;
95
96 TfLiteStatus status;
97 status = micro_allocator->StartModelAllocation(
98 model, all_ops_resolver, &node_and_registration, &eval_tensors);
99 TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
100 if (status != kTfLiteOk) return 1;
101
102 status = micro_allocator->FinishModelAllocation(model, eval_tensors,
103 &scratch_buffer_handles);
104 TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
105 if (status != kTfLiteOk) return 1;
106
107 size_t num_ops = model->subgraphs()->Get(0)->operators()->size();
108 tflite::RecordedAllocation recorded_allocation =
109 micro_allocator->GetRecordedAllocation(
110 tflite::RecordedAllocationType::kNodeAndRegistrationArray);
111 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count, num_ops);
112 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
113 num_ops * NODE_AND_REGISTRATION_STRUCT_SIZE);
114 TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,
115 num_ops * NODE_AND_REGISTRATION_STRUCT_SIZE);
116 }
117
TF_LITE_MICRO_TEST(TestRecordsMultiTenantAllocations)118 TF_LITE_MICRO_TEST(TestRecordsMultiTenantAllocations) {
119 TfLiteEvalTensor* eval_tensors = nullptr;
120 tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
121 tflite::AllOpsResolver all_ops_resolver;
122 tflite::NodeAndRegistration* node_and_registration;
123 const tflite::Model* model = tflite::GetModel(kTestConvModelData);
124
125 // Double the arena size to allocate two models inside of it:
126 uint8_t arena[kTestConvArenaSize * 2];
127
128 TfLiteStatus status;
129
130 tflite::RecordingMicroAllocator* micro_allocator =
131 tflite::RecordingMicroAllocator::Create(arena, kTestConvArenaSize * 2,
132 tflite::GetMicroErrorReporter());
133 TF_LITE_MICRO_EXPECT_NE(micro_allocator, nullptr);
134 if (micro_allocator == nullptr) return 1;
135
136 // First allocation with the model in the arena:
137 status = micro_allocator->StartModelAllocation(
138 model, all_ops_resolver, &node_and_registration, &eval_tensors);
139 TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
140 if (status != kTfLiteOk) return 1;
141
142 status = micro_allocator->FinishModelAllocation(model, eval_tensors,
143 &scratch_buffer_handles);
144 TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
145 if (status != kTfLiteOk) return 1;
146
147 // Second allocation with the same model in the arena:
148 status = micro_allocator->StartModelAllocation(
149 model, all_ops_resolver, &node_and_registration, &eval_tensors);
150 TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
151 if (status != kTfLiteOk) return 1;
152
153 status = kTfLiteOk, micro_allocator->FinishModelAllocation(
154 model, eval_tensors, &scratch_buffer_handles);
155 TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
156 if (status != kTfLiteOk) return 1;
157
158 size_t tensors_count = tflite::testing::GetModelTensorCount(model);
159
160 tflite::RecordedAllocation recorded_allocation =
161 micro_allocator->GetRecordedAllocation(
162 tflite::RecordedAllocationType::kTfLiteEvalTensorData);
163 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count, tensors_count * 2);
164 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
165 tensors_count * TF_LITE_EVAL_TENSOR_STRUCT_SIZE * 2);
166 TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,
167 tensors_count * TF_LITE_EVAL_TENSOR_STRUCT_SIZE * 2);
168 }
169
TF_LITE_MICRO_TEST(TestRecordsPersistentTfLiteTensorData)170 TF_LITE_MICRO_TEST(TestRecordsPersistentTfLiteTensorData) {
171 const tflite::Model* model = tflite::GetModel(kTestConvModelData);
172 uint8_t arena[kTestConvArenaSize];
173
174 tflite::RecordingMicroAllocator* micro_allocator =
175 tflite::RecordingMicroAllocator::Create(arena, kTestConvArenaSize,
176 tflite::GetMicroErrorReporter());
177 TF_LITE_MICRO_EXPECT_NE(micro_allocator, nullptr);
178 if (micro_allocator == nullptr) return 1;
179
180 TfLiteTensor* tensor = micro_allocator->AllocatePersistentTfLiteTensor(
181 model, /*eval_tensors=*/nullptr, 0);
182 TF_LITE_MICRO_EXPECT_NE(tensor, nullptr);
183 if (tensor == nullptr) return 1;
184
185 tflite::RecordedAllocation recorded_allocation =
186 micro_allocator->GetRecordedAllocation(
187 tflite::RecordedAllocationType::kPersistentTfLiteTensorData);
188
189 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count, static_cast<size_t>(1));
190 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
191 TF_LITE_TENSOR_STRUCT_SIZE);
192 TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,
193 TF_LITE_TENSOR_STRUCT_SIZE);
194 }
195
TF_LITE_MICRO_TEST(TestRecordsPersistentTfLiteTensorQuantizationData)196 TF_LITE_MICRO_TEST(TestRecordsPersistentTfLiteTensorQuantizationData) {
197 const tflite::Model* model = tflite::GetModel(kTestConvModelData);
198 uint8_t arena[kTestConvArenaSize];
199
200 tflite::RecordingMicroAllocator* micro_allocator =
201 tflite::RecordingMicroAllocator::Create(arena, kTestConvArenaSize,
202 tflite::GetMicroErrorReporter());
203 TF_LITE_MICRO_EXPECT_NE(micro_allocator, nullptr);
204 if (micro_allocator == nullptr) return 1;
205
206 TfLiteTensor* tensor = micro_allocator->AllocatePersistentTfLiteTensor(
207 model, /*eval_tensors=*/nullptr, 0);
208 TF_LITE_MICRO_EXPECT_NE(tensor, nullptr);
209 if (tensor == nullptr) return 1;
210
211 // Walk the model subgraph to find all tensors with quantization params and
212 // keep a tally.
213 size_t quantized_channel_bytes = 0;
214 const tflite::Tensor* cur_tensor =
215 model->subgraphs()->Get(0)->tensors()->Get(0);
216 const tflite::QuantizationParameters* quantization_params =
217 cur_tensor->quantization();
218 if (quantization_params && quantization_params->scale() &&
219 quantization_params->scale()->size() > 0 &&
220 quantization_params->zero_point() &&
221 quantization_params->zero_point()->size() > 0) {
222 size_t num_channels = quantization_params->scale()->size();
223 quantized_channel_bytes += TfLiteIntArrayGetSizeInBytes(num_channels);
224 }
225
226 // Calculate the expected allocation bytes with subgraph quantization data:
227 size_t expected_requested_bytes =
228 TF_LITE_AFFINE_QUANTIZATION_SIZE + quantized_channel_bytes;
229
230 tflite::RecordedAllocation recorded_allocation =
231 micro_allocator->GetRecordedAllocation(
232 tflite::RecordedAllocationType::
233 kPersistentTfLiteTensorQuantizationData);
234
235 // Each quantized tensors has 2 mallocs (quant struct, zero point dimensions):
236 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count, static_cast<size_t>(2));
237 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
238 expected_requested_bytes);
239 TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,
240 expected_requested_bytes);
241 }
242
TF_LITE_MICRO_TEST(TestRecordsPersistentBufferData)243 TF_LITE_MICRO_TEST(TestRecordsPersistentBufferData) {
244 uint8_t arena[kTestConvArenaSize];
245
246 tflite::RecordingMicroAllocator* micro_allocator =
247 tflite::RecordingMicroAllocator::Create(arena, kTestConvArenaSize,
248 tflite::GetMicroErrorReporter());
249 TF_LITE_MICRO_EXPECT_NE(micro_allocator, nullptr);
250 if (micro_allocator == nullptr) return 1;
251
252 void* buffer = micro_allocator->AllocatePersistentBuffer(/*bytes=*/100);
253 TF_LITE_MICRO_EXPECT_NE(buffer, nullptr);
254 if (buffer == nullptr) return 1;
255
256 tflite::RecordedAllocation recorded_allocation =
257 micro_allocator->GetRecordedAllocation(
258 tflite::RecordedAllocationType::kPersistentBufferData);
259
260 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count, static_cast<size_t>(1));
261 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
262 static_cast<size_t>(100));
263 TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,
264 static_cast<size_t>(100));
265
266 buffer = micro_allocator->AllocatePersistentBuffer(/*bytes=*/50);
267 TF_LITE_MICRO_EXPECT_NE(buffer, nullptr);
268 if (buffer == nullptr) return 1;
269
270 recorded_allocation = micro_allocator->GetRecordedAllocation(
271 tflite::RecordedAllocationType::kPersistentBufferData);
272
273 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count, static_cast<size_t>(2));
274 TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
275 static_cast<size_t>(150));
276 TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,
277 static_cast<size_t>(150));
278 }
279
280 // TODO(b/158124094): Find a way to audit OpData allocations on
281 // cross-architectures.
282
283 TF_LITE_MICRO_TESTS_END
284