From a7bdaeba611d4aa620ab2837e1a53152716320cf Mon Sep 17 00:00:00 2001 From: Nick Kreeger Date: Fri, 9 Oct 2020 05:04:47 -0700 Subject: [PATCH] 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 | 12 ++-- .../recording_simple_memory_allocator.cc | 15 +++-- .../micro/recording_simple_memory_allocator.h | 2 +- .../recording_simple_memory_allocator_test.cc | 12 ++-- .../lite/micro/simple_memory_allocator.cc | 22 ++++--- .../lite/micro/simple_memory_allocator.h | 23 ++++--- .../micro/simple_memory_allocator_test.cc | 65 ++++++++++--------- 7 files changed, 81 insertions(+), 70 deletions(-) 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( - 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(5)); TF_LITE_MICRO_EXPECT_EQ(allocator.GetRequestedBytes(), static_cast(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(10)); TF_LITE_MICRO_EXPECT_LE(allocator.GetUsedBytes(), static_cast(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(0)); TF_LITE_MICRO_EXPECT_EQ(allocator.GetRequestedBytes(), static_cast(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(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(10), + allocator.GetHeadUsedBytes()); + + TF_LITE_MICRO_EXPECT_EQ( + kTfLiteOk, allocator.SetHeadBufferSize(/*size=*/1000, /*alignment=*/1)); + TF_LITE_MICRO_EXPECT_EQ(static_cast(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(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