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 // For memory efficiency, this TensorBuffer can possibly reuse memory from the 46 // TfLiteTensor, hence caller should ensure that the TfLiteTensor always outlive 47 // this TensorBuffer. 48 class TfLiteTensorBuffer : public BaseTfLiteTensorBuffer { 49 public: 50 // If `allow_reusing=false`, then the tensor buffer won't be reused from the 51 // TfLiteTensor. 52 explicit TfLiteTensorBuffer(const TfLiteTensor* tensor, 53 bool allow_reusing = true); 54 55 ~TfLiteTensorBuffer() override; 56 size()57 inline size_t size() const override { return len_; } 58 BufferReusedFromTfLiteTensor()59 inline bool BufferReusedFromTfLiteTensor() const { 60 return reused_buffer_from_tflite_; 61 } 62 63 // This function will check if the underlying buffer in `tensor` can be 64 // reused by the tensorflow::Tensor. If it can reuse, it will return 65 // `tensor->data.raw`, otherwise it will create new tensor buffer using 66 // tensorflow's CPU allocator. 67 // TODO(b/205153246): Also consider reusing memory to avoid copying from 68 // tensorflow::Tensor to TfLiteTensor. 69 void* MaybeAllocateTensorflowBuffer(const TfLiteTensor* tensor, 70 bool allow_reusing) const; 71 72 private: 73 size_t len_; 74 bool reused_buffer_from_tflite_; 75 }; 76 77 // A string buffer. TFLITE string tensor format is different than 78 // TF's so we need perform the conversion here. 79 class StringTfLiteTensorBuffer : public BaseTfLiteTensorBuffer { 80 public: 81 explicit StringTfLiteTensorBuffer(const TfLiteTensor* tensor); 82 83 ~StringTfLiteTensorBuffer() override; 84 size()85 inline size_t size() const override { 86 return num_strings_ * sizeof(tensorflow::tstring); 87 } 88 89 private: 90 StringTfLiteTensorBuffer(const TfLiteTensor* tensor, int num_strings); 91 92 int num_strings_; 93 }; 94 95 // Sets the `tensorflow::Tensor` content from `TfLiteTensor` object. If 96 // `allow_reusing=false`, then we explicitly disallow reusing the TF Lite 97 // tensor buffer when constructing the new tensorflow Tensor. 98 tensorflow::Status SetTfTensorFromTfLite(const TfLiteTensor* tensor, 99 tensorflow::Tensor* tf_tensor, 100 bool allow_reusing = true); 101 102 } // namespace flex 103 } // namespace tflite 104 105 #endif // TENSORFLOW_LITE_DELEGATES_FLEX_BUFFER_MAP_UTIL_H_ 106