Home | History | Annotate | Download | only in kernels
      1 /* Copyright 2017 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 <cassert>
     16 #include <cmath>
     17 #include <cstdio>
     18 #include <cstdlib>
     19 #include <iostream>
     20 #include <limits>
     21 
     22 #include "tensorflow/lite/c/builtin_op_data.h"
     23 #include "tensorflow/lite/c/c_api_internal.h"
     24 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
     25 #include "tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h"
     26 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
     27 #include "tensorflow/lite/kernels/internal/tensor.h"
     28 #include "tensorflow/lite/kernels/kernel_util.h"
     29 #include "tensorflow/lite/kernels/op_macros.h"
     30 #include "tensorflow/lite/kernels/padding.h"
     31 
     32 namespace tflite {
     33 namespace ops {
     34 namespace builtin {
     35 namespace pooling {
     36 
     37 // This file has two implementation of each pooling op.
     38 enum KernelType {
     39   kReference,
     40   kGenericOptimized,
     41 };
     42 
     43 enum PoolType {
     44   kAverage,
     45   kMax,
     46   kL2,
     47 };
     48 
     49 struct OpData {
     50   TfLitePaddingValues padding;
     51 };
     52 
     53 void* Init(TfLiteContext* context, const char* buffer, size_t length) {
     54   // This is a builtin op, so we don't use the contents in 'buffer', if any.
     55   // Instead, we allocate a new object to carry information from Prepare() to
     56   // Eval().
     57   return new OpData;
     58 }
     59 
     60 void Free(TfLiteContext* context, void* buffer) {
     61   delete reinterpret_cast<OpData*>(buffer);
     62 }
     63 
     64 template <PoolType pool_type>
     65 TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) {
     66   auto* params = reinterpret_cast<TfLitePoolParams*>(node->builtin_data);
     67   OpData* data = reinterpret_cast<OpData*>(node->user_data);
     68 
     69   TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
     70   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
     71   TfLiteTensor* output = GetOutput(context, node, 0);
     72   const TfLiteTensor* input = GetInput(context, node, 0);
     73   TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4);
     74   TF_LITE_ENSURE_EQ(context, input->type, output->type);
     75 
     76   int batches = input->dims->data[0];
     77   int height = input->dims->data[1];
     78   int width = input->dims->data[2];
     79   int channels_out = input->dims->data[3];
     80 
     81   // Matching GetWindowedOutputSize in TensorFlow.
     82   auto padding = params->padding;
     83   auto compute_out_size = [padding](int image_size, int filter_size,
     84                                     int stride) -> int {
     85     return padding == kTfLitePaddingSame
     86                ? (image_size + stride - 1) / stride
     87                : padding == kTfLitePaddingValid
     88                      ? (image_size - filter_size + stride) / stride
     89                      : 0;
     90   };
     91 
     92   int out_width =
     93       compute_out_size(width, params->filter_width, params->stride_width);
     94   int out_height =
     95       compute_out_size(height, params->filter_height, params->stride_height);
     96 
     97   data->padding.height = ComputePadding(params->stride_height, 1, height,
     98                                         params->filter_height, out_height);
     99   data->padding.width = ComputePadding(params->stride_width, 1, width,
    100                                        params->filter_width, out_width);
    101 
    102   if (input->type == kTfLiteUInt8 || input->type == kTfLiteInt8) {
    103     if (pool_type == kAverage || pool_type == kMax) {
    104       TF_LITE_ENSURE_EQ(context, input->params.scale, output->params.scale);
    105       TF_LITE_ENSURE_EQ(context, input->params.zero_point,
    106                         output->params.zero_point);
    107     }
    108     if (pool_type == kL2) {
    109       // We currently don't have a quantized implementation of L2Pool
    110       TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
    111     }
    112   }
    113 
    114   TfLiteIntArray* output_size = TfLiteIntArrayCreate(4);
    115   output_size->data[0] = batches;
    116   output_size->data[1] = out_height;
    117   output_size->data[2] = out_width;
    118   output_size->data[3] = channels_out;
    119   return context->ResizeTensor(context, output, output_size);
    120 }
    121 
    122 template <KernelType kernel_type>
    123 void AverageEvalFloat(TfLiteContext* context, TfLiteNode* node,
    124                       TfLitePoolParams* params, OpData* data,
    125                       const TfLiteTensor* input, TfLiteTensor* output) {
    126   float activation_min, activation_max;
    127   CalculateActivationRange(params->activation, &activation_min,
    128                            &activation_max);
    129 #define TF_LITE_AVERAGE_POOL(type)                                       \
    130   tflite::PoolParams op_params;                                          \
    131   op_params.stride_height = params->stride_height;                       \
    132   op_params.stride_width = params->stride_width;                         \
    133   op_params.filter_height = params->filter_height;                       \
    134   op_params.filter_width = params->filter_width;                         \
    135   op_params.padding_values.height = data->padding.height;                \
    136   op_params.padding_values.width = data->padding.width;                  \
    137   op_params.float_activation_min = activation_min;                       \
    138   op_params.float_activation_max = activation_max;                       \
    139   type::AveragePool(op_params, GetTensorShape(input),                    \
    140                     GetTensorData<float>(input), GetTensorShape(output), \
    141                     GetTensorData<float>(output))
    142   if (kernel_type == kReference) {
    143     TF_LITE_AVERAGE_POOL(reference_ops);
    144   } else {
    145     TF_LITE_AVERAGE_POOL(optimized_ops);
    146   }
    147 #undef TF_LITE_AVERAGE_POOL
    148 }
    149 
    150 template <KernelType kernel_type>
    151 void AverageEvalQuantizedUint8(TfLiteContext* context, TfLiteNode* node,
    152                                TfLitePoolParams* params, OpData* data,
    153                                const TfLiteTensor* input,
    154                                TfLiteTensor* output) {
    155   int32_t activation_min;
    156   int32_t activation_max;
    157   CalculateActivationRangeUint8(params->activation, output, &activation_min,
    158                                 &activation_max);
    159 #define TF_LITE_AVERAGE_POOL(type)                                         \
    160   tflite::PoolParams op_params;                                            \
    161   op_params.stride_height = params->stride_height;                         \
    162   op_params.stride_width = params->stride_width;                           \
    163   op_params.filter_height = params->filter_height;                         \
    164   op_params.filter_width = params->filter_width;                           \
    165   op_params.padding_values.height = data->padding.height;                  \
    166   op_params.padding_values.width = data->padding.width;                    \
    167   op_params.quantized_activation_min = activation_min;                     \
    168   op_params.quantized_activation_max = activation_max;                     \
    169   type::AveragePool(op_params, GetTensorShape(input),                      \
    170                     GetTensorData<uint8_t>(input), 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 }
    179 
    180 void AverageEvalQuantizedInt8(TfLiteContext* context, TfLiteNode* node,
    181                               TfLitePoolParams* params, OpData* data,
    182                               const TfLiteTensor* input, TfLiteTensor* output) {
    183   int32_t activation_min;
    184   int32_t activation_max;
    185   CalculateActivationRangeInt8(params->activation, output, &activation_min,
    186                                &activation_max);
    187   tflite::PoolParams op_params;
    188   op_params.stride_height = params->stride_height;
    189   op_params.stride_width = params->stride_width;
    190   op_params.filter_height = params->filter_height;
    191   op_params.filter_width = params->filter_width;
    192   op_params.padding_values.height = data->padding.height;
    193   op_params.padding_values.width = data->padding.width;
    194   op_params.quantized_activation_min = activation_min;
    195   op_params.quantized_activation_max = activation_max;
    196   reference_integer_ops::AveragePool(
    197       op_params, GetTensorShape(input), GetTensorData<int8_t>(input),
    198       GetTensorShape(output), GetTensorData<int8_t>(output));
    199 }
    200 
    201 template <KernelType kernel_type>
    202 void MaxEvalFloat(TfLiteContext* context, TfLiteNode* node,
    203                   TfLitePoolParams* params, OpData* data,
    204                   const TfLiteTensor* input, TfLiteTensor* output) {
    205   float activation_min, activation_max;
    206   CalculateActivationRange(params->activation, &activation_min,
    207                            &activation_max);
    208 #define TF_LITE_MAX_POOL(type)                                                 \
    209   tflite::PoolParams op_params;                                                \
    210   op_params.stride_height = params->stride_height;                             \
    211   op_params.stride_width = params->stride_width;                               \
    212   op_params.filter_height = params->filter_height;                             \
    213   op_params.filter_width = params->filter_width;                               \
    214   op_params.padding_values.height = data->padding.height;                      \
    215   op_params.padding_values.width = data->padding.width;                        \
    216   op_params.float_activation_min = activation_min;                             \
    217   op_params.float_activation_max = activation_max;                             \
    218   type::MaxPool(op_params, GetTensorShape(input), GetTensorData<float>(input), \
    219                 GetTensorShape(output), GetTensorData<float>(output))
    220   if (kernel_type == kReference) {
    221     TF_LITE_MAX_POOL(reference_ops);
    222   } else {
    223     TF_LITE_MAX_POOL(optimized_ops);
    224   }
    225 #undef TF_LITE_MAX_POOL
    226 }
    227 
    228 template <KernelType kernel_type>
    229 void MaxEvalQuantizedUInt8(TfLiteContext* context, TfLiteNode* node,
    230                            TfLitePoolParams* params, OpData* data,
    231                            const TfLiteTensor* input, TfLiteTensor* output) {
    232   int32_t activation_min;
    233   int32_t activation_max;
    234   CalculateActivationRangeUint8(params->activation, output, &activation_min,
    235                                 &activation_max);
    236 #define TF_LITE_MAX_POOL(type)                                         \
    237   tflite::PoolParams op_params;                                        \
    238   op_params.stride_height = params->stride_height;                     \
    239   op_params.stride_width = params->stride_width;                       \
    240   op_params.filter_height = params->filter_height;                     \
    241   op_params.filter_width = params->filter_width;                       \
    242   op_params.padding_values.height = data->padding.height;              \
    243   op_params.padding_values.width = data->padding.width;                \
    244   op_params.quantized_activation_min = activation_min;                 \
    245   op_params.quantized_activation_max = activation_max;                 \
    246   type::MaxPool(op_params, GetTensorShape(input),                      \
    247                 GetTensorData<uint8_t>(input), GetTensorShape(output), \
    248                 GetTensorData<uint8_t>(output))
    249   if (kernel_type == kReference) {
    250     TF_LITE_MAX_POOL(reference_ops);
    251   } else {
    252     TF_LITE_MAX_POOL(optimized_ops);
    253   }
    254 #undef TF_LITE_MAX_POOL
    255 }
    256 
    257 template <KernelType kernel_type>
    258 void MaxEvalQuantizedInt8(TfLiteContext* context, TfLiteNode* node,
    259                           TfLitePoolParams* params, OpData* data,
    260                           const TfLiteTensor* input, TfLiteTensor* output) {
    261   int32_t activation_min;
    262   int32_t activation_max;
    263   CalculateActivationRangeInt8(params->activation, output, &activation_min,
    264                                &activation_max);
    265 #define TF_LITE_MAX_POOL(type)                                        \
    266   tflite::PoolParams op_params;                                       \
    267   op_params.stride_height = params->stride_height;                    \
    268   op_params.stride_width = params->stride_width;                      \
    269   op_params.filter_height = params->filter_height;                    \
    270   op_params.filter_width = params->filter_width;                      \
    271   op_params.padding_values.height = data->padding.height;             \
    272   op_params.padding_values.width = data->padding.width;               \
    273   op_params.quantized_activation_min = activation_min;                \
    274   op_params.quantized_activation_max = activation_max;                \
    275   type::MaxPool(op_params, GetTensorShape(input),                     \
    276                 GetTensorData<int8_t>(input), GetTensorShape(output), \
    277                 GetTensorData<int8_t>(output))
    278   TF_LITE_MAX_POOL(reference_integer_ops);
    279 #undef TF_LITE_MAX_POOL
    280 }
    281 
    282 template <KernelType kernel_type>
    283 void L2EvalFloat(TfLiteContext* context, TfLiteNode* node,
    284                  TfLitePoolParams* params, OpData* data,
    285                  const TfLiteTensor* input, TfLiteTensor* output) {
    286   float activation_min, activation_max;
    287   CalculateActivationRange(params->activation, &activation_min,
    288                            &activation_max);
    289 #define TF_LITE_L2_POOL(type)                                                 \
    290   tflite::PoolParams op_params;                                               \
    291   op_params.stride_height = params->stride_height;                            \
    292   op_params.stride_width = params->stride_width;                              \
    293   op_params.filter_height = params->filter_height;                            \
    294   op_params.filter_width = params->filter_width;                              \
    295   op_params.padding_values.height = data->padding.height;                     \
    296   op_params.padding_values.width = data->padding.width;                       \
    297   op_params.float_activation_min = activation_min;                            \
    298   op_params.float_activation_max = activation_max;                            \
    299   type::L2Pool(op_params, GetTensorShape(input), GetTensorData<float>(input), \
    300                GetTensorShape(output), GetTensorData<float>(output))
    301   if (kernel_type == kReference) {
    302     TF_LITE_L2_POOL(reference_ops);
    303   } else {
    304     TF_LITE_L2_POOL(optimized_ops);
    305   }
    306 #undef TF_LITE_L2_POOL
    307 }
    308 
    309 #undef TF_LITE_KERNEL_TYPE_DISPATCH
    310 
    311 template <KernelType kernel_type>
    312 TfLiteStatus AverageEval(TfLiteContext* context, TfLiteNode* node) {
    313   auto* params = reinterpret_cast<TfLitePoolParams*>(node->builtin_data);
    314   OpData* data = reinterpret_cast<OpData*>(node->user_data);
    315 
    316   TfLiteTensor* output = GetOutput(context, node, 0);
    317   const TfLiteTensor* input = GetInput(context, node, 0);
    318   switch (input->type) {  // Already know in/out types are same.
    319     case kTfLiteFloat32:
    320       AverageEvalFloat<kernel_type>(context, node, params, data, input, output);
    321       break;
    322     case kTfLiteUInt8:
    323       AverageEvalQuantizedUint8<kernel_type>(context, node, params, data, input,
    324                                              output);
    325       break;
    326     case kTfLiteInt8:
    327       AverageEvalQuantizedInt8(context, node, params, data, input, output);
    328       break;
    329     default:
    330       context->ReportError(context, "Type %d not currently supported.",
    331                            input->type);
    332       return kTfLiteError;
    333   }
    334   return kTfLiteOk;
    335 }
    336 
    337 template <KernelType kernel_type>
    338 TfLiteStatus MaxEval(TfLiteContext* context, TfLiteNode* node) {
    339   auto* params = reinterpret_cast<TfLitePoolParams*>(node->builtin_data);
    340   OpData* data = reinterpret_cast<OpData*>(node->user_data);
    341 
    342   TfLiteTensor* output = GetOutput(context, node, 0);
    343   const TfLiteTensor* input = GetInput(context, node, 0);
    344   switch (input->type) {  // Already know in/out types are same.
    345     case kTfLiteFloat32:
    346       MaxEvalFloat<kernel_type>(context, node, params, data, input, output);
    347       break;
    348     case kTfLiteUInt8:
    349       MaxEvalQuantizedUInt8<kernel_type>(context, node, params, data, input,
    350                                          output);
    351       break;
    352     case kTfLiteInt8:
    353       MaxEvalQuantizedInt8<kernel_type>(context, node, params, data, input,
    354                                         output);
    355       break;
    356     default:
    357       context->ReportError(context, "Type %d not currently supported.",
    358                            input->type);
    359       return kTfLiteError;
    360   }
    361   return kTfLiteOk;
    362 }
    363 
    364 template <KernelType kernel_type>
    365 TfLiteStatus L2Eval(TfLiteContext* context, TfLiteNode* node) {
    366   auto* params = reinterpret_cast<TfLitePoolParams*>(node->builtin_data);
    367   OpData* data = reinterpret_cast<OpData*>(node->user_data);
    368 
    369   TfLiteTensor* output = GetOutput(context, node, 0);
    370   const TfLiteTensor* input = GetInput(context, node, 0);
    371   switch (input->type) {  // Already know in/out types are same.
    372     case kTfLiteFloat32:
    373       L2EvalFloat<kernel_type>(context, node, params, data, input, output);
    374       break;
    375     case kTfLiteUInt8:
    376     // We don't have a quantized implementation, so just fall through to the
    377     // 'default' case.
    378     default:
    379       context->ReportError(context, "Type %d not currently supported.",
    380                            input->type);
    381       return kTfLiteError;
    382   }
    383   return kTfLiteOk;
    384 }
    385 
    386 }  // namespace pooling
    387 
    388 TfLiteRegistration* Register_AVERAGE_POOL_REF() {
    389   static TfLiteRegistration r = {pooling::Init, pooling::Free,
    390                                  pooling::GenericPrepare<pooling::kAverage>,
    391                                  pooling::AverageEval<pooling::kReference>};
    392   return &r;
    393 }
    394 
    395 TfLiteRegistration* Register_MAX_POOL_REF() {
    396   static TfLiteRegistration r = {pooling::Init, pooling::Free,
    397                                  pooling::GenericPrepare<pooling::kMax>,
    398                                  pooling::MaxEval<pooling::kReference>};
    399   return &r;
    400 }
    401 
    402 TfLiteRegistration* Register_L2_POOL_REF() {
    403   static TfLiteRegistration r = {pooling::Init, pooling::Free,
    404                                  pooling::GenericPrepare<pooling::kL2>,
    405                                  pooling::L2Eval<pooling::kReference>};
    406   return &r;
    407 }
    408 
    409 TfLiteRegistration* Register_AVERAGE_POOL_GENERIC_OPT() {
    410   static TfLiteRegistration r = {
    411       pooling::Init, pooling::Free, pooling::GenericPrepare<pooling::kAverage>,
    412       pooling::AverageEval<pooling::kGenericOptimized>};
    413   return &r;
    414 }
    415 
    416 TfLiteRegistration* Register_MAX_POOL_GENERIC_OPT() {
    417   static TfLiteRegistration r = {pooling::Init, pooling::Free,
    418                                  pooling::GenericPrepare<pooling::kMax>,
    419                                  pooling::MaxEval<pooling::kGenericOptimized>};
    420   return &r;
    421 }
    422 
    423 TfLiteRegistration* Register_L2_POOL_GENERIC_OPT() {
    424   static TfLiteRegistration r = {pooling::Init, pooling::Free,
    425                                  pooling::GenericPrepare<pooling::kL2>,
    426                                  pooling::L2Eval<pooling::kGenericOptimized>};
    427   return &r;
    428 }
    429 
    430 TfLiteRegistration* Register_AVERAGE_POOL_2D() {
    431   return Register_AVERAGE_POOL_GENERIC_OPT();
    432 }
    433 
    434 TfLiteRegistration* Register_MAX_POOL_2D() {
    435   return Register_MAX_POOL_GENERIC_OPT();
    436 }
    437 
    438 TfLiteRegistration* Register_L2_POOL_2D() {
    439   return Register_L2_POOL_GENERIC_OPT();
    440 }
    441 
    442 }  // namespace builtin
    443 }  // namespace ops
    444 }  // namespace tflite
    445