1 /* Copyright 2017 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/contrib/lite/arena_planner.h" 16 17 #include <cstdarg> 18 19 #include <gmock/gmock.h> 20 #include <gtest/gtest.h> 21 #include "tensorflow/contrib/lite/testing/util.h" 22 #include "tensorflow/core/platform/logging.h" 23 24 namespace tflite { 25 namespace { 26 27 // A simple op to be used in tests, as syntactic sugar. 28 class TestOp { 29 public: 30 TestOp(std::initializer_list<int> inputs, std::initializer_list<int> outputs, 31 std::initializer_list<int> temporaries) 32 : inputs_(inputs), outputs_(outputs), temporaries_(temporaries) {} 33 34 const std::vector<int>& inputs() const { return inputs_; } 35 const std::vector<int>& outputs() const { return outputs_; } 36 const std::vector<int>& temporaries() const { return temporaries_; } 37 38 private: 39 std::vector<int> inputs_; 40 std::vector<int> outputs_; 41 std::vector<int> temporaries_; 42 }; 43 44 // A test graph where inputs are processed by the given nodes to produce 45 // outputs. 46 class TestGraph { 47 public: 48 TestGraph(std::initializer_list<int> inputs, 49 std::initializer_list<TestOp> nodes, 50 std::initializer_list<int> outputs) 51 : inputs_(inputs), outputs_(outputs) { 52 int max_tensor_index = 0; 53 54 for (int t : inputs) { 55 max_tensor_index = std::max(max_tensor_index, t); 56 } 57 for (int t : outputs) { 58 max_tensor_index = std::max(max_tensor_index, t); 59 } 60 for (const auto& node : nodes) { 61 auto int_array = [](const std::vector<int>& x) { 62 TfLiteIntArray* lite = TfLiteIntArrayCreate(x.size()); 63 for (size_t i = 0; i < x.size(); i++) lite->data[i] = x[i]; 64 return lite; 65 }; 66 67 nodes_.push_back(TfLiteNode()); 68 nodes_.back().inputs = int_array(node.inputs()); 69 for (int t : node.inputs()) { 70 max_tensor_index = std::max(max_tensor_index, t); 71 } 72 nodes_.back().outputs = int_array(node.outputs()); 73 for (int t : node.outputs()) { 74 max_tensor_index = std::max(max_tensor_index, t); 75 } 76 nodes_.back().temporaries = int_array(node.temporaries()); 77 for (int t : node.temporaries()) { 78 max_tensor_index = std::max(max_tensor_index, t); 79 } 80 } 81 82 for (int i = 0; i <= max_tensor_index; ++i) { 83 tensors_.push_back(TfLiteTensor()); 84 // Set some default values for allocation_type and bytes, which are the 85 // only fields used by the arena planner. 86 tensors_.back().allocation_type = kTfLiteArenaRw; 87 tensors_.back().bytes = (i + 1) * 3; 88 } 89 } 90 91 ~TestGraph() { 92 for (auto node : nodes_) { 93 TfLiteIntArrayFree(node.inputs); 94 TfLiteIntArrayFree(node.outputs); 95 TfLiteIntArrayFree(node.temporaries); 96 } 97 } 98 99 const std::vector<TfLiteNode>& nodes() { return nodes_; } 100 std::vector<TfLiteTensor>* tensors() { return &tensors_; } 101 const std::vector<int>& inputs() { return inputs_; } 102 const std::vector<int>& outputs() { return outputs_; } 103 104 private: 105 std::vector<TfLiteNode> nodes_; 106 std::vector<TfLiteTensor> tensors_; 107 std::vector<int> inputs_; 108 std::vector<int> outputs_; 109 }; 110 111 // The GraphInfo for a TestGraph. 112 class TestGraphInfo : public GraphInfo { 113 public: 114 explicit TestGraphInfo(TestGraph* graph) : graph_(graph) {} 115 116 size_t num_tensors() const override { return graph_->tensors()->size(); } 117 TfLiteTensor* tensor(size_t index) override { 118 return &graph_->tensors()->at(index); 119 } 120 size_t num_nodes() const override { return graph_->nodes().size(); } 121 const TfLiteNode& node(size_t index) const override { 122 return graph_->nodes()[index]; 123 } 124 const std::vector<int>& inputs() const override { return graph_->inputs(); } 125 const std::vector<int>& outputs() const override { return graph_->outputs(); } 126 127 private: 128 TestGraph* graph_; 129 }; 130 131 void ReportError(TfLiteContext* context, const char* format, ...) { 132 const size_t kBufferSize = 1024; 133 char temp_buffer[kBufferSize]; 134 135 va_list args; 136 va_start(args, format); 137 vsnprintf(temp_buffer, kBufferSize, format, args); 138 va_end(args); 139 140 LOG(INFO) << temp_buffer; 141 } 142 143 class ArenaPlannerTest : public ::testing::Test { 144 protected: 145 void SetGraph(TestGraph* graph) { 146 graph_ = graph; 147 context_.ReportError = ReportError; 148 planner_.reset(new ArenaPlanner( 149 &context_, std::unique_ptr<GraphInfo>(new TestGraphInfo(graph)))); 150 CHECK(planner_->ResetAllocations() == kTfLiteOk); 151 CHECK(planner_->PlanAllocations() == kTfLiteOk); 152 } 153 154 void Execute(int start, int end) { 155 CHECK(planner_->ExecuteAllocations(start, end) == kTfLiteOk); 156 } 157 158 // Returns the actual offset of a given tensor, relative to the start of its 159 // arena. 160 int64_t GetOffset(int tensor_index) { 161 const TfLiteTensor& tensor = (*graph_->tensors())[tensor_index]; 162 return reinterpret_cast<int64_t>(tensor.data.raw) - 163 planner_->BasePointer(tensor.allocation_type); 164 } 165 166 // Returns the first aligned offset after a given tensor. 167 int64_t GetOffsetAfter(int tensor_index) { 168 const TfLiteTensor& tensor = (*graph_->tensors())[tensor_index]; 169 int64_t offset = GetOffset(tensor_index) + tensor.bytes; 170 // We must make sure the offset is aligned to kDefaultArenaAlignment. 171 if (offset % 4 != 0) { 172 offset += 4 - offset % 4; 173 } 174 return offset; 175 }; 176 177 TfLiteContext context_; 178 TestGraph* graph_; 179 std::unique_ptr<ArenaPlanner> planner_; 180 }; 181 182 TEST_F(ArenaPlannerTest, EmptyGraph) { 183 TestGraph graph({}, {}, {}); 184 SetGraph(&graph); 185 Execute(0, 10); 186 } 187 188 TEST_F(ArenaPlannerTest, GraphWithNoOps) { 189 TestGraph graph({0, 10}, {}, {5, 11}); 190 SetGraph(&graph); 191 Execute(0, 10); 192 EXPECT_EQ(GetOffset(0), 0); 193 EXPECT_EQ(GetOffset(10), GetOffsetAfter(0)); 194 // The outputs are never allocated because they are not connected to any 195 // inputs. 196 EXPECT_TRUE((*graph.tensors())[5].data.raw == nullptr); 197 EXPECT_TRUE((*graph.tensors())[11].data.raw == nullptr); 198 } 199 200 TEST_F(ArenaPlannerTest, GraphWithOneOp) { 201 TestGraph graph({1}, {{{1}, {2}, {}}}, {2}); 202 SetGraph(&graph); 203 Execute(0, 10); 204 EXPECT_EQ(GetOffset(1), 0); 205 EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); 206 } 207 208 TEST_F(ArenaPlannerTest, ZeroSizedTensors) { 209 TestGraph graph({1}, {{{1}, {2}, {}}}, {2}); 210 (*graph.tensors())[1].bytes = 0; 211 SetGraph(&graph); 212 // TODO(ahentz): this is currently broken because the arena finds two 213 // allocations with the same offset and returns an error. 214 ASSERT_FALSE(planner_->ExecuteAllocations(0, 10) == kTfLiteOk); 215 // EXPECT_EQ(GetOffset(1), 0); 216 // EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); 217 } 218 219 TEST_F(ArenaPlannerTest, SimpleGraph) { 220 TestGraph graph({0, 1}, 221 { 222 /* in, out, tmp */ 223 {{0, 1}, {2}, {}}, // First op 224 {{2, 0}, {4, 5}, {}}, // Second op 225 {{4, 5}, {3}, {}} // Third op 226 }, 227 {3}); 228 SetGraph(&graph); 229 Execute(0, 10); 230 231 // Alloc(+) and dealloc(-) order: +0 +1 +2 -1 +4 +5 -2 -0 +3 -4 -5 232 EXPECT_EQ(GetOffset(0), 0); 233 EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); 234 EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); 235 EXPECT_EQ(GetOffset(4), GetOffsetAfter(2)); 236 EXPECT_EQ(GetOffset(5), GetOffsetAfter(4)); 237 EXPECT_EQ(GetOffset(3), 0); 238 } 239 240 TEST_F(ArenaPlannerTest, SimpleGraphWithTemporary) { 241 TestGraph graph({0, 1}, 242 { 243 /* in, out, tmp */ 244 {{0, 1}, {2}, {}}, // First op 245 {{2, 0}, {4}, {5}}, // Second op, with temporary 246 {{4}, {3}, {}} // Third op 247 }, 248 {3}); 249 SetGraph(&graph); 250 Execute(0, 10); 251 252 // Alloc(+) and dealloc(-) order: +0 +1 +2 -1 +5 +4 -2 -0 -5 +3 -4 253 EXPECT_EQ(GetOffset(0), 0); 254 EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); 255 EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); 256 EXPECT_EQ(GetOffset(5), GetOffsetAfter(2)); 257 EXPECT_EQ(GetOffset(4), GetOffsetAfter(5)); 258 EXPECT_EQ(GetOffset(3), 0); 259 } 260 261 TEST_F(ArenaPlannerTest, SimpleGraphWithOptionals) { 262 TestGraph graph({0, -1, 1}, 263 { 264 /* in, out, tmp */ 265 {{0, 1}, {2}, {}}, // First op 266 {{2, 0}, {4, 5}, {}}, // Second op 267 {{4, -1, 5}, {3}, {}} // Third op, with optional 268 }, 269 {3}); 270 SetGraph(&graph); 271 Execute(0, 10); 272 273 // Alloc(+) and dealloc(-) order: +0 +1 +2 -1 +4 +5 -2 -0 +3 -4 -5 274 EXPECT_EQ(GetOffset(0), 0); 275 EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); 276 EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); 277 EXPECT_EQ(GetOffset(4), GetOffsetAfter(2)); 278 EXPECT_EQ(GetOffset(5), GetOffsetAfter(4)); 279 EXPECT_EQ(GetOffset(3), 0); 280 } 281 282 TEST_F(ArenaPlannerTest, SimpleGraphWithLargeTensor) { 283 TestGraph graph({0, -1, 1}, 284 { 285 /* in, out, tmp */ 286 {{0, 1}, {2}, {}}, // First op 287 {{2, 0}, {4}, {5}}, // Second op, with temporary 288 {{4, -1}, {3}, {}} // Third op, with optional 289 }, 290 {3}); 291 292 // Make #1 very large so its vacancy can be filled with #5 and #4. 293 (*graph.tensors())[1].bytes = 40; 294 295 SetGraph(&graph); 296 Execute(0, 10); 297 298 // Alloc(+) and dealloc(-) order: +0 +1 +2 -1 +5 +4 -2 -0 -5 +3 -4 299 EXPECT_EQ(GetOffset(0), 0); 300 EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); 301 EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); 302 EXPECT_EQ(GetOffset(5), GetOffsetAfter(0)); 303 EXPECT_EQ(GetOffset(4), GetOffsetAfter(5)); 304 EXPECT_EQ(GetOffset(3), 0); 305 } 306 307 TEST_F(ArenaPlannerTest, SimpleGraphWithPersistentTensor) { 308 TestGraph graph({0, -1, 1}, 309 { 310 /* in, out, tmp */ 311 {{0, 1}, {2}, {}}, // First op 312 {{2, 0}, {4}, {5}}, // Second op, with temporary 313 {{4, -1}, {3}, {}} // Third op, with optional 314 }, 315 {3}); 316 317 // Make #1 persistent so it goes into its own arena. 318 (*graph.tensors())[1].allocation_type = kTfLiteArenaRwPersistent; 319 320 SetGraph(&graph); 321 Execute(0, 10); 322 323 // Make sure #0 and #1 were given different memory locations (because they 324 // will both have offset=0, in different arenas.) 325 EXPECT_NE((*graph.tensors())[0].data.raw, (*graph.tensors())[1].data.raw); 326 327 // Alloc(+) and dealloc(-) order: +0 +1 +2 -1 +5 +4 -2 -0 -5 +3 -4 328 EXPECT_EQ(GetOffset(0), 0); 329 EXPECT_EQ(GetOffset(1), 0); 330 EXPECT_EQ(GetOffset(2), GetOffsetAfter(0)); 331 EXPECT_EQ(GetOffset(5), GetOffsetAfter(2)); 332 EXPECT_EQ(GetOffset(4), GetOffsetAfter(5)); 333 EXPECT_EQ(GetOffset(3), 0); 334 } 335 336 TEST_F(ArenaPlannerTest, SimpleGraphWithDynamicTensor) { 337 TestGraph graph({0, -1, 1}, 338 { 339 /* in, out, tmp */ 340 {{0, 1}, {2}, {}}, // First op 341 {{2, 0}, {4}, {5}}, // Second op, with temporary 342 {{4, -1}, {3}, {}} // Third op, with optional 343 }, 344 {3}); 345 346 // Make #1 dynaic so it does not get allocated. 347 (*graph.tensors())[1].allocation_type = kTfLiteDynamic; 348 349 SetGraph(&graph); 350 Execute(0, 10); 351 352 EXPECT_EQ((*graph.tensors())[1].data.raw, nullptr); 353 354 // Alloc(+) and dealloc(-) order: +0 +1 +2 -1 +5 +4 -2 -0 -5 +3 -4 355 EXPECT_EQ(GetOffset(0), 0); 356 EXPECT_EQ(GetOffset(2), GetOffsetAfter(0)); 357 EXPECT_EQ(GetOffset(5), GetOffsetAfter(2)); 358 EXPECT_EQ(GetOffset(4), GetOffsetAfter(5)); 359 EXPECT_EQ(GetOffset(3), 0); 360 } 361 362 TEST_F(ArenaPlannerTest, LargerGraphAndStepwiseAllocation) { 363 TestGraph graph({0, 1}, 364 { 365 /* in, out, tmp */ 366 {{0, 1}, {2, 3}, {}}, 367 {{2, 0}, {4, 5}, {6}}, 368 {{1, -1}, {7}, {}}, 369 {{7, 3}, {8}, {9}}, 370 {{4, 5, 8}, {10}, {}}, 371 }, 372 {10}); 373 SetGraph(&graph); 374 375 auto is_unallocated = [&](int tensor_index) { 376 return (*graph.tensors())[tensor_index].data.raw == nullptr; 377 }; 378 379 // The allocation plan is made at the beginning and is independent of 380 // the execution steps. Here's the allocation order: 381 // Op0: +0 +1 +2 +3 382 // Op1: +6 +4 +5 -6 -0 -2 383 // Op2: +7 -1 384 // Op3: +9 +8 -9 -3 -7 385 // Op4: +10 -4 -5 -8 386 387 Execute(0, 0); 388 EXPECT_EQ(GetOffset(0), 0); 389 EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); 390 EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); 391 EXPECT_EQ(GetOffset(3), GetOffsetAfter(2)); 392 EXPECT_TRUE(is_unallocated(6)); 393 EXPECT_TRUE(is_unallocated(4)); 394 EXPECT_TRUE(is_unallocated(5)); 395 EXPECT_TRUE(is_unallocated(7)); 396 EXPECT_TRUE(is_unallocated(9)); 397 EXPECT_TRUE(is_unallocated(8)); 398 EXPECT_TRUE(is_unallocated(10)); 399 400 Execute(1, 1); 401 EXPECT_EQ(GetOffset(0), 0); 402 EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); 403 EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); 404 EXPECT_EQ(GetOffset(3), GetOffsetAfter(2)); 405 EXPECT_EQ(GetOffset(6), GetOffsetAfter(3)); 406 EXPECT_EQ(GetOffset(4), GetOffsetAfter(6)); 407 EXPECT_EQ(GetOffset(5), GetOffsetAfter(4)); 408 EXPECT_TRUE(is_unallocated(7)); 409 EXPECT_TRUE(is_unallocated(9)); 410 EXPECT_TRUE(is_unallocated(8)); 411 EXPECT_TRUE(is_unallocated(10)); 412 413 Execute(2, 2); 414 EXPECT_EQ(GetOffset(0), 0); 415 EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); 416 EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); 417 EXPECT_EQ(GetOffset(3), GetOffsetAfter(2)); 418 EXPECT_EQ(GetOffset(6), GetOffsetAfter(3)); 419 EXPECT_EQ(GetOffset(4), GetOffsetAfter(6)); 420 EXPECT_EQ(GetOffset(5), GetOffsetAfter(4)); 421 // Here's an interesting allocation. Even though #6 requires only 21 bytes, 422 // its deallocation freed up 24 bytes due to the alignment requirements in 423 // the arena. That means we can fit #7 in the same space! 424 EXPECT_EQ(GetOffset(7), GetOffsetAfter(3)); 425 EXPECT_TRUE(is_unallocated(9)); 426 EXPECT_TRUE(is_unallocated(8)); 427 EXPECT_TRUE(is_unallocated(10)); 428 429 Execute(3, 3); 430 EXPECT_EQ(GetOffset(0), 0); 431 EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); 432 EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); 433 EXPECT_EQ(GetOffset(3), GetOffsetAfter(2)); 434 EXPECT_EQ(GetOffset(6), GetOffsetAfter(3)); 435 EXPECT_EQ(GetOffset(4), GetOffsetAfter(6)); 436 EXPECT_EQ(GetOffset(5), GetOffsetAfter(4)); 437 EXPECT_EQ(GetOffset(7), GetOffsetAfter(3)); 438 // The deallocation of #0, #1 and #2 freed up 24 bytes but that's not enough 439 // for #9, so it goes at the end. 440 EXPECT_EQ(GetOffset(9), GetOffsetAfter(5)); 441 EXPECT_EQ(GetOffset(8), GetOffsetAfter(9)); 442 EXPECT_TRUE(is_unallocated(10)); 443 444 Execute(4, 4); 445 EXPECT_EQ(GetOffset(0), 0); 446 EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); 447 EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); 448 EXPECT_EQ(GetOffset(3), GetOffsetAfter(2)); 449 EXPECT_EQ(GetOffset(6), GetOffsetAfter(3)); 450 EXPECT_EQ(GetOffset(4), GetOffsetAfter(6)); 451 EXPECT_EQ(GetOffset(5), GetOffsetAfter(4)); 452 EXPECT_EQ(GetOffset(7), GetOffsetAfter(3)); 453 EXPECT_EQ(GetOffset(9), GetOffsetAfter(5)); 454 EXPECT_EQ(GetOffset(8), GetOffsetAfter(9)); 455 // There's just enough space at the beginning for #10 due to the 456 // deallocation of #0, #1, #2 and #3 (total 36 bytes, #10 needs 457 // only 33.) 458 EXPECT_EQ(GetOffset(10), 0); 459 } 460 461 } // namespace 462 } // namespace tflite 463 464 int main(int argc, char** argv) { 465 ::tflite::LogToStderr(); 466 ::testing::InitGoogleTest(&argc, argv); 467 return RUN_ALL_TESTS(); 468 } 469
