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