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

@@ -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 = [

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();

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<flatbuffers::Offset<tflite::Buffer>>* 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);

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

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_

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 <cstdint>
+
+#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

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<SimpleMemoryAllocator*>(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);

tensorflow/lite/micro/simple_memory_allocator.h

@@ -20,6 +20,7 @@ limitations under the License.
 #include <cstdint>
 
 #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
 

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;