Drop various buffer pointer getters in SimpleMemoryAllocator.

The API for SimpleMemoryAllocator should just simply set the size of the head buffer and return a pointer to the start of that buffer. The current APIs exposed on this class are confusing and can easily be mixed up. This change drops those getters and relies on publicly facing methods to expose functionality. Additionaly, the head APIs are renamed to make more sense to what they do - manage the single head buffer allocation.

PiperOrigin-RevId: 336273280
Change-Id: Ibd4218c40962946b90633ca55169595057ea46c3

tensorflow/lite/micro/micro_allocator.cc

@@ -760,7 +760,7 @@ TfLiteStatus MicroAllocator::FinishPrepareNodeAllocations(int node_id) {
 
   // Ensure that the head is re-adjusted to allow for another at-most
   // kMaxScratchBuffersPerOp scratch buffer requests in the next operator:
-  TF_LITE_ENSURE_STATUS(memory_allocator_->SetHeadSize(
+  TF_LITE_ENSURE_STATUS(memory_allocator_->SetHeadBufferSize(
       sizeof(internal::ScratchBufferRequest) *
           (scratch_buffer_request_count_ + kMaxScratchBuffersPerOp),
       alignof(internal::ScratchBufferRequest)));
@@ -1116,7 +1116,7 @@ TfLiteStatus MicroAllocator::CommitStaticMemoryPlan(
   }
   // Commit the plan.
   TF_LITE_ENSURE_STATUS(CommitPlan(error_reporter_, &planner,
-                                   memory_allocator_->GetBufferHead(),
+                                   memory_allocator_->GetHeadBuffer(),
                                    allocation_info, allocation_info_count));
   head_usage = planner.GetMaximumMemorySize();
 
@@ -1132,8 +1132,8 @@ TfLiteStatus MicroAllocator::CommitStaticMemoryPlan(
   // The head is used for storing scratch buffer allocations before finalizing a
   // memory plan in this function. Ensure that the head is set to the largest
   // memory plan sent through the allocator:
-  TF_LITE_ENSURE_STATUS(
+  TF_LITE_ENSURE_STATUS(memory_allocator_->SetHeadBufferSize(
-      memory_allocator_->SetHeadSize(max_head_buffer_usage_, kBufferAlignment));
+      max_head_buffer_usage_, kBufferAlignment));
   return kTfLiteOk;
 }
 
@@ -1164,7 +1164,7 @@ TfLiteStatus MicroAllocator::InitScratchBufferData() {
   // All requests will be stored in the head section. Each kernel is allowed at
   // most kMaxScratchBuffersPerOp requests. Adjust the head to reserve at most
   // that many requests to begin:
-  TF_LITE_ENSURE_STATUS(memory_allocator_->SetHeadSize(
+  TF_LITE_ENSURE_STATUS(memory_allocator_->SetHeadBufferSize(
       sizeof(internal::ScratchBufferRequest) * kMaxScratchBuffersPerOp,
       alignof(internal::ScratchBufferRequest)));
 
@@ -1173,7 +1173,7 @@ TfLiteStatus MicroAllocator::InitScratchBufferData() {
 
 internal::ScratchBufferRequest* MicroAllocator::GetScratchBufferRequests() {
   return reinterpret_cast<internal::ScratchBufferRequest*>(
-      AlignPointerUp(memory_allocator_->GetBufferHead(),
+      AlignPointerUp(memory_allocator_->GetHeadBuffer(),
                      alignof(internal::ScratchBufferRequest)));
 }
 

tensorflow/lite/micro/recording_simple_memory_allocator.cc

@@ -57,12 +57,13 @@ size_t RecordingSimpleMemoryAllocator::GetAllocatedCount() const {
   return alloc_count_;
 }
 
-TfLiteStatus RecordingSimpleMemoryAllocator::SetHeadSize(size_t size,
+TfLiteStatus RecordingSimpleMemoryAllocator::SetHeadBufferSize(
-                                                         size_t alignment) {
+    size_t size, size_t alignment) {
-  const uint8_t* previous_head = GetHead();
+  const uint8_t* previous_head = head();
-  TfLiteStatus status = SimpleMemoryAllocator::SetHeadSize(size, alignment);
+  TfLiteStatus status =
+      SimpleMemoryAllocator::SetHeadBufferSize(size, alignment);
   if (status == kTfLiteOk) {
-    used_bytes_ += GetHead() - previous_head;
+    used_bytes_ += head() - previous_head;
     requested_head_bytes_ = size;
   }
   return status;
@@ -70,10 +71,10 @@ TfLiteStatus RecordingSimpleMemoryAllocator::SetHeadSize(size_t size,
 
 uint8_t* RecordingSimpleMemoryAllocator::AllocateFromTail(size_t size,
                                                           size_t alignment) {
-  const uint8_t* previous_tail = GetTail();
+  const uint8_t* previous_tail = tail();
   uint8_t* result = SimpleMemoryAllocator::AllocateFromTail(size, alignment);
   if (result != nullptr) {
-    used_bytes_ += previous_tail - GetTail();
+    used_bytes_ += previous_tail - tail();
     requested_tail_bytes_ += size;
     alloc_count_++;
   }

tensorflow/lite/micro/recording_simple_memory_allocator.h

@@ -47,7 +47,7 @@ class RecordingSimpleMemoryAllocator : public SimpleMemoryAllocator {
   // Returns the number of alloc calls from the head or tail.
   size_t GetAllocatedCount() const;
 
-  TfLiteStatus SetHeadSize(size_t size, size_t alignment) override;
+  TfLiteStatus SetHeadBufferSize(size_t size, size_t alignment) override;
   uint8_t* AllocateFromTail(size_t size, size_t alignment) override;
 
  private:

tensorflow/lite/micro/recording_simple_memory_allocator_test.cc

@@ -83,8 +83,8 @@ TF_LITE_MICRO_TEST(TestRecordsHeadSizeAdjustment) {
   tflite::RecordingSimpleMemoryAllocator allocator(micro_test::reporter, arena,
                                                    arena_size);
 
-  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk,
+  TF_LITE_MICRO_EXPECT_EQ(
-                          allocator.SetHeadSize(/*size=*/5, /*alignment=*/1));
+      kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/5, /*alignment=*/1));
   TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), static_cast<size_t>(5));
   TF_LITE_MICRO_EXPECT_EQ(allocator.GetRequestedBytes(),
                           static_cast<size_t>(5));
@@ -107,8 +107,8 @@ TF_LITE_MICRO_TEST(TestRecordsMisalignedHeadSizeAdjustments) {
   tflite::RecordingSimpleMemoryAllocator allocator(micro_test::reporter, arena,
                                                    arena_size);
 
-  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk,
+  TF_LITE_MICRO_EXPECT_EQ(
-                          allocator.SetHeadSize(/*size=*/10, /*alignment=*/12));
+      kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/10, /*alignment=*/12));
   // Validate used bytes in 8 byte range that can included alignment of 12:
   TF_LITE_MICRO_EXPECT_GE(allocator.GetUsedBytes(), static_cast<size_t>(10));
   TF_LITE_MICRO_EXPECT_LE(allocator.GetUsedBytes(), static_cast<size_t>(20));
@@ -125,8 +125,8 @@ TF_LITE_MICRO_TEST(TestDoesNotRecordFailedTailAllocations) {
   tflite::RecordingSimpleMemoryAllocator allocator(micro_test::reporter, arena,
                                                    arena_size);
 
-  TF_LITE_MICRO_EXPECT_EQ(
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteError, allocator.SetHeadBufferSize(
-      kTfLiteError, allocator.SetHeadSize(/*size=*/2048, /*alignment=*/1));
+                                            /*size=*/2048, /*alignment=*/1));
   TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), static_cast<size_t>(0));
   TF_LITE_MICRO_EXPECT_EQ(allocator.GetRequestedBytes(),
                           static_cast<size_t>(0));

tensorflow/lite/micro/simple_memory_allocator.cc

@@ -60,11 +60,13 @@ SimpleMemoryAllocator* SimpleMemoryAllocator::Create(
 
 SimpleMemoryAllocator::~SimpleMemoryAllocator() {}
 
-TfLiteStatus SimpleMemoryAllocator::SetHeadSize(size_t size, size_t alignment) {
+TfLiteStatus SimpleMemoryAllocator::SetHeadBufferSize(size_t size,
+                                                      size_t alignment) {
   if (head_ != temp_) {
-    TF_LITE_REPORT_ERROR(error_reporter_,
+    TF_LITE_REPORT_ERROR(
-                         "Internal error: SetHeadSize() needs to be called "
+        error_reporter_,
-                         "after ResetTempAllocations().");
+        "Internal error: SetHeadBufferSize() needs to be called "
+        "after ResetTempAllocations().");
     return kTfLiteError;
   }
 
@@ -73,7 +75,7 @@ TfLiteStatus SimpleMemoryAllocator::SetHeadSize(size_t size, size_t alignment) {
   if (available_memory < size) {
     TF_LITE_REPORT_ERROR(
         error_reporter_,
-        "Failed to adjust head size. Requested: %u, available %u, missing: %u",
+        "Failed to set head size. Requested: %u, available %u, missing: %u",
         size, available_memory, size - available_memory);
     return kTfLiteError;
   }
@@ -116,11 +118,7 @@ uint8_t* SimpleMemoryAllocator::AllocateTemp(size_t size, size_t alignment) {
 
 void SimpleMemoryAllocator::ResetTempAllocations() { temp_ = head_; }
 
-uint8_t* SimpleMemoryAllocator::GetHead() const { return head_; }
+uint8_t* SimpleMemoryAllocator::GetHeadBuffer() const { return buffer_head_; }
-
-uint8_t* SimpleMemoryAllocator::GetBufferHead() const { return buffer_head_; }
-
-uint8_t* SimpleMemoryAllocator::GetTail() const { return tail_; }
 
 size_t SimpleMemoryAllocator::GetHeadUsedBytes() const {
   return head_ - buffer_head_;
@@ -144,4 +142,8 @@ size_t SimpleMemoryAllocator::GetBufferSize() const {
   return buffer_tail_ - buffer_head_;
 }
 
+uint8_t* SimpleMemoryAllocator::head() const { return head_; }
+
+uint8_t* SimpleMemoryAllocator::tail() const { return tail_; }
+
 }  // namespace tflite

tensorflow/lite/micro/simple_memory_allocator.h

@@ -49,7 +49,7 @@ class SimpleMemoryAllocator {
   // head section). This call will fail if a chain of allocations through
   // AllocateTemp() have not been cleaned up with a call to
   // ResetTempAllocations().
-  virtual TfLiteStatus SetHeadSize(size_t size, size_t alignment);
+  virtual TfLiteStatus SetHeadBufferSize(size_t size, size_t alignment);
 
   // Allocates memory starting at the tail of the arena (highest address and
   // moving downwards).
@@ -69,16 +69,14 @@ class SimpleMemoryAllocator {
   // arena (lowest address).
   virtual void ResetTempAllocations();
 
-  // TODO(b/169834500): Consider renaming or dropping the head methods.
+  // Returns a pointer to the buffer currently assigned to the head section.
-  // GetBufferHead() will return the start of the head section. The GetHead()
+  // This buffer is set by calling SetHeadSize().
-  // method currently returns the end address of the head section. This can
+  uint8_t* GetHeadBuffer() const;
-  // easily lead to misuse by placing things at the end of the head section by
-  // calling GetHead().
-  uint8_t* GetHead() const;
-  uint8_t* GetBufferHead() const;
-  uint8_t* GetTail() const;
 
+  // Returns the size of the head section in bytes.
   size_t GetHeadUsedBytes() const;
+
+  // Returns the size of all allocations in the tail section in bytes.
   size_t GetTailUsedBytes() const;
 
   // Returns the number of bytes available with a given alignment. This number
@@ -89,6 +87,13 @@ class SimpleMemoryAllocator {
   // account any temporary allocations.
   size_t GetUsedBytes() const;
 
+ protected:
+  // Returns a pointer to the current end of the head buffer.
+  uint8_t* head() const;
+
+  // Returns a pointer to the current end of the tail buffer.
+  uint8_t* tail() const;
+
  private:
   size_t GetBufferSize() const;
 

tensorflow/lite/micro/simple_memory_allocator_test.cc

@@ -28,17 +28,20 @@ TF_LITE_MICRO_TEST(TestEnsureHeadSizeSimpleAlignment) {
   tflite::SimpleMemoryAllocator allocator(micro_test::reporter, arena,
                                           arena_size);
 
-  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk,
+  TF_LITE_MICRO_EXPECT_EQ(
-                          allocator.SetHeadSize(/*size=*/100, /*alignment=*/1));
+      kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/100, /*alignment=*/1));
-  TF_LITE_MICRO_EXPECT(arena + 100 == allocator.GetHead());
+  TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(100),
-
+                          allocator.GetHeadUsedBytes());
-  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk,
-                          allocator.SetHeadSize(/*size=*/10, /*alignment=*/1));
-  TF_LITE_MICRO_EXPECT(arena + 10 == allocator.GetHead());
 
   TF_LITE_MICRO_EXPECT_EQ(
-      kTfLiteOk, allocator.SetHeadSize(/*size=*/1000, /*alignment=*/1));
+      kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/10, /*alignment=*/1));
-  TF_LITE_MICRO_EXPECT(arena + 1000 == allocator.GetHead());
+  TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(10),
+                          allocator.GetHeadUsedBytes());
+
+  TF_LITE_MICRO_EXPECT_EQ(
+      kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/1000, /*alignment=*/1));
+  TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(1000),
+                          allocator.GetHeadUsedBytes());
 }
 
 TF_LITE_MICRO_TEST(TestAdjustHeadSizeMisalignment) {
@@ -49,22 +52,22 @@ TF_LITE_MICRO_TEST(TestAdjustHeadSizeMisalignment) {
 
   // First head adjustment of 100 bytes (aligned 12):
   TF_LITE_MICRO_EXPECT_EQ(
-      kTfLiteOk, allocator.SetHeadSize(/*size=*/100, /*alignment=*/12));
+      kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/100, /*alignment=*/12));
 
   // Offset alignment of 12 can lead to allocation within 8 byte range of
   // requested bytes based to arena alignment at runtime:
-  TF_LITE_MICRO_EXPECT_GE(allocator.GetHead(), arena + 100);
+  TF_LITE_MICRO_EXPECT_GE(allocator.GetHeadUsedBytes(), 100);
-  TF_LITE_MICRO_EXPECT_LE(allocator.GetHead(), arena + 100 + 11);
+  TF_LITE_MICRO_EXPECT_LE(allocator.GetHeadUsedBytes(), 100 + 11);
-
-  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk,
-                          allocator.SetHeadSize(/*size=*/10, /*alignment=*/12));
-  TF_LITE_MICRO_EXPECT_GE(allocator.GetHead(), arena + 10);
-  TF_LITE_MICRO_EXPECT_LE(allocator.GetHead(), arena + 100 + 11);
 
   TF_LITE_MICRO_EXPECT_EQ(
-      kTfLiteOk, allocator.SetHeadSize(/*size=*/1000, /*alignment=*/12));
+      kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/10, /*alignment=*/12));
-  TF_LITE_MICRO_EXPECT_GE(allocator.GetHead(), arena + 1000);
+  TF_LITE_MICRO_EXPECT_GE(allocator.GetHeadUsedBytes(), 10);
-  TF_LITE_MICRO_EXPECT_LE(allocator.GetHead(), arena + 1000 + 11);
+  TF_LITE_MICRO_EXPECT_LE(allocator.GetHeadUsedBytes(), 100 + 11);
+
+  TF_LITE_MICRO_EXPECT_EQ(
+      kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/1000, /*alignment=*/12));
+  TF_LITE_MICRO_EXPECT_GE(allocator.GetHeadUsedBytes(), 1000);
+  TF_LITE_MICRO_EXPECT_LE(allocator.GetHeadUsedBytes(), 1000 + 11);
 }
 
 TF_LITE_MICRO_TEST(TestAdjustHeadSizeMisalignedHandlesCorrectBytesAvailable) {
@@ -75,7 +78,7 @@ TF_LITE_MICRO_TEST(TestAdjustHeadSizeMisalignedHandlesCorrectBytesAvailable) {
 
   // First head adjustment of 100 bytes (aligned 12):
   TF_LITE_MICRO_EXPECT_EQ(
-      kTfLiteOk, allocator.SetHeadSize(/*size=*/100, /*alignment=*/12));
+      kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/100, /*alignment=*/12));
 
   // allocator.GetAvailableMemory() should also report the actual amount of
   // memory available based on a requested offset (12):
@@ -84,15 +87,15 @@ TF_LITE_MICRO_TEST(TestAdjustHeadSizeMisalignedHandlesCorrectBytesAvailable) {
   TF_LITE_MICRO_EXPECT_LE(aligned_available_bytes, arena_size - 100);
   TF_LITE_MICRO_EXPECT_GE(aligned_available_bytes, arena_size - 100 - 24);
 
-  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk,
+  TF_LITE_MICRO_EXPECT_EQ(
-                          allocator.SetHeadSize(/*size=*/10, /*alignment=*/12));
+      kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/10, /*alignment=*/12));
   aligned_available_bytes = allocator.GetAvailableMemory(/*alignment=*/12);
 
   TF_LITE_MICRO_EXPECT_LE(aligned_available_bytes, arena_size - 10);
   TF_LITE_MICRO_EXPECT_GE(aligned_available_bytes, arena_size - 10 - 24);
 
   TF_LITE_MICRO_EXPECT_EQ(
-      kTfLiteOk, allocator.SetHeadSize(/*size=*/1000, /*alignment=*/12));
+      kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/1000, /*alignment=*/12));
   aligned_available_bytes = allocator.GetAvailableMemory(/*alignment=*/12);
   TF_LITE_MICRO_EXPECT_LE(aligned_available_bytes, arena_size - 1000);
   TF_LITE_MICRO_EXPECT_GE(aligned_available_bytes, arena_size - 1000 - 24);
@@ -105,8 +108,8 @@ TF_LITE_MICRO_TEST(TestGetAvailableMemory) {
                                           arena_size);
 
   constexpr size_t allocation_size = 100;
-  allocator.SetHeadSize(/*size=*/allocation_size,
+  allocator.SetHeadBufferSize(/*size=*/allocation_size,
-                        /*alignment=*/1);
+                              /*alignment=*/1);
   allocator.AllocateFromTail(/*size=*/allocation_size,
                              /*alignment=*/1);
 
@@ -143,8 +146,8 @@ TF_LITE_MICRO_TEST(TestGetUsedBytes) {
   TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), static_cast<size_t>(0));
 
   constexpr size_t allocation_size = 100;
-  allocator.SetHeadSize(/*size=*/allocation_size,
+  allocator.SetHeadBufferSize(/*size=*/allocation_size,
-                        /*alignment=*/1);
+                              /*alignment=*/1);
   allocator.AllocateFromTail(/*size=*/allocation_size,
                              /*alignment=*/1);
 
@@ -253,16 +256,16 @@ TF_LITE_MICRO_TEST(TestEnsureHeadSizeWithoutResettingTemp) {
 
   // Adjustment to head should fail since temp allocation was not followed by a
   // call to ResetTempAllocations().
-  TF_LITE_MICRO_EXPECT_EQ(kTfLiteError, allocator.SetHeadSize(100, 1));
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteError, allocator.SetHeadBufferSize(100, 1));
 
   allocator.ResetTempAllocations();
 
   // Reduce head size back to zero.
-  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.SetHeadSize(0, 1));
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.SetHeadBufferSize(0, 1));
 
   // The most recent head allocation should be in the same location as the
   // original temp allocation pointer.
-  TF_LITE_MICRO_EXPECT(temp == allocator.GetHead());
+  TF_LITE_MICRO_EXPECT(temp == allocator.GetHeadBuffer());
 }
 
 TF_LITE_MICRO_TESTS_END