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 <unistd.h>
     16 #include <cassert>
     17 #include <cmath>
     18 #include <cstdio>
     19 #include <cstdlib>
     20 #include <iostream>
     21 #include <limits>
     22 
     23 #include "tensorflow/contrib/lite/builtin_op_data.h"
     24 #include "tensorflow/contrib/lite/context.h"
     25 #include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h"
     26 #include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h"
     27 #include "tensorflow/contrib/lite/kernels/internal/tensor.h"
     28 #include "tensorflow/contrib/lite/kernels/kernel_util.h"
     29 #include "tensorflow/contrib/lite/kernels/op_macros.h"
     30 #include "tensorflow/contrib/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   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 computeOutSize = [padding](int imageSize, int filterSize,
     84                                   int stride) -> int {
     85     return padding == kTfLitePaddingSame
     86                ? (imageSize + stride - 1) / stride
     87                : padding == kTfLitePaddingValid
     88                      ? (imageSize - filterSize + stride) / stride
     89                      : 0;
     90   };
     91 
     92   int outWidth =
     93       computeOutSize(width, params->filter_width, params->stride_width);
     94   int outHeight =
     95       computeOutSize(height, params->filter_height, params->stride_height);
     96 
     97   data->padding.height = ComputePadding(params->stride_height, height,
     98                                         params->filter_height, outHeight);
     99   data->padding.width = ComputePadding(params->stride_width, width,
    100                                        params->filter_width, outWidth);
    101 
    102   if (input->type == kTfLiteUInt8) {
    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* outputSize = TfLiteIntArrayCreate(4);
    115   outputSize->data[0] = batches;
    116   outputSize->data[1] = outHeight;
    117   outputSize->data[2] = outWidth;
    118   outputSize->data[3] = channels_out;
    119   return context->ResizeTensor(context, output, outputSize);
    120 }
    121 
    122 template <KernelType kernel_type>
    123 void AverageEvalFloat(TfLiteContext* context, TfLiteNode* node,
    124                       TfLitePoolParams* params, OpData* data,
    125                       TfLiteTensor* input, TfLiteTensor* output) {
    126   float activation_min, activation_max;
    127   CalculateActivationRangeFloat(params->activation, &activation_min,
    128                                 &activation_max);
    129 #define TF_LITE_AVERAGE_POOL(type)                                             \
    130   type::AveragePool(                                                           \
    131       GetTensorData<float>(input), GetTensorDims(input), params->stride_width, \
    132       params->stride_height, data->padding.width, data->padding.height,        \
    133       params->filter_width, params->filter_height, activation_min,             \
    134       activation_max, GetTensorData<float>(output), GetTensorDims(output))
    135   if (kernel_type == kReference) {
    136     TF_LITE_AVERAGE_POOL(reference_ops);
    137   } else {
    138     TF_LITE_AVERAGE_POOL(optimized_ops);
    139   }
    140 #undef TF_LITE_AVERAGE_POOL
    141 }
    142 
    143 template <KernelType kernel_type>
    144 void AverageEvalQuantized(TfLiteContext* context, TfLiteNode* node,
    145                           TfLitePoolParams* params, OpData* data,
    146                           TfLiteTensor* input, TfLiteTensor* output) {
    147   int32_t activation_min;
    148   int32_t activation_max;
    149   CalculateActivationRangeUint8(params->activation, output, &activation_min,
    150                                 &activation_max);
    151 #define TF_LITE_AVERAGE_POOL(type)                                       \
    152   type::AveragePool(GetTensorData<uint8_t>(input), GetTensorDims(input), \
    153                     params->stride_width, params->stride_height,         \
    154                     data->padding.width, data->padding.height,           \
    155                     params->filter_width, params->filter_height,         \
    156                     activation_min, activation_max,                      \
    157                     GetTensorData<uint8_t>(output), GetTensorDims(output))
    158   if (kernel_type == kReference) {
    159     TF_LITE_AVERAGE_POOL(reference_ops);
    160   } else {
    161     TF_LITE_AVERAGE_POOL(optimized_ops);
    162   }
    163 #undef TF_LITE_AVERAGE_POOL
    164 }
    165 
    166 template <KernelType kernel_type>
    167 void MaxEvalFloat(TfLiteContext* context, TfLiteNode* node,
    168                   TfLitePoolParams* params, OpData* data, TfLiteTensor* input,
    169                   TfLiteTensor* output) {
    170   float activation_min, activation_max;
    171   CalculateActivationRangeFloat(params->activation, &activation_min,
    172                                 &activation_max);
    173 #define TF_LITE_MAX_POOL(type)                                                 \
    174   type::MaxPool(                                                               \
    175       GetTensorData<float>(input), GetTensorDims(input), params->stride_width, \
    176       params->stride_height, data->padding.width, data->padding.height,        \
    177       params->filter_width, params->filter_height, activation_min,             \
    178       activation_max, GetTensorData<float>(output), GetTensorDims(output))
    179   if (kernel_type == kReference) {
    180     TF_LITE_MAX_POOL(reference_ops);
    181   } else {
    182     TF_LITE_MAX_POOL(optimized_ops);
    183   }
    184 #undef TF_LITE_MAX_POOL
    185 }
    186 
    187 template <KernelType kernel_type>
    188 void MaxEvalQuantized(TfLiteContext* context, TfLiteNode* node,
    189                       TfLitePoolParams* params, OpData* data,
    190                       TfLiteTensor* input, TfLiteTensor* output) {
    191   int32_t activation_min;
    192   int32_t activation_max;
    193   CalculateActivationRangeUint8(params->activation, output, &activation_min,
    194                                 &activation_max);
    195 #define TF_LITE_MAX_POOL(type)                                               \
    196   type::MaxPool(GetTensorData<uint8_t>(input), GetTensorDims(input),         \
    197                 params->stride_width, params->stride_height,                 \
    198                 data->padding.width, data->padding.height,                   \
    199                 params->filter_width, params->filter_height, activation_min, \
    200                 activation_max, GetTensorData<uint8_t>(output),              \
    201                 GetTensorDims(output))
    202   if (kernel_type == kReference) {
    203     TF_LITE_MAX_POOL(reference_ops);
    204   } else {
    205     TF_LITE_MAX_POOL(optimized_ops);
    206   }
    207 #undef TF_LITE_MAX_POOL
    208 }
    209 
    210 template <KernelType kernel_type>
    211 void L2EvalFloat(TfLiteContext* context, TfLiteNode* node,
    212                  TfLitePoolParams* params, OpData* data, TfLiteTensor* input,
    213                  TfLiteTensor* output) {
    214   float activation_min, activation_max;
    215   CalculateActivationRangeFloat(params->activation, &activation_min,
    216                                 &activation_max);
    217 #define TF_LITE_L2_POOL(type)                                                  \
    218   type::L2Pool(                                                                \
    219       GetTensorData<float>(input), GetTensorDims(input), params->stride_width, \
    220       params->stride_height, data->padding.width, data->padding.height,        \
    221       params->filter_width, params->filter_height, activation_min,             \
    222       activation_max, GetTensorData<float>(output), GetTensorDims(output))
    223   if (kernel_type == kReference) {
    224     TF_LITE_L2_POOL(reference_ops);
    225   } else {
    226     TF_LITE_L2_POOL(optimized_ops);
    227   }
    228 #undef TF_LITE_L2_POOL
    229 }
    230 
    231 #undef TF_LITE_KERNEL_TYPE_DISPATCH
    232 
    233 template <KernelType kernel_type>
    234 TfLiteStatus AverageEval(TfLiteContext* context, TfLiteNode* node) {
    235   auto* params = reinterpret_cast<TfLitePoolParams*>(node->builtin_data);
    236   OpData* data = reinterpret_cast<OpData*>(node->user_data);
    237 
    238   TfLiteTensor* output = GetOutput(context, node, 0);
    239   TfLiteTensor* input = GetInput(context, node, 0);
    240   switch (input->type) {  // Already know in/out types are same.
    241     case kTfLiteFloat32:
    242       AverageEvalFloat<kernel_type>(context, node, params, data, input, output);
    243       break;
    244     case kTfLiteUInt8:
    245       AverageEvalQuantized<kernel_type>(context, node, params, data, input,
    246                                         output);
    247       break;
    248     default:
    249       context->ReportError(context, "Type not currently supported.");
    250       return kTfLiteError;
    251   }
    252   return kTfLiteOk;
    253 }
    254 
    255 template <KernelType kernel_type>
    256 TfLiteStatus MaxEval(TfLiteContext* context, TfLiteNode* node) {
    257   auto* params = reinterpret_cast<TfLitePoolParams*>(node->builtin_data);
    258   OpData* data = reinterpret_cast<OpData*>(node->user_data);
    259 
    260   TfLiteTensor* output = GetOutput(context, node, 0);
    261   TfLiteTensor* input = GetInput(context, node, 0);
    262   switch (input->type) {  // Already know in/out types are same.
    263     case kTfLiteFloat32:
    264       MaxEvalFloat<kernel_type>(context, node, params, data, input, output);
    265       break;
    266     case kTfLiteUInt8:
    267       MaxEvalQuantized<kernel_type>(context, node, params, data, input, output);
    268       break;
    269     default:
    270       context->ReportError(context, "Type not currently supported.");
    271       return kTfLiteError;
    272   }
    273   return kTfLiteOk;
    274 }
    275 
    276 template <KernelType kernel_type>
    277 TfLiteStatus L2Eval(TfLiteContext* context, TfLiteNode* node) {
    278   auto* params = reinterpret_cast<TfLitePoolParams*>(node->builtin_data);
    279   OpData* data = reinterpret_cast<OpData*>(node->user_data);
    280 
    281   TfLiteTensor* output = GetOutput(context, node, 0);
    282   TfLiteTensor* input = GetInput(context, node, 0);
    283   switch (input->type) {  // Already know in/out types are same.
    284     case kTfLiteFloat32:
    285       L2EvalFloat<kernel_type>(context, node, params, data, input, output);
    286       break;
    287     case kTfLiteUInt8:
    288     // We don't have a quantized implementation, so just fall through to the
    289     // 'default' case.
    290     default:
    291       context->ReportError(context, "Type not currently supported.");
    292       return kTfLiteError;
    293   }
    294   return kTfLiteOk;
    295 }
    296 
    297 }  // namespace pooling
    298 
    299 TfLiteRegistration* Register_AVERAGE_POOL_REF() {
    300   static TfLiteRegistration r = {pooling::Init, pooling::Free,
    301                                  pooling::GenericPrepare<pooling::kAverage>,
    302                                  pooling::AverageEval<pooling::kReference>};
    303   return &r;
    304 }
    305 
    306 TfLiteRegistration* Register_MAX_POOL_REF() {
    307   static TfLiteRegistration r = {pooling::Init, pooling::Free,
    308                                  pooling::GenericPrepare<pooling::kMax>,
    309                                  pooling::MaxEval<pooling::kReference>};
    310   return &r;
    311 }
    312 
    313 TfLiteRegistration* Register_L2_POOL_REF() {
    314   static TfLiteRegistration r = {pooling::Init, pooling::Free,
    315                                  pooling::GenericPrepare<pooling::kL2>,
    316                                  pooling::L2Eval<pooling::kReference>};
    317   return &r;
    318 }
    319 
    320 TfLiteRegistration* Register_AVERAGE_POOL_GENERIC_OPT() {
    321   static TfLiteRegistration r = {
    322       pooling::Init, pooling::Free, pooling::GenericPrepare<pooling::kAverage>,
    323       pooling::AverageEval<pooling::kGenericOptimized>};
    324   return &r;
    325 }
    326 
    327 TfLiteRegistration* Register_MAX_POOL_GENERIC_OPT() {
    328   static TfLiteRegistration r = {pooling::Init, pooling::Free,
    329                                  pooling::GenericPrepare<pooling::kMax>,
    330                                  pooling::MaxEval<pooling::kGenericOptimized>};
    331   return &r;
    332 }
    333 
    334 TfLiteRegistration* Register_L2_POOL_GENERIC_OPT() {
    335   static TfLiteRegistration r = {pooling::Init, pooling::Free,
    336                                  pooling::GenericPrepare<pooling::kL2>,
    337                                  pooling::L2Eval<pooling::kGenericOptimized>};
    338   return &r;
    339 }
    340 
    341 TfLiteRegistration* Register_AVERAGE_POOL_2D() {
    342   return Register_AVERAGE_POOL_GENERIC_OPT();
    343 }
    344 
    345 TfLiteRegistration* Register_MAX_POOL_2D() {
    346   return Register_MAX_POOL_GENERIC_OPT();
    347 }
    348 
    349 TfLiteRegistration* Register_L2_POOL_2D() {
    350   return Register_L2_POOL_GENERIC_OPT();
    351 }
    352 
    353 }  // namespace builtin
    354 }  // namespace ops
    355 }  // namespace tflite
    356