1 /* Copyright 2021 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 #include "tensorflow/lite/simple_planner.h"
16
17 #include <algorithm>
18 #include <cstdarg>
19 #include <initializer_list>
20 #include <memory>
21 #include <utility>
22 #include <vector>
23
24 #include <gtest/gtest.h>
25 #include "tensorflow/core/platform/logging.h"
26 #include "tensorflow/lite/c/common.h"
27 #include "tensorflow/lite/graph_info.h"
28 #include "tensorflow/lite/testing/util.h"
29
30 namespace tflite {
31 namespace {
32
33 // A simple op to be used in tests, as syntactic sugar.
34 class TestOp {
35 public:
TestOp(std::initializer_list<int> inputs,std::initializer_list<int> outputs,std::initializer_list<int> temporaries)36 TestOp(std::initializer_list<int> inputs, std::initializer_list<int> outputs,
37 std::initializer_list<int> temporaries)
38 : inputs_(inputs), outputs_(outputs), temporaries_(temporaries) {}
39
inputs() const40 const std::vector<int>& inputs() const { return inputs_; }
outputs() const41 const std::vector<int>& outputs() const { return outputs_; }
temporaries() const42 const std::vector<int>& temporaries() const { return temporaries_; }
43
44 private:
45 std::vector<int> inputs_;
46 std::vector<int> outputs_;
47 std::vector<int> temporaries_;
48 };
49
50 // A test graph where inputs are processed by the given nodes to produce
51 // outputs.
52 class TestGraph {
53 public:
TestGraph(std::initializer_list<int> inputs,std::initializer_list<TestOp> nodes,std::initializer_list<int> outputs)54 TestGraph(std::initializer_list<int> inputs,
55 std::initializer_list<TestOp> nodes,
56 std::initializer_list<int> outputs)
57 : inputs_(inputs), outputs_(outputs) {
58 int max_tensor_index = 0;
59
60 for (int t : inputs) {
61 max_tensor_index = std::max(max_tensor_index, t);
62 }
63 for (int t : outputs) {
64 max_tensor_index = std::max(max_tensor_index, t);
65 }
66 for (const auto& node : nodes) {
67 auto int_array = [](const std::vector<int>& x) {
68 TfLiteIntArray* lite = TfLiteIntArrayCreate(x.size());
69 for (size_t i = 0; i < x.size(); i++) lite->data[i] = x[i];
70 return lite;
71 };
72
73 nodes_.push_back(TfLiteNode());
74 nodes_.back().inputs = int_array(node.inputs());
75 for (int t : node.inputs()) {
76 max_tensor_index = std::max(max_tensor_index, t);
77 }
78 nodes_.back().outputs = int_array(node.outputs());
79 for (int t : node.outputs()) {
80 max_tensor_index = std::max(max_tensor_index, t);
81 }
82 nodes_.back().temporaries = int_array(node.temporaries());
83 for (int t : node.temporaries()) {
84 max_tensor_index = std::max(max_tensor_index, t);
85 }
86 }
87
88 for (int i = 0; i <= max_tensor_index; ++i) {
89 tensors_.push_back(TfLiteTensor());
90 // Set some default values for allocation_type and bytes, which are the
91 // only fields used by the arena planner.
92 tensors_.back().allocation_type = kTfLiteArenaRw;
93 tensors_.back().bytes = (i + 1) * 3;
94 }
95 }
96
~TestGraph()97 ~TestGraph() {
98 for (auto node : nodes_) {
99 TfLiteIntArrayFree(node.inputs);
100 TfLiteIntArrayFree(node.outputs);
101 TfLiteIntArrayFree(node.temporaries);
102 }
103 }
104
nodes()105 const std::vector<TfLiteNode>& nodes() { return nodes_; }
tensors()106 std::vector<TfLiteTensor>* tensors() { return &tensors_; }
inputs()107 const std::vector<int>& inputs() { return inputs_; }
outputs()108 const std::vector<int>& outputs() { return outputs_; }
variables()109 const std::vector<int>& variables() { return variables_; }
110
SetVariables(const std::vector<int> & variables)111 void SetVariables(const std::vector<int>& variables) {
112 variables_ = variables;
113 }
114
Swap(TestGraph * other)115 void Swap(TestGraph* other) {
116 std::swap(nodes_, other->nodes_);
117 std::swap(tensors_, other->tensors_);
118 std::swap(inputs_, other->inputs_);
119 std::swap(outputs_, other->outputs_);
120 std::swap(variables_, other->variables_);
121 }
122
123 private:
124 std::vector<TfLiteNode> nodes_;
125 std::vector<TfLiteTensor> tensors_;
126 std::vector<int> inputs_;
127 std::vector<int> outputs_;
128 std::vector<int> variables_;
129 };
130
131 // The GraphInfo for a TestGraph.
132 class TestGraphInfo : public GraphInfo {
133 public:
TestGraphInfo(TestGraph * graph)134 explicit TestGraphInfo(TestGraph* graph) : graph_(graph) {}
135
num_tensors() const136 size_t num_tensors() const override { return graph_->tensors()->size(); }
tensor(size_t index)137 TfLiteTensor* tensor(size_t index) override {
138 return &graph_->tensors()->at(index);
139 }
num_execution_nodes() const140 size_t num_execution_nodes() const override { return graph_->nodes().size(); }
num_total_nodes() const141 size_t num_total_nodes() const override { return graph_->nodes().size(); }
node(size_t index) const142 const TfLiteNode& node(size_t index) const override {
143 return graph_->nodes()[index];
144 }
node_index(size_t index) const145 size_t node_index(size_t index) const override { return index; }
inputs() const146 const std::vector<int>& inputs() const override { return graph_->inputs(); }
outputs() const147 const std::vector<int>& outputs() const override { return graph_->outputs(); }
variables() const148 const std::vector<int>& variables() const override {
149 return graph_->variables();
150 }
151
152 private:
153 TestGraph* graph_;
154 };
155
ReportError(TfLiteContext * context,const char * format,...)156 void ReportError(TfLiteContext* context, const char* format, ...) {
157 const size_t kBufferSize = 1024;
158 char temp_buffer[kBufferSize];
159
160 va_list args;
161 va_start(args, format);
162 vsnprintf(temp_buffer, kBufferSize, format, args);
163 va_end(args);
164
165 LOG(INFO) << temp_buffer;
166 }
167
168 class SimplePlannerTest : public ::testing::Test {
169 protected:
SetGraph(TestGraph * graph,bool preserve_all_tensors=false)170 void SetGraph(TestGraph* graph, bool preserve_all_tensors = false) {
171 graph_ = graph;
172 context_.ReportError = ReportError;
173 planner_ = std::make_unique<SimplePlanner>(
174 &context_, std::unique_ptr<GraphInfo>(new TestGraphInfo(graph)));
175 CHECK(planner_->ResetAllocations() == kTfLiteOk);
176 CHECK(planner_->PlanAllocations() == kTfLiteOk);
177 }
178
SwapGraph(TestGraph * graph)179 void SwapGraph(TestGraph* graph) {
180 graph_->Swap(graph);
181 CHECK(planner_->PlanAllocations() == kTfLiteOk);
182 }
183
Execute(int start,int end)184 void Execute(int start, int end) {
185 CHECK(planner_->ExecuteAllocations(start, end) == kTfLiteOk);
186 }
187
ReleaseNonPersistentMemory()188 void ReleaseNonPersistentMemory() {
189 CHECK(planner_->ReleaseNonPersistentMemory() == kTfLiteOk);
190 }
191
AcquireNonPersistentMemory()192 void AcquireNonPersistentMemory() {
193 CHECK(planner_->AcquireNonPersistentMemory() == kTfLiteOk);
194 }
195
ResetAllocationsAfter(int node)196 void ResetAllocationsAfter(int node) {
197 CHECK(planner_->ResetAllocationsAfter(node) == kTfLiteOk);
198 }
199
HasNonPersistentMemory()200 bool HasNonPersistentMemory() {
201 return planner_ && planner_->HasNonPersistentMemory();
202 }
203
204 // Returns if the given tensor is allocated or not.
IsAllocated(int tensor_index)205 bool IsAllocated(int tensor_index) {
206 return (*graph_->tensors())[tensor_index].data.raw != nullptr;
207 }
208
209 TfLiteContext context_;
210 TestGraph* graph_;
211 std::unique_ptr<SimplePlanner> planner_;
212 };
213
TEST_F(SimplePlannerTest,EmptyGraph)214 TEST_F(SimplePlannerTest, EmptyGraph) {
215 TestGraph graph({}, {}, {});
216 SetGraph(&graph);
217 Execute(0, 10);
218 }
219
TEST_F(SimplePlannerTest,GraphWithNoOps)220 TEST_F(SimplePlannerTest, GraphWithNoOps) {
221 TestGraph graph({0, 10}, {}, {5, 11});
222 SetGraph(&graph);
223 Execute(0, 10);
224 // The outputs are never allocated because they are not connected to any
225 // inputs.
226 EXPECT_FALSE(IsAllocated(5));
227 EXPECT_FALSE(IsAllocated(11));
228 }
229
TEST_F(SimplePlannerTest,ZeroSizedTensors)230 TEST_F(SimplePlannerTest, ZeroSizedTensors) {
231 TestGraph graph({1}, {{{1}, {2}, {}}}, {2});
232 (*graph.tensors())[1].bytes = 0;
233 SetGraph(&graph);
234 ASSERT_EQ(planner_->ExecuteAllocations(0, 10), kTfLiteOk);
235 EXPECT_FALSE(IsAllocated(1));
236 EXPECT_TRUE(IsAllocated(2));
237 }
238
TEST_F(SimplePlannerTest,SimpleGraph)239 TEST_F(SimplePlannerTest, SimpleGraph) {
240 TestGraph graph({0, 1},
241 {
242 /* in, out, tmp */
243 {{0, 1}, {2}, {}}, // First op
244 {{2, 0}, {4, 5}, {}}, // Second op
245 {{4, 5}, {3}, {}} // Third op
246 },
247 {3});
248 SetGraph(&graph);
249 Execute(0, 10);
250
251 EXPECT_TRUE(IsAllocated(1));
252 EXPECT_TRUE(IsAllocated(2));
253 EXPECT_TRUE(IsAllocated(3));
254 EXPECT_TRUE(IsAllocated(4));
255 EXPECT_TRUE(IsAllocated(5));
256 }
257
TEST_F(SimplePlannerTest,SimpleGraphInputsPreserved)258 TEST_F(SimplePlannerTest, SimpleGraphInputsPreserved) {
259 TestGraph graph({0, 1},
260 {
261 /* in, out, tmp */
262 {{0, 1}, {2}, {}}, // First op
263 {{2, 0}, {4, 5}, {}}, // Second op
264 {{4, 5}, {3}, {}} // Third op
265 },
266 {3});
267 SetGraph(&graph);
268 Execute(0, 10);
269
270 EXPECT_TRUE(IsAllocated(1));
271 EXPECT_TRUE(IsAllocated(2));
272 EXPECT_TRUE(IsAllocated(3));
273 EXPECT_TRUE(IsAllocated(4));
274 EXPECT_TRUE(IsAllocated(5));
275 }
276
TEST_F(SimplePlannerTest,SimpleGraphWithTemporary)277 TEST_F(SimplePlannerTest, SimpleGraphWithTemporary) {
278 TestGraph graph({0, 1},
279 {
280 /* in, out, tmp */
281 {{0, 1}, {2}, {}}, // First op
282 {{2, 0}, {4}, {5}}, // Second op, with temporary
283 {{4}, {3}, {}} // Third op
284 },
285 {3});
286 SetGraph(&graph);
287 Execute(0, 10);
288
289 EXPECT_TRUE(IsAllocated(1));
290 EXPECT_TRUE(IsAllocated(2));
291 EXPECT_TRUE(IsAllocated(3));
292 EXPECT_TRUE(IsAllocated(4));
293 EXPECT_TRUE(IsAllocated(5));
294 }
295
TEST_F(SimplePlannerTest,SimpleGraphWithResetAllocationsAfter)296 TEST_F(SimplePlannerTest, SimpleGraphWithResetAllocationsAfter) {
297 TestGraph graph({0, 1},
298 {
299 /* in, out, tmp */
300 {{0, 1}, {2}, {}}, // First op
301 {{2, 0}, {4}, {5}}, // Second op, with temporary
302 {{4}, {3}, {}} // Third op
303 },
304 {3});
305 SetGraph(&graph);
306 Execute(0, 10);
307
308 EXPECT_TRUE(IsAllocated(2));
309 EXPECT_TRUE(IsAllocated(3));
310 EXPECT_TRUE(IsAllocated(4));
311 EXPECT_TRUE(IsAllocated(5));
312 // Reset allocations after the first node
313 ResetAllocationsAfter(0);
314
315 EXPECT_TRUE(IsAllocated(0));
316 EXPECT_TRUE(IsAllocated(1));
317 EXPECT_TRUE(IsAllocated(2));
318 EXPECT_FALSE(IsAllocated(3));
319 EXPECT_FALSE(IsAllocated(4));
320 EXPECT_FALSE(IsAllocated(5));
321 }
322
TEST_F(SimplePlannerTest,SimpleGraphWithPersistentResetAllocationsAfter)323 TEST_F(SimplePlannerTest, SimpleGraphWithPersistentResetAllocationsAfter) {
324 TestGraph graph({0, 1},
325 {
326 /* in, out, tmp */
327 {{0, 1}, {2}, {}}, // First op
328 {{2, 0}, {4}, {5}}, // Second op, with temporary
329 {{4}, {3}, {}} // Third op
330 },
331 {3});
332 // Make the tensor #5 persistent.
333 (*graph.tensors())[5].allocation_type = kTfLiteArenaRwPersistent;
334 SetGraph(&graph);
335 Execute(0, 10);
336
337 // Save the pointer of the persistent temporary tensor #5.
338 void* tensor5_ptr = (*graph.tensors())[5].data.raw;
339
340 // Reset allocations after the first node
341 ResetAllocationsAfter(0);
342
343 EXPECT_TRUE(IsAllocated(0));
344 EXPECT_TRUE(IsAllocated(1));
345 EXPECT_TRUE(IsAllocated(2));
346 EXPECT_FALSE(IsAllocated(3));
347 EXPECT_FALSE(IsAllocated(4));
348 EXPECT_TRUE(IsAllocated(5));
349
350 // Second run
351 Execute(0, 10);
352
353 // Check if the persistent pointer isn't changed.
354 EXPECT_TRUE(tensor5_ptr == (*graph.tensors())[5].data.raw);
355 }
356
357 } // namespace
358 } // namespace tflite
359