MicroAllocator

class MicroAllocator

Allocator responsible for allocating memory for all intermediate tensors necessary to invoke a model.

The lifetime of the model, tensor arena and error reporter must be at least as long as that of the allocator object, since the allocator needs them to be accessible during its entire lifetime.

The MicroAllocator simply plans out additional allocations that are required to standup a model for inference in TF Micro. This class currently relies on an additional allocator - SingleArenaBufferAllocator - for all allocations from an arena. These allocations are divided into head (non-persistent) and tail (persistent) regions:

Memory layout to help understand how it works This information could change in the future version. ************** .memory_allocator->GetBuffer() Tensors/Scratch buffers (head) ************** .head_watermark unused memory ************** .memory_allocator->GetBuffer() + ->GetMaxBufferSize()

->GetDataSize() persistent area (tail) ************** .memory_allocator->GetBuffer() + ->GetMaxBufferSize()

static MicroAllocator * tflite::MicroAllocator::Create(
    uint8_t *tensor_arena,
    size_t arena_size,
    MemoryPlannerType memory_planner_type = MemoryPlannerType::kGreedy
)

Creates a MicroAllocator instance from a given tensor arena. This arena will be managed by the created instance. The GreedyMemoryPlanner will by default be used and created on the arena. Note: Please use alignas(16) to make sure tensor_arena is 16 bytes aligned, otherwise some head room will be wasted. TODO(b/157615197): Cleanup constructor + factory usage.

Parameters

Parameters of Create
Name	Type	Default
`tensor_arena`	`uint8_t *`	Required
`arena_size`	`size_t`	Required
`memory_planner_type`	`MemoryPlannerType`	MemoryPlannerType::kGreedy

static MicroAllocator * tflite::MicroAllocator::Create(
    uint8_t *tensor_arena,
    size_t arena_size,
    MicroMemoryPlanner *memory_planner
)

Creates a MicroAllocator instance from a given tensor arena and a given MemoryPlanner. This arena will be managed by the created instance. Note: Please use alignas(16) to make sure tensor_arena is 16 bytes aligned, otherwise some head room will be wasted.

Parameters

Parameters of Create
Name	Type	Default
`tensor_arena`	`uint8_t *`	Required
`arena_size`	`size_t`	Required
`memory_planner`	`MicroMemoryPlanner *`	Required

static MicroAllocator * tflite::MicroAllocator::Create(
    SingleArenaBufferAllocator *memory_allocator,
    MicroMemoryPlanner *memory_planner
)

Creates a MicroAllocator instance using the provided SingleArenaBufferAllocator instance and the MemoryPlanner. This allocator instance will use the SingleArenaBufferAllocator instance to manage allocations internally.

Parameters

Parameters of Create
Name	Type	Default	Description
`memory_allocator`	`SingleArenaBufferAllocator *`	Required
`memory_planner`	`MicroMemoryPlanner *`	Required

static MicroAllocator * tflite::MicroAllocator::Create(
    uint8_t *persistent_tensor_arena,
    size_t persistent_arena_size,
    uint8_t *non_persistent_tensor_arena,
    size_t non_persistent_arena_size,
    MemoryPlannerType memory_planner_type = MemoryPlannerType::kGreedy
)

Parameters

Parameters of Create
Name	Type	Default
`persistent_tensor_arena`	`uint8_t *`	Required
`persistent_arena_size`	`size_t`	Required
`non_persistent_tensor_arena`	`uint8_t *`	Required
`non_persistent_arena_size`	`size_t`	Required
`memory_planner_type`	`MemoryPlannerType`	MemoryPlannerType::kGreedy

static size_t tflite::MicroAllocator::GetDefaultTailUsage(bool is_memory_planner_given)

Returns the fixed amount of memory overhead of MicroAllocator.

Parameters

Parameters of GetDefaultTailUsage
Name	Type	Default	Description
`is_memory_planner_given`	`bool`	Required

bool tflite::MicroAllocator::preserves_all_tensor() const

Returns True if the MicroAllocator uses a LinearMemoryPlanner(is compatible with the PerserveAllTensors flag / feature ) and False otherwise.

SubgraphAllocations * tflite::MicroAllocator::StartModelAllocation(const Model *model)

Allocates internal resources required for model inference for each subgraph from the arena.

This method will run through the flatbuffer data supplied in the model to properly allocate tensor, node, and op registration data. This method is expected to be followed with a call to FinishModelAllocation() Returns a pointer to an array of SubgraphAllocations (also stored in the tail of the arena) where each index corresponds to a different subgraph in the model. Return value is nullptr if the allocations failed.

Parameters

Parameters of StartModelAllocation
Name	Type	Default	Description
`model`	`const Model *`	Required

TfLiteStatus tflite::MicroAllocator::FinishModelAllocation(
    const Model *model,
    SubgraphAllocations *subgraph_allocations,
    ScratchBufferHandle **scratch_buffer_handles
)

Finish allocating internal resources required for model inference.

-Plan the memory for activation tensors and scratch buffers. -Update eval tensors for each subgraph based on planned offsets. -Allocate scratch buffer handles array and update based on planned offsets.

This method should be called after assigning model resources in StartModelAllocation(). The subgraph_allocations pointer should be the value passed into this class during StartModelAllocation(). Scratch buffer handles are stored in the out-param scratch_buffer_handles array which is allocated in this method. This value will be used in GetScratchBuffer call to retrieve scratch buffers.

Parameters

Parameters of FinishModelAllocation
Name	Type	Default
`model`	`const Model *`	Required
`subgraph_allocations`	`SubgraphAllocations *`	Required
`scratch_buffer_handles`	`ScratchBufferHandle **`	Required

virtual TfLiteTensor * tflite::MicroAllocator::AllocatePersistentTfLiteTensor(
    const Model *model,
    const SubgraphAllocations *subgraph_allocations,
    int tensor_index,
    int subgraph_index
)

Allocates a TfLiteTensor struct and populates the returned value with properties from the model flatbuffer. This struct is allocated from persistent arena memory is only guaranteed for the lifetime of the application. The eval_tensors pointer should be the value passed into this class during StartModelAllocation() and contains the source-of-truth for buffers.

Parameters

Parameters of AllocatePersistentTfLiteTensor
Name	Type	Default
`model`	`const Model *`	Required
`subgraph_allocations`	`const SubgraphAllocations *`	Required
`tensor_index`	`int`	Required
`subgraph_index`	`int`	Required

virtual TfLiteTensor * tflite::MicroAllocator::AllocateTempTfLiteTensor(
    const Model *model,
    const SubgraphAllocations *subgraph_allocations,
    int tensor_index,
    int subgraph_index
)

Allocates a TfLiteTensor struct and populates the returned value with properties from the model flatbuffer. This struct is allocated from temporary arena memory is only guaranteed until a call is made to ResetTempAllocations(). Subgraph_allocations contains the array of TfLiteEvalTensors. If the newly allocated temp at the specified subgraph and tensor index is already present int the TfLiteEvalTensor array, its data buffer will be reused.

Parameters

Parameters of AllocateTempTfLiteTensor
Name	Type	Default
`model`	`const Model *`	Required
`subgraph_allocations`	`const SubgraphAllocations *`	Required
`tensor_index`	`int`	Required
`subgraph_index`	`int`	Required

virtual void tflite::MicroAllocator::DeallocateTempTfLiteTensor(TfLiteTensor *)

Parameters

Parameters of DeallocateTempTfLiteTensor
Type	Default	Description
`TfLiteTensor *`	Required

virtual uint8_t * tflite::MicroAllocator::AllocateTempBuffer(size_t size, size_t alignment)

Returns a pointer to a buffer from the temporary arena memory and is only guaranteed until a call is made to ResetTempAllocations().

Parameters

Parameters of AllocateTempBuffer
Name	Type	Default	Description
`size`	`size_t`	Required
`alignment`	`size_t`	Required

virtual void tflite::MicroAllocator::DeallocateTempBuffer(uint8_t *buffer)

Signals that the temporary buffer no longer needed.

Parameters

Parameters of DeallocateTempBuffer
Name	Type	Default	Description
`buffer`	`uint8_t *`	Required

virtual TfLiteStatus tflite::MicroAllocator::ResetTempAllocations()

Resets all temporary allocations. This method should be called after a chain of temp allocations (e.g. chain of TfLiteTensor objects via AllocateTfLiteTensor()).

virtual bool tflite::MicroAllocator::IsAllTempDeallocated()

Returns true if all temporary buffers including temp TfLiteTensor are already deallocated.

virtual void * tflite::MicroAllocator::AllocatePersistentBuffer(size_t bytes)

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.

Parameters

Parameters of AllocatePersistentBuffer
Name	Type	Default	Description
`bytes`	`size_t`	Required

TfLiteStatus tflite::MicroAllocator::RequestScratchBufferInArena(size_t bytes, int subgraph_idx, int *buffer_idx)

Register a scratch buffer of size bytes for Node with node_id. This method only requests a buffer with a given size to be used after a model has finished allocation via FinishModelAllocation(). All requested buffers will be accessible by the out-param in that method.

Parameters

Parameters of RequestScratchBufferInArena
Name	Type	Default
`bytes`	`size_t`	Required
`subgraph_idx`	`int`	Required
`buffer_idx`	`int *`	Required

TfLiteStatus tflite::MicroAllocator::FinishPrepareNodeAllocations(int node_id)

Finish allocating a specific NodeAndRegistration prepare block (kernel entry for a model) with a given node ID. This call ensures that any scratch buffer requests and temporary allocations are handled and ready for the next node prepare block.

Parameters

Parameters of FinishPrepareNodeAllocations
Name	Type	Default	Description
`node_id`	`int`	Required

size_t tflite::MicroAllocator::used_bytes() const

Returns the arena usage in bytes, only available after FinishModelAllocation. Otherwise, it will return 0.

TfLiteBridgeBuiltinDataAllocator * tflite::MicroAllocator::GetBuiltinDataAllocator()

tflite::MicroAllocator::MicroAllocator(SingleArenaBufferAllocator *memory_allocator, MicroMemoryPlanner *memory_planner)

Parameters

Parameters of MicroAllocator
Name	Type	Default	Description
`memory_allocator`	`SingleArenaBufferAllocator *`	Required
`memory_planner`	`MicroMemoryPlanner *`	Required

tflite::MicroAllocator::MicroAllocator(
    IPersistentBufferAllocator *persistent_buffer_allocator,
    INonPersistentBufferAllocator *non_persistent_buffer_allocator,
    MicroMemoryPlanner *memory_planner
)

Parameters

Parameters of MicroAllocator
Name	Type	Default
`persistent_buffer_allocator`	`IPersistentBufferAllocator *`	Required
`non_persistent_buffer_allocator`	`INonPersistentBufferAllocator *`	Required
`memory_planner`	`MicroMemoryPlanner *`	Required

virtual tflite::MicroAllocator::~MicroAllocator()

virtual TfLiteStatus tflite::MicroAllocator::AllocateNodeAndRegistrations(
    const Model *model,
    SubgraphAllocations *subgraph_allocations
)

Allocates an array in the arena to hold pointers to the node and registration pointers required to represent the inference graph of the model.

Parameters

Parameters of AllocateNodeAndRegistrations
Name	Type	Default	Description
`model`	`const Model *`	Required
`subgraph_allocations`	`SubgraphAllocations *`	Required

virtual TfLiteStatus tflite::MicroAllocator::AllocateTfLiteEvalTensors(
    const Model *model,
    SubgraphAllocations *subgraph_allocations
)

Allocates the list of persistent TfLiteEvalTensors that are used for the “eval” phase of model inference. These structs will be the source of truth for all tensor buffers.

Parameters

Parameters of AllocateTfLiteEvalTensors
Name	Type	Default	Description
`model`	`const Model *`	Required
`subgraph_allocations`	`SubgraphAllocations *`	Required

virtual TfLiteStatus tflite::MicroAllocator::AllocateVariables(
    const SubGraph *subgraph,
    TfLiteEvalTensor *eval_tensors,
    const int32_t *offline_planner_offsets
)

Allocates persistent tensor buffers for variable tensors in the subgraph. Online and offline variable tensors are handled differently hence the offline_planner_offsets parameter is needed.

Parameters

Parameters of AllocateVariables
Name	Type	Default
`subgraph`	`const SubGraph *`	Required
`eval_tensors`	`TfLiteEvalTensor *`	Required
`offline_planner_offsets`	`const int32_t *`	Required

virtual TfLiteTensor * tflite::MicroAllocator::AllocatePersistentTfLiteTensorInternal()

Allocate and return a persistent TfLiteTensor. TODO(b/162311891): Drop this method when the interpreter has an API for accessing TfLiteEvalTensor structs.

virtual TfLiteStatus tflite::MicroAllocator::PopulateTfLiteTensorFromFlatbuffer(
    const Model *model,
    TfLiteTensor *tensor,
    int tensor_index,
    int subgraph_idx,
    bool allocate_temp
)

Populates a TfLiteTensor struct with data from the model flatbuffer. Any quantization data is allocated from either the tail (persistent) or temp sections of the arena based on the allocation flag.

Parameters

Parameters of PopulateTfLiteTensorFromFlatbuffer
Name	Type	Default
`model`	`const Model *`	Required
`tensor`	`TfLiteTensor *`	Required
`tensor_index`	`int`	Required
`subgraph_idx`	`int`	Required
`allocate_temp`	`bool`	Required