From 2673c7c0a4a922971920ac74be1c6e4fa54c7f5e Mon Sep 17 00:00:00 2001 From: Meghna Natraj Date: Fri, 24 Jan 2020 11:10:55 -0800 Subject: [PATCH] Make GetMaximumMemorySize return size_t - Fix bug 147871342 PiperOrigin-RevId: 291405835 Change-Id: I9dfe02861b6156db9d33100bd50825be17ef639e --- .../lite/micro/memory_planner/greedy_memory_planner.cc | 8 +++++--- .../lite/micro/memory_planner/greedy_memory_planner.h | 2 +- .../lite/micro/memory_planner/linear_memory_planner.cc | 2 +- .../lite/micro/memory_planner/linear_memory_planner.h | 6 +++--- tensorflow/lite/micro/memory_planner/memory_planner.h | 2 +- tensorflow/lite/micro/micro_allocator.cc | 8 ++------ 6 files changed, 13 insertions(+), 15 deletions(-) diff --git a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc index 061ec0a7a3e..d066ba884d5 100644 --- a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc +++ b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc @@ -240,17 +240,19 @@ void GreedyMemoryPlanner::CalculateOffsetsIfNeeded() { } } -int GreedyMemoryPlanner::GetMaximumMemorySize() { +size_t GreedyMemoryPlanner::GetMaximumMemorySize() { CalculateOffsetsIfNeeded(); if (buffer_count_ == 0) { return 0; } ListEntry* entry = &buffers_sorted_by_offset_[0]; - int max_size = 0; + size_t max_size = 0; while (entry) { BufferRequirements* requirements = &requirements_[entry->requirements_index]; - const int current_size = entry->offset + requirements->size; + // TODO(b/148246793): Update all size and offset variables types from + // int to size_t + const size_t current_size = entry->offset + requirements->size; if (current_size > max_size) { max_size = current_size; } diff --git a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h index 74352fad237..f2c77ed94f3 100644 --- a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h +++ b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h @@ -61,7 +61,7 @@ class GreedyMemoryPlanner : public MemoryPlanner { // Returns the high-water mark of used memory. This is the minimum size of a // memory arena you'd need to allocate to hold these buffers. - int GetMaximumMemorySize() override; + size_t GetMaximumMemorySize() override; // How many buffers have been recorded. int GetBufferCount() override; diff --git a/tensorflow/lite/micro/memory_planner/linear_memory_planner.cc b/tensorflow/lite/micro/memory_planner/linear_memory_planner.cc index 391e7ad5458..2e2e19a76ec 100644 --- a/tensorflow/lite/micro/memory_planner/linear_memory_planner.cc +++ b/tensorflow/lite/micro/memory_planner/linear_memory_planner.cc @@ -34,7 +34,7 @@ TfLiteStatus LinearMemoryPlanner::AddBuffer( return kTfLiteOk; } -int LinearMemoryPlanner::GetMaximumMemorySize() { return next_free_offset_; } +size_t LinearMemoryPlanner::GetMaximumMemorySize() { return next_free_offset_; } int LinearMemoryPlanner::GetBufferCount() { return current_buffer_count_; } diff --git a/tensorflow/lite/micro/memory_planner/linear_memory_planner.h b/tensorflow/lite/micro/memory_planner/linear_memory_planner.h index 28a2b1bf996..4d77e778237 100644 --- a/tensorflow/lite/micro/memory_planner/linear_memory_planner.h +++ b/tensorflow/lite/micro/memory_planner/linear_memory_planner.h @@ -31,16 +31,16 @@ class LinearMemoryPlanner : public MemoryPlanner { TfLiteStatus AddBuffer(tflite::ErrorReporter* error_reporter, int size, int first_time_used, int last_time_used) override; - int GetMaximumMemorySize() override; + size_t GetMaximumMemorySize() override; int GetBufferCount() override; TfLiteStatus GetOffsetForBuffer(tflite::ErrorReporter* error_reporter, int buffer_index, int* offset) override; private: static constexpr int kMaxBufferCount = 1024; - int buffer_offsets_[kMaxBufferCount]; + size_t buffer_offsets_[kMaxBufferCount]; int current_buffer_count_; - int next_free_offset_; + size_t next_free_offset_; TF_LITE_REMOVE_VIRTUAL_DELETE }; diff --git a/tensorflow/lite/micro/memory_planner/memory_planner.h b/tensorflow/lite/micro/memory_planner/memory_planner.h index 9d5cd08468b..4670d59208d 100644 --- a/tensorflow/lite/micro/memory_planner/memory_planner.h +++ b/tensorflow/lite/micro/memory_planner/memory_planner.h @@ -58,7 +58,7 @@ class MemoryPlanner { int last_time_used) = 0; // The largest contguous block of memory that's needed to hold the layout. - virtual int GetMaximumMemorySize() = 0; + virtual size_t GetMaximumMemorySize() = 0; // How many buffers have been added to the planner. virtual int GetBufferCount() = 0; // Calculated layout offset for the N-th buffer added to the planner. diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index fb876f98a4a..886dbb46dcf 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -502,12 +502,8 @@ TfLiteStatus MicroAllocator::FinishTensorAllocation() { // the tensor info array, which will be released. size_t actual_available_arena_size = arena_size - memory_allocator_->GetDataSize(); - // Make sure we have enough room. - // TODO(b/147871342): make GetMaximumMemorySize return size_t. - // int is more than enough to hold arena_size since we're only dealing with - // at most several megabytes memory. - if (planner.GetMaximumMemorySize() > - static_cast(actual_available_arena_size)) { + // Make sure we have enough arena size. + if (planner.GetMaximumMemorySize() > actual_available_arena_size) { error_reporter_->Report( "Arena size is too small for activation buffers. Needed %d but only " "%d was available.",