Track persistent buffer allocations in the RecordingMicroAllocator class.

Persistent buffer allocations in the tail section are currently not tracked. This fix ensures that there is visibility into the size and number of these allocations.

PiperOrigin-RevId: 336275255
Change-Id: I70fc1b9ff8542845663fe6119a10c6b10866ff8e

tensorflow/lite/micro/memory_arena_threshold_test.cc

@@ -49,15 +49,16 @@ constexpr int kKeywordModelNodeAndRegistrationCount = 15;
 // Run this test with '--copt=-DTF_LITE_STATIC_MEMORY' to get optimized memory
 // runtime values:
 #ifdef TF_LITE_STATIC_MEMORY
-constexpr int kKeywordModelTotalSize = 14128;
-constexpr int kKeywordModelTailSize = 13456;
+constexpr int kKeywordModelTotalSize = 14160;
+constexpr int kKeywordModelTailSize = 13488;
 #else
-constexpr int kKeywordModelTotalSize = 14496;
-constexpr int kKeywordModelTailSize = 13824;
+constexpr int kKeywordModelTotalSize = 14512;
+constexpr int kKeywordModelTailSize = 13840;
 #endif
 constexpr int kKeywordModelHeadSize = 672;
 constexpr int kKeywordModelTfLiteTensorVariableBufferDataSize = 10240;
 constexpr int kKeywordModelOpRuntimeDataSize = 148;
+constexpr int kKeywordModelPersistentBufferDataSize = 532;
 
 constexpr int kTestConvModelArenaSize = 12 * 1024;
 uint8_t test_conv_tensor_arena[kTestConvModelArenaSize];
@@ -68,14 +69,15 @@ constexpr int kTestConvModelNodeAndRegistrationCount = 7;
 // NOTE: These values are measured on x86-64:
 // TODO(b/158651472): Consider auditing these values on non-64 bit systems.
 #ifdef TF_LITE_STATIC_MEMORY
-constexpr int kTestConvModelTotalSize = 9568;
-constexpr int kTestConvModelTailSize = 1824;
+constexpr int kTestConvModelTotalSize = 9584;
+constexpr int kTestConvModelTailSize = 1840;
 #else
-constexpr int kTestConvModelTotalSize = 9728;
-constexpr int kTestConvModelTailSize = 1984;
+constexpr int kTestConvModelTotalSize = 9760;
+constexpr int kTestConvModelTailSize = 2016;
 #endif
 constexpr int kTestConvModelHeadSize = 7744;
 constexpr int kTestConvModelOpRuntimeDataSize = 136;
+constexpr int kTestConvModelPersistentBufferDataSize = 648;
 
 struct ModelAllocationThresholds {
   size_t tensor_count = 0;
@@ -85,6 +87,7 @@ struct ModelAllocationThresholds {
   size_t tail_alloc_size = 0;
   size_t tensor_variable_buffer_data_size = 0;
   size_t op_runtime_data_size = 0;
+  size_t persistent_buffer_data = 0;
 };
 
 void EnsureAllocatedSizeThreshold(const char* allocation_type, size_t actual,
@@ -148,6 +151,13 @@ void ValidateModelAllocationThresholds(
                                      kPersistentTfLiteTensorQuantizationData)
           .used_bytes,
       0);
+  EnsureAllocatedSizeThreshold(
+      "PersistentBufferData",
+      allocator
+          .GetRecordedAllocation(
+              tflite::RecordedAllocationType::kPersistentBufferData)
+          .used_bytes,
+      thresholds.persistent_buffer_data);
   EnsureAllocatedSizeThreshold(
       "NodeAndRegistration",
       allocator
@@ -195,6 +205,7 @@ TF_LITE_MICRO_TEST(TestKeywordModelMemoryThreshold) {
   thresholds.tensor_variable_buffer_data_size =
       kKeywordModelTfLiteTensorVariableBufferDataSize;
   thresholds.op_runtime_data_size = kKeywordModelOpRuntimeDataSize;
+  thresholds.persistent_buffer_data = kKeywordModelPersistentBufferDataSize;
 
   ValidateModelAllocationThresholds(interpreter.GetMicroAllocator(),
                                     thresholds);
@@ -216,6 +227,7 @@ TF_LITE_MICRO_TEST(TestConvModelMemoryThreshold) {
   thresholds.head_alloc_size = kTestConvModelHeadSize;
   thresholds.tail_alloc_size = kTestConvModelTailSize;
   thresholds.op_runtime_data_size = kTestConvModelOpRuntimeDataSize;
+  thresholds.persistent_buffer_data = kTestConvModelPersistentBufferDataSize;
 
   ValidateModelAllocationThresholds(interpreter.GetMicroAllocator(),
                                     thresholds);

tensorflow/lite/micro/micro_allocator.h

@@ -165,7 +165,7 @@ class MicroAllocator {
   // Allocates persistent buffer which has the same life time as the allocator.
   // The memory is immediately available and is allocated from the tail of the
   // arena.
-  void* AllocatePersistentBuffer(size_t bytes);
+  virtual void* AllocatePersistentBuffer(size_t bytes);
 
   // Register a scratch buffer of size `bytes` for Node with `node_id`.
   // This method only requests a buffer with a given size to be used after a

tensorflow/lite/micro/recording_micro_allocator.cc

@@ -54,6 +54,8 @@ RecordedAllocation RecordingMicroAllocator::GetRecordedAllocation(
       return recorded_persistent_tflite_tensor_data_;
     case RecordedAllocationType::kPersistentTfLiteTensorQuantizationData:
       return recorded_persistent_tflite_tensor_quantization_data_;
+    case RecordedAllocationType::kPersistentBufferData:
+      return recorded_persistent_buffer_data_;
     case RecordedAllocationType::kTfLiteTensorVariableBufferData:
       return recorded_tflite_tensor_variable_buffer_data_;
     case RecordedAllocationType::kNodeAndRegistrationArray:
@@ -91,6 +93,8 @@ void RecordingMicroAllocator::PrintAllocations() const {
   PrintRecordedAllocation(
       RecordedAllocationType::kPersistentTfLiteTensorQuantizationData,
       "Persistent TfLiteTensor quantization data", "allocations");
+  PrintRecordedAllocation(RecordedAllocationType::kPersistentBufferData,
+                          "Persistent buffer data", "allocations");
   PrintRecordedAllocation(
       RecordedAllocationType::kTfLiteTensorVariableBufferData,
       "TfLiteTensor variable buffer data", "allocations");
@@ -101,6 +105,14 @@ void RecordingMicroAllocator::PrintAllocations() const {
                           "Operator runtime data", "OpData structs");
 }
 
+void* RecordingMicroAllocator::AllocatePersistentBuffer(size_t bytes) {
+  RecordedAllocation allocations = SnapshotAllocationUsage();
+  void* buffer = MicroAllocator::AllocatePersistentBuffer(bytes);
+  RecordAllocationUsage(allocations, recorded_persistent_buffer_data_);
+
+  return buffer;
+}
+
 void RecordingMicroAllocator::PrintRecordedAllocation(
     RecordedAllocationType allocation_type, const char* allocation_name,
     const char* allocation_description) const {

tensorflow/lite/micro/recording_micro_allocator.h

@@ -29,6 +29,7 @@ enum class RecordedAllocationType {
   kTfLiteEvalTensorData,
   kPersistentTfLiteTensorData,
   kPersistentTfLiteTensorQuantizationData,
+  kPersistentBufferData,
   kTfLiteTensorVariableBufferData,
   kNodeAndRegistrationArray,
   kOpData,
@@ -67,6 +68,8 @@ class RecordingMicroAllocator : public MicroAllocator {
   // defined in RecordedAllocationType.
   void PrintAllocations() const;
 
+  void* AllocatePersistentBuffer(size_t bytes) override;
+
  protected:
   TfLiteStatus AllocateNodeAndRegistrations(
       const Model* model,
@@ -109,6 +112,7 @@ class RecordingMicroAllocator : public MicroAllocator {
   RecordedAllocation recorded_tflite_eval_tensor_data_ = {};
   RecordedAllocation recorded_persistent_tflite_tensor_data_ = {};
   RecordedAllocation recorded_persistent_tflite_tensor_quantization_data_ = {};
+  RecordedAllocation recorded_persistent_buffer_data_ = {};
   RecordedAllocation recorded_tflite_tensor_variable_buffer_data_ = {};
   RecordedAllocation recorded_node_and_registration_array_data_ = {};
   RecordedAllocation recorded_op_data_ = {};

tensorflow/lite/micro/recording_micro_allocator_test.cc

@@ -239,6 +239,43 @@ TF_LITE_MICRO_TEST(TestRecordsPersistentTfLiteTensorQuantizationData) {
                           expected_requested_bytes);
 }
 
+TF_LITE_MICRO_TEST(TestRecordsPersistentBufferData) {
+  uint8_t arena[kTestConvArenaSize];
+
+  tflite::RecordingMicroAllocator* micro_allocator =
+      tflite::RecordingMicroAllocator::Create(arena, kTestConvArenaSize,
+                                              micro_test::reporter);
+  TF_LITE_MICRO_EXPECT_NE(micro_allocator, nullptr);
+  if (micro_allocator == nullptr) return 1;
+
+  void* buffer = micro_allocator->AllocatePersistentBuffer(/*bytes=*/100);
+  TF_LITE_MICRO_EXPECT_NE(buffer, nullptr);
+  if (buffer == nullptr) return 1;
+
+  tflite::RecordedAllocation recorded_allocation =
+      micro_allocator->GetRecordedAllocation(
+          tflite::RecordedAllocationType::kPersistentBufferData);
+
+  TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count, static_cast<size_t>(1));
+  TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
+                          static_cast<size_t>(100));
+  TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,
+                          static_cast<size_t>(100));
+
+  buffer = micro_allocator->AllocatePersistentBuffer(/*bytes=*/50);
+  TF_LITE_MICRO_EXPECT_NE(buffer, nullptr);
+  if (buffer == nullptr) return 1;
+
+  recorded_allocation = micro_allocator->GetRecordedAllocation(
+      tflite::RecordedAllocationType::kPersistentBufferData);
+
+  TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count, static_cast<size_t>(2));
+  TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
+                          static_cast<size_t>(150));
+  TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,
+                          static_cast<size_t>(150));
+}
+
 // TODO(b/158124094): Find a way to audit OpData allocations on
 // cross-architectures.