Make GetMaximumMemorySize return size_t - Fix bug 147871342

PiperOrigin-RevId: 291405835
Change-Id: I9dfe02861b6156db9d33100bd50825be17ef639e

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;
     }

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;

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_; }
 

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
 };

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.

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<int>(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.",