• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #ifndef TENSORFLOW_LITE_ARENA_PLANNER_H_
16 #define TENSORFLOW_LITE_ARENA_PLANNER_H_
17 
18 #include <cstdint>
19 #include <memory>
20 #include <vector>
21 
22 #include "tensorflow/lite/c/common.h"
23 #include "tensorflow/lite/graph_info.h"
24 #include "tensorflow/lite/memory_planner.h"
25 #include "tensorflow/lite/simple_memory_arena.h"
26 #include "tensorflow/lite/util.h"
27 
28 namespace tflite {
29 
30 constexpr const int kDefaultArenaAlignment = 64;
31 struct AllocationInfo;
32 
33 // A memory planner that makes all the allocations using arenas.
34 //
35 // Before a model is executed by the interpreter, this class determines when
36 // each tensor needs to be allocated and deallocated, and preallocates all the
37 // necessary memory (the PlanAllocations phase). It then assigns portions of
38 // this memory buffer to each tensor (the ExecuteAllocations phase). Tensors may
39 // share some of the buffer if a tensor B is to be allocated after another
40 // tensor A has been deallocated.
41 //
42 // If dynamic tensors are used the planning steps can be repeated during model
43 // execution. Since dynamic tensors don't have sizes until after the
44 // corresponding operation is executed, this class supports incremental
45 // planning.
46 class ArenaPlanner : public MemoryPlanner {
47  public:
48   // Ownership of 'context' is not taken and it must remain util the
49   // ArenaPlanner is destroyed. If 'preserve_inputs' is true the inputs to the
50   // graph will not share memory with any other tensor, effectively preserving
51   // them until the end of inference.
52   ArenaPlanner(TfLiteContext* context, std::unique_ptr<GraphInfo> graph_info,
53                bool preserve_inputs, bool preserve_intermediates,
54                int tensor_alignment);
55   ~ArenaPlanner() override;
56   ArenaPlanner(const ArenaPlanner&) = delete;
57   ArenaPlanner& operator=(const ArenaPlanner&) = delete;
58 
59   TfLiteStatus ResetAllocations() override;
60   TfLiteStatus ResetAllocationsAfter(int node) override;
61   TfLiteStatus PlanAllocations() override;
62   TfLiteStatus ExecuteAllocations(int first_node, int last_node) override;
63   TfLiteStatus ReleaseNonPersistentMemory() override;
64   TfLiteStatus AcquireNonPersistentMemory() override;
65   bool HasNonPersistentMemory() override;
66 
67   // Returns the base arena location for a given allocation type.
68   std::intptr_t BasePointer(TfLiteAllocationType type);
69 
70  private:
71   // Make sure all the arenas have reserved enough memory to store all their
72   // tensors.
73   TfLiteStatus Commit();
74 
75   // Returns vector of tensor number ordered by the following algorithm.
76   // Comparator to sort tensors for the allocation algorithm:
77   // - Tensors that have lifespan through the whole model inference time go
78   // first;
79   // - Other tensors (e.g. intermediate and temporary ones) are sorted in
80   // non-increasing order of their size. If sizes of two tensors are equal, the
81   // one that needs to be allocated earlier goes first.
82   std::vector<int32_t> CreateTensorAllocationVector(int first_node,
83                                                     int last_node);
84 
85   // Traverse the allocation queue and reserve space in the appropriate arena
86   // for all tensors affected by ops in the interval [first_node, last_node].
87   TfLiteStatus CalculateAllocations(int first_node, int last_node);
88 
89   // Assign absolute memory location to a tensor, based on its relative
90   // position inside the corresponding arena buffer.
91   TfLiteStatus ResolveTensorAllocation(int tensor_index);
92 
93   // Register an allocation for all internal (temporary) tensors of
94   // 'node_index'.
95   TfLiteStatus CalculateAllocationOfInternalTensors(int node_index);
96 
97   // Register a deallocation for all internal (temporary) tensors of
98   // 'node_index'.
99   TfLiteStatus CalculateDeallocationOfInternalTensors(int node_index);
100 
101   TfLiteContext* context_;
102   std::unique_ptr<GraphInfo> graph_info_;
103 
104   // Stores allocation data for all tensors.
105   std::vector<ArenaAllocWithUsageInterval> allocs_;
106 
107   // First node, that uses the tensor. It needs to be allocated before
108   // execution of the node's operation.
109   std::vector<int32_t> alloc_node_;
110 
111   // Last node, that uses the tensor. It can be deallocated after execution of
112   // the node's operation.
113   std::vector<int32_t> dealloc_node_;
114 
115   // Raw memory buffer that is allocated for all temporary and graph outputs
116   // that are declared kTfLiteArenaRw.
117   SimpleMemoryArena arena_;
118 
119   // Raw memory buffer that is allocated for persistent tensors that are
120   // declared as kTfLiteArenaRwPersistent.
121   SimpleMemoryArena persistent_arena_;
122 
123   // Ensure that the memory self-allocated for inputs is never reused by the
124   // allocator. This allows for example, multiple runs without getting
125   // unpredictable results.
126   bool preserve_inputs_;
127 
128   // If true, then no overlapping of memory areas is done, meaning intermediate
129   // results can be queried after running (modulo running delegates).
130   bool preserve_intermediates_;
131 
132   // Number of bytes that tensor buffers should be aligned to.
133   int tensor_alignment_;
134 };
135 
136 }  // namespace tflite
137 
138 #endif  // TENSORFLOW_LITE_ARENA_PLANNER_H_
139