Introduce a "recording" MicroAllocator class.

This new class enables TFLM to measure, audit, and report memory usage in the shared tensor arena. Users may opt into this class by simply passing this class into a MicroInterpreter instance.

PiperOrigin-RevId: 314995667
Change-Id: I6a451944d55b0498a98f1cfd54244f9008e578d2

tensorflow/lite/micro/BUILD

@@ -178,11 +178,13 @@ cc_library(
 )
 
 cc_library(
-    name = "recording_simple_memory_allocator",
+    name = "recording_allocators",
     srcs = [
+        "recording_micro_allocator.cc",
         "recording_simple_memory_allocator.cc",
     ],
     hdrs = [
+        "recording_micro_allocator.h",
         "recording_simple_memory_allocator.h",
     ],
     build_for_embedded = True,
@@ -190,6 +192,7 @@ cc_library(
     deps = [
         ":micro_compatibility",
         ":micro_framework",
+        "//tensorflow/lite/core/api",
     ],
 )
 
@@ -248,7 +251,7 @@ tflite_micro_cc_test(
     ],
     deps = [
         ":micro_framework",
-        ":recording_simple_memory_allocator",
+        ":recording_allocators",
         "//tensorflow/lite/micro/testing:micro_test",
     ],
 )
@@ -264,6 +267,20 @@ tflite_micro_cc_test(
     ],
 )
 
+tflite_micro_cc_test(
+    name = "recording_micro_allocator_test",
+    srcs = [
+        "recording_micro_allocator_test.cc",
+    ],
+    deps = [
+        ":micro_framework",
+        ":op_resolvers",
+        ":recording_allocators",
+        "//tensorflow/lite/micro/testing:micro_test",
+        "//tensorflow/lite/micro/testing:test_conv_model",
+    ],
+)
+
 tflite_micro_cc_test(
     name = "memory_helpers_test",
     srcs = [

tensorflow/lite/micro/micro_allocator.cc

@@ -404,7 +404,10 @@ TfLiteStatus InitializeTfLiteTensorFromFlatbuffer(
 MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model,
                                uint8_t* tensor_arena, size_t arena_size,
                                ErrorReporter* error_reporter)
-    : model_(model), error_reporter_(error_reporter), context_(context) {
+    : model_(model),
+      context_(context),
+      error_reporter_(error_reporter),
+      active_(false) {
   uint8_t* aligned_arena = AlignPointerUp(tensor_arena, kBufferAlignment);
   if (aligned_arena != tensor_arena) {
     TF_LITE_REPORT_ERROR(
@@ -419,7 +422,20 @@ MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model,
   // destructed as it's the root allocator.
   memory_allocator_ = SimpleMemoryAllocator::Create(
       error_reporter, aligned_arena, aligned_arena_size);
+}
 
+MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model,
+                               SimpleMemoryAllocator* memory_allocator,
+                               ErrorReporter* error_reporter)
+    : memory_allocator_(memory_allocator),
+      model_(model),
+      context_(context),
+      error_reporter_(error_reporter),
+      active_(false) {}
+
+MicroAllocator::~MicroAllocator() {}
+
+TfLiteStatus MicroAllocator::Init() {
   TfLiteStatus status = InitGraphAndContextTensorData();
   // TODO(b/147871299): Consider improving this code. A better way of handling
   // failures in the constructor is to have a static function that returns a
@@ -431,9 +447,10 @@ MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model,
   } else {
     active_ = true;
   }
+  return status;
 }
 
-TfLiteStatus MicroAllocator::InitializeFromFlatbuffer(
+TfLiteStatus MicroAllocator::PrepareFromFlatbuffer(
     const MicroOpResolver& op_resolver,
     NodeAndRegistration** node_and_registrations) {
   if (!active_) {
@@ -580,12 +597,6 @@ size_t MicroAllocator::used_bytes() const {
   if (active_) {
     return 0;
   }
-  TF_LITE_REPORT_ERROR(error_reporter_, "Total buffer usage: %d bytes",
-                       memory_allocator_->GetUsedBytes());
-  TF_LITE_REPORT_ERROR(error_reporter_, "Head usage: %d bytes",
-                       memory_allocator_->GetHeadUsedBytes());
-  TF_LITE_REPORT_ERROR(error_reporter_, "Tail usage: %d bytes",
-                       memory_allocator_->GetTailUsedBytes());
   return memory_allocator_->GetUsedBytes();
 }
 
@@ -736,4 +747,14 @@ TfLiteStatus MicroAllocator::PrepareNodeAndRegistrationDataFromFlatbuffer(
   return kTfLiteOk;
 }  // namespace tflite
 
+size_t MicroAllocator::GetTensorsCount() const {
+  return context_->tensors_size;
+}
+
+size_t MicroAllocator::GetOperatorsCount() const {
+  return subgraph_->operators()->size();
+}
+
+ErrorReporter* MicroAllocator::error_reporter() { return error_reporter_; }
+
 }  // namespace tflite

tensorflow/lite/micro/micro_allocator.h

@@ -21,6 +21,7 @@ limitations under the License.
 #include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/micro/compatibility.h"
 #include "tensorflow/lite/micro/micro_op_resolver.h"
 #include "tensorflow/lite/micro/simple_memory_allocator.h"
 #include "tensorflow/lite/schema/schema_generated.h"
@@ -85,13 +86,18 @@ class MicroAllocator {
   MicroAllocator(TfLiteContext* context, const Model* model,
                  uint8_t* tensor_arena, size_t arena_size,
                  ErrorReporter* error_reporter);
+  virtual ~MicroAllocator();
+
+  // Initializes the allocator by allocating required internal structs required
+  // to prepare the model from the flatbuffer data in PrepareFromFlatbuffer.
+  TfLiteStatus Init();
 
   // Run through the model flatbuffer data (loaded from the TfLiteModel
   // instance) to allocate nodes and registrations. We need to keep them for the
   // entire life time of the model to allow persistent tensors. This method
   // needs to be called before FinishTensorAllocation method. This method also
   // allocates any internal Op data that is required from the flatbuffer.
-  TfLiteStatus InitializeFromFlatbuffer(
+  TfLiteStatus PrepareFromFlatbuffer(
       const MicroOpResolver& op_resolver,
       NodeAndRegistration** node_and_registrations);
 
@@ -123,39 +129,52 @@ class MicroAllocator {
   size_t used_bytes() const;
 
  protected:
+  MicroAllocator(TfLiteContext* context, const Model* model,
+                 SimpleMemoryAllocator* memory_allocator,
+                 ErrorReporter* error_reporter);
+
   // Allocates an array in the arena to hold pointers to the tensors required
   // to initialize and prepare a model. These allocations are stored and
   // populated on the context.
-  TfLiteStatus AllocateTfLiteTensorArray();
+  virtual TfLiteStatus AllocateTfLiteTensorArray();
 
   // Populates content on the list of tensor pointers required to initialize and
   // prepare a model from data in the flatbuffer (loaded from the TfLiteModel
   // instance). Persistent data (e.g. quantization params) is allocated from the
   // arena.
-  TfLiteStatus PopulateTfLiteTensorArrayFromFlatbuffer();
+  virtual TfLiteStatus PopulateTfLiteTensorArrayFromFlatbuffer();
 
   // Allocates an array in the arena to hold pointers to the node and
   // registration pointers required to represent the inference graph of the
   // model.
-  TfLiteStatus AllocateNodeAndRegistrations(
+  virtual TfLiteStatus AllocateNodeAndRegistrations(
       NodeAndRegistration** node_and_registrations);
 
   // Populates node and registration pointers representing the inference graph
   // of the model from values inside the flatbuffer (loaded from the TfLiteModel
   // instance). Persistent data (e.g. operator data) is allocated from the
   // arena.
-  TfLiteStatus PrepareNodeAndRegistrationDataFromFlatbuffer(
+  virtual TfLiteStatus PrepareNodeAndRegistrationDataFromFlatbuffer(
       const MicroOpResolver& op_resolver,
       NodeAndRegistration* node_and_registrations);
 
+  // Returns the number of tensors in the model subgraph.
+  size_t GetTensorsCount() const;
+
+  // Returns the number of operators in the model subgraph.
+  size_t GetOperatorsCount() const;
+
+  ErrorReporter* error_reporter();
+
  private:
   TfLiteStatus InitGraphAndContextTensorData();
 
-  const Model* model_;
   // A simple memory allocator that always allocate from the arena tail.
   SimpleMemoryAllocator* memory_allocator_;
-  ErrorReporter* error_reporter_;
+
+  const Model* model_;
   TfLiteContext* context_;
+  ErrorReporter* error_reporter_;
   // Indicating if the allocator is ready for allocation.
   bool active_ = false;
 
@@ -167,6 +186,8 @@ class MicroAllocator {
   size_t scratch_buffer_count_ = 0;
 
   const SubGraph* subgraph_;
+
+  TF_LITE_REMOVE_VIRTUAL_DELETE
 };
 
 }  // namespace tflite

tensorflow/lite/micro/micro_allocator_test.cc

@@ -151,6 +151,7 @@ TF_LITE_MICRO_TEST(TestFinishTensorAllocation) {
   uint8_t arena[arena_size];
   tflite::MicroAllocator allocator(&context, model, arena, arena_size,
                                    micro_test::reporter);
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.Init());
   TF_LITE_MICRO_EXPECT_EQ(4, context.tensors_size);
   // Memory planning hasn't been finalized, so the used bytes is unknown.
   TF_LITE_MICRO_EXPECT_EQ(0, allocator.used_bytes());
@@ -187,6 +188,7 @@ TF_LITE_MICRO_TEST(TestAllocationForModelsWithBranches) {
   uint8_t arena[arena_size];
   tflite::MicroAllocator allocator(&context, model, arena, arena_size,
                                    micro_test::reporter);
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.Init());
   TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.FinishTensorAllocation());
 
   uint8_t* start = context.tensors[0].data.uint8;
@@ -211,6 +213,7 @@ TF_LITE_MICRO_TEST(TestFinishComplexTensorAllocation) {
   uint8_t arena[arena_size];
   tflite::MicroAllocator allocator(&context, model, arena, arena_size,
                                    micro_test::reporter);
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.Init());
   TF_LITE_MICRO_EXPECT_EQ(10, context.tensors_size);
 
   TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.FinishTensorAllocation());
@@ -238,4 +241,15 @@ TF_LITE_MICRO_TEST(TestFinishComplexTensorAllocation) {
   tflite::testing::EnsureUniqueVariableTensorBuffer(&context, 7);
 }
 
+TF_LITE_MICRO_TEST(TestDoubleInitFails) {
+  const tflite::Model* model = tflite::testing::GetComplexMockModel();
+  TfLiteContext context;
+  constexpr size_t arena_size = 2048;
+  uint8_t arena[arena_size];
+  tflite::MicroAllocator allocator(&context, model, arena, arena_size,
+                                   micro_test::reporter);
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.Init());
+  TF_LITE_MICRO_EXPECT_EQ(10, context.tensors_size);
+}
+
 TF_LITE_MICRO_TESTS_END

tensorflow/lite/micro/micro_interpreter.cc

@@ -91,6 +91,12 @@ MicroInterpreter::MicroInterpreter(const Model* model,
     initialization_status_ = kTfLiteError;
     return;
   }
+  if (allocator_.Init() != kTfLiteOk) {
+    TF_LITE_REPORT_ERROR(error_reporter,
+                         "Failed to initialize the allocator.\n");
+    initialization_status_ = kTfLiteError;
+    return;
+  }
   subgraph_ = (*subgraphs)[0];
 
   context_.impl_ = static_cast<void*>(&context_helper_);
@@ -166,7 +172,7 @@ void MicroInterpreter::CorrectTensorDataEndianness(T* data, int32_t size) {
 }
 
 TfLiteStatus MicroInterpreter::AllocateTensors() {
-  TF_LITE_ENSURE_OK(&context_, allocator_.InitializeFromFlatbuffer(
+  TF_LITE_ENSURE_OK(&context_, allocator_.PrepareFromFlatbuffer(
                                    op_resolver_, &node_and_registrations_));
 
   // Only allow AllocatePersistentBuffer in Init stage.

tensorflow/lite/micro/recording_micro_allocator.cc

@@ -0,0 +1,149 @@
+/* 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_micro_allocator.h"
+
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/micro/compatibility.h"
+#include "tensorflow/lite/micro/recording_simple_memory_allocator.h"
+
+namespace tflite {
+
+RecordingMicroAllocator::RecordingMicroAllocator(
+    TfLiteContext* context, const Model* model,
+    RecordingSimpleMemoryAllocator* recording_memory_allocator,
+    ErrorReporter* error_reporter)
+    : MicroAllocator(context, model, recording_memory_allocator,
+                     error_reporter),
+      recording_memory_allocator_(recording_memory_allocator) {}
+
+RecordedAllocation RecordingMicroAllocator::GetRecordedAllocation(
+    RecordedAllocationType allocation_type) {
+  switch (allocation_type) {
+    case RecordedAllocationType::kTfLiteTensorArray:
+      return recorded_tflite_tensor_array_data_;
+    case RecordedAllocationType::kTfLiteTensorArrayQuantizationData:
+      return recorded_tflite_tensor_array_quantization_data_;
+    case RecordedAllocationType::kNodeAndRegistrationArray:
+      return recorded_node_and_registration_array_data_;
+    case RecordedAllocationType::kOpData:
+      return recorded_op_data_;
+  }
+  TF_LITE_REPORT_ERROR(error_reporter(), "Invalid allocation type supplied: %d",
+                       allocation_type);
+  return RecordedAllocation();
+}
+
+void RecordingMicroAllocator::PrintAllocations() {
+  TF_LITE_REPORT_ERROR(
+      error_reporter(),
+      "[RecordingMicroAllocator] Arena allocation total %d bytes",
+      recording_memory_allocator_->GetUsedBytes());
+  TF_LITE_REPORT_ERROR(
+      error_reporter(),
+      "[RecordingMicroAllocator] Arena allocation head %d bytes",
+      recording_memory_allocator_->GetHeadUsedBytes());
+  TF_LITE_REPORT_ERROR(
+      error_reporter(),
+      "[RecordingMicroAllocator] Arena allocation tail %d bytes",
+      recording_memory_allocator_->GetTailUsedBytes());
+  PrintRecordedAllocation(RecordedAllocationType::kTfLiteTensorArray,
+                          "TfLiteTensor struct allocation");
+  PrintRecordedAllocation(
+      RecordedAllocationType::kTfLiteTensorArrayQuantizationData,
+      "TfLiteTensor quantization data allocations");
+  PrintRecordedAllocation(RecordedAllocationType::kNodeAndRegistrationArray,
+                          "NodeAndRegistration struct allocation");
+  PrintRecordedAllocation(RecordedAllocationType::kOpData,
+                          "Operator runtime data allocation");
+}
+
+void RecordingMicroAllocator::PrintRecordedAllocation(
+    RecordedAllocationType allocation_type, const char* allocation_name) {
+  RecordedAllocation allocation = GetRecordedAllocation(allocation_type);
+  TF_LITE_REPORT_ERROR(error_reporter(),
+                       "[RecordingMicroAllocator] '%s' used %d bytes "
+                       "(requested %d bytes %d times)",
+                       allocation_name, allocation.used_bytes,
+                       allocation.requested_bytes, allocation.count);
+}
+
+TfLiteStatus RecordingMicroAllocator::AllocateTfLiteTensorArray() {
+  SnapshotAllocationUsage(recorded_tflite_tensor_array_data_);
+
+  TfLiteStatus status = MicroAllocator::AllocateTfLiteTensorArray();
+
+  RecordAllocationUsage(recorded_tflite_tensor_array_data_);
+  recorded_tflite_tensor_array_data_.count = GetTensorsCount();
+  return status;
+}
+
+TfLiteStatus
+RecordingMicroAllocator::PopulateTfLiteTensorArrayFromFlatbuffer() {
+  SnapshotAllocationUsage(recorded_tflite_tensor_array_quantization_data_);
+
+  TfLiteStatus status =
+      MicroAllocator::PopulateTfLiteTensorArrayFromFlatbuffer();
+
+  RecordAllocationUsage(recorded_tflite_tensor_array_quantization_data_);
+  return status;
+}
+
+TfLiteStatus RecordingMicroAllocator::AllocateNodeAndRegistrations(
+    NodeAndRegistration** node_and_registrations) {
+  SnapshotAllocationUsage(recorded_node_and_registration_array_data_);
+
+  TfLiteStatus status =
+      MicroAllocator::AllocateNodeAndRegistrations(node_and_registrations);
+
+  RecordAllocationUsage(recorded_node_and_registration_array_data_);
+  recorded_node_and_registration_array_data_.count = GetOperatorsCount();
+  return status;
+}
+
+TfLiteStatus
+RecordingMicroAllocator::PrepareNodeAndRegistrationDataFromFlatbuffer(
+    const MicroOpResolver& op_resolver,
+    NodeAndRegistration* node_and_registrations) {
+  SnapshotAllocationUsage(recorded_op_data_);
+
+  TfLiteStatus status =
+      MicroAllocator::PrepareNodeAndRegistrationDataFromFlatbuffer(
+          op_resolver, node_and_registrations);
+
+  RecordAllocationUsage(recorded_op_data_);
+  return status;
+}
+
+void RecordingMicroAllocator::SnapshotAllocationUsage(
+    RecordedAllocation& recorded_allocation) {
+  recorded_allocation.requested_bytes =
+      recording_memory_allocator_->GetRequestedBytes();
+  recorded_allocation.used_bytes = recording_memory_allocator_->GetUsedBytes();
+  recorded_allocation.count = recording_memory_allocator_->GetAllocatedCount();
+}
+
+void RecordingMicroAllocator::RecordAllocationUsage(
+    RecordedAllocation& recorded_allocation) {
+  recorded_allocation.requested_bytes =
+      recording_memory_allocator_->GetRequestedBytes() -
+      recorded_allocation.requested_bytes;
+  recorded_allocation.used_bytes = recording_memory_allocator_->GetUsedBytes() -
+                                   recorded_allocation.used_bytes;
+  recorded_allocation.count = recording_memory_allocator_->GetAllocatedCount() -
+                              recorded_allocation.count;
+}
+
+}  // namespace tflite

tensorflow/lite/micro/recording_micro_allocator.h

@@ -0,0 +1,92 @@
+/* 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_ALLOCATOR_H_
+#define TENSORFLOW_LITE_MICRO_RECORDING_MICRO_ALLOCATOR_H_
+
+#include "tensorflow/lite/micro/compatibility.h"
+#include "tensorflow/lite/micro/micro_allocator.h"
+#include "tensorflow/lite/micro/recording_simple_memory_allocator.h"
+
+namespace tflite {
+
+// List of buckets currently recorded by this class. Each type keeps a list of
+// allocated information during model initialization.
+enum class RecordedAllocationType {
+  kTfLiteTensorArray,
+  kTfLiteTensorArrayQuantizationData,
+  kNodeAndRegistrationArray,
+  kOpData
+};
+
+// Container for holding information about allocation recordings by a given
+// type. Each recording contains the number of bytes requested, the actual bytes
+// allocated (can defer from requested by alignment), and the number of items
+// allocated.
+typedef struct RecordedAllocation {
+  RecordedAllocation() : requested_bytes(0), used_bytes(0), count(0) {}
+  size_t requested_bytes;
+  size_t used_bytes;
+  size_t count;
+} RecordedAllocation;
+
+// Utility subclass of MicroAllocator that records all allocations
+// inside the arena. A summary of allocations can be logged through the
+// ErrorReporter by invoking LogAllocations(). Individual allocation recordings
+// can be retrieved by type through the GetRecordedAllocation() function. This
+// class should only be used for auditing memory usage or integration testing.
+class RecordingMicroAllocator : public MicroAllocator {
+ public:
+  RecordingMicroAllocator(TfLiteContext* context, const Model* model,
+                          RecordingSimpleMemoryAllocator* memory_allocator,
+                          ErrorReporter* error_reporter);
+
+  // Returns the recorded allocations information for a given allocation type.
+  RecordedAllocation GetRecordedAllocation(
+      RecordedAllocationType allocation_type);
+
+  // Logs out through the ErrorReporter all allocation recordings by type
+  // defined in RecordedAllocationType.
+  void PrintAllocations();
+
+ protected:
+  TfLiteStatus AllocateTfLiteTensorArray() override;
+  TfLiteStatus PopulateTfLiteTensorArrayFromFlatbuffer() override;
+  TfLiteStatus AllocateNodeAndRegistrations(
+      NodeAndRegistration** node_and_registrations) override;
+  TfLiteStatus PrepareNodeAndRegistrationDataFromFlatbuffer(
+      const MicroOpResolver& op_resolver,
+      NodeAndRegistration* node_and_registrations) override;
+
+  void SnapshotAllocationUsage(RecordedAllocation& recorded_allocation);
+  void RecordAllocationUsage(RecordedAllocation& recorded_allocation);
+
+ private:
+  void PrintRecordedAllocation(RecordedAllocationType allocation_type,
+                               const char* allocation_name);
+
+  RecordingSimpleMemoryAllocator* recording_memory_allocator_;
+
+  RecordedAllocation recorded_tflite_tensor_array_data_;
+  RecordedAllocation recorded_tflite_tensor_array_quantization_data_;
+  RecordedAllocation recorded_node_and_registration_array_data_;
+  RecordedAllocation recorded_op_data_;
+
+  TF_LITE_REMOVE_VIRTUAL_DELETE
+};
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_MICRO_RECORDING_MICRO_ALLOCATOR_H_

tensorflow/lite/micro/recording_micro_allocator_test.cc

@@ -0,0 +1,185 @@
+/* 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_micro_allocator.h"
+
+#include "tensorflow/lite/micro/all_ops_resolver.h"
+#include "tensorflow/lite/micro/test_helpers.h"
+#include "tensorflow/lite/micro/testing/micro_test.h"
+#include "tensorflow/lite/micro/testing/test_conv_model.h"
+
+#define TF_LITE_TENSOR_STRUCT_SIZE sizeof(TfLiteTensor)
+#define TF_LITE_AFFINE_QUANTIZATION_SIZE sizeof(TfLiteAffineQuantization)
+#define NODE_AND_REGISTRATION_STRUCT_SIZE sizeof(tflite::NodeAndRegistration)
+
+// TODO(b/158303868): Move tests into anonymous namespace.
+namespace {
+
+constexpr int kTestConvArenaSize = 1024 * 12;
+
+}  // namespace
+
+TF_LITE_MICRO_TESTS_BEGIN
+
+TF_LITE_MICRO_TEST(TestRecordedValuesDefaultToZero) {
+  TfLiteContext context;
+  const tflite::Model* model = tflite::testing::GetSimpleMockModel();
+  constexpr size_t arena_size = 1024;
+  uint8_t arena[arena_size];
+
+  tflite::RecordingSimpleMemoryAllocator memory_allocator(micro_test::reporter,
+                                                          arena, arena_size);
+  tflite::RecordingMicroAllocator micro_allocator(
+      &context, model, &memory_allocator, micro_test::reporter);
+
+  tflite::RecordedAllocation recorded_allocation;
+
+  recorded_allocation = micro_allocator.GetRecordedAllocation(
+      tflite::RecordedAllocationType::kTfLiteTensorArray);
+  TF_LITE_MICRO_EXPECT_EQ(0, recorded_allocation.requested_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(0, recorded_allocation.used_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(0, recorded_allocation.count);
+
+  recorded_allocation = micro_allocator.GetRecordedAllocation(
+      tflite::RecordedAllocationType::kTfLiteTensorArrayQuantizationData);
+  TF_LITE_MICRO_EXPECT_EQ(0, recorded_allocation.requested_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(0, recorded_allocation.used_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(0, recorded_allocation.count);
+
+  recorded_allocation = micro_allocator.GetRecordedAllocation(
+      tflite::RecordedAllocationType::kNodeAndRegistrationArray);
+  TF_LITE_MICRO_EXPECT_EQ(0, recorded_allocation.requested_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(0, recorded_allocation.used_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(0, recorded_allocation.count);
+
+  recorded_allocation = micro_allocator.GetRecordedAllocation(
+      tflite::RecordedAllocationType::kOpData);
+  TF_LITE_MICRO_EXPECT_EQ(0, recorded_allocation.requested_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(0, recorded_allocation.used_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(0, recorded_allocation.count);
+}
+
+TF_LITE_MICRO_TEST(TestRecordsTfLiteTensorArrayData) {
+  TfLiteContext context;
+  const tflite::Model* model = tflite::GetModel(kTestConvModelData);
+  uint8_t arena[kTestConvArenaSize];
+  tflite::RecordingSimpleMemoryAllocator memory_allocator(
+      micro_test::reporter, arena, kTestConvArenaSize);
+  tflite::RecordingMicroAllocator allocator(&context, model, &memory_allocator,
+                                            micro_test::reporter);
+  TfLiteStatus status = allocator.Init();
+
+  // TODO(b/158102673): Ugly workaround for not having fatal test assertions:
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, status);
+  if (status != kTfLiteOk) return 1;
+
+  tflite::RecordedAllocation recorded_allocation =
+      allocator.GetRecordedAllocation(
+          tflite::RecordedAllocationType::kTfLiteTensorArray);
+  TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count, context.tensors_size);
+  TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
+                          context.tensors_size * TF_LITE_TENSOR_STRUCT_SIZE);
+  TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,
+                          context.tensors_size * TF_LITE_TENSOR_STRUCT_SIZE);
+}
+
+TF_LITE_MICRO_TEST(TestRecordsTensorArrayQuantizationData) {
+  TfLiteContext context;
+  const tflite::Model* model = tflite::GetModel(kTestConvModelData);
+  uint8_t arena[kTestConvArenaSize];
+  tflite::RecordingSimpleMemoryAllocator memory_allocator(
+      micro_test::reporter, arena, kTestConvArenaSize);
+  tflite::RecordingMicroAllocator allocator(&context, model, &memory_allocator,
+                                            micro_test::reporter);
+  TfLiteStatus status = allocator.Init();
+
+  // TODO(b/158102673): Ugly workaround for not having fatal test assertions:
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, status);
+  if (status != kTfLiteOk) return 1;
+
+  // Walk the model subgraph to find all tensors with quantization params and
+  // keep a tally.
+  size_t quantized_tensor_count = 0;
+  size_t quantized_channel_bytes = 0;
+  for (size_t i = 0; i < context.tensors_size; ++i) {
+    const tflite::Tensor* cur_tensor =
+        model->subgraphs()->Get(0)->tensors()->Get(i);
+    const tflite::QuantizationParameters* quantization_params =
+        cur_tensor->quantization();
+    if (quantization_params && quantization_params->scale() &&
+        quantization_params->scale()->size() > 0 &&
+        quantization_params->zero_point() &&
+        quantization_params->zero_point()->size() > 0) {
+      quantized_tensor_count++;
+      size_t num_channels = quantization_params->scale()->size();
+      quantized_channel_bytes += TfLiteIntArrayGetSizeInBytes(num_channels);
+      quantized_channel_bytes += TfLiteFloatArrayGetSizeInBytes(num_channels);
+    }
+  }
+
+  // Calculate the expected allocation bytes with subgraph quantization data:
+  size_t expected_requested_bytes =
+      quantized_tensor_count * TF_LITE_AFFINE_QUANTIZATION_SIZE +
+      quantized_channel_bytes;
+
+  tflite::RecordedAllocation recorded_allocation =
+      allocator.GetRecordedAllocation(
+          tflite::RecordedAllocationType::kTfLiteTensorArrayQuantizationData);
+
+  // Each quantized tensors has 3 mallocs (quant struct, scale dimensions, zero
+  // point dimensions):
+  TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count,
+                          quantized_tensor_count * 3);
+  TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
+                          expected_requested_bytes);
+  TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,
+                          expected_requested_bytes);
+}
+
+TF_LITE_MICRO_TEST(TestRecordsNodeAndRegistrationArrayData) {
+  TfLiteContext context;
+  const tflite::Model* model = tflite::GetModel(kTestConvModelData);
+  uint8_t arena[kTestConvArenaSize];
+  tflite::RecordingSimpleMemoryAllocator memory_allocator(
+      micro_test::reporter, arena, kTestConvArenaSize);
+  tflite::RecordingMicroAllocator allocator(&context, model, &memory_allocator,
+                                            micro_test::reporter);
+  TfLiteStatus status = allocator.Init();
+
+  // TODO(b/158102673): Ugly workaround for not having fatal test assertions:
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, status);
+  if (status != kTfLiteOk) return 1;
+
+  tflite::AllOpsResolver ops_resolver;
+  tflite::NodeAndRegistration* node_and_registrations;
+  TF_LITE_MICRO_EXPECT_EQ(
+      kTfLiteOk,
+      allocator.PrepareFromFlatbuffer(ops_resolver, &node_and_registrations));
+
+  size_t num_ops = model->subgraphs()->Get(0)->operators()->size();
+  tflite::RecordedAllocation recorded_allocation =
+      allocator.GetRecordedAllocation(
+          tflite::RecordedAllocationType::kNodeAndRegistrationArray);
+  TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count, num_ops);
+  TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
+                          num_ops * NODE_AND_REGISTRATION_STRUCT_SIZE);
+  TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,
+                          num_ops * NODE_AND_REGISTRATION_STRUCT_SIZE);
+}
+
+// TODO(b/158124094): Find a way to audit OpData allocations on
+// cross-architectures.
+
+TF_LITE_MICRO_TESTS_END

tensorflow/lite/micro/tools/make/Makefile

@@ -116,6 +116,7 @@ MICROLITE_CC_BASE_SRCS := \
 $(wildcard tensorflow/lite/micro/*.cc) \
 $(wildcard tensorflow/lite/micro/kernels/*.cc) \
 $(wildcard tensorflow/lite/micro/memory_planner/*.cc) \
+$(wildcard tensorflow/lite/micro/testing/*model.cc) \
 tensorflow/lite/c/common.c \
 tensorflow/lite/core/api/error_reporter.cc \
 tensorflow/lite/core/api/flatbuffer_conversions.cc \

tensorflow/lite/micro/tools/make/targets/stm32f4_makefile.inc

@@ -58,10 +58,12 @@ ifeq ($(TARGET), stm32f4)
   MICROLITE_CC_SRCS := $(filter-out $(EXCLUDED_SRCS), $(MICROLITE_CC_SRCS))
   TEST_SCRIPT := tensorflow/lite/micro/testing/test_stm32f4_binary.sh
   # TODO, non working tests.. the micro_speech example partly works
+  # TODO(b/158324045): Examine why some tests fail here.
   EXCLUDED_TESTS := \
     tensorflow/lite/micro/micro_interpreter_test.cc \
     tensorflow/lite/micro/micro_allocator_test.cc \
     tensorflow/lite/micro/memory_helpers_test.cc \
+    tensorflow/lite/micro/recording_micro_allocator_test.cc \
     tensorflow/lite/micro/kernels/logistic_test.cc \
     tensorflow/lite/micro/kernels/logical_test.cc \
     tensorflow/lite/micro/kernels/maximum_minimum_test.cc \