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 #ifndef TENSORFLOW_LITE_ARENA_PLANNER_H_ 16 #define TENSORFLOW_LITE_ARENA_PLANNER_H_ 17 18 #include <cstdint> 19 #include <memory> 20 #include <vector> 21 22 #include "tensorflow/lite/c/common.h" 23 #include "tensorflow/lite/graph_info.h" 24 #include "tensorflow/lite/memory_planner.h" 25 #include "tensorflow/lite/simple_memory_arena.h" 26 #include "tensorflow/lite/util.h" 27 28 namespace tflite { 29 30 constexpr const int kDefaultArenaAlignment = 64; 31 struct AllocationInfo; 32 33 // A memory planner that makes all the allocations using arenas. 34 // 35 // Before a model is executed by the interpreter, this class determines when 36 // each tensor needs to be allocated and deallocated, and preallocates all the 37 // necessary memory (the PlanAllocations phase). It then assigns portions of 38 // this memory buffer to each tensor (the ExecuteAllocations phase). Tensors may 39 // share some of the buffer if a tensor B is to be allocated after another 40 // tensor A has been deallocated. 41 // 42 // If dynamic tensors are used the planning steps can be repeated during model 43 // execution. Since dynamic tensors don't have sizes until after the 44 // corresponding operation is executed, this class supports incremental 45 // planning. 46 class ArenaPlanner : public MemoryPlanner { 47 public: 48 // Ownership of 'context' is not taken and it must remain util the 49 // ArenaPlanner is destroyed. The inputs to the graph will not share 50 // memory with any other tensor, effectively preserving them until the end 51 // of inference. 52 ArenaPlanner(TfLiteContext* context, std::unique_ptr<GraphInfo> graph_info, 53 bool preserve_all_tensors, int tensor_alignment, 54 int subgraph_index = 0); 55 ~ArenaPlanner() override; 56 ArenaPlanner(const ArenaPlanner&) = delete; 57 ArenaPlanner& operator=(const ArenaPlanner&) = delete; 58 59 TfLiteStatus ResetAllocations() override; 60 TfLiteStatus ResetAllocationsAfter(int node) override; 61 TfLiteStatus PlanAllocations() override; 62 TfLiteStatus ExecuteAllocations(int first_node, int last_node) override; 63 TfLiteStatus ReleaseNonPersistentMemory() override; 64 TfLiteStatus AcquireNonPersistentMemory() override; 65 bool HasNonPersistentMemory() override; 66 void DumpDebugInfo(const std::vector<int>& execution_plan) const override; 67 void GetAllocInfo(size_t* arena_size, 68 size_t* arena_persist_size) const override; 69 70 // Returns the base arena location for a given allocation type. 71 std::intptr_t BasePointer(TfLiteAllocationType type); 72 73 private: 74 // Make sure all the arenas have reserved enough memory to store all their 75 // tensors. 76 TfLiteStatus Commit(); 77 78 // Returns vector of tensor number ordered by the following algorithm. 79 // Comparator to sort tensors for the allocation algorithm: 80 // - Tensors that have lifespan through the whole model inference time go 81 // first; 82 // - Other tensors (e.g. intermediate and temporary ones) are sorted in 83 // non-increasing order of their size. If sizes of two tensors are equal, the 84 // one that needs to be allocated earlier goes first. 85 std::vector<int32_t> CreateTensorAllocationVector(int first_node, 86 int last_node); 87 88 // Traverse the allocation queue and reserve space in the appropriate arena 89 // for all tensors affected by ops in the interval [first_node, last_node]. 90 TfLiteStatus CalculateAllocations(int first_node, int last_node); 91 92 // Assign absolute memory location to a tensor, based on its relative 93 // position inside the corresponding arena buffer. 94 TfLiteStatus ResolveTensorAllocation(int tensor_index); 95 96 // Register an allocation for all internal (temporary) tensors of 97 // 'node_index'. 98 TfLiteStatus CalculateAllocationOfInternalTensors(int node_index); 99 100 // Register a deallocation for all internal (temporary) tensors of 101 // 'node_index'. 102 TfLiteStatus CalculateDeallocationOfInternalTensors(int node_index); 103 104 TfLiteContext* context_; 105 std::unique_ptr<GraphInfo> graph_info_; 106 107 // Stores allocation data for all tensors. 108 std::vector<ArenaAllocWithUsageInterval> allocs_; 109 110 // First node, that uses the tensor. It needs to be allocated before 111 // execution of the node's operation. 112 std::vector<int32_t> alloc_node_; 113 114 // Last node, that uses the tensor. It can be deallocated after execution of 115 // the node's operation. 116 std::vector<int32_t> dealloc_node_; 117 118 // Raw memory buffer that is allocated for all temporary and graph outputs 119 // that are declared kTfLiteArenaRw. 120 SimpleMemoryArena arena_; 121 122 // Raw memory buffer that is allocated for persistent tensors that are 123 // declared as kTfLiteArenaRwPersistent. 124 SimpleMemoryArena persistent_arena_; 125 126 // If true, then no overlapping of memory areas is done, meaning intermediate 127 // tensors and temporary tensors can be queried after running. 128 // (modulo running delegates) 129 bool preserve_all_tensors_; 130 131 // Number of bytes that tensor buffers should be aligned to. 132 int tensor_alignment_; 133 }; 134 135 } // namespace tflite 136 137 #endif // TENSORFLOW_LITE_ARENA_PLANNER_H_ 138