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_DELEGATES_FLEX_BUFFER_MAP_UTIL_H_ 16 #define TENSORFLOW_LITE_DELEGATES_FLEX_BUFFER_MAP_UTIL_H_ 17 18 #include "tensorflow/core/framework/allocation_description.pb.h" 19 #include "tensorflow/core/framework/tensor.h" 20 #include "tensorflow/core/platform/status.h" 21 #include "tensorflow/lite/c/common.h" 22 23 namespace tflite { 24 namespace flex { 25 26 // A tensor buffer that is allocated, deallocated and populated by TF Lite. 27 class BaseTfLiteTensorBuffer : public tensorflow::TensorBuffer { 28 using tensorflow::TensorBuffer::TensorBuffer; 29 root_buffer()30 inline TensorBuffer* root_buffer() override { return this; } 31 32 void FillAllocationDescription( 33 tensorflow::AllocationDescription* proto) const override; 34 35 // Prevents input forwarding from mutating this buffer. OwnsMemory()36 inline bool OwnsMemory() const override { return false; } 37 38 protected: 39 void LogAllocation(); 40 void LogDeallocation(); 41 }; 42 43 // A tensor buffer for most data types. Numeric types have exactly the same 44 // representation in TFLITE and TF, so we just need use memcpy(). 45 class TfLiteTensorBuffer : public BaseTfLiteTensorBuffer { 46 public: 47 explicit TfLiteTensorBuffer(const TfLiteTensor* tensor); 48 49 ~TfLiteTensorBuffer() override; 50 size()51 inline size_t size() const override { return len_; } 52 53 private: 54 size_t len_; 55 }; 56 57 // A string buffer. TFLITE string tensor format is different than 58 // TF's so we need perform the conversion here. 59 class StringTfLiteTensorBuffer : public BaseTfLiteTensorBuffer { 60 public: 61 explicit StringTfLiteTensorBuffer(const TfLiteTensor* tensor); 62 63 ~StringTfLiteTensorBuffer() override; 64 size()65 inline size_t size() const override { 66 return num_strings_ * sizeof(tensorflow::tstring); 67 } 68 69 private: 70 StringTfLiteTensorBuffer(const TfLiteTensor* tensor, int num_strings); 71 72 int num_strings_; 73 }; 74 75 // Sets the `tensorflow::Tensor` content from `TfLiteTensor` object. 76 tensorflow::Status SetTfTensorFromTfLite(const TfLiteTensor* tensor, 77 tensorflow::Tensor* tf_tensor); 78 79 } // namespace flex 80 } // namespace tflite 81 82 #endif // TENSORFLOW_LITE_DELEGATES_FLEX_BUFFER_MAP_UTIL_H_ 83