experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

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

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

PiperOrigin-RevId: 335550726
Change-Id: Ia472939d216b950b234e9192fb60206f4a247c91
This commit is contained in:
Nick Kreeger 2020-10-05 19:21:12 -07:00 committed by TensorFlower Gardener
parent 083a88f29d
commit ba92a24c57
1 changed files with 72 additions and 83 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

155
tensorflow/lite/micro/micro_allocator.cc
View File

@ -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<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.
{
// 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<AllocationInfo*>(
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