diff options
Diffstat (limited to 'tensorflow/contrib/lite/arena_planner.h')
| -rw-r--r-- | tensorflow/contrib/lite/arena_planner.h | 17 |
1 files changed, 16 insertions, 1 deletions
diff --git a/tensorflow/contrib/lite/arena_planner.h b/tensorflow/contrib/lite/arena_planner.h index 1d84950e91..55003cf4e9 100644 --- a/tensorflow/contrib/lite/arena_planner.h +++ b/tensorflow/contrib/lite/arena_planner.h @@ -25,6 +25,10 @@ limitations under the License. namespace tflite { +// Memory allocation tuning +constexpr const int kDefaultArenaAlignment = 64; +constexpr const int kDefaultTensorAlignment = 64; + struct AllocationInfo; // A memory planner that makes all the allocations using arenas. @@ -47,7 +51,8 @@ class ArenaPlanner : public MemoryPlanner { // graph will not share memory with any other tensor, effectively preserving // them until the end of inference. ArenaPlanner(TfLiteContext* context, std::unique_ptr<GraphInfo> graph_info, - bool preserve_inputs); + bool preserve_inputs, bool preserve_intermediates, + int tensor_alignment = kDefaultTensorAlignment); ~ArenaPlanner() override; ArenaPlanner(const ArenaPlanner&) = delete; ArenaPlanner& operator=(const ArenaPlanner&) = delete; @@ -104,7 +109,17 @@ class ArenaPlanner : public MemoryPlanner { // declared as kTfLiteArenaRwPersistent. SimpleMemoryArena persistent_arena_; + // Ensure that the memory self-allocated for inputs is never reused by the + // allocator. This allows for example, multiple runs without getting + // unpredictable results. bool preserve_inputs_; + + // If true, then no overlapping of memory areas is done, meaning intermediates + // results can be queried after running (modulo running delegates). + bool preserve_intermediates_; + + // Number of bytes that tensor buffers should be aligned to. + int tensor_alignment_; }; } // namespace tflite |