1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15#include "tensorflow/lite/kernels/internal/optimized/integer_ops/pooling.h"
16
17#include <stddef.h>
18#include <stdint.h>
19
20#include <cstdlib>
21
22#include "tensorflow/lite/c/builtin_op_data.h"
23#include "tensorflow/lite/c/common.h"
24#include "tensorflow/lite/kernels/internal/compatibility.h"
25#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
26#include "tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h"
27#include "tensorflow/lite/kernels/internal/reference/pooling.h"
28#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
29#include "tensorflow/lite/kernels/internal/tensor.h"
30#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
31#include "tensorflow/lite/kernels/internal/types.h"
32#include "tensorflow/lite/kernels/kernel_util.h"
33#include "tensorflow/lite/kernels/padding.h"
34
35namespace tflite {
36namespace ops {
37namespace builtin {
38namespace pooling {
39
40// This file has two implementation of each pooling op.
41enum KernelType {
42 kReference,
43 kGenericOptimized,
44};
45
46enum PoolType {
47 kAverage,
48 kMax,
49 kL2,
50};
51
52struct OpData {
53 TfLitePaddingValues padding;
54};
55
56void* Init(TfLiteContext* context, const char* buffer, size_t length) {
57 // This is a builtin op, so we don't use the contents in 'buffer', if any.
58 // Instead, we allocate a new object to carry information from Prepare() to
59 // Eval().
60 return new OpData;
61}
62
63void Free(TfLiteContext* context, void* buffer) {
64 delete reinterpret_cast<OpData*>(buffer);
65}
66
67template <PoolType pool_type>
68TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) {
69 auto* params = reinterpret_cast<TfLitePoolParams*>(node->builtin_data);
70 OpData* data = reinterpret_cast<OpData*>(node->user_data);
71
72 TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
73 TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
74 TfLiteTensor* output;
75 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
76 const TfLiteTensor* input;
77 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
78 TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4);
79 TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
80
81 int batches = input->dims->data[0];
82 int height = input->dims->data[1];
83 int width = input->dims->data[2];
84 int channels_out = input->dims->data[3];
85
86 // Matching GetWindowedOutputSize in TensorFlow.
87 auto padding = params->padding;
88 int out_width, out_height;
89
90 // Prevent division by 0 in optimized pooling implementations
91 TF_LITE_ENSURE(context, params->stride_height > 0);
92 TF_LITE_ENSURE(context, params->stride_width > 0);
93
94 data->padding = ComputePaddingHeightWidth(
95 params->stride_height, params->stride_width, 1, 1, height, width,
96 params->filter_height, params->filter_width, padding, &out_height,
97 &out_width);
98
99 if (input->type == kTfLiteUInt8 || input->type == kTfLiteInt8) {
100 if (pool_type == kAverage || pool_type == kMax) {
101 TFLITE_DCHECK_LE(std::abs(input->params.scale - output->params.scale),
102 1.0e-6);
103 TFLITE_DCHECK_EQ(input->params.zero_point, output->params.zero_point);
104 }
105 if (pool_type == kL2) {
106 // We currently don't have a quantized implementation of L2Pool
107 TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
108 }
109 }
110
111 TfLiteIntArray* output_size = TfLiteIntArrayCreate(4);
112 output_size->data[0] = batches;
113 output_size->data[1] = out_height;
114 output_size->data[2] = out_width;
115 output_size->data[3] = channels_out;
116 return context->ResizeTensor(context, output, output_size);
117}
118
119template <KernelType kernel_type>
120TfLiteStatus AverageEvalFloat(TfLiteContext* context, TfLiteNode* node,
121 TfLitePoolParams* params, OpData* data,
122 const TfLiteTensor* input, TfLiteTensor* output) {
123 float activation_min, activation_max;
124 CalculateActivationRange(params->activation, &activation_min,
125 &activation_max);
126#define TF_LITE_AVERAGE_POOL(type) \
127 tflite::PoolParams op_params; \
128 op_params.stride_height = params->stride_height; \
129 op_params.stride_width = params->stride_width; \
130 op_params.filter_height = params->filter_height; \
131 op_params.filter_width = params->filter_width; \
132 op_params.padding_values.height = data->padding.height; \
133 op_params.padding_values.width = data->padding.width; \
134 op_params.float_activation_min = activation_min; \
135 op_params.float_activation_max = activation_max; \
136 TF_LITE_ENSURE(context, type::AveragePool(op_params, GetTensorShape(input), \
137 GetTensorData<float>(input), \
138 GetTensorShape(output), \
139 GetTensorData<float>(output)))
140 if (kernel_type == kReference) {
141 TF_LITE_AVERAGE_POOL(reference_ops);
142 } else {
143 TF_LITE_AVERAGE_POOL(optimized_ops);
144 }
145#undef TF_LITE_AVERAGE_POOL
146 return kTfLiteOk;
147}
148
149template <KernelType kernel_type>
150TfLiteStatus AverageEvalQuantizedUint8(TfLiteContext* context, TfLiteNode* node,
151 TfLitePoolParams* params, OpData* data,
152 const TfLiteTensor* input,
153 TfLiteTensor* output) {
154 int32_t activation_min;
155 int32_t activation_max;
156 (void)CalculateActivationRangeQuantized(context, params->activation, output,
157 &activation_min, &activation_max);
158#define TF_LITE_AVERAGE_POOL(type) \
159 tflite::PoolParams op_params; \
160 op_params.stride_height = params->stride_height; \
161 op_params.stride_width = params->stride_width; \
162 op_params.filter_height = params->filter_height; \
163 op_params.filter_width = params->filter_width; \
164 op_params.padding_values.height = data->padding.height; \
165 op_params.padding_values.width = data->padding.width; \
166 op_params.quantized_activation_min = activation_min; \
167 op_params.quantized_activation_max = activation_max; \
168 TF_LITE_ENSURE(context, type::AveragePool(op_params, GetTensorShape(input), \
169 GetTensorData<uint8_t>(input), \
170 GetTensorShape(output), \
171 GetTensorData<uint8_t>(output)))
172 if (kernel_type == kReference) {
173 TF_LITE_AVERAGE_POOL(reference_ops);
174 } else {
175 TF_LITE_AVERAGE_POOL(optimized_ops);
176 }
177#undef TF_LITE_AVERAGE_POOL
178 return kTfLiteOk;
179}
180
181template <KernelType kernel_type>
182TfLiteStatus AverageEvalQuantizedInt8(TfLiteContext* context, TfLiteNode* node,
183 TfLitePoolParams* params, OpData* data,
184 const TfLiteTensor* input,
185 TfLiteTensor* output) {
186 int32_t activation_min;
187 int32_t activation_max;
188
189 (void)CalculateActivationRangeQuantized(context, params->activation, output,
190 &activation_min, &activation_max);
191#define TF_LITE_AVERAGE_POOL(type) \
192 tflite::PoolParams op_params; \
193 op_params.stride_height = params->stride_height; \
194 op_params.stride_width = params->stride_width; \
195 op_params.filter_height = params->filter_height; \
196 op_params.filter_width = params->filter_width; \
197 op_params.padding_values.height = data->padding.height; \
198 op_params.padding_values.width = data->padding.width; \
199 op_params.quantized_activation_min = activation_min; \
200 op_params.quantized_activation_max = activation_max; \
201 TF_LITE_ENSURE(context, type::AveragePool(op_params, GetTensorShape(input), \
202 GetTensorData<int8_t>(input), \
203 GetTensorShape(output), \
204 GetTensorData<int8_t>(output)))
205 if (kernel_type == kReference) {
206 TF_LITE_AVERAGE_POOL(reference_integer_ops);
207 } else {
208 TF_LITE_AVERAGE_POOL(optimized_integer_ops);
209 }
210#undef TF_LITE_AVERAGE_POOL
211 return kTfLiteOk;
212}
213
214template <KernelType kernel_type>
215TfLiteStatus AverageEvalQuantizedInt16(TfLiteContext* context, TfLiteNode* node,
216 TfLitePoolParams* params, OpData* data,
217 const TfLiteTensor* input,
218 TfLiteTensor* output) {
219 int32_t activation_min;
220 int32_t activation_max;
221 CalculateActivationRangeQuantized(context, params->activation, output,
222 &activation_min, &activation_max);
223#define TF_LITE_AVERAGE_POOL(type) \
224 tflite::PoolParams op_params; \
225 op_params.stride_height = params->stride_height; \
226 op_params.stride_width = params->stride_width; \
227 op_params.filter_height = params->filter_height; \
228 op_params.filter_width = params->filter_width; \
229 op_params.padding_values.height = data->padding.height; \
230 op_params.padding_values.width = data->padding.width; \
231 op_params.quantized_activation_min = activation_min; \
232 op_params.quantized_activation_max = activation_max; \
233 TF_LITE_ENSURE(context, type::AveragePool(op_params, GetTensorShape(input), \
234 GetTensorData<int16_t>(input), \
235 GetTensorShape(output), \
236 GetTensorData<int16_t>(output)))
237 TF_LITE_AVERAGE_POOL(reference_integer_ops);
238#undef TF_LITE_AVERAGE_POOL
239 return kTfLiteOk;
240}
241
242template <KernelType kernel_type>
243void MaxEvalFloat(TfLiteContext* context, TfLiteNode* node,
244 TfLitePoolParams* params, OpData* data,
245 const TfLiteTensor* input, TfLiteTensor* output) {
246 float activation_min, activation_max;
247 CalculateActivationRange(params->activation, &activation_min,
248 &activation_max);
249#define TF_LITE_MAX_POOL(type) \
250 tflite::PoolParams op_params; \
251 op_params.stride_height = params->stride_height; \
252 op_params.stride_width = params->stride_width; \
253 op_params.filter_height = params->filter_height; \
254 op_params.filter_width = params->filter_width; \
255 op_params.padding_values.height = data->padding.height; \
256 op_params.padding_values.width = data->padding.width; \
257 op_params.float_activation_min = activation_min; \
258 op_params.float_activation_max = activation_max; \
259 type::MaxPool(op_params, GetTensorShape(input), GetTensorData<float>(input), \
260 GetTensorShape(output), GetTensorData<float>(output))
261 if (kernel_type == kReference) {
262 TF_LITE_MAX_POOL(reference_ops);
263 } else {
264 TF_LITE_MAX_POOL(optimized_ops);
265 }
266#undef TF_LITE_MAX_POOL
267}
268
269template <KernelType kernel_type>
270void MaxEvalQuantizedUInt8(TfLiteContext* context, TfLiteNode* node,
271 TfLitePoolParams* params, OpData* data,
272 const TfLiteTensor* input, TfLiteTensor* output) {
273 int32_t activation_min;
274 int32_t activation_max;
275 (void)CalculateActivationRangeQuantized(context, params->activation, output,
276 &activation_min, &activation_max);
277#define TF_LITE_MAX_POOL(type) \
278 tflite::PoolParams op_params; \
279 op_params.stride_height = params->stride_height; \
280 op_params.stride_width = params->stride_width; \
281 op_params.filter_height = params->filter_height; \
282 op_params.filter_width = params->filter_width; \
283 op_params.padding_values.height = data->padding.height; \
284 op_params.padding_values.width = data->padding.width; \
285 op_params.quantized_activation_min = activation_min; \
286 op_params.quantized_activation_max = activation_max; \
287 type::MaxPool(op_params, GetTensorShape(input), \
288 GetTensorData<uint8_t>(input), GetTensorShape(output), \
289 GetTensorData<uint8_t>(output))
290 if (kernel_type == kReference) {
291 TF_LITE_MAX_POOL(reference_ops);
292 } else {
293 TF_LITE_MAX_POOL(optimized_ops);
294 }
295#undef TF_LITE_MAX_POOL
296}
297
298template <KernelType kernel_type>
299void MaxEvalQuantizedInt8(TfLiteContext* context, TfLiteNode* node,
300 TfLitePoolParams* params, OpData* data,
301 const TfLiteTensor* input, TfLiteTensor* output) {
302 int32_t activation_min;
303 int32_t activation_max;
304 (void)CalculateActivationRangeQuantized(context, params->activation, output,
305 &activation_min, &activation_max);
306#define TF_LITE_MAX_POOL(type) \
307 tflite::PoolParams op_params; \
308 op_params.stride_height = params->stride_height; \
309 op_params.stride_width = params->stride_width; \
310 op_params.filter_height = params->filter_height; \
311 op_params.filter_width = params->filter_width; \
312 op_params.padding_values.height = data->padding.height; \
313 op_params.padding_values.width = data->padding.width; \
314 op_params.quantized_activation_min = activation_min; \
315 op_params.quantized_activation_max = activation_max; \
316 type::MaxPool(op_params, GetTensorShape(input), \
317 GetTensorData<int8_t>(input), GetTensorShape(output), \
318 GetTensorData<int8_t>(output))
319 if (kernel_type == kReference) {
320 TF_LITE_MAX_POOL(reference_integer_ops);
321 } else {
322 TF_LITE_MAX_POOL(optimized_integer_ops);
323 }
324#undef TF_LITE_MAX_POOL
325}
326
327template <KernelType kernel_type>
328void MaxEvalQuantizedInt16(TfLiteContext* context, TfLiteNode* node,
329 TfLitePoolParams* params, OpData* data,
330 const TfLiteTensor* input, TfLiteTensor* output) {
331 int32_t activation_min;
332 int32_t activation_max;
333 CalculateActivationRangeQuantized(context, params->activation, output,
334 &activation_min, &activation_max);
335#define TF_LITE_MAX_POOL(type) \
336 tflite::PoolParams op_params; \
337 op_params.stride_height = params->stride_height; \
338 op_params.stride_width = params->stride_width; \
339 op_params.filter_height = params->filter_height; \
340 op_params.filter_width = params->filter_width; \
341 op_params.padding_values.height = data->padding.height; \
342 op_params.padding_values.width = data->padding.width; \
343 op_params.quantized_activation_min = activation_min; \
344 op_params.quantized_activation_max = activation_max; \
345 type::MaxPool(op_params, GetTensorShape(input), \
346 GetTensorData<int16_t>(input), GetTensorShape(output), \
347 GetTensorData<int16_t>(output))
348 TF_LITE_MAX_POOL(reference_integer_ops);
349#undef TF_LITE_MAX_POOL
350}
351
352template <KernelType kernel_type>
353void L2EvalFloat(TfLiteContext* context, TfLiteNode* node,
354 TfLitePoolParams* params, OpData* data,
355 const TfLiteTensor* input, TfLiteTensor* output) {
356 float activation_min, activation_max;
357 CalculateActivationRange(params->activation, &activation_min,
358 &activation_max);
359#define TF_LITE_L2_POOL(type) \
360 tflite::PoolParams op_params; \
361 op_params.stride_height = params->stride_height; \
362 op_params.stride_width = params->stride_width; \
363 op_params.filter_height = params->filter_height; \
364 op_params.filter_width = params->filter_width; \
365 op_params.padding_values.height = data->padding.height; \
366 op_params.padding_values.width = data->padding.width; \
367 op_params.float_activation_min = activation_min; \
368 op_params.float_activation_max = activation_max; \
369 type::L2Pool(op_params, GetTensorShape(input), GetTensorData<float>(input), \
370 GetTensorShape(output), GetTensorData<float>(output))
371 if (kernel_type == kReference) {
372 TF_LITE_L2_POOL(reference_ops);
373 } else {
374 TF_LITE_L2_POOL(optimized_ops);
375 }
376#undef TF_LITE_L2_POOL
377}
378
379#undef TF_LITE_KERNEL_TYPE_DISPATCH
380
381template <KernelType kernel_type>
382TfLiteStatus AverageEval(TfLiteContext* context, TfLiteNode* node) {
383 auto* params = reinterpret_cast<TfLitePoolParams*>(node->builtin_data);
384 OpData* data = reinterpret_cast<OpData*>(node->user_data);
385
386 TfLiteTensor* output;
387 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
388 const TfLiteTensor* input;
389 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
390 switch (input->type) { // Already know in/out types are same.
391 case kTfLiteFloat32:
392 return AverageEvalFloat<kernel_type>(context, node, params, data, input,
393 output);
394 case kTfLiteUInt8:
395 return AverageEvalQuantizedUint8<kernel_type>(context, node, params, data,
396 input, output);
397 case kTfLiteInt8:
398 return AverageEvalQuantizedInt8<kernel_type>(context, node, params, data,
399 input, output);
400 case kTfLiteInt16:
401 return AverageEvalQuantizedInt16<kernel_type>(context, node, params, data,
402 input, output);
403 default:
404 TF_LITE_KERNEL_LOG(context, "Type %s not currently supported.",
405 TfLiteTypeGetName(input->type));
406 return kTfLiteError;
407 }
408 return kTfLiteOk;
409}
410
411template <KernelType kernel_type>
412TfLiteStatus MaxEval(TfLiteContext* context, TfLiteNode* node) {
413 auto* params = reinterpret_cast<TfLitePoolParams*>(node->builtin_data);
414 OpData* data = reinterpret_cast<OpData*>(node->user_data);
415
416 TfLiteTensor* output;
417 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
418 const TfLiteTensor* input;
419 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
420 switch (input->type) { // Already know in/out types are same.
421 case kTfLiteFloat32:
422 MaxEvalFloat<kernel_type>(context, node, params, data, input, output);
423 break;
424 case kTfLiteUInt8:
425 MaxEvalQuantizedUInt8<kernel_type>(context, node, params, data, input,
426 output);
427 break;
428 case kTfLiteInt8:
429 MaxEvalQuantizedInt8<kernel_type>(context, node, params, data, input,
430 output);
431 break;
432 case kTfLiteInt16:
433 MaxEvalQuantizedInt16<kernel_type>(context, node, params, data, input,
434 output);
435 break;
436 default:
437 TF_LITE_KERNEL_LOG(context, "Type %s not currently supported.",
438 TfLiteTypeGetName(input->type));
439 return kTfLiteError;
440 }
441 return kTfLiteOk;
442}
443
444template <KernelType kernel_type>
445TfLiteStatus L2Eval(TfLiteContext* context, TfLiteNode* node) {
446 auto* params = reinterpret_cast<TfLitePoolParams*>(node->builtin_data);
447 OpData* data = reinterpret_cast<OpData*>(node->user_data);
448
449 TfLiteTensor* output;
450 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
451 const TfLiteTensor* input;
452 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
453 switch (input->type) { // Already know in/out types are same.
454 case kTfLiteFloat32:
455 L2EvalFloat<kernel_type>(context, node, params, data, input, output);
456 break;
457 case kTfLiteUInt8:
458 // We don't have a quantized implementation, so just fall through to the
459 // 'default' case.
460 default:
461 TF_LITE_KERNEL_LOG(context, "Type %d not currently supported.",
462 input->type);
463 return kTfLiteError;
464 }
465 return kTfLiteOk;
466}
467
468} // namespace pooling
469
470TfLiteRegistration* Register_AVERAGE_POOL_REF() {
471 static TfLiteRegistration r = {pooling::Init, pooling::Free,
472 pooling::GenericPrepare<pooling::kAverage>,
473 pooling::AverageEval<pooling::kReference>};
474 return &r;
475}
476
477TfLiteRegistration* Register_MAX_POOL_REF() {
478 static TfLiteRegistration r = {pooling::Init, pooling::Free,
479 pooling::GenericPrepare<pooling::kMax>,
480 pooling::MaxEval<pooling::kReference>};
481 return &r;
482}
483
484TfLiteRegistration* Register_L2_POOL_REF() {
485 static TfLiteRegistration r = {pooling::Init, pooling::Free,
486 pooling::GenericPrepare<pooling::kL2>,
487 pooling::L2Eval<pooling::kReference>};
488 return &r;
489}
490
491TfLiteRegistration* Register_AVERAGE_POOL_GENERIC_OPT() {
492 static TfLiteRegistration r = {
493 pooling::Init, pooling::Free, pooling::GenericPrepare<pooling::kAverage>,
494 pooling::AverageEval<pooling::kGenericOptimized>};
495 return &r;
496}
497
498TfLiteRegistration* Register_MAX_POOL_GENERIC_OPT() {
499 static TfLiteRegistration r = {pooling::Init, pooling::Free,
500 pooling::GenericPrepare<pooling::kMax>,
501 pooling::MaxEval<pooling::kGenericOptimized>};
502 return &r;
503}
504
505TfLiteRegistration* Register_L2_POOL_GENERIC_OPT() {
506 static TfLiteRegistration r = {pooling::Init, pooling::Free,
507 pooling::GenericPrepare<pooling::kL2>,
508 pooling::L2Eval<pooling::kGenericOptimized>};
509 return &r;
510}
511
512TfLiteRegistration* Register_AVERAGE_POOL_2D() {
513 return Register_AVERAGE_POOL_GENERIC_OPT();
514}
515
516TfLiteRegistration* Register_MAX_POOL_2D() {
517 return Register_MAX_POOL_GENERIC_OPT();
518}
519
520TfLiteRegistration* Register_L2_POOL_2D() {
521 return Register_L2_POOL_GENERIC_OPT();
522}
523
524} // namespace builtin
525} // namespace ops
526} // namespace tflite
527