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 "flatbuffers/flexbuffers.h" // TF:flatbuffers
16 #include "tensorflow/lite/c/builtin_op_data.h"
17 #include "tensorflow/lite/c/c_api_internal.h"
18 #include "tensorflow/lite/core/subgraph.h"
19 #include "tensorflow/lite/kernels/kernel_util.h"
20
21 namespace tflite {
22 namespace ops {
23 namespace custom {
24 namespace if_kernel {
25
26 struct OpData {
27 int then_subgraph_index;
28 int else_subgraph_index;
29 };
30
Init(TfLiteContext * context,const char * buffer,size_t length)31 void* Init(TfLiteContext* context, const char* buffer, size_t length) {
32 auto* op_data = new OpData;
33 const uint8_t* buffer_t = reinterpret_cast<const uint8_t*>(buffer);
34 const flexbuffers::Map& m = flexbuffers::GetRoot(buffer_t, length).AsMap();
35 op_data->then_subgraph_index = m["then_subgraph_index"].AsInt32();
36 op_data->else_subgraph_index = m["else_subgraph_index"].AsInt32();
37 return op_data;
38 }
39
Free(TfLiteContext * context,void * buffer)40 void Free(TfLiteContext* context, void* buffer) {
41 delete reinterpret_cast<OpData*>(buffer);
42 }
43
Prepare(TfLiteContext * context,TfLiteNode * node)44 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
45 const OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
46
47 TF_LITE_ENSURE(context, node->inputs->size > 0);
48
49 // The first input is the condition.
50 const TfLiteTensor* cond = GetInput(context, node, 0);
51 // Currently only bool is supported.
52 // TODO(ycling): Support other types since TensorFlow also support
53 // non-bool types as condition.
54 TF_LITE_ENSURE_EQ(context, cond->type, kTfLiteBool);
55 TF_LITE_ENSURE_EQ(context, NumElements(cond), 1);
56
57 // The first input of the node is the condition. The rest of inputs are
58 // passed to the branch subgraphs. Therefore, the number of subgraph inputs
59 // will be the number of node inputs - 1.
60 int num_inputs = node->inputs->size - 1;
61 int num_outputs = node->outputs->size;
62
63 Subgraph* this_subgraph = reinterpret_cast<Subgraph*>(context->impl_);
64 auto* subgraphs = this_subgraph->GetSubgraphs();
65 TF_LITE_ENSURE(context, op_data->then_subgraph_index < subgraphs->size());
66 TF_LITE_ENSURE(context, op_data->else_subgraph_index < subgraphs->size());
67
68 Subgraph* then_subgraph = (*subgraphs)[op_data->then_subgraph_index].get();
69 Subgraph* else_subgraph = (*subgraphs)[op_data->else_subgraph_index].get();
70
71 for (auto* subgraph : {then_subgraph, else_subgraph}) {
72 TF_LITE_ENSURE_EQ(context, num_inputs, subgraph->inputs().size());
73 TF_LITE_ENSURE_EQ(context, num_outputs, subgraph->outputs().size());
74 }
75
76 bool has_dynamic_output_tensors = false;
77 for (auto* subgraph : {then_subgraph, else_subgraph}) {
78 for (int i = 0; i < num_inputs; ++i) {
79 // The first input of the node is the condition. The indices of the inputs
80 // passed to the subgraphs are offset by 1.
81 const TfLiteTensor* input = GetInput(context, node, i + 1);
82 std::vector<int> dims(input->dims->data,
83 input->dims->data + input->dims->size);
84 subgraph->ResizeInputTensor(i, dims);
85 TfLiteTensor* subgraph_input = subgraph->tensor(subgraph->inputs()[i]);
86 TF_LITE_ENSURE_EQ(context, input->type, subgraph_input->type);
87 }
88 // Note: The `Prepare` function is responsible to run `AllocateTensors` on
89 // both subgraphs. It's intentionally not to break out of the loop when
90 // finding a dynamic output tensor.
91 TF_LITE_ENSURE_OK(context, subgraph->AllocateTensors());
92 has_dynamic_output_tensors |= subgraph->HasDynamicTensors();
93 }
94
95 if (!has_dynamic_output_tensors) {
96 for (int i = 0; i < num_outputs; ++i) {
97 TfLiteTensor* then_output =
98 then_subgraph->tensor(then_subgraph->outputs()[i]);
99 TfLiteTensor* else_output =
100 else_subgraph->tensor(else_subgraph->outputs()[i]);
101 // If the 2 subgraphs have static but different output shapes, the output
102 // tensors of the IF op have dynamic sizes.
103 if (!TfLiteIntArrayEqual(then_output->dims, else_output->dims)) {
104 has_dynamic_output_tensors = true;
105 break;
106 }
107 }
108 }
109
110 for (int i = 0; i < num_outputs; ++i) {
111 TfLiteTensor* output = GetOutput(context, node, i);
112 if (has_dynamic_output_tensors) {
113 SetTensorToDynamic(output);
114 } else {
115 // When there's no dynamic output tensors, the 2 subgraph has exactly
116 // the same static sized outputs.
117 TfLiteTensor* then_output =
118 then_subgraph->tensor(then_subgraph->outputs()[i]);
119 TfLiteIntArray* output_size = TfLiteIntArrayCopy(then_output->dims);
120 TF_LITE_ENSURE_OK(context,
121 context->ResizeTensor(context, output, output_size));
122 }
123 }
124
125 return kTfLiteOk;
126 }
127
Eval(TfLiteContext * context,TfLiteNode * node)128 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
129 const OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
130
131 const TfLiteTensor* cond = GetInput(context, node, 0);
132 bool cond_value = cond->data.b[0];
133
134 Subgraph* this_subgraph = reinterpret_cast<Subgraph*>(context->impl_);
135 auto* subgraphs = this_subgraph->GetSubgraphs();
136
137 // Currently we copy the input / output between the subgraphs. This isn't
138 // optimized yet.
139 // TODO(b/120234921): Optimize and avoid copying tensors between subgraphs.
140 int active_branch_subgraph_index =
141 cond_value ? op_data->then_subgraph_index : op_data->else_subgraph_index;
142 Subgraph& active_branch_subgraph =
143 *(*subgraphs)[active_branch_subgraph_index];
144 for (int i = 0; i < active_branch_subgraph.inputs().size(); ++i) {
145 const TfLiteTensor* input = GetInput(context, node, i + 1);
146 TfLiteTensor* subgraph_input =
147 active_branch_subgraph.tensor(active_branch_subgraph.inputs()[i]);
148 TF_LITE_ENSURE_EQ(context, input->bytes, subgraph_input->bytes);
149 memcpy(subgraph_input->data.raw, input->data.raw, input->bytes);
150 }
151
152 // Note: It's guaranteed that the subgraphs' `AllocateTensors` are called
153 // in `Prepare`, so we don't need to do it here again.
154 TF_LITE_ENSURE_OK(context, active_branch_subgraph.Invoke());
155
156 for (int tensor_index : active_branch_subgraph.outputs()) {
157 active_branch_subgraph.EnsureTensorDataIsReadable(tensor_index);
158 }
159
160 bool has_dynamic_output_tensors = false;
161 for (int i = 0; i < node->outputs->size; ++i) {
162 TfLiteTensor* output = GetOutput(context, node, i);
163 if (IsDynamicTensor(output)) {
164 has_dynamic_output_tensors = true;
165 break;
166 }
167 }
168
169 if (has_dynamic_output_tensors) {
170 for (int i = 0; i < node->outputs->size; ++i) {
171 TfLiteTensor* output = GetOutput(context, node, i);
172 TfLiteTensor* subgraph_output =
173 active_branch_subgraph.tensor(active_branch_subgraph.outputs()[i]);
174 TfLiteIntArray* output_size = TfLiteIntArrayCopy(subgraph_output->dims);
175 TF_LITE_ENSURE_OK(context,
176 context->ResizeTensor(context, output, output_size));
177 }
178 }
179
180 for (int i = 0; i < active_branch_subgraph.outputs().size(); ++i) {
181 const TfLiteTensor* subgraph_output =
182 active_branch_subgraph.tensor(active_branch_subgraph.outputs()[i]);
183 TfLiteTensor* output = GetOutput(context, node, i);
184 TF_LITE_ENSURE_EQ(context, output->bytes, subgraph_output->bytes);
185 memcpy(output->data.raw, subgraph_output->data.raw, output->bytes);
186 }
187 return kTfLiteOk;
188 }
189
190 } // namespace if_kernel
191
Register_IF()192 TfLiteRegistration* Register_IF() {
193 static TfLiteRegistration r = {if_kernel::Init, if_kernel::Free,
194 if_kernel::Prepare, if_kernel::Eval};
195 return &r;
196 }
197
198 } // namespace custom
199 } // namespace ops
200 } // namespace tflite
201