text: // Copyright (c) 2022 PaddlePaddle 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 "paddle/phi/kernels/psroi_pool_kernel.h"
  
  #include <algorithm>
  
  #include "paddle/phi/core/kernel_registry.h"
  
  namespace phi {
  
  template <typename T, typename Context>
  void PsroiPoolKernel(const Context& ctx,
                       const DenseTensor& x,
                       const DenseTensor& rois,
                       const paddle::optional<DenseTensor>& rois_num,
                       int pooled_height,
                       int pooled_width,
                       int output_channels,
                       float spatial_scale,
                       DenseTensor* out) {
    auto in_dims = x.dims();
    int batch_size = in_dims[0];
    int input_channels = in_dims[1];
    int height = in_dims[2];
    int width = in_dims[3];
    int rois_num_t = rois.dims()[0];
  
    PADDLE_ENFORCE_EQ(input_channels,
                      output_channels * pooled_height * pooled_width,
                      errors::InvalidArgument(
                          "the channels of input "
                          "X should equal the product of "
                          "output_channels x pooled_height x pooled_width"));
  
    auto in_stride = stride(in_dims);
    auto out_stride = stride(out->dims());
  
    const T* input_data = x.data<T>();
  
    DenseTensor rois_batch_id_list;
    rois_batch_id_list.Resize({rois_num_t});
    int* rois_batch_id_data = ctx.template Alloc<int>(&rois_batch_id_list);
  
    int rois_batch_size;
    if (rois_num.get_ptr()) {
      rois_batch_size = rois_num->numel();
      auto* rois_num_data = rois_num->data<int>();
      PADDLE_ENFORCE_EQ(
          rois_batch_size,
          batch_size,
          errors::InvalidArgument(
              "The batch size of rois and the batch size of images "
              " must be the same. But received the batch size of rois is %d, "
              "and the batch size of images is %d",
              rois_batch_size,
              batch_size));
      int rois_num_count = 0;
      for (int i = 0; i < rois_batch_size; ++i) {
        rois_num_count += rois_num_data[i];
      }
      PADDLE_ENFORCE_EQ(
          rois_num_count,
          rois_num_t,
          errors::InvalidArgument(
              "the rois_num from input and RoisNum must be the same"));
      int start = 0;
      for (int n = 0; n < rois_batch_size; ++n) {
        for (int i = start; i < start + rois_num_data[n]; ++i) {
          rois_batch_id_data[i] = n;
        }
        start += rois_num_data[n];
      }
    } else {
      auto rois_lod = rois.lod().back();
      rois_batch_size = rois_lod.size() - 1;
      PADDLE_ENFORCE_EQ(
          rois_batch_size,
          batch_size,
          errors::InvalidArgument("the rois_batch_size and input(X) "
                                  "batch_size should be the same."));
      int rois_num_with_lod = rois_lod[rois_batch_size];
      PADDLE_ENFORCE_EQ(rois_num_with_lod,
                        rois_num_t,
                        errors::InvalidArgument(
                            "the rois_num from input and lod must be the same"));
      // calculate batch id index for each roi according to LoD
      for (int n = 0; n < rois_batch_size; ++n) {
        for (size_t i = rois_lod[n]; i < rois_lod[n + 1]; ++i) {
          rois_batch_id_data[i] = n;
        }
      }
    }
    T* output_data = ctx.template Alloc<T>(out);
    const T* input_rois = rois.data<T>();
  
    // calculate psroipooling, parallel processing can be implemented per ROI
    for (int n = 0; n < rois_num_t; ++n) {
      // set roi batch id
      int roi_batch_id = rois_batch_id_data[n];
  
      // [start, end) interval for spatial sampling
      const T* offset_input_rois = input_rois + n * 4;
      T roi_start_w = static_cast<T>(round(offset_input_rois[0])) * spatial_scale;
      T roi_start_h = static_cast<T>(round(offset_input_rois[1])) * spatial_scale;
      T roi_end_w =
          static_cast<T>(round(offset_input_rois[2]) + 1.) * spatial_scale;
      T roi_end_h =
          static_cast<T>(round(offset_input_rois[3]) + 1.) * spatial_scale;
      // Force too small rois to be 1 x 1
      T roi_height = std::max(roi_end_h - roi_start_h, (T)0.1);  // avoid 0
      T roi_width = std::max(roi_end_w - roi_start_w, (T)0.1);
  
      // Compute bin size w and h at input feature map
      T bin_size_h = roi_height / static_cast<T>(pooled_height);
      T bin_size_w = roi_width / static_cast<T>(pooled_width);
  
      // calculate each pixel of the output feature map.
      int out_roi_offset = n * out_stride[0];
      for (int c = 0; c < output_channels; ++c) {
        // per category
        int out_plane_offset = out_roi_offset + c * out_stride[1];
        for (int ph = 0; ph < pooled_height; ++ph) {
          int out_row_offset = out_plane_offset + ph * out_stride[2];
          for (int pw = 0; pw < pooled_width; ++pw) {
            // calculate w and h at input feature map
            int hstart = floor(static_cast<T>(ph) * bin_size_h + roi_start_h);
            int wstart = floor(static_cast<T>(pw) * bin_size_w + roi_start_w);
            int hend = ceil(static_cast<T>(ph + 1) * bin_size_h + roi_start_h);
            int wend = ceil(static_cast<T>(pw + 1) * bin_size_w + roi_start_w);
            //  Add roi offsets and clip to input boundaries
            hstart = std::min(std::max(hstart, 0), height);
            wstart = std::min(std::max(wstart, 0), width);
            hend = std::min(std::max(hend, 0), height);
            wend = std::min(std::max(wend, 0), width);
  
            int output_index = out_row_offset + pw;
            int input_channel = (c * pooled_height + ph) * pooled_width + pw;
            int input_plane_offset =
                roi_batch_id * in_stride[0] + input_channel * in_stride[1];
            const T* offset_input_data = input_data + input_plane_offset;
            T out_sum = 0.;
            bool is_empty = (hend <= hstart) || (wend <= wstart);
            for (int ih = hstart; ih < hend; ++ih) {
              for (int iw = wstart; iw < wend; ++iw) {
                int input_index = ih * in_stride[2] + iw;
                out_sum += offset_input_data[input_index];
              }
            }
            T bin_area = (hend - hstart) * (wend - wstart);
            output_data[output_index] = is_empty ? 0. : out_sum / bin_area;
          }
        }
      }
    }
  }
  
  }  // namespace phi
  
  PD_REGISTER_KERNEL(
      psroi_pool, CPU, ALL_LAYOUT, phi::PsroiPoolKernel, float, double) {
    kernel->InputAt(2).SetDataType(
        paddle::experimental::CppTypeToDataType<int>::Type());
  }
  