compute/cker/include/cker/operation/BatchToSpaceND.h

   1 /*
   2  * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
   3  * Copyright 2017 The TensorFlow Authors. All Rights Reserved.
   4  *
   5  * Licensed under the Apache License, Version 2.0 (the "License");
   6  * you may not use this file except in compliance with the License.
   7  * You may obtain a copy of the License at
   8  *
   9  *      http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 #ifndef __NNFW_CKER_BATCH_TO_SPACE_ND_H__
  19 #define __NNFW_CKER_BATCH_TO_SPACE_ND_H__
  20
  21 #include "cker/Shape.h"
  22
  23 #define UNUSED(x) ((void)(x))
  24
  25 namespace nnfw
  26 {
  27 namespace cker
  28 {
  29
  30 // Helper methods for BatchToSpaceND.
  31 // `spatial_index_dim` specifies post-crop offset index in this spatial
  32 // dimension, i.e. spatial offset introduced by flattening batch to spatial
  33 // dimension minus the crop size at beginning. `block_shape_dim` is the block
  34 // size in current dimension. `input_dim` and `output_dim` are input and output
  35 // size of BatchToSpaceND operation in current dimension.
  36 // Output start index is inclusive and end index is exclusive.
  37 inline void GetIndexRange(int spatial_index_dim, int block_shape_dim, int input_dim, int output_dim,
  38                           int *start_index, int *end_index)
  39 {
  40   // (*start_index) * block_shape_dim is effectively rounded up to the next
  41   // multiple of block_shape_dim by the integer division.
  42   *start_index = std::max(0, (-spatial_index_dim + block_shape_dim - 1) / block_shape_dim);
  43   // Similarly, (*end_index) * block_shape_dim is rounded up too (note that
  44   // end_index is exclusive).
  45   *end_index =
  46       std::min(input_dim, (output_dim - spatial_index_dim + block_shape_dim - 1) / block_shape_dim);
  47 }
  48
  49 template <typename T>
  50 inline void BatchToSpaceND(const Shape &unextended_input1_shape, const T *input1_data,
  51                            const int32_t *block_shape_data, const int32_t *crops_data,
  52                            const Shape &unextended_output_shape, T *output_data)
  53 {
  54   auto input_dim = unextended_input1_shape.DimensionsCount();
  55   auto output_dim = unextended_output_shape.DimensionsCount();
  56
  57   assert(input_dim == 3 || input_dim == 4);
  58   assert(input_dim == output_dim);
  59
  60   UNUSED(input_dim);
  61   UNUSED(output_dim);
  62
  63   // Extends the input/output shape from 3D to 4D if needed, NHC -> NH1C.
  64   auto extend_shape = [](const Shape &shape) {
  65     if (shape.DimensionsCount() == 4)
  66     {
  67       return shape;
  68     }
  69     Shape new_shape(4, 1);
  70     new_shape.SetDim(0, shape.Dims(0));
  71     new_shape.SetDim(1, shape.Dims(1));
  72     new_shape.SetDim(3, shape.Dims(2));
  73     return new_shape;
  74   };
  75   const Shape input1_shape = extend_shape(unextended_input1_shape);
  76   const Shape output_shape = extend_shape(unextended_output_shape);
  77
  78   const int32_t output_width = output_shape.Dims(2);
  79   const int32_t output_height = output_shape.Dims(1);
  80   const int32_t output_batch_size = output_shape.Dims(0);
  81
  82   const int32_t depth = input1_shape.Dims(3);
  83   const int32_t input_width = input1_shape.Dims(2);
  84   const int32_t input_height = input1_shape.Dims(1);
  85   const int32_t input_batch_size = input1_shape.Dims(0);
  86
  87   const int32_t block_shape_height = block_shape_data[0];
  88   const int32_t block_shape_width = block_shape_data[1];
  89
  90   const int32_t crops_top = crops_data[0];
  91   const int32_t crops_left = crops_data[2];
  92
  93   for (int in_batch = 0; in_batch < input_batch_size; ++in_batch)
  94   {
  95     const int out_batch = in_batch % output_batch_size;
  96     const int spatial_offset = in_batch / output_batch_size;
  97
  98     int in_h_start = 0;
  99     int in_h_end = 0;
 100     // GetIndexRange ensures start and end indices are in [0, output_height).
 101     GetIndexRange(spatial_offset / block_shape_width - crops_top, block_shape_height, input_height,
 102                   output_height, &in_h_start, &in_h_end);
 103
 104     for (int in_h = in_h_start; in_h < in_h_end; ++in_h)
 105     {
 106       const int out_h = in_h * block_shape_height + spatial_offset / block_shape_width - crops_top;
 107       assert(out_h >= 0);
 108       assert(out_h < output_height);
 109
 110       int in_w_start = 0;
 111       int in_w_end = 0;
 112       // GetIndexRange ensures start and end indices are in [0, output_width).
 113       GetIndexRange(spatial_offset % block_shape_width - crops_left, block_shape_width, input_width,
 114                     output_width, &in_w_start, &in_w_end);
 115
 116       for (int in_w = in_w_start; in_w < in_w_end; ++in_w)
 117       {
 118         const int out_w =
 119             in_w * block_shape_width + spatial_offset % block_shape_width - crops_left;
 120         assert(out_w >= 0);
 121         assert(out_w < output_width);
 122         T *out = output_data + Offset(output_shape, out_batch, out_h, out_w, 0);
 123         const T *in = input1_data + Offset(input1_shape, in_batch, in_h, in_w, 0);
 124         memcpy(out, in, depth * sizeof(T));
 125       }
 126     }
 127   }
 128 }
 129
 130 } // namespace cker
 131 } // namespace nnfw
 132
 133 #endif // __NNFW_CKER_BATCH_TO_SPACE_ND_H__