Add a memory threshold allocation test for the Keyword model.

This new test ensures that TF Micro does not regress current allocations (on x86-64 systems) for a canonical model. As RAM reduction changes are introduced, the values in this test can be updated from the console log of this test.

Current output for the keyword model:
Testing TestKeywordModelMemoryThreshold
[RecordingMicroAllocator] Arena allocation total 21440 bytes
[RecordingMicroAllocator] Arena allocation head 672 bytes
[RecordingMicroAllocator] Arena allocation tail 20768 bytes
[RecordingMicroAllocator] 'TfLiteTensor struct allocation' used 6048 bytes (requested 6048 bytes 54 times)
[RecordingMicroAllocator] 'TfLiteTensor quantization data allocations' used 2160 bytes (requested 2160 bytes 162 times)
[RecordingMicroAllocator] 'NodeAndRegistration struct allocation' used 1200 bytes (requested 1200 bytes 15 times)
[RecordingMicroAllocator] 'Operator runtime data allocation' used 148 bytes (requested 148 bytes 13 times)

PiperOrigin-RevId: 315958032
Change-Id: I226f6a01aa555970805388632559241a41ff8342

tensorflow/lite/micro/BUILD

@@ -205,6 +205,7 @@ cc_library(
     ],
     hdrs = [
         "recording_micro_allocator.h",
+        "recording_micro_interpreter.h",
         "recording_simple_memory_allocator.h",
     ],
     build_for_embedded = True,
@@ -361,3 +362,16 @@ tflite_micro_cc_test(
         "//tensorflow/lite/micro/testing:micro_test",
     ],
 )
+
+tflite_micro_cc_test(
+    name = "memory_arena_threshold_test",
+    srcs = [
+        "memory_arena_threshold_test.cc",
+    ],
+    deps = [
+        ":op_resolvers",
+        ":recording_allocators",
+        "//tensorflow/lite/micro/benchmarks:keyword_scrambled_model_data",
+        "//tensorflow/lite/micro/testing:micro_test",
+    ],
+)

tensorflow/lite/micro/benchmarks/BUILD

@@ -1,5 +1,10 @@
 licenses(["notice"])  # Apache 2.0
 
+package_group(
+    name = "micro_top_level",
+    packages = ["//tensorflow/lite/micro"],
+)
+
 cc_library(
     name = "micro_benchmark",
     hdrs = [
@@ -47,6 +52,9 @@ cc_library(
     hdrs = [
         "keyword_scrambled_model_data.h",
     ],
+    visibility = [
+        ":micro_top_level",
+    ],
 )
 
 cc_binary(

tensorflow/lite/micro/memory_arena_threshold_test.cc

@@ -0,0 +1,110 @@
+/* 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 <stdint.h>
+
+#include "tensorflow/lite/micro/all_ops_resolver.h"
+#include "tensorflow/lite/micro/benchmarks/keyword_scrambled_model_data.h"
+#include "tensorflow/lite/micro/recording_micro_allocator.h"
+#include "tensorflow/lite/micro/recording_micro_interpreter.h"
+#include "tensorflow/lite/micro/testing/micro_test.h"
+
+/**
+ * Tests to ensure arena memory allocation does not regress by more than 3%.
+ */
+
+namespace {
+
+// Ensure memory doesn't expand more that 3%:
+constexpr float kAllocationThreshold = 0.03;
+const bool kIs64BitSystem = sizeof(void*) == 8;
+
+constexpr int kKeywordTensorArenaSize = 22 * 1024;
+uint8_t tensor_arena[kKeywordTensorArenaSize];
+
+constexpr int kKeywordModelTensorCount = 54;
+constexpr int kKeywordModelNodeAndRegistrationCount = 15;
+
+// NOTE: These values are measured on x86-64:
+// TODO(b/158651472): Consider auditing these values on non-64 bit systems.
+constexpr int kKeywordModelTotalSize = 21440;
+constexpr int kKeywordModelHeadSize = 672;
+constexpr int kKeywordModelTailSize = 20768;
+constexpr int kKeywordModelTfLiteTensorQuantizationDataSize = 2160;
+constexpr int kKeywordModelOpRuntimeDataSize = 148;
+
+void EnsureAllocatedSizeThreshold(size_t actual, size_t expected) {
+  // TODO(b/158651472): Better auditing of non-64 bit systems:
+  if (kIs64BitSystem) {
+    // 64-bit systems should check floor and ceiling to catch memory savings:
+    TF_LITE_MICRO_EXPECT_NEAR(actual, expected, kAllocationThreshold);
+  } else {
+    // Non-64 bit systems should just expect allocation does not exceed the
+    // ceiling:
+    TF_LITE_MICRO_EXPECT_LE(actual, expected + expected * kAllocationThreshold);
+  }
+}
+
+}  // namespace
+
+TF_LITE_MICRO_TESTS_BEGIN
+
+TF_LITE_MICRO_TEST(TestKeywordModelMemoryThreshold) {
+  tflite::AllOpsResolver all_ops_resolver;
+  tflite::RecordingMicroInterpreter interpreter(
+      tflite::GetModel(g_keyword_scrambled_model_data), &all_ops_resolver,
+      tensor_arena, kKeywordTensorArenaSize, micro_test::reporter);
+
+  interpreter.AllocateTensors();
+
+  const tflite::RecordingMicroAllocator& allocator =
+      interpreter.GetMicroAllocator();
+  allocator.PrintAllocations();
+
+  EnsureAllocatedSizeThreshold(
+      allocator.GetSimpleMemoryAllocator()->GetUsedBytes(),
+      kKeywordModelTotalSize);
+  EnsureAllocatedSizeThreshold(
+      allocator.GetSimpleMemoryAllocator()->GetHeadUsedBytes(),
+      kKeywordModelHeadSize);
+  EnsureAllocatedSizeThreshold(
+      allocator.GetSimpleMemoryAllocator()->GetTailUsedBytes(),
+      kKeywordModelTailSize);
+  EnsureAllocatedSizeThreshold(
+      allocator
+          .GetRecordedAllocation(
+              tflite::RecordedAllocationType::kTfLiteTensorArray)
+          .used_bytes,
+      sizeof(TfLiteTensor) * kKeywordModelTensorCount);
+  EnsureAllocatedSizeThreshold(
+      allocator
+          .GetRecordedAllocation(tflite::RecordedAllocationType::
+                                     kTfLiteTensorArrayQuantizationData)
+          .used_bytes,
+      kKeywordModelTfLiteTensorQuantizationDataSize);
+  EnsureAllocatedSizeThreshold(
+      allocator
+          .GetRecordedAllocation(
+              tflite::RecordedAllocationType::kNodeAndRegistrationArray)
+          .used_bytes,
+      sizeof(tflite::NodeAndRegistration) *
+          kKeywordModelNodeAndRegistrationCount);
+  EnsureAllocatedSizeThreshold(
+      allocator.GetRecordedAllocation(tflite::RecordedAllocationType::kOpData)
+          .used_bytes,
+      kKeywordModelOpRuntimeDataSize);
+}
+
+TF_LITE_MICRO_TESTS_END

tensorflow/lite/micro/micro_allocator.cc

@@ -718,7 +718,9 @@ TfLiteStatus MicroAllocator::PrepareNodeAndRegistrationDataFromFlatbuffer(
   return kTfLiteOk;
 }
 
-ErrorReporter* MicroAllocator::error_reporter() { return error_reporter_; }
+ErrorReporter* MicroAllocator::error_reporter() const {
+  return error_reporter_;
+}
 
 TfLiteStatus MicroAllocator::InitGraphAndContextTensorData(
     const Model* model, TfLiteContext* context, const SubGraph* subgraph) {

tensorflow/lite/micro/micro_allocator.h

@@ -172,7 +172,7 @@ class MicroAllocator {
       const MicroOpResolver& op_resolver,
       NodeAndRegistration* node_and_registrations);
 
-  ErrorReporter* error_reporter();
+  ErrorReporter* error_reporter() const;
 
  private:
   // Initializes the graph and allocates TfLiteContext tensor data.

tensorflow/lite/micro/micro_interpreter.h

@@ -159,6 +159,10 @@ class MicroInterpreter {
   // arena_used_bytes() + 16.
   size_t arena_used_bytes() const { return allocator_.used_bytes(); }
 
+ protected:
+  const MicroAllocator& allocator() const { return allocator_; }
+  const TfLiteContext& context() const { return context_; }
+
  private:
   // TODO(b/158263161): Consider switching to Create() function to enable better
   // error reporting during initialization.

tensorflow/lite/micro/recording_micro_allocator.cc

@@ -29,18 +29,23 @@ RecordingMicroAllocator::RecordingMicroAllocator(
       recording_memory_allocator_(recording_memory_allocator) {}
 
 RecordingMicroAllocator* RecordingMicroAllocator::Create(
-    RecordingSimpleMemoryAllocator* memory_allocator,
+    uint8_t* tensor_arena, size_t arena_size, ErrorReporter* error_reporter) {
-    ErrorReporter* error_reporter) {
+  TFLITE_DCHECK(error_reporter != nullptr);
-  TFLITE_DCHECK(memory_allocator != nullptr);
+
-  uint8_t* allocator_buffer = memory_allocator->AllocateFromTail(
+  RecordingSimpleMemoryAllocator* simple_memory_allocator =
+      RecordingSimpleMemoryAllocator::Create(error_reporter, tensor_arena,
+                                             arena_size);
+  TFLITE_DCHECK(simple_memory_allocator != nullptr);
+
+  uint8_t* allocator_buffer = simple_memory_allocator->AllocateFromTail(
       sizeof(RecordingMicroAllocator), alignof(RecordingMicroAllocator));
   RecordingMicroAllocator* allocator = new (allocator_buffer)
-      RecordingMicroAllocator(memory_allocator, error_reporter);
+      RecordingMicroAllocator(simple_memory_allocator, error_reporter);
   return allocator;
 }
 
 RecordedAllocation RecordingMicroAllocator::GetRecordedAllocation(
-    RecordedAllocationType allocation_type) {
+    RecordedAllocationType allocation_type) const {
   switch (allocation_type) {
     case RecordedAllocationType::kTfLiteTensorArray:
       return recorded_tflite_tensor_array_data_;
@@ -56,7 +61,12 @@ RecordedAllocation RecordingMicroAllocator::GetRecordedAllocation(
   return RecordedAllocation();
 }
 
-void RecordingMicroAllocator::PrintAllocations() {
+const RecordingSimpleMemoryAllocator*
+RecordingMicroAllocator::GetSimpleMemoryAllocator() const {
+  return recording_memory_allocator_;
+}
+
+void RecordingMicroAllocator::PrintAllocations() const {
   TF_LITE_REPORT_ERROR(
       error_reporter(),
       "[RecordingMicroAllocator] Arena allocation total %d bytes",
@@ -70,24 +80,27 @@ void RecordingMicroAllocator::PrintAllocations() {
       "[RecordingMicroAllocator] Arena allocation tail %d bytes",
       recording_memory_allocator_->GetTailUsedBytes());
   PrintRecordedAllocation(RecordedAllocationType::kTfLiteTensorArray,
-                          "TfLiteTensor struct allocation");
+                          "TfLiteTensor struct allocation", "tensors");
   PrintRecordedAllocation(
       RecordedAllocationType::kTfLiteTensorArrayQuantizationData,
-      "TfLiteTensor quantization data allocations");
+      "TfLiteTensor quantization data allocations", "allocations");
   PrintRecordedAllocation(RecordedAllocationType::kNodeAndRegistrationArray,
-                          "NodeAndRegistration struct allocation");
+                          "NodeAndRegistration struct allocation",
+                          "NodeAndRegistration structs");
   PrintRecordedAllocation(RecordedAllocationType::kOpData,
-                          "Operator runtime data allocation");
+                          "Operator runtime data allocation", "OpData structs");
 }
 
 void RecordingMicroAllocator::PrintRecordedAllocation(
-    RecordedAllocationType allocation_type, const char* allocation_name) {
+    RecordedAllocationType allocation_type, const char* allocation_name,
+    const char* allocation_description) const {
   RecordedAllocation allocation = GetRecordedAllocation(allocation_type);
-  TF_LITE_REPORT_ERROR(error_reporter(),
+  TF_LITE_REPORT_ERROR(
-                       "[RecordingMicroAllocator] '%s' used %d bytes "
+      error_reporter(),
-                       "(requested %d bytes %d times)",
+      "[RecordingMicroAllocator] '%s' used %d bytes with alignment overhead "
-                       allocation_name, allocation.used_bytes,
+      "(requested %d bytes for %d %s)",
-                       allocation.requested_bytes, allocation.count);
+      allocation_name, allocation.used_bytes, allocation.requested_bytes,
+      allocation.count, allocation_description);
 }
 
 TfLiteStatus RecordingMicroAllocator::AllocateTfLiteTensorArray(

tensorflow/lite/micro/recording_micro_allocator.h

@@ -51,17 +51,19 @@ typedef struct RecordedAllocation {
 // auditing memory usage or integration testing.
 class RecordingMicroAllocator : public MicroAllocator {
  public:
-  static RecordingMicroAllocator* Create(
+  static RecordingMicroAllocator* Create(uint8_t* tensor_arena,
-      RecordingSimpleMemoryAllocator* memory_allocator,
+                                         size_t arena_size,
-      ErrorReporter* error_reporter);
+                                         ErrorReporter* error_reporter);
 
   // Returns the recorded allocations information for a given allocation type.
   RecordedAllocation GetRecordedAllocation(
-      RecordedAllocationType allocation_type);
+      RecordedAllocationType allocation_type) const;
+
+  const RecordingSimpleMemoryAllocator* GetSimpleMemoryAllocator() const;
 
   // Logs out through the ErrorReporter all allocation recordings by type
   // defined in RecordedAllocationType.
-  void PrintAllocations();
+  void PrintAllocations() const;
 
  protected:
   TfLiteStatus AllocateTfLiteTensorArray(TfLiteContext* context,
@@ -85,9 +87,10 @@ class RecordingMicroAllocator : public MicroAllocator {
                           ErrorReporter* error_reporter);
 
   void PrintRecordedAllocation(RecordedAllocationType allocation_type,
-                               const char* allocation_name);
+                               const char* allocation_name,
+                               const char* allocation_description) const;
 
-  RecordingSimpleMemoryAllocator* recording_memory_allocator_;
+  const RecordingSimpleMemoryAllocator* recording_memory_allocator_;
 
   RecordedAllocation recorded_tflite_tensor_array_data_;
   RecordedAllocation recorded_tflite_tensor_array_quantization_data_;

tensorflow/lite/micro/recording_micro_allocator_test.cc

@@ -39,11 +39,9 @@ TF_LITE_MICRO_TEST(TestRecordsTfLiteTensorArrayData) {
   tflite::NodeAndRegistration* node_and_registration;
   const tflite::Model* model = tflite::GetModel(kTestConvModelData);
   uint8_t arena[kTestConvArenaSize];
-  tflite::RecordingSimpleMemoryAllocator memory_allocator(
-      micro_test::reporter, arena, kTestConvArenaSize);
 
   tflite::RecordingMicroAllocator* micro_allocator =
-      tflite::RecordingMicroAllocator::Create(&memory_allocator,
+      tflite::RecordingMicroAllocator::Create(arena, kTestConvArenaSize,
                                               micro_test::reporter);
   TF_LITE_MICRO_EXPECT_NE(nullptr, micro_allocator);
   TF_LITE_MICRO_EXPECT_GE(kTfLiteOk, micro_allocator->StartModelAllocation(
@@ -68,11 +66,9 @@ TF_LITE_MICRO_TEST(TestRecordsTensorArrayQuantizationData) {
   tflite::NodeAndRegistration* node_and_registration;
   const tflite::Model* model = tflite::GetModel(kTestConvModelData);
   uint8_t arena[kTestConvArenaSize];
-  tflite::RecordingSimpleMemoryAllocator memory_allocator(
-      micro_test::reporter, arena, kTestConvArenaSize);
 
   tflite::RecordingMicroAllocator* micro_allocator =
-      tflite::RecordingMicroAllocator::Create(&memory_allocator,
+      tflite::RecordingMicroAllocator::Create(arena, kTestConvArenaSize,
                                               micro_test::reporter);
   TF_LITE_MICRO_EXPECT_NE(nullptr, micro_allocator);
   TF_LITE_MICRO_EXPECT_GE(kTfLiteOk, micro_allocator->StartModelAllocation(
@@ -126,11 +122,9 @@ TF_LITE_MICRO_TEST(TestRecordsNodeAndRegistrationArrayData) {
   tflite::NodeAndRegistration* node_and_registration;
   const tflite::Model* model = tflite::GetModel(kTestConvModelData);
   uint8_t arena[kTestConvArenaSize];
-  tflite::RecordingSimpleMemoryAllocator memory_allocator(
-      micro_test::reporter, arena, kTestConvArenaSize);
 
   tflite::RecordingMicroAllocator* micro_allocator =
-      tflite::RecordingMicroAllocator::Create(&memory_allocator,
+      tflite::RecordingMicroAllocator::Create(arena, kTestConvArenaSize,
                                               micro_test::reporter);
   TF_LITE_MICRO_EXPECT_NE(nullptr, micro_allocator);
   TF_LITE_MICRO_EXPECT_GE(kTfLiteOk, micro_allocator->StartModelAllocation(

tensorflow/lite/micro/recording_micro_interpreter.h

@@ -0,0 +1,58 @@
+/* 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_MICRO_INTERPRETER_H_
+#define TENSORFLOW_LITE_MICRO_RECORDING_MICRO_INTERPRETER_H_
+
+#include "tensorflow/lite/micro/micro_interpreter.h"
+#include "tensorflow/lite/micro/recording_micro_allocator.h"
+
+namespace tflite {
+
+// Utility subclass that enables internal recordings of the MicroInterpreter.
+// This class should be used to audit and analyze memory arena usage for a given
+// model and interpreter.
+//
+// After construction and the first Invoke() or AllocateTensors() call - the
+// memory usage is recorded and available through the GetMicroAllocator()
+// function. See RecordingMicroAlloctor for more details on what is currently
+// recorded from arena allocations.
+//
+// It is recommended for users to increase the tensor arena size by at least 1kb
+// to ensure enough additional memory is available for internal recordings.
+class RecordingMicroInterpreter : public MicroInterpreter {
+ public:
+  RecordingMicroInterpreter(const Model* model,
+                            const MicroOpResolver* op_resolver,
+                            uint8_t* tensor_arena, size_t tensor_arena_size,
+                            ErrorReporter* error_reporter)
+      : MicroInterpreter(model, op_resolver,
+                         RecordingMicroAllocator::Create(
+                             tensor_arena, tensor_arena_size, error_reporter),
+                         error_reporter),
+        recording_micro_allocator_(
+            static_cast<const RecordingMicroAllocator&>(allocator())) {}
+
+  const RecordingMicroAllocator& GetMicroAllocator() const {
+    return recording_micro_allocator_;
+  }
+
+ private:
+  const RecordingMicroAllocator& recording_micro_allocator_;
+};
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_MICRO_RECORDING_MICRO_INTERPRETER_H_

tensorflow/lite/micro/recording_simple_memory_allocator.cc

@@ -15,6 +15,10 @@ limitations under the License.
 
 #include "tensorflow/lite/micro/recording_simple_memory_allocator.h"
 
+#include <new>
+
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
 namespace tflite {
 
 RecordingSimpleMemoryAllocator::RecordingSimpleMemoryAllocator(
@@ -26,6 +30,20 @@ RecordingSimpleMemoryAllocator::RecordingSimpleMemoryAllocator(
 
 RecordingSimpleMemoryAllocator::~RecordingSimpleMemoryAllocator() {}
 
+RecordingSimpleMemoryAllocator* RecordingSimpleMemoryAllocator::Create(
+    ErrorReporter* error_reporter, uint8_t* buffer_head, size_t buffer_size) {
+  TFLITE_DCHECK(error_reporter != nullptr);
+  TFLITE_DCHECK(buffer_head != nullptr);
+  RecordingSimpleMemoryAllocator tmp =
+      RecordingSimpleMemoryAllocator(error_reporter, buffer_head, buffer_size);
+
+  uint8_t* allocator_buffer =
+      tmp.AllocateFromTail(sizeof(RecordingSimpleMemoryAllocator),
+                           alignof(RecordingSimpleMemoryAllocator));
+  // Use the default copy constructor to populate internal states.
+  return new (allocator_buffer) RecordingSimpleMemoryAllocator(tmp);
+}
+
 size_t RecordingSimpleMemoryAllocator::GetRequestedBytes() const {
   return requested_bytes_;
 }

tensorflow/lite/micro/recording_simple_memory_allocator.h

@@ -32,6 +32,10 @@ class RecordingSimpleMemoryAllocator : public SimpleMemoryAllocator {
   // functions.
   ~RecordingSimpleMemoryAllocator() override;
 
+  static RecordingSimpleMemoryAllocator* Create(ErrorReporter* error_reporter,
+                                                uint8_t* buffer_head,
+                                                size_t buffer_size);
+
   // Returns the number of bytes requested from the head or tail.
   size_t GetRequestedBytes() const;
 

tensorflow/lite/micro/tools/make/Makefile

@@ -111,13 +111,14 @@ $(wildcard tensorflow/lite/micro/memory_planner/*test.cc)
 
 # TODO(b/152645559): move all benchmarks to benchmarks directory.
 MICROLITE_BENCHMARK_SRCS := \
-$(wildcard tensorflow/lite/micro/benchmarks/*.cc)
+$(wildcard tensorflow/lite/micro/benchmarks/*benchmark.cc)
 
 MICROLITE_TEST_HDRS := \
 $(wildcard tensorflow/lite/micro/testing/*.h)
 
 MICROLITE_CC_BASE_SRCS := \
 $(wildcard tensorflow/lite/micro/*.cc) \
+$(wildcard tensorflow/lite/micro/benchmarks/*model_data.cc) \
 $(wildcard tensorflow/lite/micro/kernels/*.cc) \
 $(wildcard tensorflow/lite/micro/memory_planner/*.cc) \
 $(wildcard tensorflow/lite/micro/testing/*model.cc) \
@@ -135,6 +136,7 @@ MICROLITE_CC_SRCS := $(filter-out $(MICROLITE_BENCHMARK_SRCS), $(MICROLITE_CC_SR
 
 MICROLITE_CC_HDRS := \
 $(wildcard tensorflow/lite/micro/*.h) \
+$(wildcard tensorflow/lite/micro/benchmarks/*model_data.h) \
 $(wildcard tensorflow/lite/micro/kernels/*.h) \
 $(wildcard tensorflow/lite/micro/memory_planner/*.h) \
 LICENSE \