xmos-ai-tools 1.3.2.dev180__py3-none-macosx_10_15_universal2.whl

xmos_ai_tools/runtime/include/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h

@@ -0,0 +1,170 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_GREEDY_MEMORY_PLANNER_H_
+#define TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_GREEDY_MEMORY_PLANNER_H_
+
+#include "tensorflow/lite/micro/compatibility.h"
+#include "tensorflow/lite/micro/memory_planner/micro_memory_planner.h"
+
+namespace tflite_micro {
+
+constexpr int kOnlinePlannedBuffer = -1;
+
+// A memory planner that uses a greedy algorithm to arrange buffers in memory
+// to minimize the overall arena size needed.
+//
+// The algorithm works like this:
+//  - The client enters the buffer information through AddBuffer().
+//  - When a function like GetOffsetForBuffer() is called, the
+//    CalculateOffsetsIfNeeded() method is invoked.
+//  - If an up to date plan is not already present, one will be calculated.
+//  - The buffers are sorted in descending order of size.
+//  - The largest buffer is placed at offset zero.
+//  - The rest of the buffers are looped through in descending size order.
+//  - The other buffers that need to be in memory at the same time are found.
+//  - The first gap between simultaneously active buffers that the current
+//    buffer fits into will be used.
+//  - If no large-enough gap is found, the current buffer is placed after the
+//    last buffer that's simultaneously active.
+//  - This continues until all buffers are placed, and the offsets stored.
+//
+// This is not guaranteed to produce the best placement, since that's an
+// NP-Complete problem, but in practice it should produce one that's decent.
+class GreedyMemoryPlanner : public MicroMemoryPlanner {
+ public:
+  GreedyMemoryPlanner();
+  ~GreedyMemoryPlanner() override;
+
+  // You need to pass in an area of memory to be used for planning. The client
+  // should ensure the validity of the memory when it needs to use this object.
+  // This memory isn't owned by this object, so management should be handled by
+  // the client. This is so it can be stack or globally allocated if necessary
+  // on devices without dynamic memory allocation. How many buffers can be
+  // planned for will depend on the size of this scratch memory, so you should
+  // enlarge it if you see an error when calling AddBuffer(). The memory can be
+  // reused once you're done with the planner, as long as you copy the
+  // calculated offsets to another location. Each buffer requires about 36 bytes
+  // of scratch.
+  TfLiteStatus Init(unsigned char* scratch_buffer,
+                    int scratch_buffer_size) override;
+
+  // Record details of a buffer we want to place.
+  TfLiteStatus AddBuffer(int size, int first_time_used,
+                         int last_time_used) override;
+
+  // Record details of an offline planned buffer offset we want to place.
+  // offline_offset is the buffer offset from the start of the arena.
+  TfLiteStatus AddBuffer(int size, int first_time_used, int last_time_used,
+                         int offline_offset) override;
+
+  // Returns the high-water mark of used memory. This is the minimum size of a
+  // memory arena you'd need to allocate to hold these buffers.
+  size_t GetMaximumMemorySize() override;
+
+  // How many buffers have been recorded.
+  int GetBufferCount() override;
+
+  // Where a given buffer should be placed in the memory arena.
+  // This information is stored in the memory arena itself, so once the arena
+  // is used for inference, it will be overwritten.
+  TfLiteStatus GetOffsetForBuffer(int buffer_index, int* offset) override;
+
+  // Prints an ascii-art diagram of the buffer layout plan.
+  void PrintMemoryPlan() override;
+
+  // Debug method to check whether any buffer allocations are overlapping. This
+  // is an O(N^2) complexity operation, so only use for testing.
+  bool DoAnyBuffersOverlap();
+
+  // Used to store a list of buffers ordered by their offset.
+  struct ListEntry {
+    int offset;
+    int requirements_index;
+    int next_entry_index;
+  };
+
+  // Number of bytes required in order to plan a buffer.
+  static size_t per_buffer_size() {
+    const int per_buffer_size =
+        sizeof(BufferRequirements) +  // requirements_
+        sizeof(int) +                 // buffer_sizes_sorted_
+        sizeof(int) +                 // buffer_ids_sorted_
+        sizeof(ListEntry) +           // buffers_sorted_by_offset_
+        sizeof(int);                  // buffer_offsets_;
+    return per_buffer_size;
+  }
+
+  // Returns False because the GreedyMemoryPlanner doesn't preserves all tensors
+  // after invocation. Do to the fact that tensors that tensor data for tensors
+  // that aren't being used during a phase of invocation are overwritten.
+  bool preserves_all_tensors() const override { return false; }
+
+  TF_LITE_REMOVE_VIRTUAL_DELETE
+
+ private:
+  // Whether a buffer is active in a given time range.
+  bool DoesEntryOverlapInTime(const ListEntry* entry, const int first_time_used,
+                              const int last_time_used) const;
+
+  // Walks the list to return the next buffer that is active in a given time
+  // range, or a null pointer if there are none.
+  ListEntry* NextSimultaneouslyActiveBuffer(const ListEntry* start,
+                                            const int first_time_used,
+                                            const int last_time_used);
+
+  // If there isn't an up to date plan, calculate a new one.
+  void CalculateOffsetsIfNeeded();
+
+  // How many buffers we can plan for, based on the arena size we're given in
+  // the constructor.
+  int max_buffer_count_;
+
+  // The number of buffers added so far.
+  int buffer_count_;
+
+  // Records the client-provided information about each buffer.
+  struct BufferRequirements {
+    int size;
+    int offline_offset;
+    int first_time_used;
+    int last_time_used;
+  };
+
+  // Working arrays used during the layout algorithm.
+  BufferRequirements* requirements_;
+  // buffer_sizes_sorted_ and buffer_ids_sorted_ are sorted according to:
+  //   {
+  //     offline planned buffers,
+  //     online planned buffers sorted by size
+  //   }
+  int* buffer_sizes_sorted_;
+  int* buffer_ids_sorted_;
+  ListEntry* buffers_sorted_by_offset_;
+  int next_free_entry_;    // Index of the next free entry of
+                           // buffers_sorted_by_offset_
+  int first_entry_index_;  // Index of the first entry (smallest offset) of
+                           // buffers_sorted_by_offset_
+
+  // Stores the outcome of the plan, the location of each buffer in the arena.
+  int* buffer_offsets_;
+
+  // Whether buffers have been added since the last plan was calculated.
+  bool need_to_calculate_offsets_;
+};
+
+}  // namespace tflite_micro
+
+#endif  // TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_GREEDY_MEMORY_PLANNER_H_

xmos_ai_tools/runtime/include/tensorflow/lite/micro/memory_planner/linear_memory_planner.h

@@ -0,0 +1,53 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_LINEAR_MEMORY_PLANNER_H_
+#define TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_LINEAR_MEMORY_PLANNER_H_
+
+#include "tensorflow/lite/micro/compatibility.h"
+#include "tensorflow/lite/micro/memory_planner/micro_memory_planner.h"
+
+namespace tflite_micro {
+
+// The simplest possible memory planner that just lays out all buffers at
+// increasing offsets without trying to reuse memory.
+class LinearMemoryPlanner : public MicroMemoryPlanner {
+ public:
+  LinearMemoryPlanner();
+  ~LinearMemoryPlanner() override;
+
+  TfLiteStatus AddBuffer(int size, int first_time_used,
+                         int last_time_used) override;
+
+  size_t GetMaximumMemorySize() override;
+  int GetBufferCount() override;
+  TfLiteStatus GetOffsetForBuffer(int buffer_index, int* offset) override;
+
+  // Returns True because the LinearMemoryPlanner preserves all tensors after
+  // invocation.
+  bool preserves_all_tensors() const override { return true; }
+
+ private:
+  static constexpr int kMaxBufferCount = 1024;
+  size_t buffer_offsets_[kMaxBufferCount];
+  int current_buffer_count_;
+  size_t next_free_offset_;
+
+  TF_LITE_REMOVE_VIRTUAL_DELETE
+};
+
+}  // namespace tflite_micro
+
+#endif  // TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_LINEAR_MEMORY_PLANNER_H_

xmos_ai_tools/runtime/include/tensorflow/lite/micro/memory_planner/memory_plan_struct.h

@@ -0,0 +1,73 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_MEMORY_PLAN_STRUCT_H_
+#define TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_MEMORY_PLAN_STRUCT_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include "tensorflow/lite/micro/micro_utils.h"
+
+namespace tflite_micro {
+
+// This is an experimental feature and subjected to change.
+// More description is available at
+// tensorflow/lite/micro/docs/offline_memory_plan.md.
+
+// Describes a buffer's layout inside an arena. This struct should be kept as
+// small as possible for memory footprint sensitive applications and should use
+// only primitive fields, making it easy to adjust offline.
+struct BufferDescriptor {
+  // Starting offset inside an arena for this buffer.
+  // Offset is the minimum information needed for the buffer.  The user knows
+  // the model and the size of each buffer in order to lay out a valid buffer
+  // plan.
+  int32_t offset;
+};
+
+// A structure describing the lay out of buffers inside an arena.
+struct BufferPlan {
+  // Number of buffers described in this plan.
+  int32_t buffer_count;
+
+  // Each element describes one buffer.
+  // Buffer index is implicit by the order of AddBuffer() call.
+  // Specifically, indices of activation tensors are 0 … N-1 where N is the
+  // number of activation tensors.
+  // The rest are based on the order of OP requests.
+  //
+  // This is a flexible array member and should ideally be
+  // arena_entries[]; However, in order to support a variety
+  // of compilers (and without needing to add ifdef's), we
+  // are implementing the flexible array member with an array of
+  // length 1 as the last member of the struct. When the size of a BufferPlan
+  // is needed, use the provided SizeOfBufferPlan(buffer_count) that
+  // accounts for this implemenatation caveat.
+  BufferDescriptor buffer_plan_entries[1];
+};
+
+// Returns size of a BufferPlan given a buffer count. This size is compile time
+// known if buffer_count is a compile time constant.
+constexpr size_t SizeOfBufferPlan(int32_t buffer_count) {
+  // Minus 1 because a BufferPlan struct have a BufferDescriptor already.
+  // Max to provide a lower bound for the corner case of buffer_count = 0.
+  return sizeof(BufferPlan) +
+         sizeof(BufferDescriptor) * Max(buffer_count - 1, 0);
+}
+
+}  // namespace tflite_micro
+
+#endif  // TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_MEMORY_PLAN_STRUCT_H_

xmos_ai_tools/runtime/include/tensorflow/lite/micro/memory_planner/micro_memory_planner.h

@@ -0,0 +1,95 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_MICRO_MICRO_MEMORY_PLANNER_MEMORY_PLANNER_H_
+#define TENSORFLOW_LITE_MICRO_MICRO_MEMORY_PLANNER_MEMORY_PLANNER_H_
+
+#include "tensorflow/lite/c/common.h"
+
+namespace tflite_micro {
+
+// Interface class for planning the layout of memory buffers during the
+// execution of a graph.
+// It's designed to be used by a client that iterates in any order through the
+// buffers it wants to lay out, and then calls the getter functions for
+// information about the calculated layout. For example:
+//
+// SomeMemoryPlanner planner;
+// planner.AddBuffer(100, 0, 1);  // Buffer 0
+// planner.AddBuffer(50, 2, 3);   // Buffer 1
+// planner.AddBuffer(50, 2, 3);   // Buffer 2
+//
+// int offset0;
+// TF_EXPECT_OK(planner.GetOffsetForBuffer(0, &offset0));
+// int offset1;
+// TF_EXPECT_OK(planner.GetOffsetForBuffer(1, &offset1));
+// int offset2;
+// TF_EXPECT_OK(planner.GetOffsetForBuffer(2, &offset2));
+// const int arena_size_needed = planner.GetMaximumMemorySize();
+//
+// The goal is for applications to be able to experiment with different layout
+// strategies without changing their client code, by swapping out classes that
+// implement this interface.=
+class MicroMemoryPlanner {
+ public:
+  MicroMemoryPlanner() {}
+  virtual ~MicroMemoryPlanner() {}
+
+  // Pass information about a buffer's size and lifetime to the layout
+  // algorithm. The order this is called implicitly assigns an index to the
+  // result, so the buffer information that's passed into the N-th call of
+  // this method will be used as the buffer_index argument to
+  // GetOffsetForBuffer().
+  virtual TfLiteStatus AddBuffer(int size, int first_time_used,
+                                 int last_time_used) = 0;
+
+  // Record details of an offline planned buffer offset we want to place.
+  // offline_offset is the buffer offset from the start of the arena.
+  // This is to support offline memory planning from the flatbuffer metadata.
+  // By default, it returns an error.
+  virtual TfLiteStatus AddBuffer(int size, int first_time_used,
+                                 int last_time_used, int offline_offset) {
+    return kTfLiteError;
+  }
+
+  // The largest contiguous block of memory that's needed to hold the layout.
+  virtual size_t GetMaximumMemorySize() = 0;
+  // How many buffers have been added to the planner.
+  virtual int GetBufferCount() = 0;
+  // Calculated layout offset for the N-th buffer added to the planner.
+  virtual TfLiteStatus GetOffsetForBuffer(int buffer_index, int* offset) = 0;
+
+  // Provides the scratch buffer in case that the memory planner needs it.
+  // The lifetime of scratch buffers lifetime lasts until the static memory plan
+  // is committed.
+  // The default implementation is for the memory planner that does not need
+  // scratch buffer and simply returns ok.
+  virtual TfLiteStatus Init(unsigned char* scratch_buffer,
+                            int scratch_buffer_size) {
+    return kTfLiteOk;
+  }
+
+  // Method will return True if the MicroMemoryPlanner preserves all tensors
+  // after invocation, and False if it doesn't.
+  virtual bool preserves_all_tensors() const = 0;
+
+  virtual void PrintMemoryPlan() {
+    // Default does nothing.
+  }
+};
+
+}  // namespace tflite_micro
+
+#endif  // TENSORFLOW_LITE_MICRO_MICRO_MEMORY_PLANNER_MEMORY_PLANNER_H_

xmos_ai_tools/runtime/include/tensorflow/lite/micro/memory_planner/non_persistent_buffer_planner_shim.h

@@ -0,0 +1,133 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_NON_PERSISTENT_MEMORY_PLANNER_SHIM_H__
+#define TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_NON_PERSISTENT_MEMORY_PLANNER_SHIM_H__
+
+#include "tensorflow/lite/micro/compatibility.h"
+#include "tensorflow/lite/micro/memory_planner/memory_plan_struct.h"
+#include "tensorflow/lite/micro/memory_planner/micro_memory_planner.h"
+
+namespace tflite_micro {
+
+/*   This is an experimental feature and subjected to change.
+ *
+The NonPersistentMemoryPlannerShim enables TFLM to work with an external tooling
+that can plan the offset of each non persistent buffer for the Model within the
+TFLM arena.
+
+If the NonPersistentMemoryPlannerShim is used, then the final binary does not
+have any of the symbols associated with the GreedyMemoryPlanner which results in
+a reduced memory footprint.
+
+Additionally, the offline planning of the non-persistent buffers can be used to
+have a more efficient utilization compared to the GreedyMemoryPlanner.
+
+For example, consider the following hypothetical model:
+
+A1(400)                    A2(401)
+──┬─────────┐    ┌───────────
+  │         │    │
+  │         │    │
+  │         ▼    ▼
+  │       ┌────────┐
+  │       │  OP1   │
+  │       └───┬────┘       A4(201)
+  │   A3(10)  │              │
+  │           │              │
+  │           │              │
+  │       ┌───┴────┐         │
+  │       │  OP2   │◄────────┤
+  │       └───┬────┘         │
+  │   A5(11)  │      A6(202) │
+  │           │       │      │
+  │           ▼       │      │
+  │       ┌────────┐  │      │
+  │       │  OP3   │◄─┘      │
+  │       └───┬────┘         │
+  │           │      A8(200) │
+  │   A7(12)  │        │     │
+  │           │        │     │
+  │       ┌───┴────┐◄──┘     │
+  └──────►│  OP4   │         │
+          └───┬────┘◄────────┘
+              │
+      A9(13)  │
+              ▼
+
+The GreedyMemoryPlanner will give the following memory layout that requires 1012
+bytes of scratch arena size:
+
+┌─────────────────────────────────────────┬──────────────────────────┬────────┬───────┐
+│  A2(401)                                │          A1(400)         │ A4(201)│
+A3(10)│
+└─────────────────────────────────────────┴──────────────────────────┴────────┴───────┘
+
+┌───────────┬──────┬──────┐
+│ A6(202)   │A5(11)│A7(12)│
+└───────────┴──────┴──────┘
+
+┌──────────┬───────┐
+│ A8(200)  │A9(13) │
+└──────────┴───────┘
+
+But a more efficient offline memory plan that requires only 826 bytes of scratch
+arena size can be
+
+┌──────────────────────────────────────┬─────────────────────────────┬───────┬──────┐
+│      A1(400)                         │         A2(401)             │
+A3(10)│A5(11)│
+└──────────────────────────────────────┴─────────────────────────────┴───────┴──────┘
+
+                                       ┌────────────────┬────────────┬────────┬───────┐
+                                       │A4(201)         │  A8(200)   │A9(13)
+│A7(12) │ └────────────────┴────────────┴────────┴───────┘
+
+                                                        ┌─────────────┐
+                                                        │  A6(202)    │
+                                                        └─────────────┘
+
+*/
+class NonPersistentMemoryPlannerShim : public MicroMemoryPlanner {
+ public:
+  // Does not take ownership of buffer_plan, which must refer to a valid
+  // BufferPlan that outlives this object.
+  explicit NonPersistentMemoryPlannerShim(const BufferPlan* buffer_plan);
+  ~NonPersistentMemoryPlannerShim() override;
+
+  TfLiteStatus GetOffsetForBuffer(int buffer_request_index,
+                                  int* offset) override;
+
+  TfLiteStatus AddBuffer(int size, int first_time_used,
+                         int last_time_used) override;
+  size_t GetMaximumMemorySize() override;
+  int GetBufferCount() override;
+
+  // Returns False because the NonPersistentMemoryPlannerShim doesn't preserves
+  // all tensors after invocation.
+  bool preserves_all_tensors() const override { return false; }
+
+ private:
+  const BufferPlan* buffer_plan_;  // not owned, can't be null
+
+  // The number of buffers requested so far. Used for error checking.
+  int buffer_request_count_;
+
+  TF_LITE_REMOVE_VIRTUAL_DELETE
+};
+
+}  // namespace tflite_micro
+
+#endif  // TENSORFLOW_LITE_MICRO_MEMORY_PLANNER_NON_PERSISTENT_MEMORY_PLANNER_SHIM_H__

xmos_ai_tools/runtime/include/tensorflow/lite/micro/micro_allocation_info.h

@@ -0,0 +1,138 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_MICRO_MICRO_ALLOCATION_INFO_H_
+#define TENSORFLOW_LITE_MICRO_MICRO_ALLOCATION_INFO_H_
+
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/micro/compatibility.h"
+#include "tensorflow/lite/micro/flatbuffer_utils.h"
+#include "tensorflow/lite/micro/micro_allocator.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite_micro {
+
+// Used to hold information used during allocation calculations.
+struct AllocationInfo {
+  size_t bytes;
+  void** output_ptr;
+  int first_created;
+  int last_used;
+  int32_t offline_offset;
+  bool needs_allocating;
+};
+
+// Used to hold the allocation info list and related metadata for the entire
+// graph (including subgraphs). Since all subgraphs are planned together, the
+// allocation info list contains allocations for all subgraphs. Track the offset
+// into this list for each subgraph then reserve space to track all allocations.
+//
+// The AllocationInfo list is a contiguous list of allocations across all
+// subgraphs and scratch buffers. Each element here is marked as
+// s<subgraph index>t<tensor index>. The following is a possible
+// AllocationInfo list:
+// [s0t0, s0t1, s1t0, s2t1, s1t2, s3t0, s3t1, scratch0, scratch1, scratch2]
+//
+// For this example, the subgraph offsets would be [0, 2, 5] and the scratch
+// offset would be 7.
+struct GraphAllocationInfo {
+  AllocationInfo* allocation_info;
+  size_t allocation_info_count;
+  size_t* subgraph_offsets;
+  size_t scratch_offset;
+  size_t tensor_count;
+  size_t scratch_buffer_count;
+};
+
+// A helper class to construct AllocationInfo array. This array contains the
+// lifetime of tensors / scratch_buffer and will be used to calculate the memory
+// plan. Methods need to be called in order from `Create`, Init`, `Add*`, to
+// `Finish`.
+class AllocationInfoBuilder {
+ public:
+  AllocationInfoBuilder(const Model* model,
+                        INonPersistentBufferAllocator* non_persistent_allocator)
+      : model_(model), non_persistent_allocator_(non_persistent_allocator) {}
+
+  // Check if model contains offline planned buffer offsets.
+  //  - If there's no metadata available, offline_planner_offsets is not set
+  //  - If there's metadata available, offline_planner_offsets will point to the
+  //    first offset in the metadata buffer list.
+  TfLiteStatus GetOfflinePlannedOffsets(
+      const int32_t** offline_planner_offsets);
+
+  // Allocate memory for the allocation info array as well as offsets into that
+  // array for each subgraph.
+  TfLiteStatus CreateAllocationInfo(int scratch_buffer_request_count);
+
+  // Release memory used for the allocation info array.
+  TfLiteStatus FreeAllocationInfo();
+
+  // Initialize AllocationInfo for all tensors and scratch buffers in the graph.
+  TfLiteStatus InitializeAllocationInfo(const int32_t* offline_offsets,
+                                        SubgraphAllocations* allocations);
+
+  // Mark the scope of each tensor and scratch buffer across the graph. Enter
+  // all possible subgraphs invoked by each control flow operator. This method
+  // marks the maximum lifetime of each buffer so that tensors are correctly
+  // planned for all valid invocation flows.
+  TfLiteStatus MarkAllocationLifetimes(
+      int subgraph_idx, internal::ScratchBufferRequest* scratch_buffer_request,
+      ScratchBufferHandle* scratch_buffer_handles,
+      SubgraphAllocations* allocations);
+
+  // Identify control flow operators and recursively mark all subgraphs which
+  // that operator can invoke. The lifetime of all tensors within a subgraph
+  // can only be extended. The order of subgraph invocation does not matter
+  // since subgraphs within the same control flow operator are executed
+  // within their own allocation scope (planned buffers in a subgraph cannot
+  // persist beyond the end of that subgraph's invocation).
+  TfLiteStatus MarkSubgraphLifetimesIfNecessary(
+      const Operator* op,
+      internal::ScratchBufferRequest* scratch_buffer_requests,
+      ScratchBufferHandle* scratch_buffer_handles,
+      SubgraphAllocations* allocations);
+
+  // Returns the number of allocations.
+  int AllocationCount() const { return info_.allocation_info_count; }
+
+  // Returns a pointer to the built AllocationInfo array.
+  AllocationInfo* Finish() const { return info_.allocation_info; }
+
+ private:
+  // Mark the given Allocation info as first created at the specified allocation
+  // scope count. Only the first creation must be recorded since the allocation
+  // scope count monotonically increases throughout the lifetime marking
+  // process.
+  void UpdateFirstCreated(AllocationInfo* current, int allocation_scope_count);
+
+  // Mark the given AllocationInfo as last used at the specified allocation
+  // scope
+  // count. Update the last used marker every time, since the allocation scope
+  // count monotonically increases through the lifetime marking process.
+  void UpdateLastUsed(AllocationInfo* current, int allocation_scope_count);
+
+  // Validate if a subgraph satisfies assumptions.
+  TfLiteStatus ValidateSubgraph(const SubGraph* subgraph,
+                                TfLiteEvalTensor* eval_tensors);
+
+  const tflite_micro::Model* model_ = nullptr;
+  INonPersistentBufferAllocator* non_persistent_allocator_ = nullptr;
+  GraphAllocationInfo info_;
+  int allocation_scope_count_ = 0;
+};
+
+}  // namespace tflite_micro
+
+#endif  // TENSORFLOW_LITE_MICRO_MICRO_ALLOCATION_INFO_H_