diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index 6c7fe0f2482..09e4f10ba84 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -152,17 +152,12 @@ TfLiteStatus CheckOfflinePlannedOffsets(const Model* model, // plan. Methods need to be called in order from `Init`, `Add*`, to `Finish`. class AllocationInfoBuilder { public: - AllocationInfoBuilder(ErrorReporter* reporter, - SimpleMemoryAllocator* allocator) - : reporter_(reporter), allocator_(allocator) {} - - // Initializes the builder by allocating AllocationInfo array from the - // simple memory allocator. - TfLiteStatus Init(size_t tensor_count, size_t scratch_buffer_count) { - tensor_count_ = tensor_count; - buffer_count_ = scratch_buffer_count; - return Allocate(); - } + AllocationInfoBuilder(AllocationInfo* info, size_t tensor_count, + size_t scratch_buffer_count, ErrorReporter* reporter) + : info_(info), + tensor_count_(tensor_count), + buffer_count_(scratch_buffer_count), + reporter_(reporter) {} // Check if model contains offline planned buffer offsets. // - If there's no metadata available, offline_planner_offsets is not set @@ -183,33 +178,14 @@ class AllocationInfoBuilder { // Returns a pointer to the built AllocationInfo array. const AllocationInfo* Finish() const { return info_; } - size_t Size() const { return tensor_count_ + buffer_count_; } private: - // Allocate the output AllocationInfo array from the allocator_; - TfLiteStatus Allocate(); - - ErrorReporter* reporter_ = nullptr; - SimpleMemoryAllocator* allocator_ = nullptr; + AllocationInfo* info_ = nullptr; size_t tensor_count_ = 0; size_t buffer_count_ = 0; - AllocationInfo* info_ = nullptr; + ErrorReporter* reporter_ = nullptr; }; -TfLiteStatus AllocationInfoBuilder::Allocate() { - size_t bytes = sizeof(AllocationInfo) * Size(); - info_ = reinterpret_cast( - allocator_->AllocateFromTail(bytes, alignof(AllocationInfo))); - if (info_ == nullptr) { - TF_LITE_REPORT_ERROR( - reporter_, - "Failed to allocate memory for allocation_info, %d bytes required", - bytes); - return kTfLiteError; - } - return kTfLiteOk; -} - TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph, const int32_t* offline_offsets, TfLiteEvalTensor* eval_tensors) { @@ -1075,62 +1051,75 @@ TfLiteStatus MicroAllocator::CommitStaticMemoryPlan( // 2. Add them into the planner (such as the GreedyMemoryPlanner). // 3. Static memory planning using the planner. // 4. Set tensor/buffer pointers based on the offsets from the previous step. + // // Note that AllocationInfo is only needed for creating the plan. It will be - // thrown away when the child allocator (tmp_allocator) goes out of scope. - { - // TODO(b/162595810): Use temp allocation buffer instead of a stack - // instance: - SimpleMemoryAllocator tmp_allocator(error_reporter_, - memory_allocator_->GetBufferHead(), - memory_allocator_->GetTail()); + // allocated from the temp section and cleaned up at the bottom of this + // function. - AllocationInfoBuilder builder(error_reporter_, &tmp_allocator); - TF_LITE_ENSURE_STATUS(builder.Init(subgraph->tensors()->size(), - scratch_buffer_request_count_)); + size_t allocation_info_count = + subgraph->tensors()->size() + scratch_buffer_request_count_; + size_t bytes = sizeof(AllocationInfo) * allocation_info_count; - const int32_t* offline_planner_offsets = nullptr; - TF_LITE_ENSURE_STATUS( - builder.GetOfflinePlannedOffsets(model, &offline_planner_offsets)); - TF_LITE_ENSURE_STATUS( - builder.AddTensors(subgraph, offline_planner_offsets, eval_tensors)); - - internal::ScratchBufferRequest* scratch_buffer_requests = - GetScratchBufferRequests(); - - TF_LITE_ENSURE_STATUS(builder.AddScratchBuffers(scratch_buffer_requests, - scratch_buffer_handles)); - - const AllocationInfo* allocation_info = builder.Finish(); - - // Remaining arena size that memory planner can use for calculating offsets. - size_t remaining_arena_size = - tmp_allocator.GetAvailableMemory(kBufferAlignment); - uint8_t* planner_arena = - tmp_allocator.AllocateTemp(remaining_arena_size, kBufferAlignment); - TF_LITE_ENSURE(error_reporter_, planner_arena != nullptr); - GreedyMemoryPlanner planner(planner_arena, remaining_arena_size); - TF_LITE_ENSURE_STATUS( - CreatePlan(error_reporter_, &planner, allocation_info, builder.Size())); - - size_t actual_available_arena_size = - memory_allocator_->GetAvailableMemory(kBufferAlignment); - - // Make sure we have enough arena size. - if (planner.GetMaximumMemorySize() > actual_available_arena_size) { - TF_LITE_REPORT_ERROR( - error_reporter_, - "Arena size is too small for all buffers. Needed %u but only " - "%u was available.", - planner.GetMaximumMemorySize(), actual_available_arena_size); - return kTfLiteError; - } - // Commit the plan. - TF_LITE_ENSURE_STATUS(CommitPlan(error_reporter_, &planner, - memory_allocator_->GetBufferHead(), - allocation_info, builder.Size())); - head_usage = planner.GetMaximumMemorySize(); + // Allocate an array of AllocationInfo structs from the temp section. This + // struct will be used by AllocationInfoBuilder to find buffer usage. + AllocationInfo* allocation_info = reinterpret_cast( + memory_allocator_->AllocateTemp(bytes, alignof(AllocationInfo))); + if (allocation_info == nullptr) { + TF_LITE_REPORT_ERROR( + error_reporter_, + "Failed to allocate memory for allocation_info, %d bytes required", + bytes); + return kTfLiteError; } + // Use the AllocationInfoBuilder class to help determine where buffers are + // used in the subgraph. + AllocationInfoBuilder builder(allocation_info, subgraph->tensors()->size(), + scratch_buffer_request_count_, error_reporter_); + + const int32_t* offline_planner_offsets = nullptr; + TF_LITE_ENSURE_STATUS( + builder.GetOfflinePlannedOffsets(model, &offline_planner_offsets)); + TF_LITE_ENSURE_STATUS( + builder.AddTensors(subgraph, offline_planner_offsets, eval_tensors)); + + internal::ScratchBufferRequest* scratch_buffer_requests = + GetScratchBufferRequests(); + + TF_LITE_ENSURE_STATUS(builder.AddScratchBuffers(scratch_buffer_requests, + scratch_buffer_handles)); + + // Remaining arena size that memory planner can use for calculating offsets. + size_t remaining_arena_size = + memory_allocator_->GetAvailableMemory(kBufferAlignment); + uint8_t* planner_arena = + memory_allocator_->AllocateTemp(remaining_arena_size, kBufferAlignment); + TF_LITE_ENSURE(error_reporter_, planner_arena != nullptr); + GreedyMemoryPlanner planner(planner_arena, remaining_arena_size); + TF_LITE_ENSURE_STATUS(CreatePlan(error_reporter_, &planner, allocation_info, + allocation_info_count)); + + // Reset all temp allocations used above: + memory_allocator_->ResetTempAllocations(); + + size_t actual_available_arena_size = + memory_allocator_->GetAvailableMemory(kBufferAlignment); + + // Make sure we have enough arena size. + if (planner.GetMaximumMemorySize() > actual_available_arena_size) { + TF_LITE_REPORT_ERROR( + error_reporter_, + "Arena size is too small for all buffers. Needed %u but only " + "%u was available.", + planner.GetMaximumMemorySize(), actual_available_arena_size); + return kTfLiteError; + } + // Commit the plan. + TF_LITE_ENSURE_STATUS(CommitPlan(error_reporter_, &planner, + memory_allocator_->GetBufferHead(), + allocation_info, allocation_info_count)); + head_usage = planner.GetMaximumMemorySize(); + // The head is used to store memory plans for one model at a time during the // model preparation stage, and is re-purposed to store scratch buffer handles // during model invocation. The head must be as large as the greater of the