diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc
index bea32929b02..c6850246179 100644
--- a/tensorflow/lite/micro/micro_allocator.cc
+++ b/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(
- memory_allocator_->SetHeadSize(max_head_buffer_usage_, kBufferAlignment));
+ TF_LITE_ENSURE_STATUS(memory_allocator_->SetHeadBufferSize(
+ 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)));
}
diff --git a/tensorflow/lite/micro/recording_simple_memory_allocator.cc b/tensorflow/lite/micro/recording_simple_memory_allocator.cc
index 6af82ce8469..ef30aca4949 100644
--- a/tensorflow/lite/micro/recording_simple_memory_allocator.cc
+++ b/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,
- size_t alignment) {
- const uint8_t* previous_head = GetHead();
- TfLiteStatus status = SimpleMemoryAllocator::SetHeadSize(size, alignment);
+TfLiteStatus RecordingSimpleMemoryAllocator::SetHeadBufferSize(
+ size_t size, size_t alignment) {
+ const uint8_t* previous_head = head();
+ 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_++;
}
diff --git a/tensorflow/lite/micro/recording_simple_memory_allocator.h b/tensorflow/lite/micro/recording_simple_memory_allocator.h
index de18a75f937..3526716e3b4 100644
--- a/tensorflow/lite/micro/recording_simple_memory_allocator.h
+++ b/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:
diff --git a/tensorflow/lite/micro/recording_simple_memory_allocator_test.cc b/tensorflow/lite/micro/recording_simple_memory_allocator_test.cc
index 5f356cb319f..910c991978d 100644
--- a/tensorflow/lite/micro/recording_simple_memory_allocator_test.cc
+++ b/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,
- allocator.SetHeadSize(/*size=*/5, /*alignment=*/1));
+ TF_LITE_MICRO_EXPECT_EQ(
+ 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,
- allocator.SetHeadSize(/*size=*/10, /*alignment=*/12));
+ TF_LITE_MICRO_EXPECT_EQ(
+ 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(
- kTfLiteError, allocator.SetHeadSize(/*size=*/2048, /*alignment=*/1));
+ TF_LITE_MICRO_EXPECT_EQ(kTfLiteError, allocator.SetHeadBufferSize(
+ /*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));
diff --git a/tensorflow/lite/micro/simple_memory_allocator.cc b/tensorflow/lite/micro/simple_memory_allocator.cc
index 40593734044..9ddb989f726 100644
--- a/tensorflow/lite/micro/simple_memory_allocator.cc
+++ b/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_,
- "Internal error: SetHeadSize() needs to be called "
- "after ResetTempAllocations().");
+ TF_LITE_REPORT_ERROR(
+ error_reporter_,
+ "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::GetBufferHead() const { return buffer_head_; }
-
-uint8_t* SimpleMemoryAllocator::GetTail() const { return tail_; }
+uint8_t* SimpleMemoryAllocator::GetHeadBuffer() const { return buffer_head_; }
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
diff --git a/tensorflow/lite/micro/simple_memory_allocator.h b/tensorflow/lite/micro/simple_memory_allocator.h
index 6c353c84d8a..fd055f7b269 100644
--- a/tensorflow/lite/micro/simple_memory_allocator.h
+++ b/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.
- // GetBufferHead() will return the start of the head section. The GetHead()
- // method currently returns the end address of the head section. This can
- // 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 a pointer to the buffer currently assigned to the head section.
+ // This buffer is set by calling SetHeadSize().
+ uint8_t* GetHeadBuffer() 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;
diff --git a/tensorflow/lite/micro/simple_memory_allocator_test.cc b/tensorflow/lite/micro/simple_memory_allocator_test.cc
index e2a9f972fff..a38daf278a6 100644
--- a/tensorflow/lite/micro/simple_memory_allocator_test.cc
+++ b/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,
- allocator.SetHeadSize(/*size=*/100, /*alignment=*/1));
- TF_LITE_MICRO_EXPECT(arena + 100 == allocator.GetHead());
-
- 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.SetHeadBufferSize(/*size=*/100, /*alignment=*/1));
+ TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(100),
+ allocator.GetHeadUsedBytes());
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator.SetHeadSize(/*size=*/1000, /*alignment=*/1));
- TF_LITE_MICRO_EXPECT(arena + 1000 == allocator.GetHead());
+ kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/10, /*alignment=*/1));
+ 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_LE(allocator.GetHead(), arena + 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_GE(allocator.GetHeadUsedBytes(), 100);
+ TF_LITE_MICRO_EXPECT_LE(allocator.GetHeadUsedBytes(), 100 + 11);
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator.SetHeadSize(/*size=*/1000, /*alignment=*/12));
- TF_LITE_MICRO_EXPECT_GE(allocator.GetHead(), arena + 1000);
- TF_LITE_MICRO_EXPECT_LE(allocator.GetHead(), arena + 1000 + 11);
+ kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/10, /*alignment=*/12));
+ TF_LITE_MICRO_EXPECT_GE(allocator.GetHeadUsedBytes(), 10);
+ 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,
- allocator.SetHeadSize(/*size=*/10, /*alignment=*/12));
+ TF_LITE_MICRO_EXPECT_EQ(
+ 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,
- /*alignment=*/1);
+ allocator.SetHeadBufferSize(/*size=*/allocation_size,
+ /*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,
- /*alignment=*/1);
+ allocator.SetHeadBufferSize(/*size=*/allocation_size,
+ /*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