Automated sync from github.com/tensorflow/tensorflow (#71)
diff --git a/tensorflow/lite/micro/BUILD b/tensorflow/lite/micro/BUILD
index dc06fb5..3baec20 100644
--- a/tensorflow/lite/micro/BUILD
+++ b/tensorflow/lite/micro/BUILD
@@ -24,28 +24,69 @@
)
cc_library(
+ # TODO(b/187093492): Rename to micro_interpreter.
name = "micro_framework",
srcs = [
- "micro_allocator.cc",
"micro_interpreter.cc",
+ ],
+ hdrs = [
+ "micro_interpreter.h",
+ ],
+ copts = micro_copts(),
+ deps = [
+ ":memory_helpers",
+ ":micro_allocator",
+ ":micro_error_reporter",
+ ":micro_graph",
+ ":micro_profiler",
+ ":op_resolvers",
+ "//tensorflow/lite:type_to_tflitetype",
+ "//tensorflow/lite/c:common",
+ "//tensorflow/lite/core/api",
+ "//tensorflow/lite/core/api:error_reporter",
+ "//tensorflow/lite/kernels/internal:tensor",
+ "//tensorflow/lite/schema:schema_fbs",
+ "//tensorflow/lite/schema:schema_utils",
+ "@flatbuffers//:runtime_cc",
+ ],
+)
+
+cc_library(
+ name = "micro_graph",
+ srcs = ["micro_graph.cc"],
+ hdrs = ["micro_graph.h"],
+ deps = [
+ ":memory_helpers",
+ ":micro_allocator",
+ ":micro_error_reporter",
+ ":micro_profiler",
+ "//tensorflow/lite/c:common",
+ "//tensorflow/lite/kernels/internal:compatibility",
+ "//tensorflow/lite/schema:schema_fbs",
+ "@flatbuffers//:runtime_cc",
+ ],
+)
+
+cc_library(
+ name = "micro_allocator",
+ srcs = [
+ "micro_allocator.cc",
"simple_memory_allocator.cc",
],
hdrs = [
"micro_allocator.h",
- "micro_interpreter.h",
"simple_memory_allocator.h",
],
copts = micro_copts(),
deps = [
":memory_helpers",
":micro_compatibility",
- ":micro_profiler",
- ":op_resolvers",
- "//tensorflow/lite:type_to_tflitetype",
+ ":micro_error_reporter",
"//tensorflow/lite/c:common",
"//tensorflow/lite/core/api",
+ "//tensorflow/lite/core/api:error_reporter",
+ "//tensorflow/lite/core/api:op_resolver",
"//tensorflow/lite/kernels/internal:compatibility",
- "//tensorflow/lite/kernels/internal:tensor",
"//tensorflow/lite/micro/memory_planner",
"//tensorflow/lite/micro/memory_planner:greedy_memory_planner",
"//tensorflow/lite/schema:schema_fbs",
@@ -211,6 +252,7 @@
],
copts = micro_copts(),
deps = [
+ ":micro_allocator",
":micro_compatibility",
":micro_framework",
"//tensorflow/lite/core/api",
@@ -258,7 +300,9 @@
],
deps = [
":micro_compatibility",
+ ":micro_error_reporter",
":micro_framework",
+ ":micro_profiler",
":micro_utils",
":op_resolvers",
":recording_allocators",
@@ -300,8 +344,10 @@
],
deps = [
":memory_helpers",
- ":micro_framework",
+ ":micro_allocator",
+ ":micro_error_reporter",
":test_helpers",
+ "//tensorflow/lite/c:common",
"//tensorflow/lite/micro/testing:micro_test",
"//tensorflow/lite/micro/testing:test_conv_model",
],
@@ -313,7 +359,8 @@
"recording_micro_allocator_test.cc",
],
deps = [
- ":micro_framework",
+ ":micro_allocator",
+ ":micro_error_reporter",
":op_resolvers",
":recording_allocators",
":test_helpers",
@@ -385,6 +432,7 @@
"memory_arena_threshold_test.cc",
],
deps = [
+ ":micro_error_reporter",
":op_resolvers",
":recording_allocators",
"//tensorflow/lite/micro/benchmarks:keyword_scrambled_model_data",
diff --git a/tensorflow/lite/micro/memory_arena_threshold_test.cc b/tensorflow/lite/micro/memory_arena_threshold_test.cc
index c828210..a7daf2e 100644
--- a/tensorflow/lite/micro/memory_arena_threshold_test.cc
+++ b/tensorflow/lite/micro/memory_arena_threshold_test.cc
@@ -50,13 +50,13 @@
// Run this test with '--copt=-DTF_LITE_STATIC_MEMORY' to get optimized memory
// runtime values:
#ifdef TF_LITE_STATIC_MEMORY
-constexpr int kKeywordModelTotalSize = 14384;
-constexpr int kKeywordModelTailSize = 13712;
+constexpr int kKeywordModelTotalSize = 14416;
+constexpr int kKeywordModelTailSize = 13744;
constexpr int kKeywordModelPersistentTfLiteTensorDataSize = 128;
-constexpr int kKeywordModelPersistentBufferDataSize = 572;
+constexpr int kKeywordModelPersistentBufferDataSize = 564;
#else
-constexpr int kKeywordModelTotalSize = 14832;
-constexpr int kKeywordModelTailSize = 14160;
+constexpr int kKeywordModelTotalSize = 14992;
+constexpr int kKeywordModelTailSize = 14320;
constexpr int kKeywordModelPersistentTfLiteTensorDataSize = 224;
constexpr int kKeywordModelPersistentBufferDataSize = 564;
#endif
@@ -74,13 +74,13 @@
// 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 = 9744;
-constexpr int kTestConvModelTailSize = 2000;
+constexpr int kTestConvModelTotalSize = 9792;
+constexpr int kTestConvModelTailSize = 2048;
constexpr int kTestConvModelPersistentTfLiteTensorDataSize = 128;
-constexpr int kTestConvModelPersistentBufferDataSize = 672;
+constexpr int kTestConvModelPersistentBufferDataSize = 680;
#else
-constexpr int kTestConvModelTotalSize = 10016;
-constexpr int kTestConvModelTailSize = 2272;
+constexpr int kTestConvModelTotalSize = 10112;
+constexpr int kTestConvModelTailSize = 2368;
constexpr int kTestConvModelPersistentTfLiteTensorDataSize = 224;
constexpr int kTestConvModelPersistentBufferDataSize = 680;
#endif
@@ -177,11 +177,6 @@
.used_bytes,
sizeof(tflite::NodeAndRegistration) *
thresholds.node_and_registration_count);
- EnsureAllocatedSizeThreshold(
- "OpData",
- 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(TfLiteEvalTensor) * thresholds.tensor_count +
diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc
index 8655168..6bfd0ff 100644
--- a/tensorflow/lite/micro/micro_allocator.cc
+++ b/tensorflow/lite/micro/micro_allocator.cc
@@ -29,7 +29,7 @@
#include "tensorflow/lite/micro/memory_helpers.h"
#include "tensorflow/lite/micro/memory_planner/greedy_memory_planner.h"
#include "tensorflow/lite/micro/memory_planner/memory_planner.h"
-#include "tensorflow/lite/micro/micro_op_resolver.h"
+#include "tensorflow/lite/micro/micro_error_reporter.h"
#include "tensorflow/lite/micro/simple_memory_allocator.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/schema/schema_utils.h"
@@ -608,33 +608,46 @@
return allocator;
}
-TfLiteStatus MicroAllocator::StartModelAllocation(
- const Model* model, const MicroOpResolver& op_resolver,
- NodeAndRegistration** node_and_registrations,
- TfLiteEvalTensor** eval_tensors) {
+SubgraphAllocations* MicroAllocator::StartModelAllocation(const Model* model) {
TFLITE_DCHECK(model != nullptr);
if (model_is_allocating_) {
TF_LITE_REPORT_ERROR(error_reporter_,
"MicroAllocator: Model allocation started before "
"finishing previously allocated model");
- return kTfLiteError;
+ return nullptr;
}
model_is_allocating_ = true;
- TF_LITE_ENSURE_STATUS(InitScratchBufferData());
- TF_LITE_ENSURE_STATUS(AllocateTfLiteEvalTensors(model, eval_tensors));
- TF_LITE_ENSURE_STATUS(
- AllocateNodeAndRegistrations(model, node_and_registrations));
- TF_LITE_ENSURE_STATUS(PrepareNodeAndRegistrationDataFromFlatbuffer(
- model, op_resolver, *node_and_registrations));
+ uint8_t* data_allocator_buffer = memory_allocator_->AllocateFromTail(
+ sizeof(MicroBuiltinDataAllocator), alignof(MicroBuiltinDataAllocator));
+ builtin_data_allocator_ =
+ new (data_allocator_buffer) MicroBuiltinDataAllocator(memory_allocator_);
- return kTfLiteOk;
+ if (InitScratchBufferData() != kTfLiteOk) {
+ return nullptr;
+ }
+
+ // Allocate struct to store eval tensors, nodes and registrations.
+ SubgraphAllocations* output = reinterpret_cast<SubgraphAllocations*>(
+ memory_allocator_->AllocateFromTail(
+ sizeof(SubgraphAllocations) * model->subgraphs()->size(),
+ alignof(SubgraphAllocations)));
+ if (output == nullptr) {
+ MicroPrintf("Failed to allocate memory for model metadata.");
+ return nullptr;
+ }
+
+ if (AllocateTfLiteEvalTensors(model, output) != kTfLiteOk ||
+ AllocateNodeAndRegistrations(model, output) != kTfLiteOk) {
+ return nullptr;
+ }
+ return output;
}
TfLiteStatus MicroAllocator::FinishModelAllocation(
- const Model* model, TfLiteEvalTensor* eval_tensors,
+ const Model* model, SubgraphAllocations* subgraph_allocations,
ScratchBufferHandle** scratch_buffer_handles) {
if (!model_is_allocating_) {
TF_LITE_REPORT_ERROR(error_reporter_,
@@ -643,15 +656,20 @@
return kTfLiteError;
}
- const SubGraph* subgraph = GetSubGraphFromModel(model);
- TFLITE_DCHECK(subgraph != nullptr);
+ // TODO(b/187993197): Track scratch buffers for each subgraph.
+ for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs()->size();
+ subgraph_idx++) {
+ const SubGraph* subgraph = model->subgraphs()->Get(subgraph_idx);
+ TFLITE_DCHECK(subgraph != nullptr);
- TF_LITE_ENSURE_STATUS(AllocateScratchBufferHandles(
- scratch_buffer_handles, scratch_buffer_request_count_));
- TF_LITE_ENSURE_STATUS(CommitStaticMemoryPlan(model, subgraph, eval_tensors,
- *scratch_buffer_handles));
- TF_LITE_ENSURE_STATUS(AllocateVariables(subgraph, eval_tensors));
-
+ TF_LITE_ENSURE_STATUS(AllocateScratchBufferHandles(
+ scratch_buffer_handles, scratch_buffer_request_count_));
+ TF_LITE_ENSURE_STATUS(CommitStaticMemoryPlan(
+ model, subgraph_allocations[subgraph_idx].tensors,
+ *scratch_buffer_handles, subgraph_idx));
+ TF_LITE_ENSURE_STATUS(AllocateVariables(
+ subgraph, subgraph_allocations[subgraph_idx].tensors));
+ }
model_is_allocating_ = false;
return kTfLiteOk;
}
@@ -661,6 +679,7 @@
}
TfLiteStatus MicroAllocator::RequestScratchBufferInArena(size_t bytes,
+ int subgraph_idx,
int* buffer_idx) {
// All scratch buffer requests are stored in the head section of the arena
// when a model is in the prepare phase. First align a scratch buffer request
@@ -735,153 +754,66 @@
}
TfLiteStatus MicroAllocator::AllocateNodeAndRegistrations(
- const Model* model, NodeAndRegistration** node_and_registrations) {
- TFLITE_DCHECK(node_and_registrations);
+ const Model* model, SubgraphAllocations* subgraph_allocations) {
+ TFLITE_DCHECK(subgraph_allocations != nullptr);
- const SubGraph* subgraph = GetSubGraphFromModel(model);
- TFLITE_DCHECK(subgraph != nullptr);
+ for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs()->size();
+ subgraph_idx++) {
+ const SubGraph* subgraph = model->subgraphs()->Get(subgraph_idx);
+ TFLITE_DCHECK(subgraph != nullptr);
- NodeAndRegistration* output = reinterpret_cast<NodeAndRegistration*>(
- memory_allocator_->AllocateFromTail(
- sizeof(NodeAndRegistration) * subgraph->operators()->size(),
- alignof(NodeAndRegistration)));
- if (output == nullptr) {
- TF_LITE_REPORT_ERROR(
- error_reporter_,
- "Failed to allocate memory for node_and_registrations.");
- return kTfLiteError;
- }
- *node_and_registrations = output;
- return kTfLiteOk;
-}
-
-TfLiteStatus MicroAllocator::PrepareNodeAndRegistrationDataFromFlatbuffer(
- const Model* model, const MicroOpResolver& op_resolver,
- NodeAndRegistration* node_and_registrations) {
- TFLITE_DCHECK(model != nullptr);
- TFLITE_DCHECK(node_and_registrations != nullptr);
-
- const SubGraph* subgraph = GetSubGraphFromModel(model);
- TFLITE_DCHECK(subgraph != nullptr);
-
- TfLiteStatus status = kTfLiteOk;
- auto* opcodes = model->operator_codes();
- MicroBuiltinDataAllocator builtin_data_allocator(memory_allocator_);
- for (size_t i = 0; i < subgraph->operators()->size(); ++i) {
- const auto* op = subgraph->operators()->Get(i);
- const size_t index = op->opcode_index();
- if (index >= opcodes->size()) {
- TF_LITE_REPORT_ERROR(error_reporter_,
- "Missing registration for opcode_index %d\n", index);
+ // Initialize NodeAndRegistrations for the subgraph.
+ NodeAndRegistration* output = reinterpret_cast<NodeAndRegistration*>(
+ memory_allocator_->AllocateFromTail(
+ sizeof(NodeAndRegistration) * subgraph->operators()->size(),
+ alignof(NodeAndRegistration)));
+ if (output == nullptr) {
+ TF_LITE_REPORT_ERROR(
+ error_reporter_,
+ "Failed to allocate memory for node_and_registrations.");
return kTfLiteError;
}
- auto* opcode = (*opcodes)[index];
- status =
- GetRegistrationFromOpCode(opcode, op_resolver, error_reporter_,
- &(node_and_registrations[i].registration));
- if (status != kTfLiteOk) {
- TF_LITE_REPORT_ERROR(error_reporter_,
- "Failed to get registration from op code %s\n ",
- EnumNameBuiltinOperator(GetBuiltinCode(opcode)));
- return status;
- }
- const auto* registration = node_and_registrations[i].registration;
- if (registration == nullptr) {
- TF_LITE_REPORT_ERROR(error_reporter_, "Skipping op for opcode_index %d\n",
- index);
- return kTfLiteError;
- }
- BuiltinOperator op_type =
- static_cast<BuiltinOperator>(registration->builtin_code);
-
- const char* custom_data = nullptr;
- size_t custom_data_size = 0;
- unsigned char* builtin_data = nullptr;
-
- if (op_type == BuiltinOperator_CUSTOM) {
- // Custom Ops may or may not have a non-null custom_options field.
- if (op->custom_options() != nullptr) {
- custom_data =
- reinterpret_cast<const char*>(op->custom_options()->data());
- custom_data_size = op->custom_options()->size();
- }
- } else {
- if (op->custom_options() != nullptr) {
- TF_LITE_REPORT_ERROR(
- error_reporter_,
- "Unsupported behavior: found builtin operator %s with custom "
- "options.\n",
- EnumNameBuiltinOperator(op_type));
- return kTfLiteError;
- }
-
- MicroOpResolver::BuiltinParseFunction parser =
- op_resolver.GetOpDataParser(op_type);
- if (parser == nullptr) {
- TF_LITE_REPORT_ERROR(error_reporter_, "Did not find a parser for %s",
- EnumNameBuiltinOperator(op_type));
-
- return kTfLiteError;
- }
- TF_LITE_ENSURE_STATUS(parser(op, error_reporter_, &builtin_data_allocator,
- (void**)(&builtin_data)));
- }
-
- TfLiteIntArray* inputs_array;
- TF_LITE_ENSURE_STATUS(internal::FlatBufferVectorToTfLiteTypeArray(
- memory_allocator_, error_reporter_, op->inputs(), &inputs_array));
-
- TfLiteIntArray* outputs_array;
- TF_LITE_ENSURE_STATUS(internal::FlatBufferVectorToTfLiteTypeArray(
- memory_allocator_, error_reporter_, op->outputs(), &outputs_array));
-
- TfLiteNode* node = &(node_and_registrations[i].node);
- *node = {};
- node->inputs = inputs_array;
- node->outputs = outputs_array;
- node->builtin_data = reinterpret_cast<void*>(builtin_data);
- node->custom_initial_data = custom_data;
- node->custom_initial_data_size = custom_data_size;
+ subgraph_allocations[subgraph_idx].node_and_registrations = output;
}
-
return kTfLiteOk;
}
-
TfLiteTensor* MicroAllocator::AllocatePersistentTfLiteTensor(
- const Model* model, TfLiteEvalTensor* eval_tensors, int tensor_index) {
- const SubGraph* subgraph = GetSubGraphFromModel(model);
+ const Model* model, const SubgraphAllocations* subgraph_allocations,
+ int tensor_index, int subgraph_index) {
+ const SubGraph* subgraph = model->subgraphs()->Get(subgraph_index);
TFLITE_DCHECK(subgraph != nullptr);
// This value is allocated from persistent arena space. It is guaranteed to be
// around for the lifetime of the application.
- TfLiteTensor* tensor =
- AllocatePersistentTfLiteTensorInternal(model, eval_tensors, tensor_index);
+ TfLiteTensor* tensor = AllocatePersistentTfLiteTensorInternal();
// Populate any fields from the flatbuffer, since this TfLiteTensor struct is
// allocated in the persistent section of the arena, ensure that additional
// allocations also take place in that section of the arena.
- if (PopulateTfLiteTensorFromFlatbuffer(model, subgraph, tensor, tensor_index,
- /*allocate_temp=*/false) !=
- kTfLiteOk) {
+ if (PopulateTfLiteTensorFromFlatbuffer(
+ model, tensor, tensor_index, subgraph_index,
+ /*allocate_temp=*/false) != kTfLiteOk) {
TF_LITE_REPORT_ERROR(error_reporter_,
"Failed to populate a persistent TfLiteTensor struct "
"from flatbuffer data!");
return nullptr;
}
- if (eval_tensors != nullptr) {
+ if (subgraph_allocations != nullptr) {
// Tensor buffers that are allocated at runtime (e.g. non-weight buffers)
// and not located in the flatbuffer are stored on the pre-allocated list of
// TfLiteEvalTensors structs. These structs are the source of truth, simply
// point the corresponding buffer to the new TfLiteTensor data value.
- tensor->data.data = eval_tensors[tensor_index].data.data;
+ tensor->data.data =
+ subgraph_allocations[subgraph_index].tensors[tensor_index].data.data;
}
return tensor;
}
TfLiteTensor* MicroAllocator::AllocateTempTfLiteTensor(
- const Model* model, TfLiteEvalTensor* eval_tensors, int tensor_index) {
- const SubGraph* subgraph = GetSubGraphFromModel(model);
+ const Model* model, const SubgraphAllocations* subgraph_allocations,
+ int tensor_index, int subgraph_index) {
+ const SubGraph* subgraph = model->subgraphs()->Get(subgraph_index);
TFLITE_DCHECK(subgraph != nullptr);
// This value is allocated from temporary arena space. It is guaranteed to be
@@ -894,7 +826,8 @@
// Populate any fields from the flatbuffer, since this TfLiteTensor struct is
// allocated in the temp section of the arena, ensure that additional
// allocations also take place in that section of the arena.
- if (PopulateTfLiteTensorFromFlatbuffer(model, subgraph, tensor, tensor_index,
+ if (PopulateTfLiteTensorFromFlatbuffer(model, tensor, tensor_index,
+ subgraph_index,
/*allocate_temp=*/true) != kTfLiteOk) {
TF_LITE_REPORT_ERROR(
error_reporter_,
@@ -902,12 +835,13 @@
return nullptr;
}
- if (eval_tensors != nullptr) {
+ if (subgraph_allocations != nullptr) {
// Tensor buffers that are allocated at runtime (e.g. non-weight buffers)
// and not located in the flatbuffer are stored on the pre-allocated list of
// TfLiteEvalTensors structs. These structs are the source of truth, simply
// point the corresponding buffer to the new TfLiteTensor data value.
- tensor->data.data = eval_tensors[tensor_index].data.data;
+ tensor->data.data =
+ subgraph_allocations[subgraph_index].tensors[tensor_index].data.data;
}
return tensor;
}
@@ -917,38 +851,41 @@
}
TfLiteStatus MicroAllocator::AllocateTfLiteEvalTensors(
- const Model* model, TfLiteEvalTensor** eval_tensors) {
- TFLITE_DCHECK(eval_tensors != nullptr);
+ const Model* model, SubgraphAllocations* subgraph_allocations) {
+ TFLITE_DCHECK(subgraph_allocations != nullptr);
- const SubGraph* subgraph = GetSubGraphFromModel(model);
- TFLITE_DCHECK(subgraph != nullptr);
+ for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs()->size();
+ subgraph_idx++) {
+ const SubGraph* subgraph = model->subgraphs()->Get(subgraph_idx);
+ TFLITE_DCHECK(subgraph != nullptr);
- size_t alloc_count = subgraph->tensors()->size();
- TfLiteEvalTensor* tensors =
- reinterpret_cast<TfLiteEvalTensor*>(memory_allocator_->AllocateFromTail(
- sizeof(TfLiteEvalTensor) * alloc_count, alignof(TfLiteEvalTensor)));
- if (tensors == nullptr) {
- TF_LITE_REPORT_ERROR(error_reporter_,
- "Failed to allocate memory for context->eval_tensors, "
- "%d bytes required",
- sizeof(TfLiteEvalTensor) * alloc_count);
- return kTfLiteError;
- }
-
- for (size_t i = 0; i < alloc_count; ++i) {
- TfLiteStatus status = internal::InitializeTfLiteEvalTensorFromFlatbuffer(
- memory_allocator_, *subgraph->tensors()->Get(i), model->buffers(),
- error_reporter_, &tensors[i]);
- if (status != kTfLiteOk) {
- TF_LITE_REPORT_ERROR(error_reporter_, "Failed to initialize tensor %d",
- i);
+ size_t alloc_count = subgraph->tensors()->size();
+ TfLiteEvalTensor* tensors =
+ reinterpret_cast<TfLiteEvalTensor*>(memory_allocator_->AllocateFromTail(
+ sizeof(TfLiteEvalTensor) * alloc_count, alignof(TfLiteEvalTensor)));
+ if (tensors == nullptr) {
+ TF_LITE_REPORT_ERROR(
+ error_reporter_,
+ "Failed to allocate memory for context->eval_tensors, "
+ "%d bytes required",
+ sizeof(TfLiteEvalTensor) * alloc_count);
return kTfLiteError;
}
+
+ for (size_t i = 0; i < alloc_count; ++i) {
+ TfLiteStatus status = internal::InitializeTfLiteEvalTensorFromFlatbuffer(
+ memory_allocator_, *subgraph->tensors()->Get(i), model->buffers(),
+ error_reporter_, &tensors[i]);
+ if (status != kTfLiteOk) {
+ TF_LITE_REPORT_ERROR(error_reporter_, "Failed to initialize tensor %d",
+ i);
+ return kTfLiteError;
+ }
+ }
+ subgraph_allocations[subgraph_idx].tensors = tensors;
}
- *eval_tensors = tensors;
return kTfLiteOk;
}
-
TfLiteStatus MicroAllocator::AllocateVariables(const SubGraph* subgraph,
TfLiteEvalTensor* eval_tensors) {
for (size_t i = 0; i < subgraph->tensors()->size(); ++i) {
@@ -972,20 +909,20 @@
return kTfLiteOk;
}
-TfLiteTensor* MicroAllocator::AllocatePersistentTfLiteTensorInternal(
- const Model* model, TfLiteEvalTensor* eval_tensors, int tensor_index) {
+TfLiteTensor* MicroAllocator::AllocatePersistentTfLiteTensorInternal() {
return reinterpret_cast<TfLiteTensor*>(memory_allocator_->AllocateFromTail(
sizeof(TfLiteTensor), alignof(TfLiteTensor)));
}
TfLiteStatus MicroAllocator::PopulateTfLiteTensorFromFlatbuffer(
- const Model* model, const SubGraph* subgraph, TfLiteTensor* tensor,
- int tensor_index, bool allocate_temp) {
+ const Model* model, TfLiteTensor* tensor, int tensor_index,
+ int subgraph_idx, bool allocate_temp) {
// TODO(b/162311891): This method serves as a stub to ensure quantized
// allocations in the tail can be recorded. Once the interpreter has APIs for
// accessing buffers on TfLiteEvalTensor this method can be dropped.
return internal::InitializeTfLiteTensorFromFlatbuffer(
- memory_allocator_, allocate_temp, *subgraph->tensors()->Get(tensor_index),
+ memory_allocator_, allocate_temp,
+ *model->subgraphs()->Get(subgraph_idx)->tensors()->Get(tensor_index),
model->buffers(), error_reporter_, tensor);
}
@@ -993,20 +930,9 @@
return error_reporter_;
}
-const SubGraph* MicroAllocator::GetSubGraphFromModel(const Model* model) {
- auto* subgraphs = model->subgraphs();
- if (subgraphs->size() != 1) {
- TF_LITE_REPORT_ERROR(error_reporter_,
- "Only 1 subgraph is currently supported.\n");
- return nullptr;
- }
- return (*subgraphs)[0];
-}
-
TfLiteStatus MicroAllocator::CommitStaticMemoryPlan(
- const Model* model, const SubGraph* subgraph,
- TfLiteEvalTensor* eval_tensors,
- ScratchBufferHandle* scratch_buffer_handles) {
+ const Model* model, TfLiteEvalTensor* eval_tensors,
+ ScratchBufferHandle* scratch_buffer_handles, int subgraph_idx) {
size_t head_usage = 0;
// Create static memory plan
// 1. Calculate AllocationInfo to know the lifetime of each tensor/buffer.
@@ -1018,6 +944,7 @@
// allocated from the temp section and cleaned up at the bottom of this
// function.
+ const SubGraph* subgraph = model->subgraphs()->Get(subgraph_idx);
size_t allocation_info_count =
subgraph->tensors()->size() + scratch_buffer_request_count_;
size_t bytes = sizeof(AllocationInfo) * allocation_info_count;
@@ -1142,4 +1069,15 @@
alignof(internal::ScratchBufferRequest)));
}
+TfLiteStatus MicroAllocator::FlatBufferVectorToTfLiteTypeArray(
+ const flatbuffers::Vector<int32_t>* flatbuffer_array,
+ TfLiteIntArray** result) {
+ return internal::FlatBufferVectorToTfLiteTypeArray(
+ memory_allocator_, error_reporter_, flatbuffer_array, result);
+}
+
+BuiltinDataAllocator* MicroAllocator::GetBuiltinDataAllocator() {
+ return builtin_data_allocator_;
+}
+
} // namespace tflite
diff --git a/tensorflow/lite/micro/micro_allocator.h b/tensorflow/lite/micro/micro_allocator.h
index 39a12ea..d64b35f 100644
--- a/tensorflow/lite/micro/micro_allocator.h
+++ b/tensorflow/lite/micro/micro_allocator.h
@@ -21,8 +21,8 @@
#include "flatbuffers/flatbuffers.h" // from @flatbuffers
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/api/flatbuffer_conversions.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"
@@ -75,6 +75,13 @@
uint8_t* data;
} ScratchBufferHandle;
+// Stores all per-subgraph allocations. This includes the node and registration
+// array, tensor list and scratch buffer handles for each subgraph.
+typedef struct {
+ NodeAndRegistration* node_and_registrations;
+ TfLiteEvalTensor* tensors;
+} SubgraphAllocations;
+
// Allocator responsible for allocating memory for all intermediate tensors
// necessary to invoke a model.
//
@@ -114,28 +121,31 @@
static MicroAllocator* Create(SimpleMemoryAllocator* memory_allocator,
ErrorReporter* error_reporter);
- // Begin allocating internal resources required for model inference.
+ // Allocates internal resources required for model inference for each subgraph
+ // from the arena.
+ //
// This method will run through the flatbuffer data supplied in the model to
// properly allocate tensor, node, and op registration data. This method is
- // expected to be followed with a call to FinishModelAllocation() before
- // resuming allocation with another model. All persistent tensor buffers are
- // stored in the out-param eval_tensors. This value is allocated from the
- // persistent memory arena and will be used to host runtime tensor buffers.
- TfLiteStatus StartModelAllocation(
- const Model* model, const MicroOpResolver& op_resolver,
- NodeAndRegistration** node_and_registrations,
- TfLiteEvalTensor** eval_tensors);
+ // expected to be followed with a call to FinishModelAllocation() Returns a
+ // pointer to an array of SubgraphAllocations (also stored in the tail of the
+ // arena) where each index corresponds to a different subgraph in the model.
+ // Return value is nullptr if the allocations failed.
+ SubgraphAllocations* StartModelAllocation(const Model* model);
// Finish allocating internal resources required for model inference.
- // This method will plan non-persistent buffers and commit a memory plan to
- // the 'head' section of the memory arena. All variable tensor data will also
- // be allocated. This method should be called after assigning model resources
- // in StartModelAllocation(). The eval_tensors pointer should be the value
- // passed into this class during StartModelAllocation(). Scratch buffer
- // handles are stored in the out-param `scratch_buffer_handles`. This value
- // will be used in `GetScratchBuffer` call to retrieve scratch buffers.
+ //
+ // -Plan the memory for activation tensors and scratch buffers.
+ // -Update eval tensors for each subgraph based on planned offsets.
+ // -Allocate scratch buffer handles array and update based on planned offsets.
+ //
+ // This method should be called after assigning model resources
+ // in StartModelAllocation(). The subgraph_allocations pointer should be the
+ // value passed into this class during StartModelAllocation(). Scratch buffer
+ // handles are stored in the out-param `scratch_buffer_handles` array which is
+ // allocated in this method. This value will be used in `GetScratchBuffer`
+ // call to retrieve scratch buffers.
TfLiteStatus FinishModelAllocation(
- const Model* model, TfLiteEvalTensor* eval_tensors,
+ const Model* model, SubgraphAllocations* subgraph_allocations,
ScratchBufferHandle** scratch_buffer_handles);
// Allocates a TfLiteTensor struct and populates the returned value with
@@ -145,17 +155,19 @@
// class during StartModelAllocation() and contains the source-of-truth for
// buffers.
virtual TfLiteTensor* AllocatePersistentTfLiteTensor(
- const Model* model, TfLiteEvalTensor* eval_tensors, int tensor_index);
+ const Model* model, const SubgraphAllocations* subgraph_allocations,
+ int tensor_index, int subgraph_index);
// Allocates a TfLiteTensor struct and populates the returned value with
// properties from the model flatbuffer. This struct is allocated from
// temporary arena memory is only guaranteed until a call is made to
- // ResetTempAllocations(). The eval_tensors pointer should be the value passed
- // into this class during StartModelAllocation() and contains the
- // source-of-truth for buffers.
- virtual TfLiteTensor* AllocateTempTfLiteTensor(const Model* model,
- TfLiteEvalTensor* eval_tensors,
- int tensor_index);
+ // ResetTempAllocations(). Subgraph_allocaitons contains the array of
+ // TfLiteEvalTensors. If the newly allocated temp at the specified subgraph
+ // and tensor index is already present int the TfLiteEvalTensor array, its
+ // data buffer will be re-used.
+ virtual TfLiteTensor* AllocateTempTfLiteTensor(
+ const Model* model, const SubgraphAllocations* subgraph_allocations,
+ int tensor_index, int subgraph_index);
// Resets all temporary allocations. This method should be called after a
// chain of temp allocations (e.g. chain of TfLiteTensor objects via
@@ -171,7 +183,8 @@
// This method only requests a buffer with a given size to be used after a
// model has finished allocation via FinishModelAllocation(). All requested
// buffers will be accessible by the out-param in that method.
- TfLiteStatus RequestScratchBufferInArena(size_t bytes, int* buffer_idx);
+ TfLiteStatus RequestScratchBufferInArena(size_t bytes, int subgraph_idx,
+ int* buffer_idx);
// Finish allocating a specific NodeAndRegistration prepare block (kernel
// entry for a model) with a given node ID. This call ensures that any scratch
@@ -183,6 +196,14 @@
// `FinishModelAllocation`. Otherwise, it will return 0.
size_t used_bytes() const;
+ // Converts a flatbuffer int32_t array to a TfLiteIntArray, accounting for
+ // endiannes.
+ TfLiteStatus FlatBufferVectorToTfLiteTypeArray(
+ const flatbuffers::Vector<int32_t>* flatbuffer_array,
+ TfLiteIntArray** result);
+
+ BuiltinDataAllocator* GetBuiltinDataAllocator();
+
protected:
MicroAllocator(SimpleMemoryAllocator* memory_allocator,
ErrorReporter* error_reporter);
@@ -192,23 +213,13 @@
// registration pointers required to represent the inference graph of the
// model.
virtual TfLiteStatus AllocateNodeAndRegistrations(
- const Model* model, 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.
- virtual TfLiteStatus PrepareNodeAndRegistrationDataFromFlatbuffer(
- const Model* model, const MicroOpResolver& op_resolver,
- NodeAndRegistration* node_and_registrations);
+ const Model* model, SubgraphAllocations* subgraph_allocations);
// Allocates the list of persistent TfLiteEvalTensors that are used for the
// "eval" phase of model inference. These structs will be the source of truth
- // for all tensor buffers. Allocation results are stored in the out-param
- // eval_tensors.
+ // for all tensor buffers.
virtual TfLiteStatus AllocateTfLiteEvalTensors(
- const Model* model, TfLiteEvalTensor** eval_tensors);
-
+ const Model* model, SubgraphAllocations* subgraph_allocations);
// Allocates persistent tensor buffers for variable tensors in the subgraph.
virtual TfLiteStatus AllocateVariables(const SubGraph* subgraph,
TfLiteEvalTensor* eval_tensors);
@@ -216,21 +227,19 @@
// Allocate and return a persistent TfLiteTensor.
// TODO(b/162311891): Drop this method when the interpreter has an API for
// accessing TfLiteEvalTensor structs.
- virtual TfLiteTensor* AllocatePersistentTfLiteTensorInternal(
- const Model* model, TfLiteEvalTensor* eval_tensors, int tensor_index);
+ virtual TfLiteTensor* AllocatePersistentTfLiteTensorInternal();
// Populates a TfLiteTensor struct with data from the model flatbuffer. Any
// quantization data is allocated from either the tail (persistent) or temp
// sections of the arena based on the allocation flag.
- virtual TfLiteStatus PopulateTfLiteTensorFromFlatbuffer(
- const Model* model, const SubGraph* subgraph, TfLiteTensor* tensor,
- int tensor_index, bool allocate_temp);
+ virtual TfLiteStatus PopulateTfLiteTensorFromFlatbuffer(const Model* model,
+ TfLiteTensor* tensor,
+ int tensor_index,
+ int subgraph_idx,
+ bool allocate_temp);
ErrorReporter* error_reporter() const;
- // Returns the first subgraph from the model.
- const SubGraph* GetSubGraphFromModel(const Model* model);
-
private:
// Commits a memory plan for all non-persistent buffer allocations in the
// 'head' section of the memory arena. The eval_tensors pointer is the list of
@@ -240,9 +249,8 @@
// ScratchBufferHandle structs that will point to allocated buffers also in
// the head section.
virtual TfLiteStatus CommitStaticMemoryPlan(
- const Model* model, const SubGraph* subgraph,
- TfLiteEvalTensor* eval_tensors,
- ScratchBufferHandle* scratch_buffer_handles);
+ const Model* model, TfLiteEvalTensor* eval_tensors,
+ ScratchBufferHandle* scratch_buffer_handles, int subgraph_idx);
// Allocates an array of ScratchBufferHandle structs in the tail section for a
// given number of handles.
@@ -261,6 +269,9 @@
// A simple memory allocator that always allocate from the arena tail or head.
SimpleMemoryAllocator* memory_allocator_;
+ // Allocator used to allocate persistent builtin data.
+ BuiltinDataAllocator* builtin_data_allocator_;
+
ErrorReporter* error_reporter_;
bool model_is_allocating_;
diff --git a/tensorflow/lite/micro/micro_allocator_test.cc b/tensorflow/lite/micro/micro_allocator_test.cc
index 53bc55f..1c856ef 100644
--- a/tensorflow/lite/micro/micro_allocator_test.cc
+++ b/tensorflow/lite/micro/micro_allocator_test.cc
@@ -17,6 +17,7 @@
#include <cstdint>
+#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/micro/memory_helpers.h"
#include "tensorflow/lite/micro/micro_error_reporter.h"
#include "tensorflow/lite/micro/simple_memory_allocator.h"
@@ -82,19 +83,29 @@
}
void VerifyMockTensor(const Model* model, MicroAllocator* allocator,
- TfLiteEvalTensor* eval_tensors, int tensor_idx,
+ SubgraphAllocations* subgraph_allocations, int tensor_idx,
bool is_variable = false) {
- VerifyMockTfLiteTensor(allocator->AllocatePersistentTfLiteTensor(
- model, eval_tensors, tensor_idx),
- is_variable);
- VerifyMockTfLiteEvalTensor(&eval_tensors[tensor_idx]);
+ for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs()->size();
+ subgraph_idx++) {
+ VerifyMockTfLiteTensor(
+ allocator->AllocatePersistentTfLiteTensor(model, subgraph_allocations,
+ tensor_idx, subgraph_idx),
+ is_variable);
+ VerifyMockTfLiteEvalTensor(
+ &subgraph_allocations[subgraph_idx].tensors[tensor_idx]);
+ }
}
void VerifyMockWeightTensor(const Model* model, MicroAllocator* allocator,
- TfLiteEvalTensor* eval_tensors, int tensor_idx) {
- VerifyMockWeightTfLiteTensor(allocator->AllocatePersistentTfLiteTensor(
- model, eval_tensors, tensor_idx));
- VerifyMockWeightTfLiteEvalTensor(&eval_tensors[tensor_idx]);
+ SubgraphAllocations* subgraph_allocations,
+ int tensor_idx) {
+ for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs()->size();
+ subgraph_idx++) {
+ VerifyMockWeightTfLiteTensor(allocator->AllocatePersistentTfLiteTensor(
+ model, subgraph_allocations, tensor_idx, subgraph_idx));
+ VerifyMockWeightTfLiteEvalTensor(
+ &subgraph_allocations[subgraph_idx].tensors[tensor_idx]);
+ }
}
void EnsureUniqueVariableTensorBuffer(const Model* model,
@@ -109,11 +120,14 @@
}
void VerifyRegistrationAndNodeAllocation(
- NodeAndRegistration* node_and_registration, size_t count) {
- for (size_t i = 0; i < count; i++) {
- TF_LITE_MICRO_EXPECT_NE(nullptr, node_and_registration[i].registration);
- TF_LITE_MICRO_EXPECT_NE(nullptr, node_and_registration[i].node.inputs);
- TF_LITE_MICRO_EXPECT_NE(nullptr, node_and_registration[i].node.outputs);
+ SubgraphAllocations* subgraph_allocations, size_t count,
+ int num_subgraphs) {
+ for (int subgraph_idx = 0; subgraph_idx < num_subgraphs; subgraph_idx++) {
+ for (size_t i = 0; i < count; i++) {
+ TF_LITE_MICRO_EXPECT_NE(nullptr, &subgraph_allocations[subgraph_idx]
+ .node_and_registrations[i]
+ .registration);
+ }
}
}
@@ -236,30 +250,21 @@
TF_LITE_MICRO_TEST(TestFailsWhenModelStartsTwice) {
const tflite::Model* model = tflite::testing::GetSimpleMockModel();
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
- tflite::NodeAndRegistration* node_and_registration;
constexpr size_t arena_size = 1024;
uint8_t arena[arena_size];
tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create(
arena, arena_size, tflite::GetMicroErrorReporter());
TF_LITE_MICRO_EXPECT(nullptr != allocator);
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteError,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
+ TF_LITE_MICRO_EXPECT_NE(nullptr, allocator->StartModelAllocation(model));
+ TF_LITE_MICRO_EXPECT(nullptr == allocator->StartModelAllocation(model));
}
TF_LITE_MICRO_TEST(TestFailsWithWrongSequence) {
const tflite::Model* model = tflite::testing::GetSimpleMockModel();
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
- tflite::NodeAndRegistration* node_and_registration;
+ tflite::SubgraphAllocations* subgraph_allocations = nullptr;
constexpr size_t arena_size = 1024;
uint8_t arena[arena_size];
tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create(
@@ -268,47 +273,40 @@
// We can't finish allocation before it ever got started.
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteError, allocator->FinishModelAllocation(model, eval_tensors,
- &scratch_buffer_handles));
+ kTfLiteError, allocator->FinishModelAllocation(
+ model, subgraph_allocations, &scratch_buffer_handles));
// Start twice is not allowed.
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteError,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
+ TF_LITE_MICRO_EXPECT(nullptr != allocator->StartModelAllocation(model));
+ TF_LITE_MICRO_EXPECT(nullptr == allocator->StartModelAllocation(model));
}
TF_LITE_MICRO_TEST(TestMockModelAllocation) {
const tflite::Model* model = tflite::testing::GetSimpleMockModel();
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
- tflite::NodeAndRegistration* node_and_registration;
constexpr size_t arena_size = 1024;
uint8_t arena[arena_size];
tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create(
arena, arena_size, tflite::GetMicroErrorReporter());
TF_LITE_MICRO_EXPECT(nullptr != allocator);
+ tflite::SubgraphAllocations* subgraph_allocations =
+ allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator->FinishModelAllocation(model, eval_tensors,
+ kTfLiteOk, allocator->FinishModelAllocation(model, subgraph_allocations,
&scratch_buffer_handles));
size_t model_tensor_size = tflite::testing::GetModelTensorCount(model);
TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(4), model_tensor_size);
- tflite::testing::VerifyMockTensor(model, allocator, eval_tensors, 0);
- tflite::testing::VerifyMockWeightTensor(model, allocator, eval_tensors, 1);
- tflite::testing::VerifyMockTensor(model, allocator, eval_tensors, 2);
- tflite::testing::VerifyMockTensor(model, allocator, eval_tensors, 3);
+ tflite::testing::VerifyMockTensor(model, allocator, subgraph_allocations, 0);
+ tflite::testing::VerifyMockWeightTensor(model, allocator,
+ subgraph_allocations, 1);
+ tflite::testing::VerifyMockTensor(model, allocator, subgraph_allocations, 2);
+ tflite::testing::VerifyMockTensor(model, allocator, subgraph_allocations, 3);
+ TfLiteEvalTensor* eval_tensors = subgraph_allocations[0].tensors;
TF_LITE_MICRO_EXPECT_NE(eval_tensors[1].data.raw, eval_tensors[0].data.raw);
TF_LITE_MICRO_EXPECT_NE(eval_tensors[2].data.raw, eval_tensors[0].data.raw);
TF_LITE_MICRO_EXPECT_NE(eval_tensors[1].data.raw, eval_tensors[2].data.raw);
@@ -318,8 +316,9 @@
TF_LITE_MICRO_EXPECT_LE(allocator->used_bytes(), 856 + 100);
// SimpleMockModel has 2 operators:
- tflite::testing::VerifyRegistrationAndNodeAllocation(node_and_registration,
- /*count=*/2);
+ tflite::testing::VerifyRegistrationAndNodeAllocation(subgraph_allocations,
+ /*count=*/2,
+ /*num_subgraphs=*/1);
}
TF_LITE_MICRO_TEST(TestMultiTenantAllocation) {
@@ -333,31 +332,26 @@
tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create(
arena, arena_size, tflite::GetMicroErrorReporter());
TF_LITE_MICRO_EXPECT_NE(nullptr, allocator);
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
// Allocate for model 1. We use ComplexMockModel here to cover the code path
// allocatig variables.
const tflite::Model* model1 = tflite::testing::GetComplexMockModel();
- tflite::NodeAndRegistration* node_and_registration1;
+ tflite::SubgraphAllocations* subgraph_allocations1 =
+ allocator->StartModelAllocation(model1);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations1);
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model1, op_resolver,
- &node_and_registration1, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator->FinishModelAllocation(model1, eval_tensors,
+ kTfLiteOk, allocator->FinishModelAllocation(model1, subgraph_allocations1,
&scratch_buffer_handles));
const size_t single_model_used_bytes = allocator->used_bytes();
// Allocate for model 2.
const tflite::Model* model2 = tflite::testing::GetComplexMockModel();
- tflite::NodeAndRegistration* node_and_registration2;
+ tflite::SubgraphAllocations* subgraph_allocations2 =
+ allocator->StartModelAllocation(model2);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations2);
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model2, op_resolver,
- &node_and_registration2, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator->FinishModelAllocation(model2, eval_tensors,
+ kTfLiteOk, allocator->FinishModelAllocation(model2, subgraph_allocations2,
&scratch_buffer_handles));
// Allocation for two instances of the same model takes less memory as `head`
@@ -367,90 +361,97 @@
TF_LITE_MICRO_TEST(TestAllocationForModelsWithBranches) {
const tflite::Model* model = tflite::testing::GetSimpleModelWithBranch();
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
- tflite::NodeAndRegistration* node_and_registration;
constexpr size_t arena_size = 4096;
uint8_t arena[arena_size];
tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create(
arena, arena_size, tflite::GetMicroErrorReporter());
TF_LITE_MICRO_EXPECT_NE(nullptr, allocator);
+ tflite::SubgraphAllocations* subgraph_allocations =
+ allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator->FinishModelAllocation(model, eval_tensors,
+ kTfLiteOk, allocator->FinishModelAllocation(model, subgraph_allocations,
&scratch_buffer_handles));
- uint8_t* start = eval_tensors[0].data.uint8;
+ uint8_t* start = subgraph_allocations[0].tensors[0].data.uint8;
// Check test_helpers.cc BuildSimpleModelWithBranch for model structure.
// t0 is the first tensor, so place it in offset 0.
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[0].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[0].data.uint8 - start);
// bytes = 2 * 2 * 3 * sizeof(float32) = 48, same for other tensors.
size_t buffer_size;
- TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, tflite::TfLiteEvalTensorByteLength(
- &eval_tensors[0], &buffer_size));
+ TF_LITE_MICRO_EXPECT_EQ(
+ kTfLiteOk, tflite::TfLiteEvalTensorByteLength(
+ &subgraph_allocations[0].tensors[0], &buffer_size));
TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(48), buffer_size);
// t1 can't reuse any memory, as n0 requires both t0 and t1.
- TF_LITE_MICRO_EXPECT_EQ(96, eval_tensors[1].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 96, subgraph_allocations[0].tensors[1].data.uint8 - start);
// t2 can't reuse any memory, as n1 requires both t0 and t2. Also n2 requires
// both t1 and t2.
- TF_LITE_MICRO_EXPECT_EQ(48, eval_tensors[2].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 48, subgraph_allocations[0].tensors[2].data.uint8 - start);
// t3 reuses the same memory from t0 as t0 is not an input to any node.
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[3].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[3].data.uint8 - start);
// SimpleModelWithBranch has 3 operators:
- tflite::testing::VerifyRegistrationAndNodeAllocation(node_and_registration,
- /*count=*/3);
+ tflite::testing::VerifyRegistrationAndNodeAllocation(subgraph_allocations,
+ /*count=*/3,
+ /*num_subgraphs=*/1);
}
TF_LITE_MICRO_TEST(TestAllocationForComplexModelAllocation) {
const tflite::Model* model = tflite::testing::GetComplexMockModel();
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
- tflite::NodeAndRegistration* node_and_registration;
constexpr size_t arena_size = 2048;
uint8_t arena[arena_size];
tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create(
arena, arena_size, tflite::GetMicroErrorReporter());
TF_LITE_MICRO_EXPECT(nullptr != allocator);
+ tflite::SubgraphAllocations* subgraph_allocations =
+ allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator->FinishModelAllocation(model, eval_tensors,
+ kTfLiteOk, allocator->FinishModelAllocation(model, subgraph_allocations,
&scratch_buffer_handles));
size_t model_tensor_size = tflite::testing::GetModelTensorCount(model);
TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(10), model_tensor_size);
// NOTE: Tensor indexes match the values in GetComplexMockModel().
- tflite::testing::VerifyMockTensor(model, allocator, eval_tensors, 0);
- tflite::testing::VerifyMockTensor(model, allocator, eval_tensors, 1,
+ tflite::testing::VerifyMockTensor(model, allocator, subgraph_allocations, 0);
+ tflite::testing::VerifyMockTensor(model, allocator, subgraph_allocations, 1,
/*is_variable=*/true);
- tflite::testing::VerifyMockWeightTensor(model, allocator, eval_tensors, 2);
- tflite::testing::VerifyMockTensor(model, allocator, eval_tensors, 3);
- tflite::testing::VerifyMockTensor(model, allocator, eval_tensors, 4,
+ tflite::testing::VerifyMockWeightTensor(model, allocator,
+ subgraph_allocations, 2);
+ tflite::testing::VerifyMockTensor(model, allocator, subgraph_allocations, 3);
+ tflite::testing::VerifyMockTensor(model, allocator, subgraph_allocations, 4,
/*is_variable=*/true);
- tflite::testing::VerifyMockWeightTensor(model, allocator, eval_tensors, 5);
- tflite::testing::VerifyMockTensor(model, allocator, eval_tensors, 6);
- tflite::testing::VerifyMockTensor(model, allocator, eval_tensors, 7,
+ tflite::testing::VerifyMockWeightTensor(model, allocator,
+ subgraph_allocations, 5);
+ tflite::testing::VerifyMockTensor(model, allocator, subgraph_allocations, 6);
+ tflite::testing::VerifyMockTensor(model, allocator, subgraph_allocations, 7,
/*is_variable=*/true);
- tflite::testing::VerifyMockWeightTensor(model, allocator, eval_tensors, 8);
- tflite::testing::VerifyMockTensor(model, allocator, eval_tensors, 9);
+ tflite::testing::VerifyMockWeightTensor(model, allocator,
+ subgraph_allocations, 8);
+ tflite::testing::VerifyMockTensor(model, allocator, subgraph_allocations, 9);
// // Ensure that variable tensors have unique address
- tflite::testing::EnsureUniqueVariableTensorBuffer(model, eval_tensors, 1);
- tflite::testing::EnsureUniqueVariableTensorBuffer(model, eval_tensors, 4);
- tflite::testing::EnsureUniqueVariableTensorBuffer(model, eval_tensors, 7);
+ tflite::testing::EnsureUniqueVariableTensorBuffer(
+ model, subgraph_allocations[0].tensors, 1);
+ tflite::testing::EnsureUniqueVariableTensorBuffer(
+ model, subgraph_allocations[0].tensors, 4);
+ tflite::testing::EnsureUniqueVariableTensorBuffer(
+ model, subgraph_allocations[0].tensors, 7);
// ComplexMockModel has 3 operators:
- tflite::testing::VerifyRegistrationAndNodeAllocation(node_and_registration,
- /*count=*/3);
+ tflite::testing::VerifyRegistrationAndNodeAllocation(subgraph_allocations,
+ /*count=*/3,
+ /*num_subgraphs=*/1);
}
TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) {
@@ -458,7 +459,6 @@
int subgraph = 0;
constexpr int number_tensors = 4;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
- tflite::NodeAndRegistration* node_and_registration;
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
number_tensors] = {version, subgraph,
number_tensors, // header
@@ -484,7 +484,6 @@
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
number_tensors, metadata_buffer, node_list, number_connections);
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
constexpr size_t arena_size = 4096;
@@ -492,33 +491,36 @@
tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create(
arena, arena_size, tflite::GetMicroErrorReporter());
+ tflite::SubgraphAllocations* subgraph_allocations =
+ allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator->FinishModelAllocation(model, eval_tensors,
+ kTfLiteOk, allocator->FinishModelAllocation(model, subgraph_allocations,
&scratch_buffer_handles));
// Since all of the tensors are online planned and the model structure is
// identical to that in TestAllocationForModelsWithBranches,
// the offsets be should identical to that test.
- uint8_t* start = eval_tensors[0].data.uint8;
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[0].data.uint8 - start);
+ uint8_t* start = subgraph_allocations[0].tensors[0].data.uint8;
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[0].data.uint8 - start);
size_t buffer_size;
- TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, tflite::TfLiteEvalTensorByteLength(
- &eval_tensors[0], &buffer_size));
+ TF_LITE_MICRO_EXPECT_EQ(
+ kTfLiteOk, tflite::TfLiteEvalTensorByteLength(
+ &subgraph_allocations[0].tensors[0], &buffer_size));
TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(48), buffer_size);
- TF_LITE_MICRO_EXPECT_EQ(96, eval_tensors[1].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(48, eval_tensors[2].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[3].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 96, subgraph_allocations[0].tensors[1].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 48, subgraph_allocations[0].tensors[2].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[3].data.uint8 - start);
}
TF_LITE_MICRO_TEST(OfflinePlannerBasic) {
constexpr int number_tensors = 4;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
- tflite::NodeAndRegistration* node_and_registration;
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
number_tensors] = {1, 0, number_tensors,
/*t0=*/0,
@@ -537,32 +539,33 @@
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
number_tensors, metadata_buffer, node_list, number_connections);
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
constexpr size_t arena_size = 4096;
uint8_t arena[arena_size];
tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create(
arena, arena_size, tflite::GetMicroErrorReporter());
+ tflite::SubgraphAllocations* subgraph_allocations =
+ allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator->FinishModelAllocation(model, eval_tensors,
+ kTfLiteOk, allocator->FinishModelAllocation(model, subgraph_allocations,
&scratch_buffer_handles));
- uint8_t* start = eval_tensors[0].data.uint8;
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[0].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(48, eval_tensors[1].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[2].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(48, eval_tensors[3].data.uint8 - start);
+ uint8_t* start = subgraph_allocations[0].tensors[0].data.uint8;
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[0].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 48, subgraph_allocations[0].tensors[1].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[2].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 48, subgraph_allocations[0].tensors[3].data.uint8 - start);
}
TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) {
constexpr int number_tensors = 4;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
- tflite::NodeAndRegistration* node_and_registration;
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
number_tensors] = {/*version=*/1,
/*subgraph=*/0,
@@ -583,33 +586,34 @@
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
number_tensors, metadata_buffer, node_list, number_connections);
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
constexpr size_t arena_size = 4096;
uint8_t arena[arena_size];
tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create(
arena, arena_size, tflite::GetMicroErrorReporter());
+ tflite::SubgraphAllocations* subgraph_allocations =
+ allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator->FinishModelAllocation(model, eval_tensors,
+ kTfLiteOk, allocator->FinishModelAllocation(model, subgraph_allocations,
&scratch_buffer_handles));
- uint8_t* start = eval_tensors[0].data.uint8;
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[0].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[1].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(48, eval_tensors[2].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[3].data.uint8 - start);
+ uint8_t* start = subgraph_allocations[0].tensors[0].data.uint8;
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[0].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[1].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 48, subgraph_allocations[0].tensors[2].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[3].data.uint8 - start);
// TF_LITE_MICRO_EXPECT_EQ(static_cast<size_t>(48), context.tensors[0].bytes);
}
TF_LITE_MICRO_TEST(OfflinePlannerOfflineOnline) {
constexpr int number_tensors = 5;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
- tflite::NodeAndRegistration* node_and_registration;
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
number_tensors] = {/*version=*/1,
/*subgraph=*/0,
@@ -635,27 +639,30 @@
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
number_tensors, metadata_buffer, node_list, number_connections);
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
constexpr size_t arena_size = 4096;
uint8_t arena[arena_size];
tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create(
arena, arena_size, tflite::GetMicroErrorReporter());
+ tflite::SubgraphAllocations* subgraph_allocations =
+ allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator->FinishModelAllocation(model, eval_tensors,
+ kTfLiteOk, allocator->FinishModelAllocation(model, subgraph_allocations,
&scratch_buffer_handles));
- uint8_t* start = eval_tensors[0].data.uint8;
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[0].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(48, eval_tensors[1].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(96, eval_tensors[2].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(48, eval_tensors[4].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[3].data.uint8 - start);
+ uint8_t* start = subgraph_allocations[0].tensors[0].data.uint8;
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[0].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 48, subgraph_allocations[0].tensors[1].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 96, subgraph_allocations[0].tensors[2].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 48, subgraph_allocations[0].tensors[4].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[3].data.uint8 - start);
}
TF_LITE_MICRO_TEST(TestAllocatePersistentTfLiteTensor) {
@@ -667,13 +674,15 @@
TF_LITE_MICRO_EXPECT_NE(allocator, nullptr);
TfLiteTensor* tensor1 = allocator->AllocatePersistentTfLiteTensor(
- model, /*eval_tensors=*/nullptr, /*tensor_index=*/1);
+ model, /*subgraph_allocations=*/nullptr, /*tensor_index=*/1,
+ /*subgraph_index=*/0);
TF_LITE_MICRO_EXPECT_NE(tensor1, nullptr);
TF_LITE_MICRO_EXPECT_NE(tensor1->quantization.params, nullptr);
TF_LITE_MICRO_EXPECT_FALSE(tensor1->is_variable);
TfLiteTensor* tensor2 = allocator->AllocatePersistentTfLiteTensor(
- model, /*eval_tensors=*/nullptr, /*tensor_index=*/2);
+ model, /*subgraph_allocations=*/nullptr, /*tensor_index=*/2,
+ /*subgraph_index=*/0);
TF_LITE_MICRO_EXPECT_NE(tensor2, nullptr);
TF_LITE_MICRO_EXPECT_NE(tensor2->quantization.params, nullptr);
TF_LITE_MICRO_EXPECT_FALSE(tensor2->is_variable);
@@ -692,7 +701,8 @@
TF_LITE_MICRO_EXPECT_NE(allocator, nullptr);
TfLiteTensor* tensor1 = allocator->AllocateTempTfLiteTensor(
- model, /*eval_tensors=*/nullptr, /*tensor_index=*/1);
+ model, /*subgraph_allocations=*/nullptr, /*tensor_index=*/1,
+ /*subgraph_index=*/0);
TF_LITE_MICRO_EXPECT_NE(tensor1, nullptr);
}
@@ -705,11 +715,13 @@
TF_LITE_MICRO_EXPECT_NE(allocator, nullptr);
TfLiteTensor* tensor1 = allocator->AllocateTempTfLiteTensor(
- model, /*eval_tensors=*/nullptr, /*tensor_index=*/1);
+ model, /*subgraph_allocations=*/nullptr, /*tensor_index=*/1,
+ /*subgraph_index=*/0);
TF_LITE_MICRO_EXPECT_NE(tensor1, nullptr);
TfLiteTensor* tensor2 = allocator->AllocateTempTfLiteTensor(
- model, /*eval_tensors=*/nullptr, /*tensor_index=*/2);
+ model, /*subgraph_allocations=*/nullptr, /*tensor_index=*/2,
+ /*subgraph_index=*/0);
TF_LITE_MICRO_EXPECT_NE(tensor2, nullptr);
// The address of tensor2 should be higher than the address of tensor1
@@ -726,13 +738,15 @@
TF_LITE_MICRO_EXPECT(allocator != nullptr);
TfLiteTensor* tensor1 = allocator->AllocateTempTfLiteTensor(
- model, /*eval_tensors=*/nullptr, /*tensor_index=*/1);
+ model, /*subgraph_allocations=*/nullptr, /*tensor_index=*/1,
+ /*subgraph_index=*/0);
TF_LITE_MICRO_EXPECT(tensor1 != nullptr);
allocator->ResetTempAllocations();
TfLiteTensor* tensor2 = allocator->AllocateTempTfLiteTensor(
- model, /*eval_tensors=*/nullptr, /*tensor_index=*/2);
+ model, /*subgraph_allocations=*/nullptr, /*tensor_index=*/2,
+ /*subgraph_index=*/0);
TF_LITE_MICRO_EXPECT(tensor2 != nullptr);
// The address of tensor2 should be equal than the address of tensor1 since
@@ -743,7 +757,6 @@
TF_LITE_MICRO_TEST(TestOperatorInputsNotInSubgraphInputs) {
constexpr int number_tensors = 5;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
- tflite::NodeAndRegistration* node_and_registration;
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
number_tensors] = {/*version=*/1,
/*subgraph=*/0,
@@ -770,33 +783,35 @@
number_tensors, metadata_buffer, node_list, number_connections,
/*Only first tensor (t0) is in subgraph input list=*/1);
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
constexpr size_t arena_size = 4096;
uint8_t arena[arena_size];
tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create(
arena, arena_size, tflite::GetMicroErrorReporter());
+ tflite::SubgraphAllocations* subgraph_allocations =
+ allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator->FinishModelAllocation(model, eval_tensors,
+ kTfLiteOk, allocator->FinishModelAllocation(model, subgraph_allocations,
&scratch_buffer_handles));
- uint8_t* start = eval_tensors[0].data.uint8;
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[0].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[1].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[2].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(48, eval_tensors[3].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[4].data.uint8 - start);
+ uint8_t* start = subgraph_allocations[0].tensors[0].data.uint8;
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[0].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[1].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[2].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 48, subgraph_allocations[0].tensors[3].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[4].data.uint8 - start);
}
TF_LITE_MICRO_TEST(TestTypicalFirstOpAndSecondOpWithScratchTensors) {
constexpr int number_tensors = 6;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
- tflite::NodeAndRegistration* node_and_registration;
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
number_tensors] = {/*version=*/1,
/*subgraph=*/0,
@@ -826,28 +841,32 @@
number_tensors, metadata_buffer, node_list, number_connections,
/*Only first tensor (t0) is in subgraph input list=*/1);
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
constexpr size_t arena_size = 4096;
uint8_t arena[arena_size];
tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create(
arena, arena_size, tflite::GetMicroErrorReporter());
+ tflite::SubgraphAllocations* subgraph_allocations =
+ allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- allocator->StartModelAllocation(model, op_resolver,
- &node_and_registration, &eval_tensors));
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk, allocator->FinishModelAllocation(model, eval_tensors,
+ kTfLiteOk, allocator->FinishModelAllocation(model, subgraph_allocations,
&scratch_buffer_handles));
- uint8_t* start = eval_tensors[0].data.uint8;
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[0].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[1].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[2].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[3].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(48, eval_tensors[4].data.uint8 - start);
- TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[5].data.uint8 - start);
+ uint8_t* start = subgraph_allocations[0].tensors[0].data.uint8;
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[0].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[1].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[2].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[3].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 48, subgraph_allocations[0].tensors[4].data.uint8 - start);
+ TF_LITE_MICRO_EXPECT_EQ(
+ 0, subgraph_allocations[0].tensors[5].data.uint8 - start);
}
TF_LITE_MICRO_TESTS_END
diff --git a/tensorflow/lite/micro/micro_graph.cc b/tensorflow/lite/micro/micro_graph.cc
new file mode 100644
index 0000000..7080d9d
--- /dev/null
+++ b/tensorflow/lite/micro/micro_graph.cc
@@ -0,0 +1,247 @@
+/* Copyright 2021 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/micro_graph.h"
+
+#include "flatbuffers/flatbuffers.h" // from @flatbuffers
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/micro/memory_helpers.h"
+#include "tensorflow/lite/micro/micro_error_reporter.h"
+#include "tensorflow/lite/micro/micro_profiler.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+#ifndef TF_LITE_STRIP_ERROR_STRINGS
+const char* OpNameFromRegistration(const TfLiteRegistration* registration) {
+ if (registration->builtin_code == BuiltinOperator_CUSTOM) {
+ return registration->custom_name;
+ } else {
+ return EnumNameBuiltinOperator(BuiltinOperator(registration->builtin_code));
+ }
+}
+#endif // !defined(TF_LITE_STRIP_ERROR_STRINGS)
+
+} // namespace
+
+MicroGraph::MicroGraph(TfLiteContext* context, const Model* model,
+ MicroAllocator* allocator)
+ : context_(context),
+ model_(model),
+ allocator_(allocator),
+ current_subgraph_index_(0) {
+ if (model != nullptr) {
+ subgraphs_ = model->subgraphs();
+ }
+}
+
+MicroGraph::~MicroGraph() {}
+
+TfLiteStatus MicroGraph::InitSubgraphs() {
+ int previous_subgraph_idx = current_subgraph_index_;
+
+ for (size_t subgraph_idx = 0; subgraph_idx < subgraphs_->size();
+ subgraph_idx++) {
+ current_subgraph_index_ = subgraph_idx;
+
+ const SubGraph* subgraph = (*subgraphs_)[subgraph_idx];
+ for (size_t i = 0; i < subgraph->operators()->size(); ++i) {
+ TfLiteNode* node =
+ &(subgraph_allocations_[subgraph_idx].node_and_registrations[i].node);
+ const TfLiteRegistration* registration =
+ subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[i]
+ .registration;
+ size_t init_data_size;
+ const char* init_data;
+ if (registration->builtin_code == BuiltinOperator_CUSTOM) {
+ init_data = reinterpret_cast<const char*>(node->custom_initial_data);
+ init_data_size = node->custom_initial_data_size;
+ } else {
+ init_data = reinterpret_cast<const char*>(node->builtin_data);
+ init_data_size = 0;
+ }
+ if (registration->init) {
+ node->user_data =
+ registration->init(context_, init_data, init_data_size);
+ }
+ }
+ }
+ current_subgraph_index_ = previous_subgraph_idx;
+
+ return kTfLiteOk;
+}
+
+TfLiteStatus MicroGraph::PrepareSubgraphs() {
+ int previous_subgraph_idx = current_subgraph_index_;
+
+ for (size_t subgraph_idx = 0; subgraph_idx < subgraphs_->size();
+ subgraph_idx++) {
+ current_subgraph_index_ = subgraph_idx;
+
+ const SubGraph* subgraph = (*subgraphs_)[subgraph_idx];
+ for (size_t i = 0; i < subgraph->operators()->size(); ++i) {
+ TfLiteNode* node =
+ &(subgraph_allocations_[subgraph_idx].node_and_registrations[i].node);
+ const TfLiteRegistration* registration =
+ subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[i]
+ .registration;
+ if (registration->prepare != nullptr) {
+ TfLiteStatus prepare_status = registration->prepare(context_, node);
+ if (prepare_status != kTfLiteOk) {
+ MicroPrintf("Node %s (number %df) failed to prepare with status %d",
+ OpNameFromRegistration(registration), i, prepare_status);
+ return kTfLiteError;
+ }
+ }
+ allocator_->FinishPrepareNodeAllocations(/*node_id=*/i);
+ }
+ }
+ current_subgraph_index_ = previous_subgraph_idx;
+
+ return kTfLiteOk;
+}
+
+TfLiteStatus MicroGraph::FreeSubgraphs() {
+ int previous_subgraph_idx = current_subgraph_index_;
+
+ for (size_t subgraph_idx = 0; subgraph_idx < subgraphs_->size();
+ subgraph_idx++) {
+ current_subgraph_index_ = subgraph_idx;
+ const SubGraph* subgraph = (*subgraphs_)[subgraph_idx];
+ for (size_t i = 0; i < subgraph->operators()->size(); ++i) {
+ TfLiteNode* node =
+ &(subgraph_allocations_[subgraph_idx].node_and_registrations[i].node);
+ const TfLiteRegistration* registration =
+ subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[i]
+ .registration;
+ // registration is allocated outside the interpreter, so double check to
+ // make sure it's not nullptr;
+ if (registration != nullptr && registration->free != nullptr) {
+ registration->free(context_, node->user_data);
+ }
+ }
+ }
+ current_subgraph_index_ = previous_subgraph_idx;
+
+ return kTfLiteOk;
+}
+
+TfLiteStatus MicroGraph::InvokeSubgraph(int subgraph_idx) {
+ int previous_subgraph_idx = current_subgraph_index_;
+ current_subgraph_index_ = subgraph_idx;
+
+ if (static_cast<size_t>(subgraph_idx) >= subgraphs_->size()) {
+ MicroPrintf("Accessing subgraph %d but only %d subgraphs found",
+ subgraph_idx, subgraphs_->size());
+ return kTfLiteError;
+ }
+ const SubGraph* subgraph = (*subgraphs_)[subgraph_idx];
+
+ for (size_t i = 0; i < subgraph->operators()->size(); ++i) {
+ TfLiteNode* node =
+ &(subgraph_allocations_[subgraph_idx].node_and_registrations[i].node);
+ const TfLiteRegistration* registration = subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[i]
+ .registration;
+
+// This ifdef is needed (even though ScopedMicroProfiler itself is a no-op with
+// -DTF_LITE_STRIP_ERROR_STRINGS) because the function OpNameFromRegistration is
+// only defined for builds with the error strings.
+#if !defined(TF_LITE_STRIP_ERROR_STRINGS)
+ ScopedMicroProfiler scoped_profiler(
+ OpNameFromRegistration(registration),
+ reinterpret_cast<MicroProfiler*>(context_->profiler));
+#endif
+
+ TFLITE_DCHECK(registration->invoke);
+ TfLiteStatus invoke_status = registration->invoke(context_, node);
+
+ // All TfLiteTensor structs used in the kernel are allocated from temp
+ // memory in the allocator. This creates a chain of allocations in the
+ // temp section. The call below resets the chain of allocations to
+ // prepare for the next call.
+ allocator_->ResetTempAllocations();
+
+ if (invoke_status == kTfLiteError) {
+ MicroPrintf("Node %s (number %d) failed to invoke with status %d",
+ OpNameFromRegistration(registration), i, invoke_status);
+ return kTfLiteError;
+ } else if (invoke_status != kTfLiteOk) {
+ return invoke_status;
+ }
+ }
+ current_subgraph_index_ = previous_subgraph_idx;
+ return kTfLiteOk;
+}
+
+TfLiteStatus MicroGraph::ResetVariableTensors() {
+ for (size_t subgraph_idx = 0; subgraph_idx < subgraphs_->size();
+ subgraph_idx++) {
+ const SubGraph* subgraph = (*subgraphs_)[subgraph_idx];
+ for (size_t i = 0; i < subgraph->tensors()->size(); ++i) {
+ auto* tensor = subgraph->tensors()->Get(i);
+ if (tensor->is_variable()) {
+ size_t buffer_size;
+ TF_LITE_ENSURE_STATUS(TfLiteEvalTensorByteLength(
+ &subgraph_allocations_[subgraph_idx].tensors[i], &buffer_size));
+
+ int value = 0;
+ if (tensor->type() == tflite::TensorType_INT8) {
+ value = tensor->quantization()->zero_point()->Get(0);
+ }
+ memset(subgraph_allocations_[subgraph_idx].tensors[i].data.raw, value,
+ buffer_size);
+ }
+ }
+ }
+
+ return kTfLiteOk;
+}
+
+int MicroGraph::NumSubgraphs() { return model_->subgraphs()->size(); }
+
+void MicroGraph::SetSubgraphAllocations(
+ SubgraphAllocations* subgraph_allocations) {
+ subgraph_allocations_ = subgraph_allocations;
+}
+
+size_t MicroGraph::NumSubgraphInputs(int subgraph_idx) {
+ return model_->subgraphs()->Get(subgraph_idx)->inputs()->size();
+}
+
+TfLiteEvalTensor* MicroGraph::GetSubgraphInput(int subgraph_idx,
+ int input_idx) {
+ int tensor_idx =
+ model_->subgraphs()->Get(subgraph_idx)->inputs()->Get(input_idx);
+ return &subgraph_allocations_[subgraph_idx].tensors[tensor_idx];
+}
+
+size_t MicroGraph::NumSubgraphOutputs(int subgraph_idx) {
+ return model_->subgraphs()->Get(subgraph_idx)->outputs()->size();
+}
+
+TfLiteEvalTensor* MicroGraph::GetSubgraphOutput(int subgraph_idx,
+ int output_idx) {
+ int tensor_idx =
+ model_->subgraphs()->Get(subgraph_idx)->outputs()->Get(output_idx);
+ return &subgraph_allocations_[subgraph_idx].tensors[tensor_idx];
+}
+
+} // namespace tflite
diff --git a/tensorflow/lite/micro/micro_graph.h b/tensorflow/lite/micro/micro_graph.h
new file mode 100644
index 0000000..62316c7
--- /dev/null
+++ b/tensorflow/lite/micro/micro_graph.h
@@ -0,0 +1,97 @@
+/* Copyright 2021 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_MICRO_GRAPH_H_
+#define TENSORFLOW_LITE_MICRO_MICRO_GRAPH_H_
+
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/micro/micro_allocator.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+
+// Abstracts the details of interacting with the tflite::Model.
+//
+// Provides methods to access, initialize, prepare, invoke and free any
+// subgraph in the tflite::Graph.
+class MicroGraph {
+ public:
+ // The lifetime of the context, model, and allocator must be at least as long
+ // as that of the graph object, since the this class may need to access them
+ // at any time.
+ MicroGraph(TfLiteContext* context, const Model* model,
+ MicroAllocator* allocator);
+ virtual ~MicroGraph();
+
+ // Sets up builtin data and calls TfLiteRegistration->Init for every operator
+ // in every subgraph in the model.
+ virtual TfLiteStatus InitSubgraphs();
+
+ // Calls TfLiteRegistration->Prepare for every operator in every subgraph in
+ // the model.
+ virtual TfLiteStatus PrepareSubgraphs();
+
+ // Calls TfLiteRegistration->Free for every operator in every subgraph in the
+ // model.
+ virtual TfLiteStatus FreeSubgraphs();
+
+ // Calls TfLiteRegistration->Invoke for every operator in a single subgraph in
+ // the model.
+ virtual TfLiteStatus InvokeSubgraph(int subgraph_idx);
+
+ // Zeros out all variable tensors in all subgraphs in the model.
+ virtual TfLiteStatus ResetVariableTensors();
+
+ // Number of tensor inputs to a specified subgraph in the model.
+ virtual size_t NumSubgraphInputs(int subgraph_idx);
+
+ // Get the specified input tensor of a specified subgraph in the model.
+ virtual TfLiteEvalTensor* GetSubgraphInput(int subgraph_idx, int input_idx);
+
+ // Number of tensor outputs to a specified subgraph in the model.
+ virtual size_t NumSubgraphOutputs(int subgraph_idx);
+
+ // Get the specified output tensor of a specified subgraph in the model.
+ virtual TfLiteEvalTensor* GetSubgraphOutput(int subgraph_idx, int output_idx);
+
+ // Number of subgraphs in the model.
+ virtual int NumSubgraphs();
+
+ // Hook to pass in subgraph allocations tracked within the interpreter,
+ // allowing MicroGraph to init / prepare / invoke subgraphs in the model.
+ void SetSubgraphAllocations(SubgraphAllocations* subgraph_allocations);
+
+ // Get the current subgraph index. Within an on operator, this is guaranteed
+ // to be the subgraph of that operator.
+ int GetCurrentSubgraphIndex() { return current_subgraph_index_; }
+
+ // Gets the list of alloctions for each subgraph. This is the source of truth
+ // for all per-subgraph allocation data.
+ SubgraphAllocations* GetAllocations() { return subgraph_allocations_; }
+
+ private:
+ TfLiteContext* context_;
+ const Model* model_;
+ MicroAllocator* allocator_;
+ SubgraphAllocations* subgraph_allocations_ = nullptr;
+ int current_subgraph_index_;
+ const flatbuffers::Vector<flatbuffers::Offset<SubGraph>>* subgraphs_;
+
+ TF_LITE_REMOVE_VIRTUAL_DELETE
+};
+
+} // namespace tflite
+
+#endif // TENSORFLOW_LITE_MICRO_MICRO_GRAPH_H_
diff --git a/tensorflow/lite/micro/micro_interpreter.cc b/tensorflow/lite/micro/micro_interpreter.cc
index e05a35a..5b7ba0b 100644
--- a/tensorflow/lite/micro/micro_interpreter.cc
+++ b/tensorflow/lite/micro/micro_interpreter.cc
@@ -28,21 +28,9 @@
#include "tensorflow/lite/micro/micro_op_resolver.h"
#include "tensorflow/lite/micro/micro_profiler.h"
#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/schema/schema_utils.h"
namespace tflite {
-namespace {
-
-#ifndef TF_LITE_STRIP_ERROR_STRINGS
-const char* OpNameFromRegistration(const TfLiteRegistration* registration) {
- if (registration->builtin_code == BuiltinOperator_CUSTOM) {
- return registration->custom_name;
- } else {
- return EnumNameBuiltinOperator(BuiltinOperator(registration->builtin_code));
- }
-}
-#endif // !defined(TF_LITE_STRIP_ERROR_STRINGS)
-
-} // namespace
MicroInterpreter::MicroInterpreter(const Model* model,
const MicroOpResolver& op_resolver,
@@ -55,9 +43,10 @@
error_reporter_(error_reporter),
allocator_(*MicroAllocator::Create(tensor_arena, tensor_arena_size,
error_reporter)),
+
+ graph_(&context_, model, &allocator_),
tensors_allocated_(false),
initialization_status_(kTfLiteError),
- eval_tensors_(nullptr),
input_tensors_(nullptr),
output_tensors_(nullptr) {
Init(profiler);
@@ -72,103 +61,148 @@
op_resolver_(op_resolver),
error_reporter_(error_reporter),
allocator_(*allocator),
+ graph_(&context_, model, allocator),
tensors_allocated_(false),
initialization_status_(kTfLiteError),
- eval_tensors_(nullptr),
input_tensors_(nullptr),
output_tensors_(nullptr) {
Init(profiler);
}
MicroInterpreter::~MicroInterpreter() {
- if (node_and_registrations_ != nullptr) {
- for (size_t i = 0; i < subgraph_->operators()->size(); ++i) {
- TfLiteNode* node = &(node_and_registrations_[i].node);
- const TfLiteRegistration* registration =
- node_and_registrations_[i].registration;
- // registration is allocated outside the interpreter, so double check to
- // make sure it's not nullptr;
- if (registration != nullptr && registration->free != nullptr) {
- registration->free(&context_, node->user_data);
- }
- }
+ if (graph_.GetAllocations() != nullptr) {
+ graph_.FreeSubgraphs();
}
}
void MicroInterpreter::Init(MicroProfiler* profiler) {
- const flatbuffers::Vector<flatbuffers::Offset<SubGraph>>* subgraphs =
- model_->subgraphs();
- if (subgraphs->size() != 1) {
- TF_LITE_REPORT_ERROR(error_reporter_,
- "Only 1 subgraph is currently supported.\n");
- initialization_status_ = kTfLiteError;
- return;
- }
- subgraph_ = (*subgraphs)[0];
-
context_.impl_ = static_cast<void*>(this);
context_.ReportError = ReportOpError;
context_.GetTensor = GetTensor;
+ context_.ReportError = ReportOpError;
+ context_.GetTensor = GetTensor;
context_.GetEvalTensor = GetEvalTensor;
- context_.recommended_num_threads = 1;
context_.profiler = profiler;
initialization_status_ = kTfLiteOk;
}
+TfLiteStatus MicroInterpreter::PrepareNodeAndRegistrationDataFromFlatbuffer() {
+ for (int subgraph_idx = 0; subgraph_idx < graph_.NumSubgraphs();
+ subgraph_idx++) {
+ const SubGraph* subgraph = model_->subgraphs()->Get(subgraph_idx);
+ TFLITE_DCHECK(subgraph != nullptr);
+
+ auto* opcodes = model_->operator_codes();
+ BuiltinDataAllocator* builtin_data_allocator =
+ allocator_.GetBuiltinDataAllocator();
+ for (size_t i = 0; i < subgraph->operators()->size(); ++i) {
+ const auto* op = subgraph->operators()->Get(i);
+ const size_t index = op->opcode_index();
+ if (index >= opcodes->size()) {
+ MicroPrintf("Missing registration for opcode_index %d\n", index);
+ return kTfLiteError;
+ }
+ const auto* opcode = opcodes->Get(index);
+ TfLiteStatus status =
+ GetRegistrationFromOpCode(opcode, op_resolver_, error_reporter_,
+ &(graph_.GetAllocations()[subgraph_idx]
+ .node_and_registrations[i]
+ .registration));
+ if (status != kTfLiteOk) {
+ MicroPrintf("Failed to get registration from op code %s\n ",
+ EnumNameBuiltinOperator(GetBuiltinCode(opcode)));
+ return status;
+ }
+ const auto* registration = graph_.GetAllocations()[subgraph_idx]
+ .node_and_registrations[i]
+ .registration;
+ if (registration == nullptr) {
+ MicroPrintf("Skipping op for opcode_index %d\n", index);
+ return kTfLiteError;
+ }
+ BuiltinOperator op_type =
+ static_cast<BuiltinOperator>(registration->builtin_code);
+
+ const char* custom_data = nullptr;
+ size_t custom_data_size = 0;
+ unsigned char* builtin_data = nullptr;
+
+ if (op_type == BuiltinOperator_CUSTOM) {
+ // Custom Ops may or may not have a non-null custom_options field.
+ if (op->custom_options() != nullptr) {
+ custom_data =
+ reinterpret_cast<const char*>(op->custom_options()->data());
+ custom_data_size = op->custom_options()->size();
+ }
+ } else {
+ if (op->custom_options() != nullptr) {
+ MicroPrintf(
+ "Unsupported behavior: found builtin operator %s with custom "
+ "options.\n",
+ EnumNameBuiltinOperator(op_type));
+ return kTfLiteError;
+ }
+
+ MicroOpResolver::BuiltinParseFunction parser =
+ op_resolver_.GetOpDataParser(op_type);
+ if (parser == nullptr) {
+ MicroPrintf("Did not find a parser for %s",
+ EnumNameBuiltinOperator(op_type));
+
+ return kTfLiteError;
+ }
+ TF_LITE_ENSURE_STATUS(parser(op, error_reporter_,
+ builtin_data_allocator,
+ (void**)(&builtin_data)));
+ }
+
+ TfLiteIntArray* inputs_array;
+ TF_LITE_ENSURE_STATUS(allocator_.FlatBufferVectorToTfLiteTypeArray(
+ op->inputs(), &inputs_array));
+
+ TfLiteIntArray* outputs_array;
+ TF_LITE_ENSURE_STATUS(allocator_.FlatBufferVectorToTfLiteTypeArray(
+ op->outputs(), &outputs_array));
+
+ TfLiteNode* node = &(
+ graph_.GetAllocations()[subgraph_idx].node_and_registrations[i].node);
+ *node = {};
+ node->inputs = inputs_array;
+ node->outputs = outputs_array;
+ node->builtin_data = reinterpret_cast<void*>(builtin_data);
+ node->custom_initial_data = custom_data;
+ node->custom_initial_data_size = custom_data_size;
+ }
+ }
+ return kTfLiteOk;
+}
+
TfLiteStatus MicroInterpreter::AllocateTensors() {
- if (allocator_.StartModelAllocation(model_, op_resolver_,
- &node_and_registrations_,
- &eval_tensors_) != kTfLiteOk) {
+ SubgraphAllocations* allocations = allocator_.StartModelAllocation(model_);
+
+ if (allocations == nullptr) {
TF_LITE_REPORT_ERROR(error_reporter_,
"Failed starting model allocation.\n");
initialization_status_ = kTfLiteError;
return kTfLiteError;
}
- context_.tensors_size = subgraph_->tensors()->size();
+ graph_.SetSubgraphAllocations(allocations);
+
+ TF_LITE_ENSURE_STATUS(PrepareNodeAndRegistrationDataFromFlatbuffer());
// Only allow AllocatePersistentBuffer in Init stage.
context_.AllocatePersistentBuffer = AllocatePersistentBuffer;
context_.RequestScratchBufferInArena = nullptr;
context_.GetScratchBuffer = nullptr;
-
- for (size_t i = 0; i < subgraph_->operators()->size(); ++i) {
- auto* node = &(node_and_registrations_[i].node);
- auto* registration = node_and_registrations_[i].registration;
- size_t init_data_size;
- const char* init_data;
- if (registration->builtin_code == BuiltinOperator_CUSTOM) {
- init_data = reinterpret_cast<const char*>(node->custom_initial_data);
- init_data_size = node->custom_initial_data_size;
- } else {
- init_data = reinterpret_cast<const char*>(node->builtin_data);
- init_data_size = 0;
- }
- if (registration->init) {
- node->user_data =
- registration->init(&context_, init_data, init_data_size);
- }
- }
+ context_.GetExecutionPlan = GetGraph;
+ graph_.InitSubgraphs();
// Both AllocatePersistentBuffer and RequestScratchBufferInArena is
// available in Prepare stage.
context_.RequestScratchBufferInArena = RequestScratchBufferInArena;
- for (size_t i = 0; i < subgraph_->operators()->size(); ++i) {
- auto* node = &(node_and_registrations_[i].node);
- auto* registration = node_and_registrations_[i].registration;
- if (registration->prepare) {
- TfLiteStatus prepare_status = registration->prepare(&context_, node);
- if (prepare_status != kTfLiteOk) {
- TF_LITE_REPORT_ERROR(
- error_reporter_,
- "Node %s (number %df) failed to prepare with status %d",
- OpNameFromRegistration(registration), i, prepare_status);
- return kTfLiteError;
- }
- }
- allocator_.FinishPrepareNodeAllocations(/*node_id=*/i);
- }
+ graph_.PrepareSubgraphs();
// Prepare is done, we're ready for Invoke. Memory allocation is no longer
// allowed. Kernels can only fetch scratch buffers via GetScratchBuffer.
@@ -176,9 +210,9 @@
context_.RequestScratchBufferInArena = nullptr;
context_.GetScratchBuffer = GetScratchBuffer;
- TF_LITE_ENSURE_OK(&context_,
- allocator_.FinishModelAllocation(model_, eval_tensors_,
- &scratch_buffer_handles_));
+ TF_LITE_ENSURE_OK(&context_, allocator_.FinishModelAllocation(
+ model_, graph_.GetAllocations(),
+ &scratch_buffer_handles_));
// TODO(b/162311891): Drop these allocations when the interpreter supports
// handling buffers from TfLiteEvalTensor.
@@ -196,7 +230,7 @@
for (size_t i = 0; i < inputs_size(); ++i) {
input_tensors_[i] = allocator_.AllocatePersistentTfLiteTensor(
- model_, eval_tensors_, inputs().Get(i));
+ model_, graph_.GetAllocations(), inputs().Get(i), 0);
if (input_tensors_[i] == nullptr) {
TF_LITE_REPORT_ERROR(error_reporter_,
"Failed to initialize input tensor %d", i);
@@ -220,7 +254,7 @@
for (size_t i = 0; i < outputs_size(); ++i) {
output_tensors_[i] = allocator_.AllocatePersistentTfLiteTensor(
- model_, eval_tensors_, outputs().Get(i));
+ model_, graph_.GetAllocations(), outputs().Get(i), 0);
if (output_tensors_[i] == nullptr) {
TF_LITE_REPORT_ERROR(error_reporter_,
"Failed to initialize output tensor %d", i);
@@ -246,41 +280,7 @@
if (!tensors_allocated_) {
TF_LITE_ENSURE_OK(&context_, AllocateTensors());
}
-
- for (size_t i = 0; i < subgraph_->operators()->size(); ++i) {
- auto* node = &(node_and_registrations_[i].node);
- auto* registration = node_and_registrations_[i].registration;
-
-// This ifdef is needed (even though ScopedMicroProfiler itself is a no-op with
-// -DTF_LITE_STRIP_ERROR_STRINGS) because the function OpNameFromRegistration is
-// only defined for builds with the error strings.
-#if !defined(TF_LITE_STRIP_ERROR_STRINGS)
- ScopedMicroProfiler scoped_profiler(
- OpNameFromRegistration(registration),
- reinterpret_cast<MicroProfiler*>(context_.profiler));
-#endif
-
- TFLITE_DCHECK(registration->invoke);
- TfLiteStatus invoke_status = registration->invoke(&context_, node);
-
- // All TfLiteTensor structs used in the kernel are allocated from temp
- // memory in the allocator. This creates a chain of allocations in the
- // temp section. The call below resets the chain of allocations to
- // prepare for the next call.
- allocator_.ResetTempAllocations();
-
- if (invoke_status == kTfLiteError) {
- TF_LITE_REPORT_ERROR(
- error_reporter_,
- "Node %s (number %d) failed to invoke with status %d",
- OpNameFromRegistration(registration), i, invoke_status);
- return kTfLiteError;
- } else if (invoke_status != kTfLiteOk) {
- return invoke_status;
- }
- }
-
- return kTfLiteOk;
+ return graph_.InvokeSubgraph(0);
}
TfLiteTensor* MicroInterpreter::input(size_t index) {
@@ -306,42 +306,27 @@
}
TfLiteStatus MicroInterpreter::ResetVariableTensors() {
- for (size_t i = 0; i < subgraph_->tensors()->size(); ++i) {
- auto* tensor = subgraph_->tensors()->Get(i);
- if (tensor->is_variable()) {
- size_t buffer_size;
- TF_LITE_ENSURE_STATUS(
- TfLiteEvalTensorByteLength(&eval_tensors_[i], &buffer_size));
-
- int value = 0;
- if (tensor->type() == tflite::TensorType_INT8) {
- value = tensor->quantization()->zero_point()->Get(0);
- }
- memset(eval_tensors_[i].data.raw, value, buffer_size);
- }
- }
-
- return kTfLiteOk;
+ return graph_.ResetVariableTensors();
}
-void* MicroInterpreter::AllocatePersistentBuffer(TfLiteContext* context,
+void* MicroInterpreter::AllocatePersistentBuffer(TfLiteContext* ctx,
size_t bytes) {
- return reinterpret_cast<MicroInterpreter*>(context->impl_)
+ return reinterpret_cast<MicroInterpreter*>(ctx->impl_)
->allocator_.AllocatePersistentBuffer(bytes);
}
-TfLiteStatus MicroInterpreter::RequestScratchBufferInArena(
- TfLiteContext* context, size_t bytes, int* buffer_idx) {
- // All scratch buffer requests are managed in the allocator. Simply route the
- // request and let the allocator manage allocations.
- return static_cast<MicroInterpreter*>(context->impl_)
- ->allocator_.RequestScratchBufferInArena(bytes, buffer_idx);
+TfLiteStatus MicroInterpreter::RequestScratchBufferInArena(TfLiteContext* ctx,
+ size_t bytes,
+ int* buffer_idx) {
+ MicroInterpreter* interpreter =
+ reinterpret_cast<MicroInterpreter*>(ctx->impl_);
+ return interpreter->allocator_.RequestScratchBufferInArena(
+ bytes, interpreter->graph_.GetCurrentSubgraphIndex(), buffer_idx);
}
-void* MicroInterpreter::GetScratchBuffer(TfLiteContext* context,
- int buffer_idx) {
+void* MicroInterpreter::GetScratchBuffer(TfLiteContext* ctx, int buffer_idx) {
MicroInterpreter* interpreter =
- static_cast<MicroInterpreter*>(context->impl_);
+ reinterpret_cast<MicroInterpreter*>(ctx->impl_);
ScratchBufferHandle* handle =
interpreter->scratch_buffer_handles_ + buffer_idx;
return handle->data;
@@ -364,14 +349,25 @@
MicroInterpreter* interpreter =
static_cast<MicroInterpreter*>(context->impl_);
return interpreter->allocator_.AllocateTempTfLiteTensor(
- interpreter->model_, interpreter->eval_tensors_, tensor_idx);
+ interpreter->model_, interpreter->graph_.GetAllocations(), tensor_idx,
+ interpreter->get_subgraph_index());
}
TfLiteEvalTensor* MicroInterpreter::GetEvalTensor(
const struct TfLiteContext* context, int tensor_idx) {
MicroInterpreter* interpreter =
- static_cast<MicroInterpreter*>(context->impl_);
- return &interpreter->eval_tensors_[tensor_idx];
+ reinterpret_cast<MicroInterpreter*>(context->impl_);
+ return &interpreter->graph_
+ .GetAllocations()[interpreter->get_subgraph_index()]
+ .tensors[tensor_idx];
+}
+
+TfLiteStatus MicroInterpreter::GetGraph(struct TfLiteContext* context,
+ TfLiteIntArray** args) {
+ MicroInterpreter* interpreter =
+ reinterpret_cast<MicroInterpreter*>(context->impl_);
+ *args = reinterpret_cast<TfLiteIntArray*>(&interpreter->graph_);
+ return kTfLiteOk;
}
} // namespace tflite
diff --git a/tensorflow/lite/micro/micro_interpreter.h b/tensorflow/lite/micro/micro_interpreter.h
index d34015a..67d5420 100644
--- a/tensorflow/lite/micro/micro_interpreter.h
+++ b/tensorflow/lite/micro/micro_interpreter.h
@@ -23,6 +23,7 @@
#include "tensorflow/lite/core/api/error_reporter.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/micro/micro_allocator.h"
+#include "tensorflow/lite/micro/micro_graph.h"
#include "tensorflow/lite/micro/micro_op_resolver.h"
#include "tensorflow/lite/micro/micro_profiler.h"
#include "tensorflow/lite/portable_type_to_tflitetype.h"
@@ -69,9 +70,11 @@
TfLiteStatus Invoke();
TfLiteTensor* input(size_t index);
- size_t inputs_size() const { return subgraph_->inputs()->Length(); }
+ size_t inputs_size() const {
+ return model_->subgraphs()->Get(0)->inputs()->size();
+ }
const flatbuffers::Vector<int32_t>& inputs() const {
- return *subgraph_->inputs();
+ return *model_->subgraphs()->Get(0)->inputs();
}
TfLiteTensor* input_tensor(size_t index) { return input(index); }
template <class T>
@@ -85,9 +88,11 @@
}
TfLiteTensor* output(size_t index);
- size_t outputs_size() const { return subgraph_->outputs()->Length(); }
+ size_t outputs_size() const {
+ return model_->subgraphs()->Get(0)->outputs()->size();
+ }
const flatbuffers::Vector<int32_t>& outputs() const {
- return *subgraph_->outputs();
+ return *model_->subgraphs()->Get(0)->outputs();
}
TfLiteTensor* output_tensor(size_t index) { return output(index); }
template <class T>
@@ -105,13 +110,16 @@
TfLiteStatus initialization_status() const { return initialization_status_; }
- size_t operators_size() const { return subgraph_->operators()->size(); }
-
- // For debugging only.
- const NodeAndRegistration node_and_registration(int node_index) const {
- return node_and_registrations_[node_index];
+ size_t operators_size() const {
+ return model_->subgraphs()->Get(0)->operators()->size();
}
+ // 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();
+
// For debugging only.
// Returns the actual used arena in bytes. This method gives the optimal arena
// size. It's only available after `AllocateTensors` has been called.
@@ -129,32 +137,34 @@
// error reporting during initialization.
void Init(MicroProfiler* profiler);
+ // Gets the current subgraph index used from within context methods.
+ int get_subgraph_index() { return graph_.GetCurrentSubgraphIndex(); }
+
// Static functions that are bound to the TfLiteContext instance:
- static void* AllocatePersistentBuffer(TfLiteContext* Context, size_t bytes);
- static TfLiteStatus RequestScratchBufferInArena(TfLiteContext* context,
+ static void* AllocatePersistentBuffer(TfLiteContext* ctx, size_t bytes);
+ static TfLiteStatus RequestScratchBufferInArena(TfLiteContext* ctx,
size_t bytes,
int* buffer_idx);
- static void* GetScratchBuffer(TfLiteContext* context, int buffer_idx);
+ static void* GetScratchBuffer(TfLiteContext* ctx, int buffer_idx);
static void ReportOpError(struct TfLiteContext* context, const char* format,
...);
static TfLiteTensor* GetTensor(const struct TfLiteContext* context,
int tensor_idx);
static TfLiteEvalTensor* GetEvalTensor(const struct TfLiteContext* context,
int tensor_idx);
-
- NodeAndRegistration* node_and_registrations_ = nullptr;
+ static TfLiteStatus GetGraph(struct TfLiteContext* context,
+ TfLiteIntArray** args);
const Model* model_;
const MicroOpResolver& op_resolver_;
ErrorReporter* error_reporter_;
TfLiteContext context_ = {};
MicroAllocator& allocator_;
+ MicroGraph graph_;
bool tensors_allocated_;
TfLiteStatus initialization_status_;
- const SubGraph* subgraph_ = nullptr;
- TfLiteEvalTensor* eval_tensors_ = nullptr;
ScratchBufferHandle* scratch_buffer_handles_ = nullptr;
// TODO(b/162311891): Clean these pointers up when this class supports buffers
diff --git a/tensorflow/lite/micro/micro_interpreter_test.cc b/tensorflow/lite/micro/micro_interpreter_test.cc
index c0f8667..6ff6233 100644
--- a/tensorflow/lite/micro/micro_interpreter_test.cc
+++ b/tensorflow/lite/micro/micro_interpreter_test.cc
@@ -302,7 +302,7 @@
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
- constexpr size_t allocator_buffer_size = 512;
+ constexpr size_t allocator_buffer_size = 528;
uint8_t allocator_buffer[allocator_buffer_size];
tflite::RecordingMicroAllocator* allocator =
diff --git a/tensorflow/lite/micro/micro_utils.h b/tensorflow/lite/micro/micro_utils.h
index b9a3121..afa50a3 100644
--- a/tensorflow/lite/micro/micro_utils.h
+++ b/tensorflow/lite/micro/micro_utils.h
@@ -19,6 +19,7 @@
#include <algorithm>
#include <cmath>
#include <cstdint>
+#include <limits>
#include "tensorflow/lite/c/common.h"
diff --git a/tensorflow/lite/micro/recording_micro_allocator.cc b/tensorflow/lite/micro/recording_micro_allocator.cc
index 6bb5297..1e32e23 100644
--- a/tensorflow/lite/micro/recording_micro_allocator.cc
+++ b/tensorflow/lite/micro/recording_micro_allocator.cc
@@ -130,58 +130,48 @@
}
TfLiteStatus RecordingMicroAllocator::AllocateNodeAndRegistrations(
- const Model* model, NodeAndRegistration** node_and_registrations) {
- RecordedAllocation allocations = SnapshotAllocationUsage();
-
- TfLiteStatus status = MicroAllocator::AllocateNodeAndRegistrations(
- model, node_and_registrations);
-
- RecordAllocationUsage(allocations,
- recorded_node_and_registration_array_data_);
- // The allocation count in SimpleMemoryAllocator will only be 1. To provide
- // better logging, decrement by 1 and add in the actual number of operators
- // used in the graph:
- // The allocation for this recording will always be 1. This is because the
- // parent class mallocs one large allocation for the number of nodes in the
- // graph (e.g. sizeof(NodeAndRegistration) * num_nodes).
- // To prevent extra overhead and potential for fragmentation, manually adjust
- // the accounting by decrementing by 1 and adding the actual number of nodes
- // used in the graph:
- recorded_node_and_registration_array_data_.count +=
- GetSubGraphFromModel(model)->operators()->size() - 1;
- return status;
-}
-
-TfLiteStatus
-RecordingMicroAllocator::PrepareNodeAndRegistrationDataFromFlatbuffer(
- const Model* model, const MicroOpResolver& op_resolver,
- NodeAndRegistration* node_and_registrations) {
+ const Model* model, SubgraphAllocations* subgraph_allocations) {
RecordedAllocation allocations = SnapshotAllocationUsage();
TfLiteStatus status =
- MicroAllocator::PrepareNodeAndRegistrationDataFromFlatbuffer(
- model, op_resolver, node_and_registrations);
-
- RecordAllocationUsage(allocations, recorded_op_data_);
+ MicroAllocator::AllocateNodeAndRegistrations(model, subgraph_allocations);
+ for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs()->size();
+ subgraph_idx++) {
+ RecordAllocationUsage(allocations,
+ recorded_node_and_registration_array_data_);
+ // The allocation count in SimpleMemoryAllocator will only be 1. To provide
+ // better logging, decrement by 1 and add in the actual number of operators
+ // used in the graph:
+ // The allocation for this recording will always be 1. This is because the
+ // parent class mallocs one large allocation for the number of nodes in the
+ // graph (e.g. sizeof(NodeAndRegistration) * num_nodes).
+ // To prevent extra overhead and potential for fragmentation, manually
+ // adjust the accounting by decrementing by 1 and adding the actual number
+ // of nodes used in the graph:
+ recorded_node_and_registration_array_data_.count +=
+ model->subgraphs()->Get(subgraph_idx)->operators()->size() - 1;
+ }
return status;
}
TfLiteStatus RecordingMicroAllocator::AllocateTfLiteEvalTensors(
- const Model* model, TfLiteEvalTensor** eval_tensors) {
+ const Model* model, SubgraphAllocations* subgraph_allocations) {
RecordedAllocation allocations = SnapshotAllocationUsage();
TfLiteStatus status =
- MicroAllocator::AllocateTfLiteEvalTensors(model, eval_tensors);
-
- RecordAllocationUsage(allocations, recorded_tflite_eval_tensor_data_);
- // The allocation for this recording will always be 1. This is because the
- // parent class mallocs one large allocation for the number of tensors in the
- // graph (e.g. sizeof(TfLiteEvalTensor) * num_tensors).
- // To prevent extra overhead and potential for fragmentation, manually adjust
- // the accounting by decrementing by 1 and adding the actual number of tensors
- // used in the graph:
- recorded_tflite_eval_tensor_data_.count +=
- GetSubGraphFromModel(model)->tensors()->size() - 1;
+ MicroAllocator::AllocateTfLiteEvalTensors(model, subgraph_allocations);
+ for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs()->size();
+ subgraph_idx++) {
+ RecordAllocationUsage(allocations, recorded_tflite_eval_tensor_data_);
+ // The allocation for this recording will always be 1. This is because the
+ // parent class mallocs one large allocation for the number of tensors in
+ // the graph (e.g. sizeof(TfLiteEvalTensor) * num_tensors). To prevent extra
+ // overhead and potential for fragmentation, manually adjust the accounting
+ // by decrementing by 1 and adding the actual number of tensors used in the
+ // graph:
+ recorded_tflite_eval_tensor_data_.count +=
+ model->subgraphs()->Get(subgraph_idx)->tensors()->size() - 1;
+ }
return status;
}
@@ -197,24 +187,24 @@
return status;
}
-TfLiteTensor* RecordingMicroAllocator::AllocatePersistentTfLiteTensorInternal(
- const Model* model, TfLiteEvalTensor* eval_tensors, int tensor_index) {
+TfLiteTensor*
+RecordingMicroAllocator::AllocatePersistentTfLiteTensorInternal() {
RecordedAllocation allocations = SnapshotAllocationUsage();
- TfLiteTensor* result = MicroAllocator::AllocatePersistentTfLiteTensorInternal(
- model, eval_tensors, tensor_index);
+ TfLiteTensor* result =
+ MicroAllocator::AllocatePersistentTfLiteTensorInternal();
RecordAllocationUsage(allocations, recorded_persistent_tflite_tensor_data_);
return result;
}
TfLiteStatus RecordingMicroAllocator::PopulateTfLiteTensorFromFlatbuffer(
- const Model* model, const SubGraph* subgraph, TfLiteTensor* tensor,
- int tensor_index, bool allocate_temp) {
+ const Model* model, TfLiteTensor* tensor, int tensor_index,
+ int subgraph_index, bool allocate_temp) {
RecordedAllocation allocations = SnapshotAllocationUsage();
TfLiteStatus status = MicroAllocator::PopulateTfLiteTensorFromFlatbuffer(
- model, subgraph, tensor, tensor_index, allocate_temp);
+ model, tensor, tensor_index, subgraph_index, allocate_temp);
RecordAllocationUsage(allocations,
recorded_persistent_tflite_tensor_quantization_data_);
diff --git a/tensorflow/lite/micro/recording_micro_allocator.h b/tensorflow/lite/micro/recording_micro_allocator.h
index 47246e1..c3f4223 100644
--- a/tensorflow/lite/micro/recording_micro_allocator.h
+++ b/tensorflow/lite/micro/recording_micro_allocator.h
@@ -72,27 +72,22 @@
protected:
TfLiteStatus AllocateNodeAndRegistrations(
- const Model* model,
- NodeAndRegistration** node_and_registrations) override;
- TfLiteStatus PrepareNodeAndRegistrationDataFromFlatbuffer(
- const Model* model, const MicroOpResolver& op_resolver,
- NodeAndRegistration* node_and_registrations) override;
+ const Model* model, SubgraphAllocations* subgraph_allocations) override;
TfLiteStatus AllocateTfLiteEvalTensors(
- const Model* model, TfLiteEvalTensor** eval_tensors) override;
+ const Model* model, SubgraphAllocations* subgraph_allocations) override;
TfLiteStatus AllocateVariables(const SubGraph* subgraph,
TfLiteEvalTensor* eval_tensors) override;
// TODO(b/162311891): Once all kernels have been updated to the new API drop
// this method. It is only used to record TfLiteTensor persistent allocations.
- TfLiteTensor* AllocatePersistentTfLiteTensorInternal(
- const Model* model, TfLiteEvalTensor* eval_tensors,
- int tensor_index) override;
+ TfLiteTensor* AllocatePersistentTfLiteTensorInternal() override;
+
// TODO(b/162311891): Once all kernels have been updated to the new API drop
// this function since all allocations for quantized data will take place in
// the temp section.
TfLiteStatus PopulateTfLiteTensorFromFlatbuffer(const Model* model,
- const SubGraph* subgraph,
TfLiteTensor* tensor,
int tensor_index,
+ int subgraph_index,
bool allocate_temp) override;
private:
@@ -115,6 +110,8 @@
RecordedAllocation recorded_persistent_buffer_data_ = {};
RecordedAllocation recorded_tflite_tensor_variable_buffer_data_ = {};
RecordedAllocation recorded_node_and_registration_array_data_ = {};
+
+ // TODO(b/187993291): Re-enable OpData allocating tracking.
RecordedAllocation recorded_op_data_ = {};
TF_LITE_REMOVE_VIRTUAL_DELETE
diff --git a/tensorflow/lite/micro/recording_micro_allocator_test.cc b/tensorflow/lite/micro/recording_micro_allocator_test.cc
index b5f4080..1733726 100644
--- a/tensorflow/lite/micro/recording_micro_allocator_test.cc
+++ b/tensorflow/lite/micro/recording_micro_allocator_test.cc
@@ -16,6 +16,7 @@
#include "tensorflow/lite/micro/recording_micro_allocator.h"
#include "tensorflow/lite/micro/all_ops_resolver.h"
+#include "tensorflow/lite/micro/micro_allocator.h"
#include "tensorflow/lite/micro/micro_error_reporter.h"
#include "tensorflow/lite/micro/test_helpers.h"
#include "tensorflow/lite/micro/testing/micro_test.h"
@@ -36,10 +37,8 @@
TF_LITE_MICRO_TESTS_BEGIN
TF_LITE_MICRO_TEST(TestRecordsTfLiteEvalTensorArrayData) {
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
tflite::AllOpsResolver all_ops_resolver;
- tflite::NodeAndRegistration* node_and_registration;
const tflite::Model* model = tflite::GetModel(kTestConvModelData);
uint8_t arena[kTestConvArenaSize];
@@ -51,14 +50,13 @@
TF_LITE_MICRO_EXPECT_NE(micro_allocator, nullptr);
if (micro_allocator == nullptr) return 1;
- TfLiteStatus status;
- status = micro_allocator->StartModelAllocation(
- model, all_ops_resolver, &node_and_registration, &eval_tensors);
- TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
- if (status != kTfLiteOk) return 1;
+ tflite::SubgraphAllocations* subgraph_allocations =
+ micro_allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
+ if (subgraph_allocations == nullptr) return 1;
- status = micro_allocator->FinishModelAllocation(model, eval_tensors,
- &scratch_buffer_handles);
+ TfLiteStatus status = micro_allocator->FinishModelAllocation(
+ model, subgraph_allocations, &scratch_buffer_handles);
TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
if (status != kTfLiteOk) return 1;
@@ -80,10 +78,8 @@
}
TF_LITE_MICRO_TEST(TestRecordsNodeAndRegistrationArrayData) {
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
tflite::AllOpsResolver all_ops_resolver;
- tflite::NodeAndRegistration* node_and_registration;
const tflite::Model* model = tflite::GetModel(kTestConvModelData);
uint8_t arena[kTestConvArenaSize];
@@ -93,14 +89,13 @@
TF_LITE_MICRO_EXPECT_NE(micro_allocator, nullptr);
if (micro_allocator == nullptr) return 1;
- TfLiteStatus status;
- status = micro_allocator->StartModelAllocation(
- model, all_ops_resolver, &node_and_registration, &eval_tensors);
- TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
- if (status != kTfLiteOk) return 1;
+ tflite::SubgraphAllocations* subgraph_allocations =
+ micro_allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
+ if (subgraph_allocations == nullptr) return 1;
- status = micro_allocator->FinishModelAllocation(model, eval_tensors,
- &scratch_buffer_handles);
+ TfLiteStatus status = micro_allocator->FinishModelAllocation(
+ model, subgraph_allocations, &scratch_buffer_handles);
TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
if (status != kTfLiteOk) return 1;
@@ -108,6 +103,7 @@
tflite::RecordedAllocation recorded_allocation =
micro_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);
@@ -116,10 +112,8 @@
}
TF_LITE_MICRO_TEST(TestRecordsMultiTenantAllocations) {
- TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
tflite::AllOpsResolver all_ops_resolver;
- tflite::NodeAndRegistration* node_and_registration;
const tflite::Model* model = tflite::GetModel(kTestConvModelData);
// Double the arena size to allocate two models inside of it:
@@ -134,24 +128,23 @@
if (micro_allocator == nullptr) return 1;
// First allocation with the model in the arena:
- status = micro_allocator->StartModelAllocation(
- model, all_ops_resolver, &node_and_registration, &eval_tensors);
- TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
- if (status != kTfLiteOk) return 1;
+ tflite::SubgraphAllocations* subgraph_allocations =
+ micro_allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
+ if (subgraph_allocations == nullptr) return 1;
- status = micro_allocator->FinishModelAllocation(model, eval_tensors,
+ status = micro_allocator->FinishModelAllocation(model, subgraph_allocations,
&scratch_buffer_handles);
TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
if (status != kTfLiteOk) return 1;
// Second allocation with the same model in the arena:
- status = micro_allocator->StartModelAllocation(
- model, all_ops_resolver, &node_and_registration, &eval_tensors);
- TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
- if (status != kTfLiteOk) return 1;
+ subgraph_allocations = micro_allocator->StartModelAllocation(model);
+ TF_LITE_MICRO_EXPECT(nullptr != subgraph_allocations);
+ if (subgraph_allocations == nullptr) return 1;
status = kTfLiteOk, micro_allocator->FinishModelAllocation(
- model, eval_tensors, &scratch_buffer_handles);
+ model, subgraph_allocations, &scratch_buffer_handles);
TF_LITE_MICRO_EXPECT_EQ(status, kTfLiteOk);
if (status != kTfLiteOk) return 1;
@@ -160,6 +153,8 @@
tflite::RecordedAllocation recorded_allocation =
micro_allocator->GetRecordedAllocation(
tflite::RecordedAllocationType::kTfLiteEvalTensorData);
+
+ // Node and tensor arrays must be allocated as well as each node and tensor.
TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count, tensors_count * 2);
TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
tensors_count * TF_LITE_EVAL_TENSOR_STRUCT_SIZE * 2);
@@ -178,7 +173,7 @@
if (micro_allocator == nullptr) return 1;
TfLiteTensor* tensor = micro_allocator->AllocatePersistentTfLiteTensor(
- model, /*eval_tensors=*/nullptr, 0);
+ model, /*subgraph_allocations=*/nullptr, 0, 0);
TF_LITE_MICRO_EXPECT_NE(tensor, nullptr);
if (tensor == nullptr) return 1;
@@ -204,7 +199,7 @@
if (micro_allocator == nullptr) return 1;
TfLiteTensor* tensor = micro_allocator->AllocatePersistentTfLiteTensor(
- model, /*eval_tensors=*/nullptr, 0);
+ model, /*subgraph_allocations=*/nullptr, 0, 0);
TF_LITE_MICRO_EXPECT_NE(tensor, nullptr);
if (tensor == nullptr) return 1;
