Home | History | Annotate | Download | only in lite 
      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_CONTRIB_LITE_SIMPLE_MEMORY_ARENA_H_
     16 #define TENSORFLOW_CONTRIB_LITE_SIMPLE_MEMORY_ARENA_H_
     17 
     18 #include <list>
     19 #include <memory>
     20 #include "tensorflow/contrib/lite/context.h"
     21 
     22 namespace tflite {
     23 
     24 // This little structure holds the offset and the size for a dynamic memory
     25 // allocation in the memory arena. When the arena is commited and the
     26 // underlying buffer is set, the alloc can be resolved into an actual memory
     27 // pointer.
     28 struct ArenaAlloc {
     29   ArenaAlloc() : offset(0), size(0) {}
     30 
     31   size_t offset;
     32   size_t size;
     33 
     34   inline bool operator<(const ArenaAlloc& other) const {
     35     return offset < other.offset;
     36   }
     37 };
     38 
     39 // This small class is responsible for allocating, deallocating and reusing
     40 // dynamic memory from a common underlying buffer. The arena can be used in
     41 // scenarios when the pattern of memory allocations and deallocations is
     42 // repetitive, e.g. running NN inference in multiple iterations.
     43 class SimpleMemoryArena {
     44  public:
     45   explicit SimpleMemoryArena(size_t arena_alignment)
     46       : commited_(false),
     47         arena_alignment_(arena_alignment),
     48         high_water_mark_(0),
     49         underlying_buffer_size_(0),
     50         allocs_() {}
     51 
     52   TfLiteStatus Allocate(TfLiteContext* context, size_t alignment, size_t size,
     53                         ArenaAlloc* new_alloc);
     54 
     55   TfLiteStatus Deallocate(TfLiteContext* context, const ArenaAlloc& alloc);
     56 
     57   inline size_t RequiredBufferSize() {
     58     // Add in a small amount of padding to reduce the chance of resize events
     59     // for small allocations.
     60     size_t padding = arena_alignment_;
     61     return arena_alignment_ + high_water_mark_ + padding;
     62   }
     63 
     64   TfLiteStatus Commit(TfLiteContext* context);
     65 
     66   TfLiteStatus ResolveAlloc(TfLiteContext* context, const ArenaAlloc& alloc,
     67                             char** output_ptr);
     68 
     69   TfLiteStatus Clear();
     70 
     71   int64_t BasePointer() const {
     72     return reinterpret_cast<int64_t>(underlying_buffer_aligned_ptr_);
     73   }
     74 
     75  private:
     76   bool commited_;
     77   size_t arena_alignment_;
     78   size_t high_water_mark_;
     79   std::unique_ptr<char[]> underlying_buffer_;
     80   size_t underlying_buffer_size_;
     81   char* underlying_buffer_aligned_ptr_;
     82   // TODO(maciekc): add list iterator to the ArenaAlloc to lookup quickly.
     83   std::list<ArenaAlloc> allocs_;
     84 };
     85 
     86 }  // namespace tflite
     87 
     88 #endif  // TENSORFLOW_CONTRIB_LITE_SIMPLE_MEMORY_ARENA_H_
     89