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 <stddef.h>
16 #include <stdint.h>
17
18 #include "tensorflow/lite/c/common.h"
19 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
20 #include "tensorflow/lite/kernels/internal/tensor.h"
21 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
22 #include "tensorflow/lite/kernels/kernel_util.h"
23
24 namespace tflite {
25 namespace ops {
26 namespace builtin {
27 namespace select {
28
29 constexpr int kInputTensorCondition = 0;
30 constexpr int kInputTensorX = 1;
31 constexpr int kInputTensorY = 2;
32 constexpr int kOutputTensor = 0;
33
34 enum KernelType {
35 kVersionOne,
36 kVersionTwo,
37 };
38
39 struct OpData {
40 bool requires_broadcast;
41 // True if input condition is scalar or input condition has rank one and
42 // matches the first dimension of other inputs.
43 bool has_low_rank_input_condition;
44 };
45
SelectInit(TfLiteContext * context,const char * buffer,size_t length)46 void* SelectInit(TfLiteContext* context, const char* buffer, size_t length) {
47 auto* data = new OpData;
48 data->requires_broadcast = false;
49 data->has_low_rank_input_condition = false;
50 return data;
51 }
52
SelectFree(TfLiteContext * context,void * buffer)53 void SelectFree(TfLiteContext* context, void* buffer) {
54 delete reinterpret_cast<OpData*>(buffer);
55 }
56
57 template <KernelType kernel_type>
SelectPrepare(TfLiteContext * context,TfLiteNode * node)58 TfLiteStatus SelectPrepare(TfLiteContext* context, TfLiteNode* node) {
59 OpData* data = reinterpret_cast<OpData*>(node->user_data);
60
61 TF_LITE_ENSURE_EQ(context, NumInputs(node), 3);
62 TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
63
64 const TfLiteTensor* input_condition;
65 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensorCondition,
66 &input_condition));
67 const TfLiteTensor* input_x;
68 TF_LITE_ENSURE_OK(context,
69 GetInputSafe(context, node, kInputTensorX, &input_x));
70 const TfLiteTensor* input_y;
71 TF_LITE_ENSURE_OK(context,
72 GetInputSafe(context, node, kInputTensorY, &input_y));
73 TfLiteTensor* output;
74 TF_LITE_ENSURE_OK(context,
75 GetOutputSafe(context, node, kOutputTensor, &output));
76
77 // Input must be bool.
78 TF_LITE_ENSURE_TYPES_EQ(context, input_condition->type, kTfLiteBool);
79 TF_LITE_ENSURE_TYPES_EQ(context, input_x->type, input_y->type);
80 output->type = input_x->type;
81
82 bool same_shape = HaveSameShapes(input_condition, input_x) &&
83 HaveSameShapes(input_x, input_y);
84 TfLiteIntArray* output_size;
85 if (!same_shape) {
86 switch (kernel_type) {
87 case kVersionOne: {
88 bool is_input_condition_scalar = NumDimensions(input_condition) == 0;
89 bool has_rank_one_input_condition =
90 NumDimensions(input_condition) == 1 &&
91 SizeOfDimension(input_condition, 0) == SizeOfDimension(input_x, 0);
92 data->has_low_rank_input_condition =
93 is_input_condition_scalar || has_rank_one_input_condition;
94 TF_LITE_ENSURE(context, data->has_low_rank_input_condition);
95
96 output_size = TfLiteIntArrayCopy(input_x->dims);
97
98 // Input tensors must have the same type and size
99 TF_LITE_ENSURE(context, HaveSameShapes(input_x, input_y));
100 break;
101 }
102 case kVersionTwo: {
103 TF_LITE_ENSURE_OK(context, CalculateShapeForBroadcast(
104 context, input_condition, input_x,
105 input_y, &output_size));
106 data->requires_broadcast = true;
107 break;
108 }
109 default:
110 return kTfLiteError;
111 }
112 } else {
113 output_size = TfLiteIntArrayCopy(input_x->dims);
114 }
115
116 return context->ResizeTensor(context, output, output_size);
117 }
118
SelectEval(TfLiteContext * context,TfLiteNode * node)119 TfLiteStatus SelectEval(TfLiteContext* context, TfLiteNode* node) {
120 OpData* data = reinterpret_cast<OpData*>(node->user_data);
121 const TfLiteTensor* input_condition;
122 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensorCondition,
123 &input_condition));
124 const TfLiteTensor* input_x;
125 TF_LITE_ENSURE_OK(context,
126 GetInputSafe(context, node, kInputTensorX, &input_x));
127 const TfLiteTensor* input_y;
128 TF_LITE_ENSURE_OK(context,
129 GetInputSafe(context, node, kInputTensorY, &input_y));
130 TfLiteTensor* output;
131 TF_LITE_ENSURE_OK(context,
132 GetOutputSafe(context, node, kOutputTensor, &output));
133
134 #define TF_LITE_SELECT(type, op) \
135 reference_ops::op(GetTensorShape(input_condition), \
136 GetTensorData<bool>(input_condition), \
137 GetTensorShape(input_x), GetTensorData<type>(input_x), \
138 GetTensorShape(input_y), GetTensorData<type>(input_y), \
139 GetTensorShape(output), GetTensorData<type>(output));
140
141 #define TF_LITE_SWITCH(type, op) \
142 switch (type) { \
143 break; \
144 case kTfLiteBool: \
145 TF_LITE_SELECT(bool, op); \
146 break; \
147 case kTfLiteFloat32: \
148 TF_LITE_SELECT(float, op); \
149 break; \
150 case kTfLiteUInt8: \
151 TF_LITE_SELECT(uint8_t, op); \
152 break; \
153 case kTfLiteInt8: \
154 TF_LITE_SELECT(int8_t, op); \
155 break; \
156 case kTfLiteInt16: \
157 TF_LITE_SELECT(int16_t, op); \
158 break; \
159 case kTfLiteInt32: \
160 TF_LITE_SELECT(int32_t, op); \
161 break; \
162 case kTfLiteInt64: \
163 TF_LITE_SELECT(int64_t, op); \
164 break; \
165 default: \
166 context->ReportError(context, \
167 "Does not support type other than bool|float|int, " \
168 "got %d", \
169 type); \
170 return kTfLiteError; \
171 }
172
173 if (data->has_low_rank_input_condition) {
174 TF_LITE_SWITCH(input_x->type, RankOneSelect);
175 } else if (data->requires_broadcast) {
176 TF_LITE_SWITCH(input_x->type, BroadcastSelect4DSlow);
177 } else {
178 TF_LITE_SWITCH(input_x->type, Select);
179 }
180
181 #undef TF_LITE_SELECT
182 #undef TF_LITE_SWITCH
183 return kTfLiteOk;
184 }
185
186 } // namespace select
187
188 // Select op selects values of 'x' if the corresponding value of 'condition' is
189 // true or the value of 'y' if false. There are valid condition input sizes:
190 //
191 // 1. Either the same shape (in which case the select is elementwise), or
192 // 2. condition must be Rank 1 and match over the first dimension, or
193 // 3. condition is scalar
Register_SELECT()194 TfLiteRegistration* Register_SELECT() {
195 static TfLiteRegistration r = {select::SelectInit, select::SelectFree,
196 select::SelectPrepare<select::kVersionOne>,
197 select::SelectEval};
198 return &r;
199 }
200
201 // SelectV2 op selects values of 'x' if the corresponding value of 'condition'
202 // is true or the value of 'y' if false. There are valid condition input sizes:
203 //
204 // 1. Either the same shape (in which case the select is elementwise), or
205 // 2. Broadcastable shapes between 'condition', 'x' and 'y'.
Register_SELECT_V2()206 TfLiteRegistration* Register_SELECT_V2() {
207 static TfLiteRegistration r = {select::SelectInit, select::SelectFree,
208 select::SelectPrepare<select::kVersionTwo>,
209 select::SelectEval};
210 return &r;
211 }
212
213 } // namespace builtin
214 } // namespace ops
215 } // namespace tflite
216