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{}; 94 95 private: 96 // Free all the all allocations. 97 void FreeAllAllocations(); 98 99 // Assign absolute memory location to a tensor. 100 TfLiteStatus ResolveTensorAllocation(int tensor_index); 101 102 TfLiteContext* context_; 103 std::unique_ptr<GraphInfo> graph_info_; 104 105 // Stores allocation data for all tensors. 106 std::vector<SimpleAlloc> allocs_; 107 108 // First node, that uses the tensor. It needs to be allocated before 109 // execution of the node's operation. 110 std::vector<int32_t> alloc_node_; 111 112 // Last node, that uses the tensor. It can be deallocated after execution of 113 // the node's operation. 114 std::vector<int32_t> dealloc_node_; 115 }; 116 117 } // namespace tflite 118 119 #endif // TENSORFLOW_LITE_SIMPLE_PLANNER_H_ 120