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/delegate.h"
16
17 #include <vector>
18
19 #include "absl/strings/str_cat.h"
20 #include "tensorflow/core/lib/core/status.h"
21 #include "tensorflow/lite/context_util.h"
22 #include "tensorflow/lite/delegates/flex/buffer_map.h"
23 #include "tensorflow/lite/delegates/flex/kernel.h"
24 #include "tensorflow/lite/delegates/flex/util.h"
25 #include "tensorflow/lite/string_util.h"
26 #include "tensorflow/lite/util.h"
27
28 namespace tflite {
29 namespace flex {
30 namespace delegate {
31
Prepare(TfLiteContext * context,TfLiteDelegate * delegate)32 TfLiteStatus Prepare(TfLiteContext* context, TfLiteDelegate* delegate) {
33 // If the TensorFlow Lite thread count is explicitly configured, use it,
34 // otherwise rely on the default TensorFlow threading behavior.
35 tensorflow::SessionOptions session_options;
36 if (context->recommended_num_threads > 0) {
37 session_options.config.set_intra_op_parallelism_threads(
38 context->recommended_num_threads);
39 }
40
41 if (!reinterpret_cast<DelegateData*>(delegate->data_)
42 ->Prepare(session_options)
43 .ok()) {
44 context->ReportError(context, "Failed to initialize TensorFlow context.");
45 return kTfLiteError;
46 }
47
48 // Get the nodes in the current execution plan. Interpreter owns this array.
49 TfLiteIntArray* plan;
50 TF_LITE_ENSURE_STATUS(context->GetExecutionPlan(context, &plan));
51
52 // Add all custom ops starting with "Flex" to list of supported nodes.
53 std::vector<int> supported_nodes;
54 for (int node_index : TfLiteIntArrayView(plan)) {
55 TfLiteNode* node;
56 TfLiteRegistration* registration;
57 TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
58 context, node_index, &node, ®istration));
59
60 if (IsFlexOp(registration->custom_name)) {
61 supported_nodes.push_back(node_index);
62 }
63 }
64
65 // Request TFLite to partition the graph and make kernels for each independent
66 // node sub set.
67 TfLiteIntArray* size_and_nodes =
68 ConvertVectorToTfLiteIntArray(supported_nodes);
69 context->ReplaceNodeSubsetsWithDelegateKernels(context, GetKernel(),
70 size_and_nodes, delegate);
71 TfLiteIntArrayFree(size_and_nodes);
72 return kTfLiteOk;
73 }
74
CopyFromBufferHandle(TfLiteContext * context,TfLiteDelegate * delegate,TfLiteBufferHandle buffer_handle,TfLiteTensor * output)75 TfLiteStatus CopyFromBufferHandle(TfLiteContext* context,
76 TfLiteDelegate* delegate,
77 TfLiteBufferHandle buffer_handle,
78 TfLiteTensor* output) {
79 BufferMap* buffer_map =
80 reinterpret_cast<DelegateData*>(delegate->data_)->GetBufferMap(context);
81
82 if (!buffer_map->HasTensor(buffer_handle)) {
83 context->ReportError(context, "Invalid tensor index %d.", buffer_handle);
84 return kTfLiteError;
85 }
86
87 tensorflow::Tensor t = buffer_map->GetTensor(buffer_handle);
88
89 if (output->type == kTfLiteString) {
90 if (t.dtype() != tensorflow::DT_STRING) {
91 context->ReportError(context,
92 "Inconsistent type for TF string tensor index %d.",
93 buffer_handle);
94 return kTfLiteError;
95 }
96 DynamicBuffer dynamic_buffer;
97
98 auto tf_data = t.flat<string>();
99 for (int i = 0; i < t.NumElements(); ++i) {
100 dynamic_buffer.AddString(tf_data(i).data(), tf_data(i).size());
101 }
102
103 dynamic_buffer.WriteToTensor(output, /*new_shape=*/nullptr);
104 return kTfLiteOk;
105 }
106
107 tensorflow::StringPiece t_data = t.tensor_data();
108
109 if (output->bytes != t_data.size()) {
110 context->ReportError(context,
111 absl::StrCat("The given ", output->bytes,
112 " bytes are not enough to store "
113 "TensorFlow's aligned buffer of size ",
114 t_data.size(), " bytes.")
115 .c_str());
116 return kTfLiteError;
117 }
118
119 memcpy(output->data.raw, t_data.data(), t_data.size());
120 return kTfLiteOk;
121 }
122
123 } // namespace delegate
124 } // namespace flex
125
126 // Corresponding weak declaration found in lite/model.cc.
127 std::unique_ptr<TfLiteDelegate, void (*)(TfLiteDelegate*)>
AcquireFlexDelegate()128 AcquireFlexDelegate() {
129 return std::unique_ptr<TfLiteDelegate, void (*)(TfLiteDelegate*)>(
130 tflite::FlexDelegate::Create().release(), [](TfLiteDelegate* delegate) {
131 delete reinterpret_cast<tflite::FlexDelegate*>(delegate);
132 });
133 }
134
Create()135 std::unique_ptr<FlexDelegate> FlexDelegate::Create() {
136 return std::unique_ptr<FlexDelegate>(new FlexDelegate());
137 }
138
FlexDelegate()139 FlexDelegate::FlexDelegate() : TfLiteDelegate(TfLiteDelegateCreate()) {
140 data_ = &delegate_data_;
141 Prepare = &flex::delegate::Prepare;
142 CopyFromBufferHandle = &flex::delegate::CopyFromBufferHandle;
143 flags = kTfLiteDelegateFlagsAllowDynamicTensors;
144 }
145
~FlexDelegate()146 FlexDelegate::~FlexDelegate() {}
147
148 } // namespace tflite
149