1 /* Copyright 2019 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 <stdint.h>
16
17 #include <memory>
18 #include <vector>
19
20 #include <gtest/gtest.h>
21 #include "tensorflow/lite/interpreter.h"
22 #include "tensorflow/lite/kernels/subgraph_test_util.h"
23 #include "tensorflow/lite/profiling/memory_info.h"
24
25 namespace tflite {
26
27 using subgraph_test_util::CheckIntTensor;
28 using subgraph_test_util::CheckScalarStringTensor;
29 using subgraph_test_util::CheckStringTensor;
30 using subgraph_test_util::ControlFlowOpTest;
31 using subgraph_test_util::FillIntTensor;
32 using subgraph_test_util::FillScalarStringTensor;
33
34 namespace {
35
36 class WhileTest : public ControlFlowOpTest {};
37
38 // The test builds a model that produces the i-th number of
39 // triangular number sequence.
TEST_F(WhileTest,TestTriangularNumberSequence)40 TEST_F(WhileTest, TestTriangularNumberSequence) {
41 const std::vector<int> expected = {1, 3, 6, 10, 15, 21, 28};
42 for (int i = 0; i < expected.size(); ++i) {
43 interpreter_ = std::make_unique<Interpreter>();
44 AddSubgraphs(2);
45 builder_->BuildLessEqualCondSubgraph(interpreter_->subgraph(1), i);
46 builder_->BuildAccumulateLoopBodySubgraph(interpreter_->subgraph(2));
47 builder_->BuildWhileSubgraph(&interpreter_->primary_subgraph());
48
49 interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {1});
50 interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {1});
51 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
52 FillIntTensor(interpreter_->tensor(interpreter_->inputs()[0]), {1});
53 FillIntTensor(interpreter_->tensor(interpreter_->inputs()[1]), {1});
54
55 // Check While BODY inputs are static tensors.
56 auto body_subgraph = interpreter_->subgraph(2);
57 TfLiteTensor* subgraph_input2 =
58 body_subgraph->tensor(body_subgraph->inputs()[1]);
59 ASSERT_EQ(subgraph_input2->allocation_type, kTfLiteArenaRw);
60
61 ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
62 TfLiteTensor* output1 = interpreter_->tensor(interpreter_->outputs()[0]);
63 CheckIntTensor(output1, {1}, {i + 1});
64 TfLiteTensor* output2 = interpreter_->tensor(interpreter_->outputs()[1]);
65 CheckIntTensor(output2, {1}, {expected[i]});
66 }
67 }
68
TEST_F(WhileTest,TestTriangularNumberSequenceWithShallowCopy)69 TEST_F(WhileTest, TestTriangularNumberSequenceWithShallowCopy) {
70 const std::vector<int> expected = {1, 3, 6, 10, 15, 21, 28};
71 for (int i = 0; i < expected.size(); ++i) {
72 interpreter_ = std::make_unique<Interpreter>();
73 AddSubgraphs(2);
74 builder_->BuildLessEqualCondSubgraph(interpreter_->subgraph(1), i);
75 builder_->BuildAccumulateLoopBodySubgraph(interpreter_->subgraph(2));
76 builder_->BuildWhileSubgraph(&interpreter_->primary_subgraph());
77
78 interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {1});
79 // Use 4MB inputs to test shallow copy.
80 interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {1000000});
81 // Apply DynamicAllocationForLargeTensors option to enable shallow copy.
82 InterpreterOptions options;
83 options.OptimizeMemoryForLargeTensors(1000000);
84 ASSERT_EQ(interpreter_->ApplyOptions(&options), kTfLiteOk);
85 const size_t initial_mem_usage =
86 profiling::memory::GetMemoryUsage().max_rss_kb;
87 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
88 // Memory usage shouldn't exceed 9MB (2 x inputs + margin).
89 ASSERT_LE(
90 profiling::memory::GetMemoryUsage().max_rss_kb - initial_mem_usage,
91 9000);
92 FillIntTensor(interpreter_->tensor(interpreter_->inputs()[0]), {1});
93 const std::vector<int> input_vector(1000000, 1);
94 FillIntTensor(interpreter_->tensor(interpreter_->inputs()[1]),
95 input_vector);
96 auto body_subgraph = interpreter_->subgraph(2);
97
98 ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
99
100 // While BODY inputs are dynamic tensors with shallow copy.
101 TfLiteTensor* subgraph_input2 =
102 body_subgraph->tensor(body_subgraph->inputs()[1]);
103 ASSERT_EQ(subgraph_input2->allocation_type, kTfLiteDynamic);
104
105 TfLiteTensor* output1 = interpreter_->tensor(interpreter_->outputs()[0]);
106 CheckIntTensor(output1, {1}, {i + 1});
107 TfLiteTensor* output2 = interpreter_->tensor(interpreter_->outputs()[1]);
108 const std::vector<int> expected2(1000000, expected[i]);
109 CheckIntTensor(output2, {1000000}, expected2);
110 }
111 }
112
TEST_F(WhileTest,TestPadLoop)113 TEST_F(WhileTest, TestPadLoop) {
114 interpreter_ = std::make_unique<Interpreter>();
115 AddSubgraphs(2);
116 builder_->BuildLessEqualCondSubgraph(interpreter_->subgraph(1), 3);
117 builder_->BuildPadLoopBodySubgraph(interpreter_->subgraph(2), {1, 2});
118 builder_->BuildWhileSubgraph(&interpreter_->primary_subgraph());
119
120 interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {1});
121 interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {2});
122 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
123
124 FillIntTensor(interpreter_->tensor(interpreter_->inputs()[0]), {1});
125 FillIntTensor(interpreter_->tensor(interpreter_->inputs()[1]), {5, 7});
126
127 ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
128 TfLiteTensor* output1 = interpreter_->tensor(interpreter_->outputs()[0]);
129 CheckIntTensor(output1, {1}, {4});
130 TfLiteTensor* output2 = interpreter_->tensor(interpreter_->outputs()[1]);
131 CheckIntTensor(output2, {11}, {0, 0, 0, 5, 7, 0, 0, 0, 0, 0, 0});
132
133 // The extra invocation serves as a regression test: There was a bug that
134 // invoking a while loop with dynamic shaped body makes the interpreter
135 // state uninvokable.
136 ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
137 }
138
TEST_F(WhileTest,TestPadLoopWithShallowCopy)139 TEST_F(WhileTest, TestPadLoopWithShallowCopy) {
140 interpreter_ = std::make_unique<Interpreter>();
141 AddSubgraphs(2);
142 builder_->BuildLessEqualCondSubgraph(interpreter_->subgraph(1), 3);
143 builder_->BuildPadLoopBodySubgraph(interpreter_->subgraph(2), {1, 2});
144 builder_->BuildWhileSubgraph(&interpreter_->primary_subgraph());
145
146 interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {1});
147 // Use 4MB inputs to test shallow copy.
148 interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {1000000});
149 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
150
151 FillIntTensor(interpreter_->tensor(interpreter_->inputs()[0]), {1});
152 std::vector<int> input_vector(1000000, 0);
153 input_vector[0] = 5;
154 input_vector[1] = 7;
155 FillIntTensor(interpreter_->tensor(interpreter_->inputs()[1]), input_vector);
156
157 ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
158 TfLiteTensor* output1 = interpreter_->tensor(interpreter_->outputs()[0]);
159 CheckIntTensor(output1, {1}, {4});
160 TfLiteTensor* output2 = interpreter_->tensor(interpreter_->outputs()[1]);
161 std::vector<int> output_vector(1000009, 0);
162 output_vector[3] = 5;
163 output_vector[4] = 7;
164 CheckIntTensor(output2, {1000009}, output_vector);
165
166 // The extra invocation serves as a regression test: There was a bug that
167 // invoking a while loop with dynamic shaped body makes the interpreter
168 // state uninvokable.
169 ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
170 }
171
TEST_F(WhileTest,TestWhileLoopWithDynamicTensor)172 TEST_F(WhileTest, TestWhileLoopWithDynamicTensor) {
173 interpreter_ = std::make_unique<Interpreter>();
174 AddSubgraphs(2);
175 builder_->BuildLessEqualCondSubgraphWithDynamicTensor(
176 interpreter_->subgraph(1), 3);
177 builder_->BuildBodySubgraphWithDynamicTensor(interpreter_->subgraph(2));
178 builder_->BuildWhileSubgraphWithDynamicTensor(
179 &interpreter_->primary_subgraph());
180
181 interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {});
182 interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {});
183 interpreter_->ResizeInputTensor(interpreter_->inputs()[2], {1});
184 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
185
186 FillScalarStringTensor(interpreter_->tensor(interpreter_->inputs()[0]), "A");
187 FillScalarStringTensor(interpreter_->tensor(interpreter_->inputs()[1]), "A");
188 FillIntTensor(interpreter_->tensor(interpreter_->inputs()[2]), {1});
189
190 ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
191 TfLiteTensor* string_output1 =
192 interpreter_->tensor(interpreter_->outputs()[0]);
193 CheckScalarStringTensor(string_output1, "A");
194 TfLiteTensor* string_output2 =
195 interpreter_->tensor(interpreter_->outputs()[1]);
196 CheckStringTensor(string_output2, {4}, {"A", "A", "A", "A"});
197 TfLiteTensor* integer_output =
198 interpreter_->tensor(interpreter_->outputs()[2]);
199 CheckIntTensor(integer_output, {1}, {4});
200
201 // The extra invocation serves as a regression test: There was a bug that
202 // invoking a while loop with dynamic shaped body makes the interpreter
203 // state uninvokable.
204 ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
205 }
206
207 } // namespace
208 } // namespace tflite
209