greedy_memory_planner.h

00001 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
00002 
00003 Licensed under the Apache License, Version 2.0 (the "License");
00004 you may not use this file except in compliance with the License.
00005 You may obtain a copy of the License at
00006 
00007     http://www.apache.org/licenses/LICENSE-2.0
00008 
00009 Unless required by applicable law or agreed to in writing, software
00010 distributed under the License is distributed on an "AS IS" BASIS,
00011 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
00012 See the License for the specific language governing permissions and
00013 limitations under the License.
00014 ==============================================================================*/
00015 
00016 #ifndef TENSORFLOW_LITE_EXPERIMENTAL_MICRO_MEMORY_PLANNER_GREEDY_MEMORY_PLANNER_H_
00017 #define TENSORFLOW_LITE_EXPERIMENTAL_MICRO_MEMORY_PLANNER_GREEDY_MEMORY_PLANNER_H_
00018 
00019 #include "tensorflow/lite/experimental/micro/memory_planner/memory_planner.h"
00020 
00021 namespace tflite {
00022 
00023 // A memory planner that uses a greedy algorithm to arrange buffers in memory
00024 // to minimize the overall arena size needed.
00025 //
00026 // The algorithm works like this:
00027 //  - The client enters the buffer information through AddBuffer().
00028 //  - When a function like GetOffsetForBuffer() is called, the
00029 //    CalculateOffsetsIfNeeded() method is invoked.
00030 //  - If an up to date plan is not already present, one will be calculated.
00031 //  - The buffers are sorted in descending order of size.
00032 //  - The largest buffer is placed at offset zero.
00033 //  - The rest of the buffers are looped through in descending size order.
00034 //  - The other buffers that need to be in memory at the same time are found.
00035 //  - The first gap between simultaneously active buffers that the current
00036 //    buffer fits into will be used.
00037 //  - If no large-enough gap is found, the current buffer is placed after the
00038 //    last buffer that's simultaneously active.
00039 //  - This continues until all buffers are placed, and the offsets stored.
00040 //
00041 // This is not guaranteed to produce the best placement, since that's an
00042 // NP-Complete problem, but in practice it should produce one that's decent.
00043 class GreedyMemoryPlanner : public MemoryPlanner {
00044  public:
00045   // You need to pass in an area of memory to be used for planning. This memory
00046   // needs to have a lifetime as long as the planner, but isn't owned by this
00047   // object, so management should be handled by the client. This is so it can be
00048   // stack or globally allocated if necessary on devices without dynamic memory
00049   // allocation. How many buffers can be planned for will depend on the size of
00050   // this scratch memory, so you should enlarge it if you see an error when
00051   // calling AddBuffer(). The memory can be reused once you're done with the
00052   // planner, as long as you copy the calculated offsets to another location.
00053   // Each buffer requires about 36 bytes of scratch.
00054   GreedyMemoryPlanner(unsigned char* scratch_buffer, int scratch_buffer_size);
00055   ~GreedyMemoryPlanner() override;
00056 
00057   // Record details of a buffer we want to place.
00058   TfLiteStatus AddBuffer(ErrorReporter* error_reporter, int size,
00059                          int first_time_used, int last_time_used) override;
00060 
00061   // Returns the high-water mark of used memory. This is the minimum size of a
00062   // memory arena you'd need to allocate to hold these buffers.
00063   int GetMaximumMemorySize() override;
00064 
00065   // How many buffers have been recorded.
00066   int GetBufferCount() override;
00067 
00068   // Where a given buffer should be placed in the memory arena.
00069   // This information is stored in the memory arena itself, so once the arena
00070   // is used for inference, it will be overwritten.
00071   TfLiteStatus GetOffsetForBuffer(ErrorReporter* error_reporter,
00072                                   int buffer_index, int* offset) override;
00073 
00074   // Prints an ascii-art diagram of the buffer layout plan.
00075   void PrintMemoryPlan(ErrorReporter* error_reporter);
00076 
00077   // Debug method to check whether any buffer allocations are overlapping. This
00078   // is an O(N^2) complexity operation, so only use for testing.
00079   bool DoAnyBuffersOverlap(ErrorReporter* error_reporter);
00080 
00081   // Used to store a list of buffers ordered by their offset.
00082   struct ListEntry {
00083     int offset;
00084     int requirements_index;
00085     int next_entry_index;
00086   };
00087 
00088  private:
00089   // Whether a buffer is active in a given time range.
00090   bool DoesEntryOverlapInTime(const ListEntry* entry, const int first_time_used,
00091                               const int last_time_used) const;
00092 
00093   // Walks the list to return the next buffer that is active in a given time
00094   // range, or a null pointer if there are none.
00095   ListEntry* NextSimultaneouslyActiveBuffer(const ListEntry* start,
00096                                             const int first_time_used,
00097                                             const int last_time_used);
00098 
00099   // If there isn't an up to date plan, calculate a new one.
00100   void CalculateOffsetsIfNeeded();
00101 
00102   // How many buffers we can plan for, based on the arena size we're given in
00103   // the constructor.
00104   int max_buffer_count_;
00105 
00106   // The number of buffers added so far.
00107   int buffer_count_;
00108 
00109   // Records the client-provided information about each buffer.
00110   struct BufferRequirements {
00111     int size;
00112     int first_time_used;
00113     int last_time_used;
00114   };
00115 
00116   // Working arrays used during the layout algorithm.
00117   BufferRequirements* requirements_;
00118   int* buffer_sizes_sorted_by_size_;
00119   int* buffer_ids_sorted_by_size_;
00120   ListEntry* buffers_sorted_by_offset_;
00121   int next_free_entry_;
00122 
00123   // Stores the outcome of the plan, the location of each buffer in the arena.
00124   int* buffer_offsets_;
00125 
00126   // Whether buffers have been added since the last plan was calculated.
00127   bool need_to_calculate_offsets_;
00128 };
00129 
00130 }  // namespace tflite
00131 
00132 #endif  // TENSORFLOW_LITE_EXPERIMENTAL_MICRO_MEMORY_PLANNER_GREEDY_MEMORY_PLANNER_H_