text: /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
  
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  
      http://www.apache.org/licenses/LICENSE-2.0
  
  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
  ==============================================================================*/
  
  #include "tensorflow/lite/delegates/gpu/cl/kernels/pooling.h"
  
  #include <string>
  
  #include "tensorflow/lite/delegates/gpu/cl/kernels/util.h"
  #include "tensorflow/lite/delegates/gpu/cl/kernels/work_group_picking.h"
  
  namespace tflite {
  namespace gpu {
  namespace cl {
  
  Pooling::Pooling(const OperationDef& definition,
                   const Pooling2DAttributes& attr)
      : GPUOperation(definition),
        stride_(attr.strides.w, attr.strides.h, 0, 0),
        padding_(-attr.padding.prepended.w, -attr.padding.prepended.h, 0, 0),
        kernel_size_(attr.kernel.w, attr.kernel.h, 0, 0),
        type_(attr.type),
        output_indices_(attr.output_indices) {}
  
  Pooling::Pooling(const OperationDef& definition,
                   const Pooling3DAttributes& attr)
      : GPUOperation(definition),
        stride_(attr.strides.w, attr.strides.h, attr.strides.d, 0),
        padding_(-attr.padding.prepended.w, -attr.padding.prepended.h,
                 -attr.padding.prepended.d, 0),
        kernel_size_(attr.kernel.w, attr.kernel.h, attr.kernel.d, 0),
        type_(attr.type),
        output_indices_(attr.output_indices) {}
  
  Pooling::Pooling(Pooling&& kernel)
      : GPUOperation(std::move(kernel)),
        stride_(kernel.stride_),
        padding_(kernel.padding_),
        kernel_size_(kernel.kernel_size_),
        type_(kernel.type_),
        output_indices_(kernel.output_indices_) {}
  
  Pooling& Pooling::operator=(Pooling&& kernel) {
    if (this != &kernel) {
      std::swap(stride_, kernel.stride_);
      std::swap(padding_, kernel.padding_);
      std::swap(kernel_size_, kernel.kernel_size_);
      std::swap(type_, kernel.type_);
      std::swap(output_indices_, kernel.output_indices_);
      GPUOperation::operator=(std::move(kernel));
    }
    return *this;
  }
  
  std::string Pooling::GetAveragePoolingKernelCode(const OperationDef& op_def,
                                                   bool stride_correction,
                                                   const CLDevice& device) {
    auto src_desc = op_def.src_tensors[0];
    src_desc.SetTextureAddressMode(GetFastestZeroMode(device));
    if (op_def.IsBatchSupported()) {
      src_desc.SetStateVar("BatchedWidth", "true");
    }
    AddSrcTensor("src_tensor", src_desc);
    auto dst_desc = op_def.dst_tensors[0];
    if (op_def.IsBatchSupported()) {
      dst_desc.SetStateVar("BatchedWidth", "true");
    }
    AddDstTensor("dst_tensor", dst_desc);
    if (op_def.dst_tensors[0].HasAxis(Axis::WIDTH)) {
      args_.AddInt("kernel_size_x");
      args_.AddInt("padding_x");
      args_.AddInt("stride_x");
    }
    if (op_def.dst_tensors[0].HasAxis(Axis::HEIGHT)) {
      args_.AddInt("kernel_size_y");
      args_.AddInt("padding_y");
      args_.AddInt("stride_y");
    }
    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
      args_.AddInt("kernel_size_z");
      args_.AddInt("padding_z");
      args_.AddInt("stride_z");
    }
  
    std::map<Axis, std::string> axis_to_src_coord = {
        {Axis::WIDTH, "x_c"},  {Axis::HEIGHT, "y_c"}, {Axis::DEPTH, "d_c"},
        {Axis::CHANNELS, "Z"}, {Axis::BATCH, "B"},
    };
  
    std::map<Axis, std::string> axis_to_dst_coord = {
        {Axis::WIDTH, "X"},    {Axis::HEIGHT, "Y"}, {Axis::DEPTH, "D"},
        {Axis::CHANNELS, "Z"}, {Axis::BATCH, "B"},
    };
  
    std::vector<std::string> src_coords;
    std::vector<std::string> dst_coords;
    for (auto axis : {Axis::WIDTH, Axis::HEIGHT, Axis::DEPTH, Axis::CHANNELS}) {
      if (op_def.dst_tensors[0].HasAxis(axis)) {
        dst_coords.push_back(axis_to_dst_coord[axis]);
      }
      if (op_def.src_tensors[0].HasAxis(axis)) {
        src_coords.push_back(axis_to_src_coord[axis]);
      }
    }
    std::string src_coord = src_coords[0];
    for (int i = 1; i < src_coords.size(); ++i) {
      src_coord += ", " + src_coords[i];
    }
    std::string dst_coord = dst_coords[0];
    for (int i = 1; i < dst_coords.size(); ++i) {
      dst_coord += ", " + dst_coords[i];
    }
  
    const bool manual_clamp =
        op_def.src_tensors[0].storage_type == TensorStorageType::BUFFER ||
        op_def.src_tensors[0].storage_type == TensorStorageType::IMAGE_BUFFER;
  
    std::string c = GetCommonDefines(op_def.precision);
    c += "__kernel void main_function(\n";
    c += "$0) {\n";
    c += "  int X = get_global_id(0);\n";
    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
      c += "  int linear_id_1 = get_global_id(1);\n";
      c += "  int Y = linear_id_1 / args.dst_tensor.Depth();\n";
      c += "  int D = linear_id_1 % args.dst_tensor.Depth();\n";
    } else {
      c += "  int Y = get_global_id(1);\n";
    }
    c += "  int Z = get_global_id(2);\n";
    c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
         "Z >= args.dst_tensor.Slices()) { \n";
    c += "    return; \n";
    c += "  } \n";
    c += "  float4 r = (float4)(0.0f);\n";
    c += "  float window_size = 0.0;\n";
    if (stride_correction) {
      c += "  int xs = " +
           GetXStrideCorrected("X", "args.src_tensor.Batch()", "args.stride_x",
                               "args.padding_x") +
           ";\n";
    } else {
      c += "  int xs = X * args.stride_x + args.padding_x;\n";
    }
    c += "  int ys = Y * args.stride_y + args.padding_y;\n";
    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
      c += "  int ds = D * args.stride_z + args.padding_z;\n";
      c += "  for (int kz = 0; kz < args.kernel_size_z; ++kz) {\n";
      c += "    int d_c = ds + kz;\n";
      c += "    if (d_c < 0 || d_c >= args.src_tensor.Depth()) continue;\n";
    }
    c += "  for (int ky = 0; ky < args.kernel_size_y; ++ky) {\n";
    c += "    int y_c = ys + ky;\n";
    c += "    bool outside_y = y_c < 0 || y_c >= args.src_tensor.Height();\n";
    c += "    for (int kx = 0; kx < args.kernel_size_x; ++kx) {\n";
    if (op_def.IsBatchSupported()) {
      c += "      int x_c = xs + kx * args.src_tensor.Batch();\n";
    } else {
      c += "      int x_c = xs + kx;\n";
    }
    c += "      bool outside = outside_y || x_c < 0 || x_c >= "
         "args.src_tensor.Width();\n";
    if (manual_clamp) {
      c += "     r += !outside ? args.src_tensor.Read<float>(" + src_coord +
           ") : "
           "(float4)(0.0f);\n";
    } else {
      c += "      r += args.src_tensor.Read<float>(" + src_coord + ");\n";
    }
    c += "        window_size += !outside ? 1.0 : 0.0;\n";
    c += "    }\n";
    c += "  }\n";
    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
      c += "  }  // Depth\n";
    }
    // If window_size==0, window covered nothing. This situation is a sign of
    // incorrectly constructed operation. NaNs are expected as output.
    c += "  FLT4 result = TO_FLT4(r / window_size);\n";
    c += "  args.dst_tensor.Write(result, " + dst_coord + ");\n";
    c += "}\n";
  
    return c;
  }
  
  std::string Pooling::GetMaxPoolingKernelCode(const OperationDef& op_def,
                                               bool stride_correction,
                                               bool output_indices) {
    auto src_desc = op_def.src_tensors[0];
    if (op_def.IsBatchSupported()) {
      src_desc.SetStateVar("BatchedWidth", "true");
    }
    AddSrcTensor("src_tensor", src_desc);
    auto dst_desc = op_def.dst_tensors[0];
    if (op_def.IsBatchSupported()) {
      dst_desc.SetStateVar("BatchedWidth", "true");
    }
    AddDstTensor("dst_tensor", dst_desc);
    if (output_indices) {
      auto dst_ind_desc = op_def.dst_tensors[1];
      if (op_def.IsBatchSupported()) {
        dst_ind_desc.SetStateVar("BatchedWidth", "true");
      }
      AddDstTensor("dst_indices", dst_ind_desc);
    }
    if (op_def.dst_tensors[0].HasAxis(Axis::WIDTH)) {
      args_.AddInt("kernel_size_x");
      args_.AddInt("padding_x");
      args_.AddInt("stride_x");
    }
    if (op_def.dst_tensors[0].HasAxis(Axis::HEIGHT)) {
      args_.AddInt("kernel_size_y");
      args_.AddInt("padding_y");
      args_.AddInt("stride_y");
    }
    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
      args_.AddInt("kernel_size_z");
      args_.AddInt("padding_z");
      args_.AddInt("stride_z");
    }
  
    std::map<Axis, std::string> axis_to_src_coord = {
        {Axis::WIDTH, "x_c"},  {Axis::HEIGHT, "y_c"}, {Axis::DEPTH, "d_c"},
        {Axis::CHANNELS, "Z"}, {Axis::BATCH, "B"},
    };
  
    std::map<Axis, std::string> axis_to_dst_coord = {
        {Axis::WIDTH, "X"},    {Axis::HEIGHT, "Y"}, {Axis::DEPTH, "D"},
        {Axis::CHANNELS, "Z"}, {Axis::BATCH, "B"},
    };
  
    std::vector<std::string> src_coords;
    std::vector<std::string> dst_coords;
    for (auto axis : {Axis::WIDTH, Axis::HEIGHT, Axis::DEPTH, Axis::CHANNELS}) {
      if (op_def.dst_tensors[0].HasAxis(axis)) {
        dst_coords.push_back(axis_to_dst_coord[axis]);
      }
      if (op_def.src_tensors[0].HasAxis(axis)) {
        src_coords.push_back(axis_to_src_coord[axis]);
      }
    }
    std::string src_coord = src_coords[0];
    for (int i = 1; i < src_coords.size(); ++i) {
      src_coord += ", " + src_coords[i];
    }
    std::string dst_coord = dst_coords[0];
    for (int i = 1; i < dst_coords.size(); ++i) {
      dst_coord += ", " + dst_coords[i];
    }
  
    std::string c = GetCommonDefines(op_def.precision);
    c += "__kernel void main_function(\n";
    c += "$0) {\n";
    c += "  int X = get_global_id(0);\n";
    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
      c += "  int linear_id_1 = get_global_id(1);\n";
      c += "  int Y = linear_id_1 / args.dst_tensor.Depth();\n";
      c += "  int D = linear_id_1 % args.dst_tensor.Depth();\n";
    } else {
      c += "  int Y = get_global_id(1);\n";
    }
    c += "  int Z = get_global_id(2);\n";
    c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
         "Z >= args.dst_tensor.Slices()) { \n";
    c += "    return; \n";
    c += "  } \n";
    c += "  FLT4 maximum = (FLT4)(-10000.0f);\n";
    if (output_indices) {
      c += "  FLT4 indexes = (FLT4)(0.0f);\n";
    }
    if (stride_correction) {
      c += "  int xs = " +
           GetXStrideCorrected("X", "args.src_tensor.Batch()", "args.stride_x",
                               "args.padding_x") +
           ";\n";
    } else {
      c += "  int xs = X * args.stride_x + args.padding_x;\n";
    }
    c += "  int ys = Y * args.stride_y + args.padding_y;\n";
    c += "  for (int ky = 0; ky < args.kernel_size_y; ++ky) {\n";
    c += "    int y_c = ys + ky;\n";
    c += "    if (y_c < 0 || y_c >= args.src_tensor.Height()) continue;\n";
    c += "    for (int kx = 0; kx < args.kernel_size_x; ++kx) {\n";
    if (op_def.IsBatchSupported()) {
      c += "      int x_c = xs + kx * args.src_tensor.Batch();\n";
    } else {
      c += "      int x_c = xs + kx;\n";
    }
    c += "      if (x_c < 0 || x_c >= args.src_tensor.Width()) continue;\n";
    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
      c += "    int ds = D * args.stride_z + args.padding_z;\n";
      c += "    for (int kz = 0; kz < args.kernel_size_z; ++kz) {\n";
      c += "    int d_c = ds + kz;\n";
      c += "      if (d_c < 0 || d_c >= args.src_tensor.Depth()) continue;\n";
    }
    c += "      FLT4 src = args.src_tensor.Read(" + src_coord + ");\n";
    if (output_indices) {
      if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
        c += "      FLT index_counter = (FLT)((ky * args.kernel_size_x + kx) * "
             "args.kernel_size_z + kz) + (FLT)(0.1f);\n";
      } else {
        c += "      FLT index_counter = (FLT)(ky * args.kernel_size_x + kx) + "
             "(FLT)(0.1f);\n";
      }
      c += "      if (src.x > maximum.x) {\n";
      c += "        indexes.x = index_counter;\n";
      c += "        maximum.x = src.x;\n";
      c += "      }\n";
      c += "      if (src.y > maximum.y) {\n";
      c += "        indexes.y = index_counter;\n";
      c += "        maximum.y = src.y;\n";
      c += "      }\n";
      c += "      if (src.z > maximum.z) {\n";
      c += "        indexes.z = index_counter;\n";
      c += "        maximum.z = src.z;\n";
      c += "      }\n";
      c += "      if (src.w > maximum.w) {\n";
      c += "        indexes.w = index_counter;\n";
      c += "        maximum.w = src.w;\n";
      c += "      }\n";
    } else {
      c += "      maximum = max(src, maximum);\n";
    }
    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
      c += "    }  // Depth\n";
    }
    c += "    }\n";
    c += "  }\n";
    c += "  args.dst_tensor.Write(maximum, " + dst_coord + ");\n";
    if (output_indices) {
      c += "  args.dst_indices.Write(indexes, " + dst_coord + ");\n";
    }
    c += "}\n";
  
    return c;
  }
  
  absl::Status Pooling::Compile(const CreationContext& creation_context) {
    std::string code;
    const bool stride_correction =
        definition_.IsBatchSupported() && stride_.x != 1;
    switch (type_) {
      case PoolingType::AVERAGE:
        code = GetAveragePoolingKernelCode(definition_, stride_correction,
                                           *creation_context.device);
        break;
      case PoolingType::MAX:
        code = GetMaxPoolingKernelCode(definition_, stride_correction,
                                       output_indices_);
        break;
      default:
        return absl::InvalidArgumentError(
            "You should create another kernel with this params");
        break;
    }
    std::string element_wise_code;
    RETURN_IF_ERROR(
        MergeOperations(linked_operations_, &args_, &element_wise_code));
    RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
                                            {{"dst_tensor", element_wise_code}},
                                            &code));
    return creation_context.cache->GetOrCreateCLKernel(
        code, "main_function", *creation_context.context,
        *creation_context.device, &kernel_);
  }
  
  absl::Status Pooling::BindArguments() {
    if (definition_.dst_tensors[0].HasAxis(Axis::WIDTH)) {
      RETURN_IF_ERROR(args_.SetInt("stride_x", stride_.x));
      RETURN_IF_ERROR(args_.SetInt("padding_x", padding_.x * src_[0]->Batch()));
      RETURN_IF_ERROR(args_.SetInt("kernel_size_x", kernel_size_.x));
    }
    if (definition_.dst_tensors[0].HasAxis(Axis::HEIGHT)) {
      RETURN_IF_ERROR(args_.SetInt("stride_y", stride_.y));
      RETURN_IF_ERROR(args_.SetInt("padding_y", padding_.y));
      RETURN_IF_ERROR(args_.SetInt("kernel_size_y", kernel_size_.y));
    }
    if (definition_.dst_tensors[0].HasAxis(Axis::DEPTH)) {
      RETURN_IF_ERROR(args_.SetInt("stride_z", stride_.z));
      RETURN_IF_ERROR(args_.SetInt("padding_z", padding_.z));
      RETURN_IF_ERROR(args_.SetInt("kernel_size_z", kernel_size_.z));
    }
    return absl::OkStatus();
  }
  
  int3 Pooling::GetGridSize() const {
    const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
    const int grid_y = dst_[0]->Height() * dst_[0]->Depth();
    const int grid_z = dst_[0]->Slices();
    return int3(grid_x, grid_y, grid_z);
  }
  
  Pooling CreatePooling(const OperationDef& definition,
                        const Pooling2DAttributes& attr) {
    return Pooling(definition, attr);
  }
  
  Pooling CreatePooling(const OperationDef& definition,
                        const Pooling3DAttributes& attr) {
    return Pooling(definition, attr);
  }
  
  }  // namespace cl
  }  // namespace gpu
  }  // namespace tflite
  