| 33 | namespace tflite { |
| 34 | |
| 35 | TfLiteStatus SimpleMemoryArena::Allocate(TfLiteContext* context, |
| 36 | size_t alignment, size_t size, |
| 37 | ArenaAlloc* new_alloc) { |
| 38 | TF_LITE_ENSURE(context, alignment <= arena_alignment_); |
| 39 | |
| 40 | if (size == 0) { |
| 41 | new_alloc->offset = 0; |
| 42 | new_alloc->size = 0; |
| 43 | return kTfLiteOk; |
| 44 | } |
| 45 | |
| 46 | size_t current_top = 0; |
| 47 | |
| 48 | if (!allocs_.empty()) { |
| 49 | auto last = allocs_.rbegin(); |
| 50 | current_top = last->offset + last->size; |
| 51 | } |
| 52 | |
| 53 | // If we don't find a better gap just allocate at the end of the buffer. |
| 54 | size_t best_offset = AlignTo(alignment, current_top); |
| 55 | size_t best_offset_fit = std::numeric_limits<size_t>::max(); |
| 56 | auto best_insertion_it = allocs_.end(); |
| 57 | |
| 58 | // Go through the sorted allocs and look at the gaps between them. |
| 59 | size_t current_offset = 0; |
| 60 | for (auto it = allocs_.begin(); it != allocs_.end(); ++it) { |
| 61 | size_t aligned_current_offset = AlignTo(alignment, current_offset); |
| 62 | // If we found a gap larger than required size, and smaller than previous |
| 63 | // best fit, take it. |
| 64 | if (aligned_current_offset + size <= it->offset && |
| 65 | it->offset - current_offset < best_offset_fit) { |
| 66 | best_offset = aligned_current_offset; |
| 67 | best_offset_fit = it->offset - current_offset; |
| 68 | best_insertion_it = it; |
| 69 | } |
| 70 | current_offset = it->offset + it->size; |
| 71 | } |
| 72 | |
| 73 | // Update the required buffer size. |
| 74 | high_water_mark_ = std::max(high_water_mark_, best_offset + size); |
| 75 | |
| 76 | new_alloc->offset = best_offset; |
| 77 | new_alloc->size = size; |
| 78 | allocs_.insert(best_insertion_it, *new_alloc); |
| 79 | |
| 80 | return kTfLiteOk; |
| 81 | } |
| 82 | |
| 83 | TfLiteStatus SimpleMemoryArena::Deallocate(TfLiteContext* context, |
| 84 | const ArenaAlloc& alloc) { |