| /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. |
| |
| Licensed under the Apache License, Version 2.0 (the "License"); |
| you may not use this file except in compliance with the License. |
| You may obtain a copy of the License at |
| |
| http://www.apache.org/licenses/LICENSE-2.0 |
| |
| Unless required by applicable law or agreed to in writing, software |
| distributed under the License is distributed on an "AS IS" BASIS, |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| See the License for the specific language governing permissions and |
| limitations under the License. |
| ==============================================================================*/ |
| |
| #include "tensorflow/lite/micro/recording_micro_allocator.h" |
| |
| #include "tensorflow/lite/core/api/error_reporter.h" |
| #include "tensorflow/lite/kernels/internal/compatibility.h" |
| #include "tensorflow/lite/micro/compatibility.h" |
| #include "tensorflow/lite/micro/recording_simple_memory_allocator.h" |
| |
| namespace tflite { |
| |
| RecordingMicroAllocator::RecordingMicroAllocator( |
| RecordingSimpleMemoryAllocator* recording_memory_allocator, |
| ErrorReporter* error_reporter) |
| : MicroAllocator(recording_memory_allocator, error_reporter), |
| recording_memory_allocator_(recording_memory_allocator) {} |
| |
| RecordingMicroAllocator* RecordingMicroAllocator::Create( |
| uint8_t* tensor_arena, size_t arena_size, ErrorReporter* error_reporter) { |
| TFLITE_DCHECK(error_reporter != nullptr); |
| |
| RecordingSimpleMemoryAllocator* simple_memory_allocator = |
| RecordingSimpleMemoryAllocator::Create(error_reporter, tensor_arena, |
| arena_size); |
| TFLITE_DCHECK(simple_memory_allocator != nullptr); |
| |
| uint8_t* allocator_buffer = simple_memory_allocator->AllocateFromTail( |
| sizeof(RecordingMicroAllocator), alignof(RecordingMicroAllocator)); |
| RecordingMicroAllocator* allocator = new (allocator_buffer) |
| RecordingMicroAllocator(simple_memory_allocator, error_reporter); |
| return allocator; |
| } |
| |
| RecordedAllocation RecordingMicroAllocator::GetRecordedAllocation( |
| RecordedAllocationType allocation_type) const { |
| switch (allocation_type) { |
| case RecordedAllocationType::kTfLiteTensorArray: |
| return recorded_tflite_tensor_array_data_; |
| case RecordedAllocationType::kTfLiteTensorArrayQuantizationData: |
| return recorded_tflite_tensor_array_quantization_data_; |
| case RecordedAllocationType::kTfLiteTensorVariableBufferData: |
| return recorded_tflite_tensor_variable_buffer_data_; |
| case RecordedAllocationType::kNodeAndRegistrationArray: |
| return recorded_node_and_registration_array_data_; |
| case RecordedAllocationType::kOpData: |
| return recorded_op_data_; |
| } |
| TF_LITE_REPORT_ERROR(error_reporter(), "Invalid allocation type supplied: %d", |
| allocation_type); |
| return RecordedAllocation(); |
| } |
| |
| const RecordingSimpleMemoryAllocator* |
| RecordingMicroAllocator::GetSimpleMemoryAllocator() const { |
| return recording_memory_allocator_; |
| } |
| |
| void RecordingMicroAllocator::PrintAllocations() const { |
| TF_LITE_REPORT_ERROR( |
| error_reporter(), |
| "[RecordingMicroAllocator] Arena allocation total %d bytes", |
| recording_memory_allocator_->GetUsedBytes()); |
| TF_LITE_REPORT_ERROR( |
| error_reporter(), |
| "[RecordingMicroAllocator] Arena allocation head %d bytes", |
| recording_memory_allocator_->GetHeadUsedBytes()); |
| TF_LITE_REPORT_ERROR( |
| error_reporter(), |
| "[RecordingMicroAllocator] Arena allocation tail %d bytes", |
| recording_memory_allocator_->GetTailUsedBytes()); |
| PrintRecordedAllocation(RecordedAllocationType::kTfLiteTensorArray, |
| "TfLiteTensor struct", "tensors"); |
| PrintRecordedAllocation( |
| RecordedAllocationType::kTfLiteTensorArrayQuantizationData, |
| "TfLiteTensor quantization data", "allocations"); |
| PrintRecordedAllocation( |
| RecordedAllocationType::kTfLiteTensorVariableBufferData, |
| "TfLiteTensor variable buffer data", "allocations"); |
| PrintRecordedAllocation(RecordedAllocationType::kNodeAndRegistrationArray, |
| "NodeAndRegistration struct", |
| "NodeAndRegistration structs"); |
| PrintRecordedAllocation(RecordedAllocationType::kOpData, |
| "Operator runtime data", "OpData structs"); |
| } |
| |
| void RecordingMicroAllocator::PrintRecordedAllocation( |
| RecordedAllocationType allocation_type, const char* allocation_name, |
| const char* allocation_description) const { |
| RecordedAllocation allocation = GetRecordedAllocation(allocation_type); |
| TF_LITE_REPORT_ERROR( |
| error_reporter(), |
| "[RecordingMicroAllocator] '%s' used %d bytes with alignment overhead " |
| "(requested %d bytes for %d %s)", |
| allocation_name, allocation.used_bytes, allocation.requested_bytes, |
| allocation.count, allocation_description); |
| } |
| |
| TfLiteStatus RecordingMicroAllocator::AllocateTfLiteTensorArray( |
| TfLiteContext* context, const SubGraph* subgraph) { |
| SnapshotAllocationUsage(recorded_tflite_tensor_array_data_); |
| |
| TfLiteStatus status = |
| MicroAllocator::AllocateTfLiteTensorArray(context, subgraph); |
| |
| RecordAllocationUsage(recorded_tflite_tensor_array_data_); |
| recorded_tflite_tensor_array_data_.count = context->tensors_size; |
| return status; |
| } |
| |
| TfLiteStatus RecordingMicroAllocator::PopulateTfLiteTensorArrayFromFlatbuffer( |
| const Model* model, TfLiteContext* context, const SubGraph* subgraph) { |
| SnapshotAllocationUsage(recorded_tflite_tensor_array_quantization_data_); |
| |
| TfLiteStatus status = MicroAllocator::PopulateTfLiteTensorArrayFromFlatbuffer( |
| model, context, subgraph); |
| |
| RecordAllocationUsage(recorded_tflite_tensor_array_quantization_data_); |
| return status; |
| } |
| |
| TfLiteStatus RecordingMicroAllocator::AllocateNodeAndRegistrations( |
| const SubGraph* subgraph, NodeAndRegistration** node_and_registrations) { |
| SnapshotAllocationUsage(recorded_node_and_registration_array_data_); |
| |
| TfLiteStatus status = MicroAllocator::AllocateNodeAndRegistrations( |
| subgraph, node_and_registrations); |
| |
| RecordAllocationUsage(recorded_node_and_registration_array_data_); |
| recorded_node_and_registration_array_data_.count = |
| subgraph->operators()->size(); |
| return status; |
| } |
| |
| TfLiteStatus |
| RecordingMicroAllocator::PrepareNodeAndRegistrationDataFromFlatbuffer( |
| const Model* model, const SubGraph* subgraph, |
| const MicroOpResolver& op_resolver, |
| NodeAndRegistration* node_and_registrations) { |
| SnapshotAllocationUsage(recorded_op_data_); |
| |
| TfLiteStatus status = |
| MicroAllocator::PrepareNodeAndRegistrationDataFromFlatbuffer( |
| model, subgraph, op_resolver, node_and_registrations); |
| |
| RecordAllocationUsage(recorded_op_data_); |
| return status; |
| } |
| |
| TfLiteStatus RecordingMicroAllocator::AllocateVariables( |
| TfLiteContext* context, const SubGraph* subgraph) { |
| SnapshotAllocationUsage(recorded_tflite_tensor_variable_buffer_data_); |
| |
| TfLiteStatus status = MicroAllocator::AllocateVariables(context, subgraph); |
| |
| RecordAllocationUsage(recorded_tflite_tensor_variable_buffer_data_); |
| return status; |
| } |
| |
| void RecordingMicroAllocator::SnapshotAllocationUsage( |
| RecordedAllocation& recorded_allocation) { |
| recorded_allocation.requested_bytes = |
| recording_memory_allocator_->GetRequestedBytes(); |
| recorded_allocation.used_bytes = recording_memory_allocator_->GetUsedBytes(); |
| recorded_allocation.count = recording_memory_allocator_->GetAllocatedCount(); |
| } |
| |
| void RecordingMicroAllocator::RecordAllocationUsage( |
| RecordedAllocation& recorded_allocation) { |
| recorded_allocation.requested_bytes = |
| recording_memory_allocator_->GetRequestedBytes() - |
| recorded_allocation.requested_bytes; |
| recorded_allocation.used_bytes = recording_memory_allocator_->GetUsedBytes() - |
| recorded_allocation.used_bytes; |
| recorded_allocation.count = recording_memory_allocator_->GetAllocatedCount() - |
| recorded_allocation.count; |
| } |
| |
| } // namespace tflite |