1 /* Copyright 2020 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 16 #ifndef TENSORFLOW_LITE_MICRO_RECORDING_MICRO_ALLOCATOR_H_ 17 #define TENSORFLOW_LITE_MICRO_RECORDING_MICRO_ALLOCATOR_H_ 18 19 #include "tensorflow/lite/micro/compatibility.h" 20 #include "tensorflow/lite/micro/micro_allocator.h" 21 #include "tensorflow/lite/micro/recording_simple_memory_allocator.h" 22 23 namespace tflite { 24 25 // List of buckets currently recorded by this class. Each type keeps a list of 26 // allocated information during model initialization. 27 // TODO(b/169834511): Add tracking for scratch buffer allocations. 28 enum class RecordedAllocationType { 29 kTfLiteEvalTensorData, 30 kPersistentTfLiteTensorData, 31 kPersistentTfLiteTensorQuantizationData, 32 kPersistentBufferData, 33 kTfLiteTensorVariableBufferData, 34 kNodeAndRegistrationArray, 35 kOpData, 36 }; 37 38 // Container for holding information about allocation recordings by a given 39 // type. Each recording contains the number of bytes requested, the actual bytes 40 // allocated (can defer from requested by alignment), and the number of items 41 // allocated. 42 struct RecordedAllocation { 43 size_t requested_bytes; 44 size_t used_bytes; 45 size_t count; 46 }; 47 48 // Utility subclass of MicroAllocator that records all allocations 49 // inside the arena. A summary of allocations can be logged through the 50 // ErrorReporter by invoking LogAllocations(). This special allocator requires 51 // an instance of RecordingSimpleMemoryAllocator to capture allocations in the 52 // head and tail. Arena allocation recording can be retrieved by type through 53 // the GetRecordedAllocation() function. This class should only be used for 54 // auditing memory usage or integration testing. 55 class RecordingMicroAllocator : public MicroAllocator { 56 public: 57 static RecordingMicroAllocator* Create(uint8_t* tensor_arena, 58 size_t arena_size, 59 ErrorReporter* error_reporter); 60 61 // Returns the recorded allocations information for a given allocation type. 62 RecordedAllocation GetRecordedAllocation( 63 RecordedAllocationType allocation_type) const; 64 65 const RecordingSimpleMemoryAllocator* GetSimpleMemoryAllocator() const; 66 67 // Logs out through the ErrorReporter all allocation recordings by type 68 // defined in RecordedAllocationType. 69 void PrintAllocations() const; 70 71 void* AllocatePersistentBuffer(size_t bytes) override; 72 73 protected: 74 TfLiteStatus AllocateNodeAndRegistrations( 75 const Model* model, 76 NodeAndRegistration** node_and_registrations) override; 77 TfLiteStatus PrepareNodeAndRegistrationDataFromFlatbuffer( 78 const Model* model, const MicroOpResolver& op_resolver, 79 NodeAndRegistration* node_and_registrations) override; 80 TfLiteStatus AllocateTfLiteEvalTensors( 81 const Model* model, TfLiteEvalTensor** eval_tensors) override; 82 TfLiteStatus AllocateVariables(const SubGraph* subgraph, 83 TfLiteEvalTensor* eval_tensors) override; 84 // TODO(b/162311891): Once all kernels have been updated to the new API drop 85 // this method. It is only used to record TfLiteTensor persistent allocations. 86 TfLiteTensor* AllocatePersistentTfLiteTensorInternal( 87 const Model* model, TfLiteEvalTensor* eval_tensors, 88 int tensor_index) override; 89 // TODO(b/162311891): Once all kernels have been updated to the new API drop 90 // this function since all allocations for quantized data will take place in 91 // the temp section. 92 TfLiteStatus PopulateTfLiteTensorFromFlatbuffer(const Model* model, 93 const SubGraph* subgraph, 94 TfLiteTensor* tensor, 95 int tensor_index, 96 bool allocate_temp) override; 97 98 private: 99 RecordingMicroAllocator(RecordingSimpleMemoryAllocator* memory_allocator, 100 ErrorReporter* error_reporter); 101 102 void PrintRecordedAllocation(RecordedAllocationType allocation_type, 103 const char* allocation_name, 104 const char* allocation_description) const; 105 106 RecordedAllocation SnapshotAllocationUsage() const; 107 void RecordAllocationUsage(const RecordedAllocation& snapshotted_allocation, 108 RecordedAllocation& recorded_allocation); 109 110 const RecordingSimpleMemoryAllocator* recording_memory_allocator_; 111 112 RecordedAllocation recorded_tflite_eval_tensor_data_ = {}; 113 RecordedAllocation recorded_persistent_tflite_tensor_data_ = {}; 114 RecordedAllocation recorded_persistent_tflite_tensor_quantization_data_ = {}; 115 RecordedAllocation recorded_persistent_buffer_data_ = {}; 116 RecordedAllocation recorded_tflite_tensor_variable_buffer_data_ = {}; 117 RecordedAllocation recorded_node_and_registration_array_data_ = {}; 118 RecordedAllocation recorded_op_data_ = {}; 119 120 TF_LITE_REMOVE_VIRTUAL_DELETE 121 }; 122 123 } // namespace tflite 124 125 #endif // TENSORFLOW_LITE_MICRO_RECORDING_MICRO_ALLOCATOR_H_ 126