Drop the stack-allocated instance of SimpleMemoryAllocator when committing the memory plan in MicroAllocator.

All allocations are now handled internally in ::CommitStaticMemoryPlan() to improve readability and tracking of allocations.

PiperOrigin-RevId: 335550726
Change-Id: Ia472939d216b950b234e9192fb60206f4a247c91

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,
+  AllocationInfoBuilder(AllocationInfo* info, size_t tensor_count,
-                        SimpleMemoryAllocator* allocator)
+                        size_t scratch_buffer_count, ErrorReporter* reporter)
-      : reporter_(reporter), allocator_(allocator) {}
+      : info_(info),
-
+        tensor_count_(tensor_count),
-  // Initializes the builder by allocating AllocationInfo array from the
+        buffer_count_(scratch_buffer_count),
-  // simple memory allocator.
+        reporter_(reporter) {}
-  TfLiteStatus Init(size_t tensor_count, size_t scratch_buffer_count) {
-    tensor_count_ = tensor_count;
-    buffer_count_ = scratch_buffer_count;
-    return Allocate();
-  }
 
   // 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_;
+  AllocationInfo* info_ = nullptr;
-  TfLiteStatus Allocate();
-
-  ErrorReporter* reporter_ = nullptr;
-  SimpleMemoryAllocator* allocator_ = 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<AllocationInfo*>(
-      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.
+  // allocated from the temp section and cleaned up at the bottom of this
-  {
+  // function.
-    // TODO(b/162595810): Use temp allocation buffer instead of a stack
-    // instance:
-    SimpleMemoryAllocator tmp_allocator(error_reporter_,
-                                        memory_allocator_->GetBufferHead(),
-                                        memory_allocator_->GetTail());
 
-    AllocationInfoBuilder builder(error_reporter_, &tmp_allocator);
+  size_t allocation_info_count =
-    TF_LITE_ENSURE_STATUS(builder.Init(subgraph->tensors()->size(),
+      subgraph->tensors()->size() + scratch_buffer_request_count_;
-                                       scratch_buffer_request_count_));
+  size_t bytes = sizeof(AllocationInfo) * allocation_info_count;
 
-    const int32_t* offline_planner_offsets = nullptr;
+  // Allocate an array of AllocationInfo structs from the temp section. This
-    TF_LITE_ENSURE_STATUS(
+  // struct will be used by AllocationInfoBuilder to find buffer usage.
-        builder.GetOfflinePlannedOffsets(model, &offline_planner_offsets));
+  AllocationInfo* allocation_info = reinterpret_cast<AllocationInfo*>(
-    TF_LITE_ENSURE_STATUS(
+      memory_allocator_->AllocateTemp(bytes, alignof(AllocationInfo)));
-        builder.AddTensors(subgraph, offline_planner_offsets, eval_tensors));
+  if (allocation_info == nullptr) {
-
+    TF_LITE_REPORT_ERROR(
-    internal::ScratchBufferRequest* scratch_buffer_requests =
+        error_reporter_,
-        GetScratchBufferRequests();
+        "Failed to allocate memory for allocation_info, %d bytes required",
-
+        bytes);
-    TF_LITE_ENSURE_STATUS(builder.AddScratchBuffers(scratch_buffer_requests,
+    return kTfLiteError;
-                                                    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();
   }
 
+  // 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