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, SubgraphAllocations* subgraph_allocations) override;
76   TfLiteStatus AllocateTfLiteEvalTensors(
77       const Model* model, SubgraphAllocations* subgraph_allocations) override;
78   TfLiteStatus AllocateVariables(const SubGraph* subgraph,
79                                  TfLiteEvalTensor* eval_tensors) override;
80   // TODO(b/162311891): Once all kernels have been updated to the new API drop
81   // this method. It is only used to record TfLiteTensor persistent allocations.
82   TfLiteTensor* AllocatePersistentTfLiteTensorInternal() override;
83 
84   // TODO(b/162311891): Once all kernels have been updated to the new API drop
85   // this function since all allocations for quantized data will take place in
86   // the temp section.
87   TfLiteStatus PopulateTfLiteTensorFromFlatbuffer(const Model* model,
88                                                   TfLiteTensor* tensor,
89                                                   int tensor_index,
90                                                   int subgraph_index,
91                                                   bool allocate_temp) override;
92 
93  private:
94   RecordingMicroAllocator(RecordingSimpleMemoryAllocator* memory_allocator,
95                           ErrorReporter* error_reporter);
96 
97   void PrintRecordedAllocation(RecordedAllocationType allocation_type,
98                                const char* allocation_name,
99                                const char* allocation_description) const;
100 
101   RecordedAllocation SnapshotAllocationUsage() const;
102   void RecordAllocationUsage(const RecordedAllocation& snapshotted_allocation,
103                              RecordedAllocation& recorded_allocation);
104 
105   const RecordingSimpleMemoryAllocator* recording_memory_allocator_;
106 
107   RecordedAllocation recorded_tflite_eval_tensor_data_ = {};
108   RecordedAllocation recorded_persistent_tflite_tensor_data_ = {};
109   RecordedAllocation recorded_persistent_tflite_tensor_quantization_data_ = {};
110   RecordedAllocation recorded_persistent_buffer_data_ = {};
111   RecordedAllocation recorded_tflite_tensor_variable_buffer_data_ = {};
112   RecordedAllocation recorded_node_and_registration_array_data_ = {};
113 
114   // TODO(b/187993291): Re-enable OpData allocating tracking.
115   RecordedAllocation recorded_op_data_ = {};
116 
117   TF_LITE_REMOVE_VIRTUAL_DELETE
118 };
119 
120 }  // namespace tflite
121 
122 #endif  // TENSORFLOW_LITE_MICRO_RECORDING_MICRO_ALLOCATOR_H_
123