Test for missing large tail allocations in the "recording" MicroAllocator.
This change ensures that any large tail allocation should be handled and recorded with the RecordingMicroAllocator. The "misc" category for the keyword model currently is at 1004 bytes on my machine. I'll start with exactly 1kb ceiling for now.
PiperOrigin-RevId: 316171598
Change-Id: I1fc7a38b7bd89266de9d386648c5ef8dc290899c
diff --git a/tensorflow/lite/micro/memory_arena_threshold_test.cc b/tensorflow/lite/micro/memory_arena_threshold_test.cc
index 37d323c..19c3d0f 100644
--- a/tensorflow/lite/micro/memory_arena_threshold_test.cc
+++ b/tensorflow/lite/micro/memory_arena_threshold_test.cc
@@ -30,6 +30,7 @@
// Ensure memory doesn't expand more that 3%:
constexpr float kAllocationThreshold = 0.03;
+constexpr float kAllocationTailMiscCeiling = 1024;
const bool kIs64BitSystem = sizeof(void*) == 8;
constexpr int kKeywordModelTensorArenaSize = 22 * 1024;
@@ -127,6 +128,17 @@
allocator.GetRecordedAllocation(tflite::RecordedAllocationType::kOpData)
.used_bytes,
thresholds.op_runtime_data_size);
+
+ // Ensure tail allocation recording is not missing any large chunks:
+ size_t tail_est_length = sizeof(TfLiteTensor) * thresholds.tensor_count +
+ thresholds.tensor_quantization_data_size +
+ thresholds.tensor_variable_buffer_data_size +
+ sizeof(tflite::NodeAndRegistration) *
+ thresholds.node_and_registration_count +
+ thresholds.op_runtime_data_size;
+
+ TF_LITE_MICRO_EXPECT_LE(thresholds.tail_alloc_size - tail_est_length,
+ kAllocationTailMiscCeiling);
}
} // namespace
