• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 #ifndef TENSORFLOW_LITE_DELEGATES_FLEX_BUFFER_MAP_H_
16 #define TENSORFLOW_LITE_DELEGATES_FLEX_BUFFER_MAP_H_
17 
18 #include <map>
19 
20 #include "tensorflow/lite/c/c_api_internal.h"
21 #include "tensorflow/core/framework/tensor.h"
22 
23 namespace tflite {
24 namespace flex {
25 
26 // Maps a TF Lite tensor index into a TensorFlow tensor.
27 //
28 // The TF Lite interpreter assigns integer indices to each of its tensors, but
29 // the Flex delegate deals in terms of TensorFlow tensors. This class maps
30 // from indices to tensors and allows the creation of new tensors to be
31 // associated with a given index.
32 class BufferMap {
33  public:
34   BufferMap();
35   ~BufferMap();
36 
37   // Returns true if the given 'tensor_index' has a corresponding
38   // tensorflow::Tensor.
39   bool HasTensor(int tensor_index) const;
40 
41   // Returns true if the given 'tensor_index' has a corresponding
42   // tensorflow::Tensor *and* the content is owned by TensorFlow (that is, the
43   // mapping was added by SetFromTensorFlow()).
44   bool IsTensorFlowTensor(int tensor_index) const;
45 
46   // Returns the tensorflow::Tensor associated with the given 'tensor_index'.
47   // Precondition: HasTensor() is true.
48   tensorflow::Tensor GetTensor(int tensor_index) const;
49 
50   // Associates the given tensorflow::Tensor with the given 'tensor_index'.
51   // Note that TensorFlow Tensors share data buffers, so this method is only a
52   // shallow copy.
53   void SetFromTensorFlow(int tensor_index, tensorflow::Tensor tensor);
54 
55   // Same as above but creates a new tensorflow::Tensor with a copy of the
56   // given TfLiteTensor's data.
57   void SetFromTfLite(int tensor_index, const TfLiteTensor* tensor);
58 
59  private:
60   // Mapping from TL Lite tensor ID to TensorFlow's Tensor. All tensors that
61   // are inputs or outputs of a subgraph will be added here, irrespective of
62   // whether their data are managed by TF Lite or TensorFlow.
63   std::map<int, tensorflow::Tensor> id_to_tensor_;
64   // A list of tensors that are completely managed by TensorFlow. Most of the
65   // time, TF Lite will populate tensors that are inputs to subgraphs, while
66   // TensorFlow will populate output tensors. Occasionally, however, an input
67   // tensor is coming from a previous subgraph and could have been populated by
68   // TensorFlow. This set keeps track of all input or output tensors that have
69   // been populated by tensorflow.
70   std::set<int> owned_by_tf_;
71 };
72 
73 }  // namespace flex
74 }  // namespace tflite
75 
76 #endif  // TENSORFLOW_LITE_DELEGATES_FLEX_BUFFER_MAP_H_
77