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
This commit is contained in:
parent
e76bff1201
commit
8c8dc2699b
@ -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<flatbuffers::Offset<Buffer>>* 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<TfLiteTensor*>(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<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;
|
||||
}
|
||||
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<BuiltinOperator>(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<const char*>(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<TfLiteIntArray*>(
|
||||
reinterpret_cast<const TfLiteIntArray*>(op->inputs()));
|
||||
TfLiteIntArray* outputs_array = const_cast<TfLiteIntArray*>(
|
||||
reinterpret_cast<const TfLiteIntArray*>(op->outputs()));
|
||||
|
||||
TfLiteNode* node = &(output[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;
|
||||
}
|
||||
*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<TfLiteTensor*>(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<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 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<BuiltinOperator>(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<const char*>(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<TfLiteIntArray*>(
|
||||
reinterpret_cast<const TfLiteIntArray*>(op->inputs()));
|
||||
TfLiteIntArray* outputs_array = const_cast<TfLiteIntArray*>(
|
||||
reinterpret_cast<const TfLiteIntArray*>(op->outputs()));
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
return kTfLiteOk;
|
||||
}
|
||||
|
||||
} // namespace tflite
|
||||
|
@ -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<flatbuffers::Offset<Buffer>>* 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.
|
||||
|
@ -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);
|
||||
|
@ -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.
|
||||
|
Loading…
Reference in New Issue
Block a user