From 8c8dc2699bc8d91345ae7ea6d38a20f15efb8f31 Mon Sep 17 00:00:00 2001 From: Nick Kreeger Date: Tue, 26 May 2020 12:12:38 -0700 Subject: [PATCH] Cleanup and refactor all allocations in MicroAllocator to function calls. This change is a precursor to adding a new memory logging MicroAllocator subclass that will enable TFLM to keep track of tensor arena tail allocations. Outside of moving all arena allocations to utility methods - I also cleaned up the organization of the methods inside of the cc file. PiperOrigin-RevId: 313242666 Change-Id: Icddcc07187419fe314bc57708170cda8cd35690a --- tensorflow/lite/micro/micro_allocator.cc | 284 ++++++++++-------- tensorflow/lite/micro/micro_allocator.h | 58 +++- tensorflow/lite/micro/micro_allocator_test.cc | 6 +- tensorflow/lite/micro/micro_interpreter.cc | 2 +- 4 files changed, 202 insertions(+), 148 deletions(-) diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index 1dd1fa4b63c..b67e158980d 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -258,7 +258,7 @@ TfLiteStatus CommitPlan(ErrorReporter* error_reporter, MemoryPlanner* planner, namespace internal { -TfLiteStatus InitializeRuntimeTensor( +TfLiteStatus InitializeTfLiteTensorFromFlatbuffer( SimpleMemoryAllocator* allocator, const tflite::Tensor& flatbuffer_tensor, const flatbuffers::Vector>* buffers, ErrorReporter* error_reporter, TfLiteTensor* result) { @@ -380,58 +380,9 @@ TfLiteStatus InitializeRuntimeTensor( } return kTfLiteOk; } + } // namespace internal -TfLiteStatus MicroAllocator::Init() { - auto* subgraphs = model_->subgraphs(); - if (subgraphs->size() != 1) { - TF_LITE_REPORT_ERROR(error_reporter_, - "Only 1 subgraph is currently supported.\n"); - return kTfLiteError; - } - subgraph_ = (*subgraphs)[0]; - - context_->tensors_size = subgraph_->tensors()->size(); - context_->tensors = - reinterpret_cast(memory_allocator_->AllocateFromTail( - sizeof(TfLiteTensor) * context_->tensors_size, - alignof(TfLiteTensor))); - if (context_->tensors == nullptr) { - TF_LITE_REPORT_ERROR( - error_reporter_, - "Failed to allocate memory for context->tensors, %d bytes required", - sizeof(TfLiteTensor) * context_->tensors_size); - return kTfLiteError; - } - - // Initialize runtime tensors in context_ using the flatbuffer. - for (size_t i = 0; i < subgraph_->tensors()->size(); ++i) { - TfLiteStatus status = internal::InitializeRuntimeTensor( - memory_allocator_, *subgraph_->tensors()->Get(i), model_->buffers(), - error_reporter_, &context_->tensors[i]); - if (status != kTfLiteOk) { - TF_LITE_REPORT_ERROR(error_reporter_, "Failed to initialize tensor %d", - i); - return kTfLiteError; - } - } - - return kTfLiteOk; -} - -size_t MicroAllocator::used_bytes() const { - if (active_) { - return 0; - } - TF_LITE_REPORT_ERROR(error_reporter_, "Total buffer usage: %d bytes", - memory_allocator_->GetUsedBytes()); - TF_LITE_REPORT_ERROR(error_reporter_, "Head usage: %d bytes", - memory_allocator_->GetHeadUsedBytes()); - TF_LITE_REPORT_ERROR(error_reporter_, "Tail usage: %d bytes", - memory_allocator_->GetTailUsedBytes()); - return memory_allocator_->GetUsedBytes(); -} - MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model, uint8_t* tensor_arena, size_t arena_size, ErrorReporter* error_reporter) @@ -450,7 +401,8 @@ MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model, // destructed as it's the root allocator. memory_allocator_ = CreateInPlaceSimpleMemoryAllocator( error_reporter, aligned_arena, aligned_arena_size); - TfLiteStatus status = Init(); + + TfLiteStatus status = InitGraphAndContextTensorData(); // TODO(b/147871299): Consider improving this code. A better way of handling // failures in the constructor is to have a static function that returns a // pointer to the class. If allocation failed, a nullptr will be returned. @@ -463,88 +415,15 @@ MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model, } } -TfLiteStatus MicroAllocator::AllocateNodeAndRegistrations( +TfLiteStatus MicroAllocator::InitializeFromFlatbuffer( const OpResolver& op_resolver, NodeAndRegistration** node_and_registrations) { if (!active_) { return kTfLiteError; } - - auto* output = reinterpret_cast( - 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; - } - 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); - size_t index = op->opcode_index(); - if (index >= opcodes->size()) { - TF_LITE_REPORT_ERROR(error_reporter_, - "Missing registration for opcode_index %d\n", index); - return kTfLiteError; - } - auto* opcode = (*opcodes)[index]; - status = GetRegistrationFromOpCode(opcode, op_resolver, error_reporter_, - &(output[i].registration)); - if (status != kTfLiteOk) { - TF_LITE_REPORT_ERROR(error_reporter_, - "Failed to get registration from op code %s\n ", - EnumNameBuiltinOperator(opcode->builtin_code())); - return status; - } - const auto* registration = output[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(registration->builtin_code); - - if (op_type != BuiltinOperator_CUSTOM && op->custom_options()) { - TF_LITE_REPORT_ERROR( - error_reporter_, - "Unsupported behavior: found builtin operator %s with custom " - "options.\n", - EnumNameBuiltinOperator(op_type)); - return kTfLiteError; - } - - const char* custom_data = nullptr; - size_t custom_data_size = 0; - unsigned char* builtin_data = nullptr; - if (op->custom_options()) { - custom_data = reinterpret_cast(op->custom_options()->data()); - custom_data_size = op->custom_options()->size(); - } else { - TF_LITE_ENSURE_STATUS(ParseOpData(op, op_type, error_reporter_, - &builtin_data_allocator, - (void**)(&builtin_data))); - } - - // Disregard const qualifier to workaround with existing API. - TfLiteIntArray* inputs_array = const_cast( - reinterpret_cast(op->inputs())); - TfLiteIntArray* outputs_array = const_cast( - reinterpret_cast(op->outputs())); - - TfLiteNode* node = &(output[i].node); - *node = {}; - node->inputs = inputs_array; - node->outputs = outputs_array; - node->builtin_data = reinterpret_cast(builtin_data); - node->custom_initial_data = custom_data; - node->custom_initial_data_size = custom_data_size; - } - *node_and_registrations = output; + TF_LITE_ENSURE_STATUS(AllocateNodeAndRegistrations(node_and_registrations)); + TF_LITE_ENSURE_STATUS(PrepareNodeAndRegistrationDataFromFlatbuffer( + op_resolver, *node_and_registrations)); return kTfLiteOk; } @@ -679,4 +558,151 @@ void* MicroAllocator::GetScratchBuffer(int buffer_idx) const { return scratch_buffer_handles_[scratch_buffer_count_ - buffer_idx - 1].data; } +size_t MicroAllocator::used_bytes() const { + if (active_) { + return 0; + } + TF_LITE_REPORT_ERROR(error_reporter_, "Total buffer usage: %d bytes", + memory_allocator_->GetUsedBytes()); + TF_LITE_REPORT_ERROR(error_reporter_, "Head usage: %d bytes", + memory_allocator_->GetHeadUsedBytes()); + TF_LITE_REPORT_ERROR(error_reporter_, "Tail usage: %d bytes", + memory_allocator_->GetTailUsedBytes()); + return memory_allocator_->GetUsedBytes(); +} + +TfLiteStatus MicroAllocator::InitGraphAndContextTensorData() { + auto* subgraphs = model_->subgraphs(); + if (subgraphs->size() != 1) { + TF_LITE_REPORT_ERROR(error_reporter_, + "Only 1 subgraph is currently supported.\n"); + return kTfLiteError; + } + subgraph_ = (*subgraphs)[0]; + + TF_LITE_ENSURE_STATUS(AllocateTfLiteTensorArray()); + TF_LITE_ENSURE_STATUS(PopulateTfLiteTensorArrayFromFlatbuffer()); + + return kTfLiteOk; +} + +TfLiteStatus MicroAllocator::AllocateTfLiteTensorArray() { + context_->tensors_size = subgraph_->tensors()->size(); + context_->tensors = + reinterpret_cast(memory_allocator_->AllocateFromTail( + sizeof(TfLiteTensor) * context_->tensors_size, + alignof(TfLiteTensor))); + if (context_->tensors == nullptr) { + TF_LITE_REPORT_ERROR( + error_reporter_, + "Failed to allocate memory for context->tensors, %d bytes required", + sizeof(TfLiteTensor) * context_->tensors_size); + return kTfLiteError; + } + return kTfLiteOk; +} + +TfLiteStatus MicroAllocator::PopulateTfLiteTensorArrayFromFlatbuffer() { + // Initialize tensors in context_ using the flatbuffer for quantization data. + for (size_t i = 0; i < subgraph_->tensors()->size(); ++i) { + TfLiteStatus status = internal::InitializeTfLiteTensorFromFlatbuffer( + memory_allocator_, *subgraph_->tensors()->Get(i), model_->buffers(), + error_reporter_, &context_->tensors[i]); + if (status != kTfLiteOk) { + TF_LITE_REPORT_ERROR(error_reporter_, "Failed to initialize tensor %d", + i); + return kTfLiteError; + } + } + return kTfLiteOk; +} + +TfLiteStatus MicroAllocator::AllocateNodeAndRegistrations( + NodeAndRegistration** node_and_registrations) { + NodeAndRegistration* output = reinterpret_cast( + 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 OpResolver& op_resolver, + NodeAndRegistration* node_and_registrations) { + 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); + 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(opcode->builtin_code())); + 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(registration->builtin_code); + + if (op_type != BuiltinOperator_CUSTOM && op->custom_options()) { + TF_LITE_REPORT_ERROR( + error_reporter_, + "Unsupported behavior: found builtin operator %s with custom " + "options.\n", + EnumNameBuiltinOperator(op_type)); + return kTfLiteError; + } + + const char* custom_data = nullptr; + size_t custom_data_size = 0; + unsigned char* builtin_data = nullptr; + if (op->custom_options()) { + custom_data = reinterpret_cast(op->custom_options()->data()); + custom_data_size = op->custom_options()->size(); + } else { + TF_LITE_ENSURE_STATUS(ParseOpData(op, op_type, error_reporter_, + &builtin_data_allocator, + (void**)(&builtin_data))); + } + + // Disregard const qualifier to workaround with existing API. + TfLiteIntArray* inputs_array = const_cast( + reinterpret_cast(op->inputs())); + TfLiteIntArray* outputs_array = const_cast( + reinterpret_cast(op->outputs())); + + TfLiteNode* node = &(node_and_registrations[i].node); + *node = {}; + node->inputs = inputs_array; + node->outputs = outputs_array; + node->builtin_data = reinterpret_cast(builtin_data); + node->custom_initial_data = custom_data; + node->custom_initial_data_size = custom_data_size; + } + + return kTfLiteOk; +} + } // namespace tflite diff --git a/tensorflow/lite/micro/micro_allocator.h b/tensorflow/lite/micro/micro_allocator.h index d05974f365a..1dd90c36a4d 100644 --- a/tensorflow/lite/micro/micro_allocator.h +++ b/tensorflow/lite/micro/micro_allocator.h @@ -30,9 +30,9 @@ namespace tflite { // Namespace used for unittests. namespace internal { -// Sets up all of the data structure members for a runtime tensor -// based on the contents of a serialized tensor. -TfLiteStatus InitializeRuntimeTensor( +// Sets up all of the data structure members for a TfLiteTensor based on the +// contents of a serialized tensor in the flatbuffer. +TfLiteStatus InitializeTfLiteTensorFromFlatbuffer( SimpleMemoryAllocator* allocator, const tflite::Tensor& flatbuffer_tensor, const flatbuffers::Vector>* buffers, ErrorReporter* error_reporter, TfLiteTensor* result); @@ -86,6 +86,15 @@ class MicroAllocator { uint8_t* tensor_arena, size_t arena_size, ErrorReporter* error_reporter); + // Run through the model flatbuffer data (loaded from the TfLiteModel + // instance) to allocate nodes and registrations. We need to keep them for the + // entire life time of the model to allow persistent tensors. This method + // needs to be called before FinishTensorAllocation method. This method also + // allocates any internal Op data that is required from the flatbuffer. + TfLiteStatus InitializeFromFlatbuffer( + const OpResolver& op_resolver, + NodeAndRegistration** node_and_registrations); + // Runs through the model and allocates all necessary input, output and // intermediate tensors. // WARNING: doing any allocation after calling this method has the risk of @@ -93,17 +102,6 @@ class MicroAllocator { // called in this class. TfLiteStatus FinishTensorAllocation(); - // Returns the arena usage in bytes, only available after - // `FinishTensorAllocation`. Otherwise, it will return 0. - size_t used_bytes() const; - - // Run through the model to allocate nodes and registrations. We need to keep - // them for the entire life time of the model to allow persistent tensors. - // This method needs to be called before FinishTensorAllocation method. - TfLiteStatus AllocateNodeAndRegistrations( - const OpResolver& op_resolver, - NodeAndRegistration** node_and_registrations); - // Allocates persistent buffer which has the same life time as the allocator. // The memory is immediately available and is allocated from the tail of the // arena. @@ -120,8 +118,38 @@ class MicroAllocator { // Returns the pointer to the planned scratch buffer. void* GetScratchBuffer(int buffer_idx) const; + // Returns the arena usage in bytes, only available after + // `FinishTensorAllocation`. Otherwise, it will return 0. + size_t used_bytes() const; + + protected: + // Allocates an array in the arena to hold pointers to the tensors required + // to initialize and prepare a model. These allocations are stored and + // populated on the context. + TfLiteStatus AllocateTfLiteTensorArray(); + + // Populates content on the list of tensor pointers required to initialize and + // prepare a model from data in the flatbuffer (loaded from the TfLiteModel + // instance). Persistent data (e.g. quantization params) is allocated from the + // arena. + TfLiteStatus PopulateTfLiteTensorArrayFromFlatbuffer(); + + // Allocates an array in the arena to hold pointers to the node and + // registration pointers required to represent the inference graph of the + // model. + TfLiteStatus AllocateNodeAndRegistrations( + NodeAndRegistration** node_and_registrations); + + // Populates node and registration pointers representing the inference graph + // of the model from values inside the flatbuffer (loaded from the TfLiteModel + // instance). Persistent data (e.g. operator data) is allocated from the + // arena. + TfLiteStatus PrepareNodeAndRegistrationDataFromFlatbuffer( + const OpResolver& op_resolver, + NodeAndRegistration* node_and_registrations); + private: - TfLiteStatus Init(); + TfLiteStatus InitGraphAndContextTensorData(); const Model* model_; // A simple memory allocator that always allocate from the arena tail. diff --git a/tensorflow/lite/micro/micro_allocator_test.cc b/tensorflow/lite/micro/micro_allocator_test.cc index 78419edbbf9..b34b2dc2866 100644 --- a/tensorflow/lite/micro/micro_allocator_test.cc +++ b/tensorflow/lite/micro/micro_allocator_test.cc @@ -77,7 +77,7 @@ TF_LITE_MICRO_TEST(TestInitializeRuntimeTensor) { TfLiteTensor allocated_tensor; TF_LITE_MICRO_EXPECT_EQ( - kTfLiteOk, tflite::internal::InitializeRuntimeTensor( + kTfLiteOk, tflite::internal::InitializeTfLiteTensorFromFlatbuffer( &simple_allocator, *tensor, buffers, micro_test::reporter, &allocated_tensor)); TF_LITE_MICRO_EXPECT_EQ(kTfLiteInt32, allocated_tensor.type); @@ -103,7 +103,7 @@ TF_LITE_MICRO_TEST(TestInitializeQuantizedTensor) { TfLiteTensor allocated_tensor; TF_LITE_MICRO_EXPECT_EQ( - kTfLiteOk, tflite::internal::InitializeRuntimeTensor( + kTfLiteOk, tflite::internal::InitializeTfLiteTensorFromFlatbuffer( &simple_allocator, *tensor, buffers, micro_test::reporter, &allocated_tensor)); TF_LITE_MICRO_EXPECT_EQ(kTfLiteInt32, allocated_tensor.type); @@ -129,7 +129,7 @@ TF_LITE_MICRO_TEST(TestMissingQuantization) { TfLiteTensor allocated_tensor; TF_LITE_MICRO_EXPECT_EQ( - kTfLiteOk, tflite::internal::InitializeRuntimeTensor( + kTfLiteOk, tflite::internal::InitializeTfLiteTensorFromFlatbuffer( &simple_allocator, *tensor, buffers, micro_test::reporter, &allocated_tensor)); TF_LITE_MICRO_EXPECT_EQ(kTfLiteInt32, allocated_tensor.type); diff --git a/tensorflow/lite/micro/micro_interpreter.cc b/tensorflow/lite/micro/micro_interpreter.cc index b46f9ecb9ea..6b78966020e 100644 --- a/tensorflow/lite/micro/micro_interpreter.cc +++ b/tensorflow/lite/micro/micro_interpreter.cc @@ -165,7 +165,7 @@ void MicroInterpreter::CorrectTensorDataEndianness(T* data, int32_t size) { } TfLiteStatus MicroInterpreter::AllocateTensors() { - TF_LITE_ENSURE_OK(&context_, allocator_.AllocateNodeAndRegistrations( + TF_LITE_ENSURE_OK(&context_, allocator_.InitializeFromFlatbuffer( op_resolver_, &node_and_registrations_)); // Only allow AllocatePersistentBuffer in Init stage.