xref: /aosp_15_r20/external/tensorflow/tensorflow/lite/simple_planner.h (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1 /* Copyright 2021 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_SIMPLE_PLANNER_H_
16 #define TENSORFLOW_LITE_SIMPLE_PLANNER_H_
17 
18 #include <cassert>
19 #include <cstdint>
20 #include <memory>
21 #include <vector>
22 
23 #include "tensorflow/lite/c/common.h"
24 #include "tensorflow/lite/graph_info.h"
25 #include "tensorflow/lite/memory_planner.h"
26 #include "tensorflow/lite/util.h"
27 
28 namespace tflite {
29 
30 // A structure to keep heap allocation records. This structure is used by
31 // SimplePlanner::allocs_.
32 struct SimpleAlloc {
SimpleAllocSimpleAlloc33   SimpleAlloc() { reset(); }
34 
35   // Size of allocation.
36   size_t size;
37   // The index of the node that first needs to use this tensor.
38   int32_t node;
39   // Allocated heap memory address of allocation.
40   char* ptr;
41 
42   // Reset member variables.
resetSimpleAlloc43   inline void reset() {
44     size = 0;
45     node = 0;
46     ptr = nullptr;
47   }
48 
49   // Allocate heap memory for a tensor with the given size and first_node
50   // information.
allocSimpleAlloc51   inline bool alloc(size_t new_size, int32_t new_first_node) {
52     if (new_size == 0) {
53       return false;
54     }
55     size = new_size;
56     node = new_first_node;
57     assert(ptr == nullptr);
58     ptr = static_cast<char*>(malloc(new_size));
59     return true;
60   }
61 
62   // Free allocated heap memory and reset member variables.
freeSimpleAlloc63   inline void free() {
64     if (ptr) {
65       ::free(ptr);
66     }
67     reset();
68   }
69 };
70 
71 // A memory planner that makes all the allocations using malloc()/free().
72 //
73 // This is simple implementation of MemoryPlanner which uses malloc()/free()
74 // instead of memory areana. This planner is designed for AddressSanitizer.
75 class SimplePlanner : public MemoryPlanner {
76  public:
77   // Ownership of 'context' is not taken and it must remain util the
78   // ArenaPlanner is destroyed. The inputs to the graph will not share
79   // memory with any other tensor, effectively preserving them until the end
80   // of inference.
81   SimplePlanner(TfLiteContext* context, std::unique_ptr<GraphInfo> graph_info);
82   ~SimplePlanner() override;
83   SimplePlanner(const SimplePlanner&) = delete;
84   SimplePlanner& operator=(const SimplePlanner&) = delete;
85 
86   TfLiteStatus ResetAllocations() override;
87   TfLiteStatus ResetAllocationsAfter(int node) override;
88   TfLiteStatus PlanAllocations() override;
89   TfLiteStatus ExecuteAllocations(int first_node, int last_node) override;
90   TfLiteStatus ReleaseNonPersistentMemory() override;
91   TfLiteStatus AcquireNonPersistentMemory() override;
HasNonPersistentMemory()92   bool HasNonPersistentMemory() override { return true; };
DumpDebugInfo(const std::vector<int> & execution_plan)93   void DumpDebugInfo(const std::vector<int>& execution_plan) const override{};
GetAllocInfo(size_t * arena_size,size_t * arena_persist_size)94   void GetAllocInfo(size_t* arena_size,
95                     size_t* arena_persist_size) const override{};
96 
97  private:
98   // Free all the all allocations.
99   void FreeAllAllocations();
100 
101   // Assign absolute memory location to a tensor.
102   TfLiteStatus ResolveTensorAllocation(int tensor_index);
103 
104   TfLiteContext* context_;
105   std::unique_ptr<GraphInfo> graph_info_;
106 
107   // Stores allocation data for all tensors.
108   std::vector<SimpleAlloc> allocs_;
109 
110   // First node, that uses the tensor. It needs to be allocated before
111   // execution of the node's operation.
112   std::vector<int32_t> alloc_node_;
113 
114   // Last node, that uses the tensor. It can be deallocated after execution of
115   // the node's operation.
116   std::vector<int32_t> dealloc_node_;
117 };
118 
119 }  // namespace tflite
120 
121 #endif  // TENSORFLOW_LITE_SIMPLE_PLANNER_H_
122