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 16 #ifndef TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_GREEDY_MEMORY_PLANNER_H_ 17 #define TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_GREEDY_MEMORY_PLANNER_H_ 18 19 #include "tensorflow/lite/micro/compatibility.h" 20 #include "tensorflow/lite/micro/memory_planner/memory_planner.h" 21 22 namespace tflite { 23 24 constexpr int kOnlinePlannedBuffer = -1; 25 26 // A memory planner that uses a greedy algorithm to arrange buffers in memory 27 // to minimize the overall arena size needed. 28 // 29 // The algorithm works like this: 30 // - The client enters the buffer information through AddBuffer(). 31 // - When a function like GetOffsetForBuffer() is called, the 32 // CalculateOffsetsIfNeeded() method is invoked. 33 // - If an up to date plan is not already present, one will be calculated. 34 // - The buffers are sorted in descending order of size. 35 // - The largest buffer is placed at offset zero. 36 // - The rest of the buffers are looped through in descending size order. 37 // - The other buffers that need to be in memory at the same time are found. 38 // - The first gap between simultaneously active buffers that the current 39 // buffer fits into will be used. 40 // - If no large-enough gap is found, the current buffer is placed after the 41 // last buffer that's simultaneously active. 42 // - This continues until all buffers are placed, and the offsets stored. 43 // 44 // This is not guaranteed to produce the best placement, since that's an 45 // NP-Complete problem, but in practice it should produce one that's decent. 46 class GreedyMemoryPlanner : public MemoryPlanner { 47 public: 48 // You need to pass in an area of memory to be used for planning. This memory 49 // needs to have a lifetime as long as the planner, but isn't owned by this 50 // object, so management should be handled by the client. This is so it can be 51 // stack or globally allocated if necessary on devices without dynamic memory 52 // allocation. How many buffers can be planned for will depend on the size of 53 // this scratch memory, so you should enlarge it if you see an error when 54 // calling AddBuffer(). The memory can be reused once you're done with the 55 // planner, as long as you copy the calculated offsets to another location. 56 // Each buffer requires about 36 bytes of scratch. 57 GreedyMemoryPlanner(unsigned char* scratch_buffer, int scratch_buffer_size); 58 ~GreedyMemoryPlanner() override; 59 60 // Record details of a buffer we want to place. 61 TfLiteStatus AddBuffer(ErrorReporter* error_reporter, int size, 62 int first_time_used, int last_time_used) override; 63 64 // Record details of an offline planned buffer offset we want to place. 65 // offline_offset is the buffer offset from the start of the arena. 66 TfLiteStatus AddBuffer(ErrorReporter* error_reporter, int size, 67 int first_time_used, int last_time_used, 68 int offline_offset); 69 70 // Returns the high-water mark of used memory. This is the minimum size of a 71 // memory arena you'd need to allocate to hold these buffers. 72 size_t GetMaximumMemorySize() override; 73 74 // How many buffers have been recorded. 75 int GetBufferCount() override; 76 77 // Where a given buffer should be placed in the memory arena. 78 // This information is stored in the memory arena itself, so once the arena 79 // is used for inference, it will be overwritten. 80 TfLiteStatus GetOffsetForBuffer(ErrorReporter* error_reporter, 81 int buffer_index, int* offset) override; 82 83 // Prints an ascii-art diagram of the buffer layout plan. 84 void PrintMemoryPlan(); 85 86 // Debug method to check whether any buffer allocations are overlapping. This 87 // is an O(N^2) complexity operation, so only use for testing. 88 bool DoAnyBuffersOverlap(ErrorReporter* error_reporter); 89 90 // Used to store a list of buffers ordered by their offset. 91 struct ListEntry { 92 int offset; 93 int requirements_index; 94 int next_entry_index; 95 }; 96 97 // Number of bytes required in order to plan a buffer. per_buffer_size()98 static size_t per_buffer_size() { 99 const int per_buffer_size = 100 sizeof(BufferRequirements) + // requirements_ 101 sizeof(int) + // buffer_sizes_sorted_ 102 sizeof(int) + // buffer_ids_sorted_ 103 sizeof(ListEntry) + // buffers_sorted_by_offset_ 104 sizeof(int); // buffer_offsets_; 105 return per_buffer_size; 106 } 107 108 private: 109 // Whether a buffer is active in a given time range. 110 bool DoesEntryOverlapInTime(const ListEntry* entry, const int first_time_used, 111 const int last_time_used) const; 112 113 // Walks the list to return the next buffer that is active in a given time 114 // range, or a null pointer if there are none. 115 ListEntry* NextSimultaneouslyActiveBuffer(const ListEntry* start, 116 const int first_time_used, 117 const int last_time_used); 118 119 // If there isn't an up to date plan, calculate a new one. 120 void CalculateOffsetsIfNeeded(); 121 122 // How many buffers we can plan for, based on the arena size we're given in 123 // the constructor. 124 int max_buffer_count_; 125 126 // The number of buffers added so far. 127 int buffer_count_; 128 129 // Records the client-provided information about each buffer. 130 struct BufferRequirements { 131 int size; 132 int offline_offset; 133 int first_time_used; 134 int last_time_used; 135 }; 136 137 // Working arrays used during the layout algorithm. 138 BufferRequirements* requirements_; 139 // buffer_sizes_sorted_ and buffer_ids_sorted_ are sorted according to: 140 // { 141 // offline planned buffers, 142 // online planned buffers sorted by size 143 // } 144 int* buffer_sizes_sorted_; 145 int* buffer_ids_sorted_; 146 ListEntry* buffers_sorted_by_offset_; 147 int next_free_entry_; // Index of the next free entry of 148 // buffers_sorted_by_offset_ 149 int first_entry_index_; // Index of the first entry (smallest offset) of 150 // buffers_sorted_by_offset_ 151 152 // Stores the outcome of the plan, the location of each buffer in the arena. 153 int* buffer_offsets_; 154 155 // Whether buffers have been added since the last plan was calculated. 156 bool need_to_calculate_offsets_; 157 158 TF_LITE_REMOVE_VIRTUAL_DELETE 159 }; 160 161 } // namespace tflite 162 163 #endif // TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_GREEDY_MEMORY_PLANNER_H_ 164
