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 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 40 void Free(TfLiteContext* context, void* buffer) { 41 delete reinterpret_cast<OpData*>(buffer); 42 } 43 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 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 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
