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 #include "tensorflow/lite/micro/simple_memory_allocator.h"
17
18 #include <cstddef>
19 #include <cstdint>
20 #include <new>
21
22 #include "tensorflow/lite/c/common.h"
23 #include "tensorflow/lite/core/api/error_reporter.h"
24 #include "tensorflow/lite/kernels/internal/compatibility.h"
25 #include "tensorflow/lite/micro/memory_helpers.h"
26
27 namespace tflite {
28
SimpleMemoryAllocator(ErrorReporter * error_reporter,uint8_t * buffer_head,uint8_t * buffer_tail)29 SimpleMemoryAllocator::SimpleMemoryAllocator(ErrorReporter* error_reporter,
30 uint8_t* buffer_head,
31 uint8_t* buffer_tail)
32 : error_reporter_(error_reporter),
33 buffer_head_(buffer_head),
34 buffer_tail_(buffer_tail),
35 head_(buffer_head),
36 tail_(buffer_tail),
37 temp_(buffer_head_) {}
38
SimpleMemoryAllocator(ErrorReporter * error_reporter,uint8_t * buffer,size_t buffer_size)39 SimpleMemoryAllocator::SimpleMemoryAllocator(ErrorReporter* error_reporter,
40 uint8_t* buffer,
41 size_t buffer_size)
42 : SimpleMemoryAllocator(error_reporter, buffer, buffer + buffer_size) {}
43
44 /* static */
Create(ErrorReporter * error_reporter,uint8_t * buffer_head,size_t buffer_size)45 SimpleMemoryAllocator* SimpleMemoryAllocator::Create(
46 ErrorReporter* error_reporter, uint8_t* buffer_head, size_t buffer_size) {
47 TFLITE_DCHECK(error_reporter != nullptr);
48 TFLITE_DCHECK(buffer_head != nullptr);
49 SimpleMemoryAllocator tmp =
50 SimpleMemoryAllocator(error_reporter, buffer_head, buffer_size);
51
52 // Allocate enough bytes from the buffer to create a SimpleMemoryAllocator.
53 // The new instance will use the current adjusted tail buffer from the tmp
54 // allocator instance.
55 uint8_t* allocator_buffer = tmp.AllocateFromTail(
56 sizeof(SimpleMemoryAllocator), alignof(SimpleMemoryAllocator));
57 // Use the default copy constructor to populate internal states.
58 return new (allocator_buffer) SimpleMemoryAllocator(tmp);
59 }
60
~SimpleMemoryAllocator()61 SimpleMemoryAllocator::~SimpleMemoryAllocator() {}
62
SetHeadBufferSize(size_t size,size_t alignment)63 TfLiteStatus SimpleMemoryAllocator::SetHeadBufferSize(size_t size,
64 size_t alignment) {
65 if (head_ != temp_) {
66 TF_LITE_REPORT_ERROR(
67 error_reporter_,
68 "Internal error: SetHeadBufferSize() needs to be called "
69 "after ResetTempAllocations().");
70 return kTfLiteError;
71 }
72
73 uint8_t* const aligned_result = AlignPointerUp(buffer_head_, alignment);
74 const size_t available_memory = tail_ - aligned_result;
75 if (available_memory < size) {
76 TF_LITE_REPORT_ERROR(
77 error_reporter_,
78 "Failed to set head size. Requested: %u, available %u, missing: %u",
79 size, available_memory, size - available_memory);
80 return kTfLiteError;
81 }
82 head_ = aligned_result + size;
83 temp_ = head_;
84
85 return kTfLiteOk;
86 }
87
AllocateFromTail(size_t size,size_t alignment)88 uint8_t* SimpleMemoryAllocator::AllocateFromTail(size_t size,
89 size_t alignment) {
90 uint8_t* const aligned_result = AlignPointerDown(tail_ - size, alignment);
91 if (aligned_result < head_) {
92 #ifndef TF_LITE_STRIP_ERROR_STRINGS
93 const size_t missing_memory = head_ - aligned_result;
94 TF_LITE_REPORT_ERROR(error_reporter_,
95 "Failed to allocate tail memory. Requested: %u, "
96 "available %u, missing: %u",
97 size, size - missing_memory, missing_memory);
98 #endif
99 return nullptr;
100 }
101 tail_ = aligned_result;
102 return aligned_result;
103 }
104
AllocateTemp(size_t size,size_t alignment)105 uint8_t* SimpleMemoryAllocator::AllocateTemp(size_t size, size_t alignment) {
106 uint8_t* const aligned_result = AlignPointerUp(temp_, alignment);
107 const size_t available_memory = tail_ - aligned_result;
108 if (available_memory < size) {
109 TF_LITE_REPORT_ERROR(error_reporter_,
110 "Failed to allocate temp memory. Requested: %u, "
111 "available %u, missing: %u",
112 size, available_memory, size - available_memory);
113 return nullptr;
114 }
115 temp_ = aligned_result + size;
116 return aligned_result;
117 }
118
ResetTempAllocations()119 void SimpleMemoryAllocator::ResetTempAllocations() { temp_ = head_; }
120
GetHeadBuffer() const121 uint8_t* SimpleMemoryAllocator::GetHeadBuffer() const { return buffer_head_; }
122
GetHeadUsedBytes() const123 size_t SimpleMemoryAllocator::GetHeadUsedBytes() const {
124 return head_ - buffer_head_;
125 }
126
GetTailUsedBytes() const127 size_t SimpleMemoryAllocator::GetTailUsedBytes() const {
128 return buffer_tail_ - tail_;
129 }
130
GetAvailableMemory(size_t alignment) const131 size_t SimpleMemoryAllocator::GetAvailableMemory(size_t alignment) const {
132 uint8_t* const aligned_temp = AlignPointerUp(temp_, alignment);
133 uint8_t* const aligned_tail = AlignPointerDown(tail_, alignment);
134 return aligned_tail - aligned_temp;
135 }
136
GetUsedBytes() const137 size_t SimpleMemoryAllocator::GetUsedBytes() const {
138 return GetBufferSize() - (tail_ - temp_);
139 }
140
GetBufferSize() const141 size_t SimpleMemoryAllocator::GetBufferSize() const {
142 return buffer_tail_ - buffer_head_;
143 }
144
head() const145 uint8_t* SimpleMemoryAllocator::head() const { return head_; }
146
tail() const147 uint8_t* SimpleMemoryAllocator::tail() const { return tail_; }
148
149 } // namespace tflite
150