1 /* Copyright 2018 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 #include "tensorflow/lite/delegates/flex/buffer_map.h"
16
17 #include "tensorflow/c/c_api_internal.h"
18 #include "tensorflow/core/framework/allocation_description.pb.h"
19 #include "tensorflow/core/framework/log_memory.h"
20 #include "tensorflow/core/framework/typed_allocator.h"
21 #include "tensorflow/lite/delegates/flex/util.h"
22 #include "tensorflow/lite/string_type.h"
23 #include "tensorflow/lite/string_util.h"
24
25 namespace tflite {
26 namespace flex {
27 namespace {
28 // A tensor buffer that is allocated, deallocated and populated by TF Lite.
29 class BaseTfLiteTensorBuffer : public tensorflow::TensorBuffer {
30 using tensorflow::TensorBuffer::TensorBuffer;
31
root_buffer()32 TensorBuffer* root_buffer() override { return this; }
FillAllocationDescription(tensorflow::AllocationDescription * proto) const33 void FillAllocationDescription(
34 tensorflow::AllocationDescription* proto) const override {
35 tensorflow::int64 rb = size();
36 proto->set_requested_bytes(rb);
37 proto->set_allocator_name(tensorflow::cpu_allocator()->Name());
38 }
39
40 // Prevents input forwarding from mutating this buffer.
OwnsMemory() const41 bool OwnsMemory() const override { return false; }
42
43 protected:
LogAllocation()44 void LogAllocation() {
45 if (tensorflow::LogMemory::IsEnabled() && data() != nullptr) {
46 tensorflow::LogMemory::RecordRawAllocation(
47 "TfLiteTensorBuffer_New",
48 tensorflow::LogMemory::EXTERNAL_TENSOR_ALLOCATION_STEP_ID, size(),
49 data(), tensorflow::cpu_allocator());
50 }
51 }
LogDeallocation()52 void LogDeallocation() {
53 if (tensorflow::LogMemory::IsEnabled() && data() != nullptr) {
54 tensorflow::LogMemory::RecordRawDeallocation(
55 "TfLiteTensorBuffer_Delete",
56 tensorflow::LogMemory::EXTERNAL_TENSOR_ALLOCATION_STEP_ID, data(),
57 tensorflow::cpu_allocator(), false);
58 }
59 }
60 };
61
62 // A tensor buffer for most data types. Numeric types have exactly the same
63 // representation in TFLITE and TF, so we just need use memcpy().
64 class TfLiteTensorBuffer : public BaseTfLiteTensorBuffer {
65 public:
TfLiteTensorBuffer(const TfLiteTensor * tensor)66 explicit TfLiteTensorBuffer(const TfLiteTensor* tensor)
67 : BaseTfLiteTensorBuffer(tensorflow::cpu_allocator()->AllocateRaw(
68 EIGEN_MAX_ALIGN_BYTES, tensor->bytes)) {
69 // TODO(ahentz): if we can guarantee that TF Lite allocated tensors with
70 // the same alignment as TensorFlow (EIGEN_MAX_ALIGN_BYTES), then we can
71 // potentially eliminate the copy below.
72 len_ = tensor->bytes;
73
74 LogAllocation();
75
76 if (data()) {
77 std::memcpy(data(), tensor->data.raw, tensor->bytes);
78 }
79 }
80
~TfLiteTensorBuffer()81 ~TfLiteTensorBuffer() override {
82 LogDeallocation();
83 tensorflow::cpu_allocator()->DeallocateRaw(data());
84 }
85
size() const86 size_t size() const override { return len_; }
87
88 private:
89 size_t len_;
90 };
91
92 // A string buffer. TFLITE string tensor format is different than
93 // TF's so we need perform the conversion here.
94 class StringTfLiteTensorBuffer : public BaseTfLiteTensorBuffer {
95 public:
StringTfLiteTensorBuffer(const TfLiteTensor * tensor)96 explicit StringTfLiteTensorBuffer(const TfLiteTensor* tensor)
97 : StringTfLiteTensorBuffer(
98 tensor, tensor->data.raw != nullptr ? GetStringCount(tensor) : 0) {}
99
~StringTfLiteTensorBuffer()100 ~StringTfLiteTensorBuffer() override {
101 LogDeallocation();
102 tensorflow::TypedAllocator::Deallocate<tensorflow::tstring>(
103 tensorflow::cpu_allocator(), static_cast<tensorflow::tstring*>(data()),
104 num_strings_);
105 }
106
size() const107 size_t size() const override {
108 return num_strings_ * sizeof(tensorflow::tstring);
109 }
110
111 private:
StringTfLiteTensorBuffer(const TfLiteTensor * tensor,int num_strings)112 StringTfLiteTensorBuffer(const TfLiteTensor* tensor, int num_strings)
113 : BaseTfLiteTensorBuffer(
114 num_strings != 0
115 ? tensorflow::TypedAllocator::Allocate<tensorflow::tstring>(
116 tensorflow::cpu_allocator(), num_strings,
117 tensorflow::AllocationAttributes())
118 : nullptr),
119 num_strings_(num_strings) {
120 LogAllocation();
121
122 if (data()) {
123 tensorflow::tstring* p = static_cast<tensorflow::tstring*>(data());
124 for (size_t i = 0; i < num_strings_; ++p, ++i) {
125 auto ref = GetString(tensor, i);
126 p->assign(ref.str, ref.len);
127 }
128 }
129 }
130
131 int num_strings_;
132 };
133
134 } // namespace
135
BufferMap()136 BufferMap::BufferMap() {}
137
~BufferMap()138 BufferMap::~BufferMap() {}
139
HasTensor(int tensor_index) const140 bool BufferMap::HasTensor(int tensor_index) const {
141 return id_to_tensor_.count(tensor_index) != 0;
142 }
143
IsTensorFlowTensor(int tensor_index) const144 bool BufferMap::IsTensorFlowTensor(int tensor_index) const {
145 return HasTensor(tensor_index) && owned_by_tf_.count(tensor_index) > 0;
146 }
147
GetTensor(int tensor_index) const148 tensorflow::Tensor BufferMap::GetTensor(int tensor_index) const {
149 return id_to_tensor_.at(tensor_index);
150 }
151
GetTensorPtr(int tensor_index) const152 const tensorflow::Tensor* BufferMap::GetTensorPtr(int tensor_index) const {
153 auto& tensor = id_to_tensor_.at(tensor_index);
154 return &tensor;
155 }
156
SetFromTfLite(int tensor_index,const TfLiteTensor * tensor)157 void BufferMap::SetFromTfLite(int tensor_index, const TfLiteTensor* tensor) {
158 tensorflow::TensorShape shape;
159 int num_dims = tensor->dims->size;
160 for (int i = 0; i < num_dims; ++i) {
161 shape.AddDim(tensor->dims->data[i]);
162 }
163 // TODO(ahentz): we assume this is a new tensor and allocate a new buffer
164 // for it. This is not always the best approach. For example, this might
165 // be a reallocation after resizing tensors. In that case it would be
166 // preferable to somehow reuse the buffer.
167 BaseTfLiteTensorBuffer* buf;
168 if (tensor->type == kTfLiteString) {
169 buf = new StringTfLiteTensorBuffer(tensor);
170 } else {
171 buf = new TfLiteTensorBuffer(tensor);
172 }
173 tensorflow::Tensor t = tensorflow::TensorCApi::MakeTensor(
174 GetTensorFlowDataType(tensor->type), shape, buf);
175 buf->Unref();
176
177 id_to_tensor_[tensor_index] = std::move(t);
178 owned_by_tf_.erase(tensor_index);
179 }
180
SetFromTensorFlow(int tensor_index,tensorflow::Tensor tensor)181 void BufferMap::SetFromTensorFlow(int tensor_index, tensorflow::Tensor tensor) {
182 id_to_tensor_[tensor_index] = std::move(tensor);
183 owned_by_tf_.insert(tensor_index);
184 }
185
186 } // namespace flex
187 } // namespace tflite
188