From 7d0ab6178803ea00b1693f8628e5914c46ebde44 Mon Sep 17 00:00:00 2001 From: Nick Kreeger Date: Fri, 29 May 2020 13:44:24 -0700 Subject: [PATCH] Add special "recording" SimpleMemoryAllocator class to help with logging tail allocations. This new helper class will enable TFLM to log and record where the allocations in the shared arena are going. A future change will use this new class in a special "recording" MicroAllocator subclass. All these logging mechanisms will be opt-in by code. PiperOrigin-RevId: 313843072 Change-Id: I3fc9205e475e89b4a3795c3cc79c31d2166da2c8 --- tensorflow/lite/micro/BUILD | 28 +++++ tensorflow/lite/micro/micro_allocator.cc | 2 +- tensorflow/lite/micro/micro_allocator_test.cc | 21 ++-- .../recording_simple_memory_allocator.cc | 65 ++++++++++ .../micro/recording_simple_memory_allocator.h | 59 +++++++++ .../recording_simple_memory_allocator_test.cc | 119 ++++++++++++++++++ .../lite/micro/simple_memory_allocator.cc | 22 ++-- .../lite/micro/simple_memory_allocator.h | 20 +-- tensorflow/lite/micro/testing/test_utils.cc | 4 +- 9 files changed, 313 insertions(+), 27 deletions(-) create mode 100644 tensorflow/lite/micro/recording_simple_memory_allocator.cc create mode 100644 tensorflow/lite/micro/recording_simple_memory_allocator.h create mode 100644 tensorflow/lite/micro/recording_simple_memory_allocator_test.cc diff --git a/tensorflow/lite/micro/BUILD b/tensorflow/lite/micro/BUILD index 36cc14512b6..ae4ba2fefd6 100644 --- a/tensorflow/lite/micro/BUILD +++ b/tensorflow/lite/micro/BUILD @@ -131,6 +131,22 @@ cc_library( ], ) +cc_library( + name = "recording_simple_memory_allocator", + srcs = [ + "recording_simple_memory_allocator.cc", + ], + hdrs = [ + "recording_simple_memory_allocator.h", + ], + build_for_embedded = True, + copts = micro_copts(), + deps = [ + ":micro_compatibility", + ":micro_framework", + ], +) + tflite_micro_cc_test( name = "micro_error_reporter_test", srcs = [ @@ -177,6 +193,18 @@ tflite_micro_cc_test( ], ) +tflite_micro_cc_test( + name = "recording_simple_memory_allocator_test", + srcs = [ + "recording_simple_memory_allocator_test.cc", + ], + deps = [ + ":micro_framework", + ":recording_simple_memory_allocator", + "//tensorflow/lite/micro/testing:micro_test", + ], +) + tflite_micro_cc_test( name = "micro_allocator_test", srcs = [ diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index 35f4bdabd20..ad26483ec3c 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -417,7 +417,7 @@ MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model, // Creates a root memory allocator managing the arena. The allocator itself // also locates in the arena buffer. This allocator doesn't need to be // destructed as it's the root allocator. - memory_allocator_ = CreateInPlaceSimpleMemoryAllocator( + memory_allocator_ = SimpleMemoryAllocator::Create( error_reporter, aligned_arena, aligned_arena_size); TfLiteStatus status = InitGraphAndContextTensorData(); diff --git a/tensorflow/lite/micro/micro_allocator_test.cc b/tensorflow/lite/micro/micro_allocator_test.cc index b34b2dc2866..013c1cefabf 100644 --- a/tensorflow/lite/micro/micro_allocator_test.cc +++ b/tensorflow/lite/micro/micro_allocator_test.cc @@ -68,8 +68,9 @@ TF_LITE_MICRO_TEST(TestInitializeRuntimeTensor) { TfLiteContext context; constexpr size_t arena_size = 1024; uint8_t arena[arena_size]; - tflite::SimpleMemoryAllocator simple_allocator(micro_test::reporter, arena, - arena_size); + tflite::SimpleMemoryAllocator* simple_allocator = + tflite::SimpleMemoryAllocator::Create(micro_test::reporter, arena, + arena_size); const tflite::Tensor* tensor = tflite::testing::Create1dFlatbufferTensor(100); const flatbuffers::Vector>* buffers = @@ -78,7 +79,7 @@ TF_LITE_MICRO_TEST(TestInitializeRuntimeTensor) { TfLiteTensor allocated_tensor; TF_LITE_MICRO_EXPECT_EQ( kTfLiteOk, tflite::internal::InitializeTfLiteTensorFromFlatbuffer( - &simple_allocator, *tensor, buffers, micro_test::reporter, + simple_allocator, *tensor, buffers, micro_test::reporter, &allocated_tensor)); TF_LITE_MICRO_EXPECT_EQ(kTfLiteInt32, allocated_tensor.type); TF_LITE_MICRO_EXPECT_EQ(1, allocated_tensor.dims->size); @@ -93,8 +94,9 @@ TF_LITE_MICRO_TEST(TestInitializeQuantizedTensor) { TfLiteContext context; constexpr size_t arena_size = 1024; uint8_t arena[arena_size]; - tflite::SimpleMemoryAllocator simple_allocator(micro_test::reporter, arena, - arena_size); + tflite::SimpleMemoryAllocator* simple_allocator = + tflite::SimpleMemoryAllocator::Create(micro_test::reporter, arena, + arena_size); const tflite::Tensor* tensor = tflite::testing::CreateQuantizedFlatbufferTensor(100); @@ -104,7 +106,7 @@ TF_LITE_MICRO_TEST(TestInitializeQuantizedTensor) { TfLiteTensor allocated_tensor; TF_LITE_MICRO_EXPECT_EQ( kTfLiteOk, tflite::internal::InitializeTfLiteTensorFromFlatbuffer( - &simple_allocator, *tensor, buffers, micro_test::reporter, + simple_allocator, *tensor, buffers, micro_test::reporter, &allocated_tensor)); TF_LITE_MICRO_EXPECT_EQ(kTfLiteInt32, allocated_tensor.type); TF_LITE_MICRO_EXPECT_EQ(1, allocated_tensor.dims->size); @@ -119,8 +121,9 @@ TF_LITE_MICRO_TEST(TestMissingQuantization) { TfLiteContext context; constexpr size_t arena_size = 1024; uint8_t arena[arena_size]; - tflite::SimpleMemoryAllocator simple_allocator(micro_test::reporter, arena, - arena_size); + tflite::SimpleMemoryAllocator* simple_allocator = + tflite::SimpleMemoryAllocator::Create(micro_test::reporter, arena, + arena_size); const tflite::Tensor* tensor = tflite::testing::CreateMissingQuantizationFlatbufferTensor(100); @@ -130,7 +133,7 @@ TF_LITE_MICRO_TEST(TestMissingQuantization) { TfLiteTensor allocated_tensor; TF_LITE_MICRO_EXPECT_EQ( kTfLiteOk, tflite::internal::InitializeTfLiteTensorFromFlatbuffer( - &simple_allocator, *tensor, buffers, micro_test::reporter, + simple_allocator, *tensor, buffers, micro_test::reporter, &allocated_tensor)); TF_LITE_MICRO_EXPECT_EQ(kTfLiteInt32, allocated_tensor.type); TF_LITE_MICRO_EXPECT_EQ(1, allocated_tensor.dims->size); diff --git a/tensorflow/lite/micro/recording_simple_memory_allocator.cc b/tensorflow/lite/micro/recording_simple_memory_allocator.cc new file mode 100644 index 00000000000..934fa260e30 --- /dev/null +++ b/tensorflow/lite/micro/recording_simple_memory_allocator.cc @@ -0,0 +1,65 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/lite/micro/recording_simple_memory_allocator.h" + +namespace tflite { + +RecordingSimpleMemoryAllocator::RecordingSimpleMemoryAllocator( + ErrorReporter* error_reporter, uint8_t* buffer_head, size_t buffer_size) + : SimpleMemoryAllocator(error_reporter, buffer_head, buffer_size), + requested_bytes_(0), + used_bytes_(0), + alloc_count_(0) {} + +RecordingSimpleMemoryAllocator::~RecordingSimpleMemoryAllocator() {} + +size_t RecordingSimpleMemoryAllocator::GetRequestedBytes() const { + return requested_bytes_; +} + +size_t RecordingSimpleMemoryAllocator::GetUsedBytes() const { + return used_bytes_; +} + +size_t RecordingSimpleMemoryAllocator::GetAllocatedCount() const { + return alloc_count_; +} + +uint8_t* RecordingSimpleMemoryAllocator::AllocateFromHead(size_t size, + size_t alignment) { + const uint8_t* previous_head = GetHead(); + uint8_t* result = SimpleMemoryAllocator::AllocateFromHead(size, alignment); + if (result != nullptr) { + used_bytes_ += GetHead() - previous_head; + requested_bytes_ += size; + alloc_count_++; + } + return result; +} + +uint8_t* RecordingSimpleMemoryAllocator::AllocateFromTail(size_t size, + size_t alignment) { + const uint8_t* previous_tail = GetTail(); + uint8_t* result = SimpleMemoryAllocator::AllocateFromTail(size, alignment); + if (result != nullptr) { + used_bytes_ += previous_tail - GetTail(); + requested_bytes_ += size; + alloc_count_++; + } + return result; +} + +} // namespace tflite diff --git a/tensorflow/lite/micro/recording_simple_memory_allocator.h b/tensorflow/lite/micro/recording_simple_memory_allocator.h new file mode 100644 index 00000000000..77edadb35be --- /dev/null +++ b/tensorflow/lite/micro/recording_simple_memory_allocator.h @@ -0,0 +1,59 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_LITE_MICRO_RECORDING_SIMPLE_MEMORY_ALLOCATOR_H_ +#define TENSORFLOW_LITE_MICRO_RECORDING_SIMPLE_MEMORY_ALLOCATOR_H_ + +#include "tensorflow/lite/micro/compatibility.h" +#include "tensorflow/lite/micro/simple_memory_allocator.h" + +namespace tflite { + +// Utility class used to log allocations of a SimpleMemoryAllocator. Should only +// be used in debug/evaluation settings or unit tests to evaluate allocation +// usage. +class RecordingSimpleMemoryAllocator : public SimpleMemoryAllocator { + public: + RecordingSimpleMemoryAllocator(ErrorReporter* error_reporter, + uint8_t* buffer_head, size_t buffer_size); + // TODO(b/157615197): Cleanup constructors/destructor and use factory + // functions. + ~RecordingSimpleMemoryAllocator() override; + + // Returns the number of bytes requested from the head or tail. + size_t GetRequestedBytes() const; + + // Returns the number of bytes actually allocated from the head or tail. This + // value will be >= to the number of requested bytes due to padding and + // alignment. + size_t GetUsedBytes() const; + + // Returns the number of alloc calls from the head or tail. + size_t GetAllocatedCount() const; + + uint8_t* AllocateFromHead(size_t size, size_t alignment) override; + uint8_t* AllocateFromTail(size_t size, size_t alignment) override; + + private: + size_t requested_bytes_; + size_t used_bytes_; + size_t alloc_count_; + + TF_LITE_REMOVE_VIRTUAL_DELETE +}; + +} // namespace tflite + +#endif // TENSORFLOW_LITE_MICRO_RECORDING_SIMPLE_MEMORY_ALLOCATOR_H_ diff --git a/tensorflow/lite/micro/recording_simple_memory_allocator_test.cc b/tensorflow/lite/micro/recording_simple_memory_allocator_test.cc new file mode 100644 index 00000000000..8fc4745a70e --- /dev/null +++ b/tensorflow/lite/micro/recording_simple_memory_allocator_test.cc @@ -0,0 +1,119 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/lite/micro/recording_simple_memory_allocator.h" + +#include + +#include "tensorflow/lite/micro/test_helpers.h" +#include "tensorflow/lite/micro/testing/micro_test.h" + +TF_LITE_MICRO_TESTS_BEGIN + +TF_LITE_MICRO_TEST(TestRecordsTailAllocations) { + constexpr size_t arena_size = 1024; + uint8_t arena[arena_size]; + tflite::RecordingSimpleMemoryAllocator allocator(micro_test::reporter, arena, + arena_size); + + uint8_t* result = allocator.AllocateFromTail(/*size=*/10, /*alignment=*/1); + TF_LITE_MICRO_EXPECT_NE(result, nullptr); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), 10); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetRequestedBytes(), 10); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetAllocatedCount(), 1); + + result = allocator.AllocateFromTail(/*size=*/20, /*alignment=*/1); + TF_LITE_MICRO_EXPECT_NE(result, nullptr); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), 30); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetRequestedBytes(), 30); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetAllocatedCount(), 2); +} + +TF_LITE_MICRO_TEST(TestRecordsMisalignedTailAllocations) { + constexpr size_t arena_size = 1024; + uint8_t arena[arena_size]; + tflite::RecordingSimpleMemoryAllocator allocator(micro_test::reporter, arena, + arena_size); + + uint8_t* result = allocator.AllocateFromTail(/*size=*/10, /*alignment=*/12); + TF_LITE_MICRO_EXPECT_NE(result, nullptr); + // Validate used bytes in 8 byte range that can included alignment of 12: + TF_LITE_MICRO_EXPECT_GE(allocator.GetUsedBytes(), 10); + TF_LITE_MICRO_EXPECT_LE(allocator.GetUsedBytes(), 20); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetRequestedBytes(), 10); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetAllocatedCount(), 1); +} + +TF_LITE_MICRO_TEST(TestDoesNotRecordFailedTailAllocations) { + constexpr size_t arena_size = 1024; + uint8_t arena[arena_size]; + tflite::RecordingSimpleMemoryAllocator allocator(micro_test::reporter, arena, + arena_size); + + uint8_t* result = allocator.AllocateFromTail(/*size=*/2048, /*alignment=*/1); + TF_LITE_MICRO_EXPECT_EQ(result, nullptr); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), 0); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetRequestedBytes(), 0); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetAllocatedCount(), 0); +} + +TF_LITE_MICRO_TEST(TestRecordsHeadAllocations) { + constexpr size_t arena_size = 1024; + uint8_t arena[arena_size]; + tflite::RecordingSimpleMemoryAllocator allocator(micro_test::reporter, arena, + arena_size); + + uint8_t* result = allocator.AllocateFromHead(/*size=*/5, /*alignment=*/1); + TF_LITE_MICRO_EXPECT_NE(result, nullptr); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), 5); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetRequestedBytes(), 5); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetAllocatedCount(), 1); + + result = allocator.AllocateFromTail(/*size=*/15, /*alignment=*/1); + TF_LITE_MICRO_EXPECT_NE(result, nullptr); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), 20); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetRequestedBytes(), 20); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetAllocatedCount(), 2); +} + +TF_LITE_MICRO_TEST(TestRecordsMisalignedHeadAllocations) { + constexpr size_t arena_size = 1024; + uint8_t arena[arena_size]; + tflite::RecordingSimpleMemoryAllocator allocator(micro_test::reporter, arena, + arena_size); + + uint8_t* result = allocator.AllocateFromHead(/*size=*/10, /*alignment=*/12); + TF_LITE_MICRO_EXPECT_NE(result, nullptr); + // Validate used bytes in 8 byte range that can included alignment of 12: + TF_LITE_MICRO_EXPECT_GE(allocator.GetUsedBytes(), 10); + TF_LITE_MICRO_EXPECT_LE(allocator.GetUsedBytes(), 20); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetRequestedBytes(), 10); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetAllocatedCount(), 1); +} + +TF_LITE_MICRO_TEST(TestDoesNotRecordFailedTailAllocations) { + constexpr size_t arena_size = 1024; + uint8_t arena[arena_size]; + tflite::RecordingSimpleMemoryAllocator allocator(micro_test::reporter, arena, + arena_size); + + uint8_t* result = allocator.AllocateFromHead(/*size=*/2048, /*alignment=*/1); + TF_LITE_MICRO_EXPECT_EQ(result, nullptr); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetUsedBytes(), 0); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetRequestedBytes(), 0); + TF_LITE_MICRO_EXPECT_EQ(allocator.GetAllocatedCount(), 0); +} + +TF_LITE_MICRO_TESTS_END diff --git a/tensorflow/lite/micro/simple_memory_allocator.cc b/tensorflow/lite/micro/simple_memory_allocator.cc index d55e7e87640..65d60161e9a 100644 --- a/tensorflow/lite/micro/simple_memory_allocator.cc +++ b/tensorflow/lite/micro/simple_memory_allocator.cc @@ -37,17 +37,23 @@ SimpleMemoryAllocator::SimpleMemoryAllocator(ErrorReporter* error_reporter, size_t buffer_size) : SimpleMemoryAllocator(error_reporter, buffer, buffer + buffer_size) {} -SimpleMemoryAllocator* CreateInPlaceSimpleMemoryAllocator( - ErrorReporter* error_reporter, uint8_t* buffer, size_t buffer_size) { +/* static */ +SimpleMemoryAllocator* SimpleMemoryAllocator::Create( + ErrorReporter* error_reporter, uint8_t* buffer_head, size_t buffer_size) { SimpleMemoryAllocator tmp = - SimpleMemoryAllocator(error_reporter, buffer, buffer_size); - SimpleMemoryAllocator* in_place_allocator = - reinterpret_cast(tmp.AllocateFromTail( - sizeof(SimpleMemoryAllocator), alignof(SimpleMemoryAllocator))); - *in_place_allocator = tmp; - return in_place_allocator; + SimpleMemoryAllocator(error_reporter, buffer_head, buffer_size); + + // Allocate enough bytes from the buffer to create a SimpleMemoryAllocator. + // The new instance will use the current adjusted tail buffer from the tmp + // allocator instance. + uint8_t* allocator_buffer = tmp.AllocateFromTail( + sizeof(SimpleMemoryAllocator), alignof(SimpleMemoryAllocator)); + return new (allocator_buffer) + SimpleMemoryAllocator(error_reporter, tmp.head_, tmp.tail_); } +SimpleMemoryAllocator::~SimpleMemoryAllocator() {} + uint8_t* SimpleMemoryAllocator::AllocateFromHead(size_t size, size_t alignment) { uint8_t* const aligned_result = AlignPointerUp(head_, alignment); diff --git a/tensorflow/lite/micro/simple_memory_allocator.h b/tensorflow/lite/micro/simple_memory_allocator.h index 5be260f9ed2..426ced032f6 100644 --- a/tensorflow/lite/micro/simple_memory_allocator.h +++ b/tensorflow/lite/micro/simple_memory_allocator.h @@ -20,6 +20,7 @@ limitations under the License. #include #include "tensorflow/lite/core/api/error_reporter.h" +#include "tensorflow/lite/micro/compatibility.h" namespace tflite { @@ -28,17 +29,25 @@ namespace tflite { // This makes it pretty wasteful, so we should use a more intelligent method. class SimpleMemoryAllocator { public: + // TODO(b/157615197): Cleanup constructors/destructor and use factory + // functions. SimpleMemoryAllocator(ErrorReporter* error_reporter, uint8_t* buffer_head, uint8_t* buffer_tail); SimpleMemoryAllocator(ErrorReporter* error_reporter, uint8_t* buffer, size_t buffer_size); + virtual ~SimpleMemoryAllocator(); + + // Creates a new SimpleMemoryAllocator from a given buffer head and size. + static SimpleMemoryAllocator* Create(ErrorReporter* error_reporter, + uint8_t* buffer_head, + size_t buffer_size); // Allocates memory starting at the head of the arena (lowest address and // moving upwards). - uint8_t* AllocateFromHead(size_t size, size_t alignment); + virtual uint8_t* AllocateFromHead(size_t size, size_t alignment); // Allocates memory starting at the tail of the arena (highest address and // moving downwards). - uint8_t* AllocateFromTail(size_t size, size_t alignment); + virtual uint8_t* AllocateFromTail(size_t size, size_t alignment); uint8_t* GetHead() const; uint8_t* GetTail() const; @@ -57,12 +66,9 @@ class SimpleMemoryAllocator { uint8_t* buffer_tail_; uint8_t* head_; uint8_t* tail_; -}; -// Allocate a SimpleMemoryAllocator from the buffer and then return the pointer -// to this allocator. -SimpleMemoryAllocator* CreateInPlaceSimpleMemoryAllocator( - ErrorReporter* error_reporter, uint8_t* buffer, size_t buffer_size); + TF_LITE_REMOVE_VIRTUAL_DELETE +}; } // namespace tflite diff --git a/tensorflow/lite/micro/testing/test_utils.cc b/tensorflow/lite/micro/testing/test_utils.cc index 62621db40d3..7d7a5554b10 100644 --- a/tensorflow/lite/micro/testing/test_utils.cc +++ b/tensorflow/lite/micro/testing/test_utils.cc @@ -119,8 +119,8 @@ int32_t F2Q32(float value, float scale) { // TODO(b/141330728): Move this method elsewhere as part clean up. void PopulateContext(TfLiteTensor* tensors, int tensors_size, ErrorReporter* error_reporter, TfLiteContext* context) { - simple_memory_allocator_ = CreateInPlaceSimpleMemoryAllocator( - error_reporter, raw_arena_, kArenaSize); + simple_memory_allocator_ = + SimpleMemoryAllocator::Create(error_reporter, raw_arena_, kArenaSize); TFLITE_DCHECK(simple_memory_allocator_ != nullptr); scratch_buffer_count_ = 0;