/* Copyright (c) 2023 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/multiclass_nms3_kernel.h"

#ifdef PADDLE_WITH_HIP
#include <hip/hip_runtime.h>
#else
#include "cuda.h"  // NOLINT
#endif
#include "paddle/phi/backends/context_pool.h"
#include "paddle/phi/common/place.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/core/tensor_utils.h"
#include "paddle/phi/kernels/funcs/concat_and_split_functor.h"
#include "paddle/phi/kernels/funcs/cub.h"
#include "paddle/phi/kernels/funcs/gather.cu.h"
#include "paddle/phi/kernels/funcs/math_function.h"
#include "paddle/phi/kernels/nonzero_kernel.h"

#define CUDA_MEM_ALIGN 256

namespace phi {

template <typename T>
struct Bbox {
  T xmin, ymin, xmax, ymax;
  Bbox(T xmin, T ymin, T xmax, T ymax)
      : xmin(xmin), ymin(ymin), xmax(xmax), ymax(ymax) {}
  Bbox() = default;
};

template <typename KeyT, typename ValueT>
size_t CalcCubSortPairsWorkspaceSize(int num_items, int num_segments) {
  size_t temp_storage_bytes = 0;
  cub::DeviceSegmentedRadixSort::SortPairsDescending(
      reinterpret_cast<void*>(NULL),
      temp_storage_bytes,
      reinterpret_cast<KeyT*>(NULL),
      reinterpret_cast<KeyT*>(NULL),
      reinterpret_cast<ValueT*>(NULL),
      reinterpret_cast<ValueT*>(NULL),
      num_items,     // # items
      num_segments,  // # segments
      reinterpret_cast<int*>(NULL),
      reinterpret_cast<int*>(NULL));
  return temp_storage_bytes;
}

template <typename T>
size_t CalcDetectionForwardBBoxDataSize(int N, int C1) {
  return static_cast<int64_t>(N) * C1 * sizeof(T);
}

template <typename T>
size_t CalcDetectionForwardBBoxPermuteSize(bool share_location, int N, int C1) {
  return share_location ? 0 : N * C1 * sizeof(T);
}

template <typename T>
size_t CalcDetectionForwardPreNMSSize(int N, int C2) {
  return static_cast<int64_t>(N) * C2 * sizeof(T);
}

template <typename T>
size_t CalcDetectionForwardPostNMSSize(int N, int num_classes, int top_k) {
  return static_cast<int64_t>(N) * num_classes * top_k * sizeof(T);
}

size_t CalcTotalWorkspaceSize(size_t* workspaces, int count) {
  size_t total = 0;
  for (int i = 0; i < count; i++) {
    total += workspaces[i];
    if (workspaces[i] % CUDA_MEM_ALIGN) {
      total += CUDA_MEM_ALIGN - (workspaces[i] % CUDA_MEM_ALIGN);
    }
  }
  return total;
}

template <typename T>
size_t CalcSortScoresPerClassWorkspaceSize(const int num,
                                           const int num_classes,
                                           const int num_preds_per_class) {
  size_t wss[4];
  const int64_t array_len =
      static_cast<int64_t>(num) * num_classes * num_preds_per_class;
  PADDLE_ENFORCE_LE_INT_MAX(array_len,
                            "num_images * num_classes * num_preds_per_class");
  wss[0] = array_len * sizeof(T);    // temp scores
  wss[1] = array_len * sizeof(int);  // temp indices
  wss[2] =
      (static_cast<int64_t>(num) * num_classes + 1) * sizeof(int);  // offsets
  const int64_t num_segments = static_cast<int64_t>(num) * num_classes;
  PADDLE_ENFORCE_LE_INT_MAX(num_segments, "num_segments");
  wss[3] = CalcCubSortPairsWorkspaceSize<T, int>(
      static_cast<int>(array_len), static_cast<int>(num_segments));

  return CalcTotalWorkspaceSize(wss, 4);
}

template <typename T>
size_t CalcSortScoresPerImageWorkspaceSize(const int num_images,
                                           const int num_items_per_image) {
  const int64_t array_len =
      static_cast<int64_t>(num_images) * num_items_per_image;
  PADDLE_ENFORCE_LE_INT_MAX(array_len, "num_images * num_items_per_image");
  size_t wss[2];
  wss[0] = (num_images + 1) * sizeof(int);  // offsets
  wss[1] = CalcCubSortPairsWorkspaceSize<T, int>(static_cast<int>(array_len),
                                                 num_images);

  return CalcTotalWorkspaceSize(wss, 2);
}

template <typename T>
size_t CalcDetectionInferenceWorkspaceSize(bool share_location,
                                           int N,
                                           int C1,
                                           int C2,
                                           int num_classes,
                                           int num_preds_per_class,
                                           int top_k) {
  size_t wss[6];
  wss[0] = CalcDetectionForwardBBoxDataSize<T>(N, C1);
  wss[1] = CalcDetectionForwardPreNMSSize<T>(N, C2);
  wss[2] = CalcDetectionForwardPreNMSSize<int>(N, C2);
  wss[3] = CalcDetectionForwardPostNMSSize<T>(N, num_classes, top_k);
  wss[4] = CalcDetectionForwardPostNMSSize<int>(N, num_classes, top_k);
  wss[5] =
      std::max(CalcSortScoresPerClassWorkspaceSize<T>(
                   N, num_classes, num_preds_per_class),
               CalcSortScoresPerImageWorkspaceSize<T>(N, num_classes * top_k));
  return CalcTotalWorkspaceSize(wss, 6);
}

// ALIGNPTR
int8_t* AlignPtr(int8_t* ptr, uintptr_t to) {
  uintptr_t addr = (uintptr_t)ptr;
  if (addr % to) {
    addr += to - addr % to;
  }
  return reinterpret_cast<int8_t*>(addr);
}

// GetNEXTWORKSPACEPTR
int8_t* GetNextWorkspacePtr(int8_t* ptr, uintptr_t previous_workspace_size) {
  uintptr_t addr = (uintptr_t)ptr;
  addr += previous_workspace_size;
  return AlignPtr(reinterpret_cast<int8_t*>(addr), CUDA_MEM_ALIGN);
}

/* ==================
 * sortScoresPerClass
 * ================== */
template <typename T_SCORE, unsigned nthds_per_cta>
__launch_bounds__(nthds_per_cta) __global__
    void PrepareSortData(const int num,
                         const int num_classes,
                         const int num_preds_per_class,
                         const int background_label_id,
                         const float confidence_threshold,
                         T_SCORE* conf_scores_gpu,
                         T_SCORE* temp_scores,
                         T_SCORE score_shift,
                         int* temp_idx,
                         int* d_offsets) {
  // Prepare scores data for sort
  const int cur_idx = blockIdx.x * nthds_per_cta + threadIdx.x;
  const int num_preds_per_batch = num_classes * num_preds_per_class;
  T_SCORE clip_val =
      T_SCORE(static_cast<float>(score_shift) + 1.f - 1.f / 1024.f);
  if (cur_idx < num_preds_per_batch) {
    const int class_idx = cur_idx / num_preds_per_class;
    for (int i = 0; i < num; i++) {
      const int target_idx = i * num_preds_per_batch + cur_idx;
      const T_SCORE score = conf_scores_gpu[target_idx];

      // "Clear" background labeled score and index
      // Because we do not care about background
      if (class_idx == background_label_id) {
        // Set scores to 0
        // Set label = -1
        // add shift of 1.0 to normalize the score values
        // to the range [1, 2).
        // add a constant shift to scores will not change the sort
        // result, but will help reduce the computation because
        // we only need to sort the mantissa part of the floating-point
        // numbers
        temp_scores[target_idx] = score_shift;
        temp_idx[target_idx] = -1;
        conf_scores_gpu[target_idx] = score_shift;
      } else {  // "Clear" scores lower than threshold
        if (static_cast<float>(score) > confidence_threshold) {
          // add shift of 1.0 to normalize the score values
          // to the range [1, 2).
          // add a constant shift to scores will not change the sort
          // result, but will help reduce the computation because
          // we only need to sort the mantissa part of the floating-point
          // numbers
          temp_scores[target_idx] = score + score_shift;
          if (static_cast<float>(score_shift) > 0.f &&
              (temp_scores[target_idx] >= clip_val))
            temp_scores[target_idx] = clip_val;
          temp_idx[target_idx] = cur_idx + i * num_preds_per_batch;
        } else {
          // Set scores to 0
          // Set label = -1
          // add shift of 1.0 to normalize the score values
          // to the range [1, 2).
          // add a constant shift to scores will not change the sort
          // result, but will help reduce the computation because
          // we only need to sort the mantissa part of the floating-point
          // numbers
          temp_scores[target_idx] = score_shift;
          temp_idx[target_idx] = -1;
          conf_scores_gpu[target_idx] = score_shift;
          // TODO(tizheng): HERE writing memory too many times
        }
      }

      if ((cur_idx % num_preds_per_class) == 0) {
        const int offset_ct = i * num_classes + cur_idx / num_preds_per_class;
        d_offsets[offset_ct] = offset_ct * num_preds_per_class;
        // set the last element in d_offset
        if (blockIdx.x == 0 && threadIdx.x == 0)
          d_offsets[num * num_classes] = num * num_preds_per_batch;
      }
    }
  }
}

template <typename T_SCORE>
void SortScoresPerClassGPU(gpuStream_t stream,
                           const int num,
                           const int num_classes,
                           const int num_preds_per_class,
                           const int background_label_id,
                           const float confidence_threshold,
                           void* conf_scores_gpu,
                           void* index_array_gpu,
                           void* workspace,
                           const int score_bits,
                           const float score_shift) {
  const int num_segments = num * num_classes;
  void* temp_scores = workspace;
  const int array_len = num * num_classes * num_preds_per_class;
  void* temp_idx = GetNextWorkspacePtr(reinterpret_cast<int8_t*>(temp_scores),
                                       array_len * sizeof(T_SCORE));
  void* d_offsets = GetNextWorkspacePtr(reinterpret_cast<int8_t*>(temp_idx),
                                        array_len * sizeof(int));
  size_t cubOffsetSize = (num_segments + 1) * sizeof(int);
  void* cubWorkspace =
      GetNextWorkspacePtr(reinterpret_cast<int8_t*>(d_offsets), cubOffsetSize);

  const int BS = 512;
  const int GS = (num_classes * num_preds_per_class + BS - 1) / BS;
  // prepare the score, index, and offsets for CUB radix sort
  // also normalize the scores to the range [1, 2)
  // so we only need to sort the mantissa of floating-point numbers
  // since their sign bit and exponential bits are identical
  // we will subtract the 1.0 shift in gatherTopDetections()
  PrepareSortData<T_SCORE, BS>
      <<<GS, BS, 0, stream>>>(num,
                              num_classes,
                              num_preds_per_class,
                              background_label_id,
                              confidence_threshold,
                              reinterpret_cast<T_SCORE*>(conf_scores_gpu),
                              reinterpret_cast<T_SCORE*>(temp_scores),
                              T_SCORE(score_shift),
                              reinterpret_cast<int*>(temp_idx),
                              reinterpret_cast<int*>(d_offsets));

  size_t temp_storage_bytes =
      CalcCubSortPairsWorkspaceSize<T_SCORE, int>(array_len, num_segments);
  size_t begin_bit = 0;
  size_t end_bit = sizeof(T_SCORE) * 8;
  if (sizeof(T_SCORE) == 2 && score_bits > 0 && score_bits <= 10) {
    // only sort score_bits in 10 mantissa bits.
    end_bit = 10;
    begin_bit = end_bit - score_bits;
  }
  cub::DeviceSegmentedRadixSort::SortPairsDescending(
      cubWorkspace,
      temp_storage_bytes,
      reinterpret_cast<T_SCORE*>(temp_scores),
      reinterpret_cast<T_SCORE*>(conf_scores_gpu),
      reinterpret_cast<int*>(temp_idx),
      reinterpret_cast<int*>(index_array_gpu),
      array_len,
      num_segments,
      reinterpret_cast<int*>(d_offsets),
      reinterpret_cast<int*>(d_offsets) + 1,
      begin_bit,
      end_bit,
      stream);
#ifdef PADDLE_WITH_HIP
  PADDLE_ENFORCE_GPU_SUCCESS(hipGetLastError());
#else
  PADDLE_ENFORCE_GPU_SUCCESS(cudaGetLastError());
#endif
}

/* ===========
 * allClassNMS
 * =========== */
template <typename T_BBOX>
__device__ float CalcBboxSize(const Bbox<T_BBOX>& bbox, const bool normalized) {
  if (static_cast<float>(bbox.xmax) < static_cast<float>(bbox.xmin) ||
      static_cast<float>(bbox.ymax) < static_cast<float>(bbox.ymin)) {
    // If bbox is invalid (e.g. xmax < xmin or ymax < ymin), return 0.
    return 0;
  } else {
    float width = static_cast<float>(bbox.xmax) - static_cast<float>(bbox.xmin);
    float height =
        static_cast<float>(bbox.ymax) - static_cast<float>(bbox.ymin);
    if (normalized) {
      return width * height;
    } else {
      // If bbox is not within range [0, 1].
      return (width + 1.f) * (height + 1.f);
    }
  }
}

template <typename T_BBOX>
__device__ void CalcIntersectBbox(const Bbox<T_BBOX>& bbox1,
                                  const Bbox<T_BBOX>& bbox2,
                                  Bbox<T_BBOX>* intersect_bbox) {
  if (bbox2.xmin > bbox1.xmax || bbox2.xmax < bbox1.xmin ||
      bbox2.ymin > bbox1.ymax || bbox2.ymax < bbox1.ymin) {
    // Return [0, 0, 0, 0] if there is no intersection.
    intersect_bbox->xmin = T_BBOX(0);
    intersect_bbox->ymin = T_BBOX(0);
    intersect_bbox->xmax = T_BBOX(0);
    intersect_bbox->ymax = T_BBOX(0);
  } else {
    intersect_bbox->xmin = max(bbox1.xmin, bbox2.xmin);
    intersect_bbox->ymin = max(bbox1.ymin, bbox2.ymin);
    intersect_bbox->xmax = min(bbox1.xmax, bbox2.xmax);
    intersect_bbox->ymax = min(bbox1.ymax, bbox2.ymax);
  }
}

template <typename T_BBOX>
__device__ Bbox<T_BBOX> GetDiagonalMinMaxSortedBox(const Bbox<T_BBOX>& bbox1) {
  Bbox<T_BBOX> result;
  result.xmin = min(bbox1.xmin, bbox1.xmax);
  result.xmax = max(bbox1.xmin, bbox1.xmax);

  result.ymin = min(bbox1.ymin, bbox1.ymax);
  result.ymax = max(bbox1.ymin, bbox1.ymax);
  return result;
}

template <typename T_BBOX>
__device__ void GetFlippedBox(const T_BBOX* bbox1,
                              bool flip_xy,
                              Bbox<T_BBOX>* result) {
  result->xmin = flip_xy ? bbox1[1] : bbox1[0];
  result->ymin = flip_xy ? bbox1[0] : bbox1[1];
  result->xmax = flip_xy ? bbox1[3] : bbox1[2];
  result->ymax = flip_xy ? bbox1[2] : bbox1[3];
}

template <typename T_BBOX>
__device__ float CalcJaccardOverlap(const Bbox<T_BBOX>& bbox1,
                                    const Bbox<T_BBOX>& bbox2,
                                    const bool normalized,
                                    const bool caffe_semantics) {
  Bbox<T_BBOX> intersect_bbox;

  Bbox<T_BBOX> localbbox1 = GetDiagonalMinMaxSortedBox(bbox1);
  Bbox<T_BBOX> localbbox2 = GetDiagonalMinMaxSortedBox(bbox2);

  CalcIntersectBbox(localbbox1, localbbox2, &intersect_bbox);

  float intersect_width, intersect_height;
  // Only when using Caffe semantics, IOU calculation adds "1" to width and
  // height if bbox is not normalized.
  // https://github.com/weiliu89/caffe/blob/ssd/src/caffe/util/bbox_util.cpp#L92-L97
  if (normalized || !caffe_semantics) {
    intersect_width = static_cast<float>(intersect_bbox.xmax) -
                      static_cast<float>(intersect_bbox.xmin);
    intersect_height = static_cast<float>(intersect_bbox.ymax) -
                       static_cast<float>(intersect_bbox.ymin);
  } else {
    intersect_width = static_cast<float>(intersect_bbox.xmax) -
                      static_cast<float>(intersect_bbox.xmin) +
                      static_cast<float>(T_BBOX(1));
    intersect_height = static_cast<float>(intersect_bbox.ymax) -
                       static_cast<float>(intersect_bbox.ymin) +
                       static_cast<float>(T_BBOX(1));
  }
  if (intersect_width > 0 && intersect_height > 0) {
    float intersect_size = intersect_width * intersect_height;
    float bbox1_size = CalcBboxSize(localbbox1, normalized);
    float bbox2_size = CalcBboxSize(localbbox2, normalized);
    return intersect_size / (bbox1_size + bbox2_size - intersect_size);
  } else {
    return 0.;
  }
}

template <typename T_SCORE, typename T_BBOX, int TSIZE>
__global__ void AllClassNMSKernel(
    const int num,
    const int num_classes,
    const int num_preds_per_class,
    const int top_k,
    const float nms_threshold,
    const bool share_location,
    const bool is_normalized,
    T_BBOX* bbox_data,  // bbox_data should be float to preserve location
                        // information
    T_SCORE* before_nms_scores,
    int* before_nms_index_array,
    T_SCORE* after_nms_scores,
    int* after_nms_index_array,
    bool flip_xy,
    const float score_shift,
    bool caffe_semantics) {
  // __shared__ bool kept_bboxinfo_flag[CAFFE_CUDA_NUM_THREADS * TSIZE];
  extern __shared__ bool kept_bboxinfo_flag[];

  for (int i = 0; i < num; i++) {
    int32_t const offset = i * num_classes * num_preds_per_class +
                           blockIdx.x * num_preds_per_class;
    // Should not write data beyond [offset, top_k).
    int32_t const max_idx = offset + top_k;
    // Should not read beyond [offset, num_preds_per_class).
    int32_t const max_read_idx = offset + min(top_k, num_preds_per_class);
    int32_t const bbox_idx_offset =
        i * num_preds_per_class * (share_location ? 1 : num_classes);

    // local thread data
    int loc_bboxIndex[TSIZE];
    Bbox<T_BBOX> loc_bbox[TSIZE];

    // initialize Bbox, Bboxinfo, kept_bboxinfo_flag
    // Eliminate shared memory RAW hazard
    __syncthreads();
#pragma unroll
    for (int t = 0; t < TSIZE; t++) {
      const int cur_idx = threadIdx.x + blockDim.x * t;
      const int item_idx = offset + cur_idx;
      // Init all output data
      if (item_idx < max_idx) {
        // Do not access data if it exceeds read boundary
        if (item_idx < max_read_idx) {
          loc_bboxIndex[t] = before_nms_index_array[item_idx];
        } else {
          loc_bboxIndex[t] = -1;
        }

        if (loc_bboxIndex[t] != -1) {
          const int bbox_data_idx =
              share_location
                  ? (loc_bboxIndex[t] % num_preds_per_class + bbox_idx_offset)
                  : loc_bboxIndex[t];
          GetFlippedBox(&bbox_data[bbox_data_idx * 4], flip_xy, &loc_bbox[t]);
          kept_bboxinfo_flag[cur_idx] = true;
        } else {
          kept_bboxinfo_flag[cur_idx] = false;
        }
      } else {
        kept_bboxinfo_flag[cur_idx] = false;
      }
    }

    // filter out overlapped boxes with lower scores
    int ref_item_idx = offset;

    int32_t ref_bbox_idx = -1;
    if (ref_item_idx < max_read_idx) {
      ref_bbox_idx =
          share_location
              ? (before_nms_index_array[ref_item_idx] % num_preds_per_class +
                 bbox_idx_offset)
              : before_nms_index_array[ref_item_idx];
    }
    while ((ref_bbox_idx != -1) && ref_item_idx < max_read_idx) {
      Bbox<T_BBOX> ref_bbox;
      GetFlippedBox(&bbox_data[ref_bbox_idx * 4], flip_xy, &ref_bbox);

      // Eliminate shared memory RAW hazard
      __syncthreads();

      for (int t = 0; t < TSIZE; t++) {
        const int cur_idx = threadIdx.x + blockDim.x * t;
        const int item_idx = offset + cur_idx;

        if ((kept_bboxinfo_flag[cur_idx]) && (item_idx > ref_item_idx)) {
          if (CalcJaccardOverlap(
                  ref_bbox, loc_bbox[t], is_normalized, caffe_semantics) >
              nms_threshold) {
            kept_bboxinfo_flag[cur_idx] = false;
          }
        }
      }
      __syncthreads();

      do {
        ref_item_idx++;
      } while (ref_item_idx < max_read_idx &&
               !kept_bboxinfo_flag[ref_item_idx - offset]);

      // Move to next valid point
      if (ref_item_idx < max_read_idx) {
        ref_bbox_idx =
            share_location
                ? (before_nms_index_array[ref_item_idx] % num_preds_per_class +
                   bbox_idx_offset)
                : before_nms_index_array[ref_item_idx];
      }
    }

    // store data
    for (int t = 0; t < TSIZE; t++) {
      const int cur_idx = threadIdx.x + blockDim.x * t;
      const int read_item_idx = offset + cur_idx;
      const int write_item_idx =
          (i * num_classes * top_k + blockIdx.x * top_k) + cur_idx;
      /*
       * If not not keeping the bbox
       * Set the score to 0
       * Set the bounding box index to -1
       */
      if (read_item_idx < max_idx) {
        after_nms_scores[write_item_idx] =
            kept_bboxinfo_flag[cur_idx]
                ? T_SCORE(before_nms_scores[read_item_idx])
                : T_SCORE(score_shift);
        after_nms_index_array[write_item_idx] =
            kept_bboxinfo_flag[cur_idx] ? loc_bboxIndex[t] : -1;
      }
    }
  }
}

template <typename T_SCORE, typename T_BBOX>
void AllClassNMSGPU(gpuStream_t stream,
                    const int num,
                    const int num_classes,
                    const int num_preds_per_class,
                    const int top_k,
                    const float nms_threshold,
                    const bool share_location,
                    const bool is_normalized,
                    void* bbox_data,
                    void* before_nms_scores,
                    void* before_nms_index_array,
                    void* after_nms_scores,
                    void* after_nms_index_array,
                    bool flip_xy,
                    const float score_shift,
                    bool caffe_semantics) {
#define P(tsize) AllClassNMSKernel<T_SCORE, T_BBOX, (tsize)>

  void (*kernel[8])(const int,
                    const int,
                    const int,
                    const int,
                    const float,
                    const bool,
                    const bool,
                    T_BBOX*,
                    T_SCORE*,
                    int*,
                    T_SCORE*,
                    int*,
                    bool,
                    const float,
                    bool) = {
      P(1),
      P(2),
      P(3),
      P(4),
      P(5),
      P(6),
      P(7),
      P(8),
  };

  const int BS = 512;
  const int GS = num_classes;
  const int t_size = (top_k + BS - 1) / BS;

  kernel[t_size - 1]<<<GS, BS, BS * t_size * sizeof(bool), stream>>>(
      num,
      num_classes,
      num_preds_per_class,
      top_k,
      nms_threshold,
      share_location,
      is_normalized,
      reinterpret_cast<T_BBOX*>(bbox_data),
      reinterpret_cast<T_SCORE*>(before_nms_scores),
      reinterpret_cast<int*>(before_nms_index_array),
      reinterpret_cast<T_SCORE*>(after_nms_scores),
      reinterpret_cast<int*>(after_nms_index_array),
      flip_xy,
      score_shift,
      caffe_semantics);

#ifdef PADDLE_WITH_HIP
  PADDLE_ENFORCE_GPU_SUCCESS(hipGetLastError());
#else
  PADDLE_ENFORCE_GPU_SUCCESS(cudaGetLastError());
#endif
}

/* ==================
 * sortScoresPerImage
 * ================== */
template <unsigned nthds_per_cta>
__launch_bounds__(nthds_per_cta) __global__
    void SetUniformOffsetsKernel(const int num_segments,
                                 const int offset,
                                 int* d_offsets) {
  const int idx = blockIdx.x * nthds_per_cta + threadIdx.x;
  if (idx <= num_segments) d_offsets[idx] = idx * offset;
}

void SetUniformOffsets(gpuStream_t stream,
                       const int num_segments,
                       const int offset,
                       int* d_offsets) {
#ifdef PADDLE_WITH_HIP
  const int BS = 256;
#else
  const int BS = 32;
#endif
  const int GS = (num_segments + 1 + BS - 1) / BS;
  SetUniformOffsetsKernel<BS>
      <<<GS, BS, 0, stream>>>(num_segments, offset, d_offsets);
}

/* ================
 * gatherNMSOutputs
 * ================ */
template <typename T_BBOX>
__device__ T_BBOX saturate(T_BBOX v) {
  return max(min(v, T_BBOX(1)), T_BBOX(0));
}

template <typename T_BBOX, typename T_SCORE, unsigned nthds_per_cta>
__launch_bounds__(nthds_per_cta) __global__
    void GatherNMSOutputsKernel(const bool share_location,
                                const int num_images,
                                const int num_preds_per_class,
                                const int num_classes,
                                const int top_k,
                                const int keep_top_k,
                                const int* indices,
                                const T_SCORE* scores,
                                const T_BBOX* bbox_data,
                                int* num_detections,
                                T_BBOX* nmsed_boxes,
                                T_BBOX* nmsed_scores,
                                T_BBOX* nmsed_classes,
                                int* nmsed_indices,
                                int* nmsed_valid_mask,
                                bool clip_boxes,
                                const T_SCORE score_shift) {
  if (keep_top_k > top_k) return;
  for (int64_t i = static_cast<int64_t>(blockIdx.x) * nthds_per_cta +
                   static_cast<int64_t>(threadIdx.x);
       i < num_images * keep_top_k;
       i += gridDim.x * nthds_per_cta) {
    const int imgId = i / keep_top_k;
    const int detId = i % keep_top_k;
    const int offset = imgId * num_classes * top_k;
    const int index = indices[offset + detId];
    const T_SCORE score = scores[offset + detId];
    if (index == -1) {
      nmsed_classes[i] = -1;
      nmsed_scores[i] = 0;
      nmsed_boxes[i * 4] = 0;
      nmsed_boxes[i * 4 + 1] = 0;
      nmsed_boxes[i * 4 + 2] = 0;
      nmsed_boxes[i * 4 + 3] = 0;
      nmsed_indices[i] = -1;
      nmsed_valid_mask[i] = 0;
    } else {
      const int bbox_offset =
          imgId * (share_location ? num_preds_per_class
                                  : (num_classes * num_preds_per_class));
      const int bbox_id =
          ((share_location ? (index % num_preds_per_class)
                           : index % (num_classes * num_preds_per_class)) +
           bbox_offset) *
          4;
      nmsed_classes[i] = (index % (num_classes * num_preds_per_class)) /
                         num_preds_per_class;  // label
      nmsed_scores[i] = score;                 // confidence score
      nmsed_scores[i] = nmsed_scores[i] - score_shift;
      const T_BBOX xMin = bbox_data[bbox_id];
      const T_BBOX yMin = bbox_data[bbox_id + 1];
      const T_BBOX xMax = bbox_data[bbox_id + 2];
      const T_BBOX yMax = bbox_data[bbox_id + 3];
      // clipped bbox xmin
      nmsed_boxes[i * 4] = clip_boxes ? saturate(xMin) : xMin;
      // clipped bbox ymin
      nmsed_boxes[i * 4 + 1] = clip_boxes ? saturate(yMin) : yMin;
      // clipped bbox xmax
      nmsed_boxes[i * 4 + 2] = clip_boxes ? saturate(xMax) : xMax;
      // clipped bbox ymax
      nmsed_boxes[i * 4 + 3] = clip_boxes ? saturate(yMax) : yMax;
      nmsed_indices[i] = bbox_id >> 2;
      nmsed_valid_mask[i] = 1;
      atomicAdd(&num_detections[i / keep_top_k], 1);
    }
  }
}

template <typename T_BBOX, typename T_SCORE>
void GatherNMSOutputsGPU(gpuStream_t stream,
                         const bool share_location,
                         const int num_images,
                         const int num_preds_per_class,
                         const int num_classes,
                         const int top_k,
                         const int keep_top_k,
                         const void* indices,
                         const void* scores,
                         const void* bbox_data,
                         void* num_detections,
                         void* nmsed_boxes,
                         void* nmsed_scores,
                         void* nmsed_classes,
                         void* nmsed_indices,
                         void* nmsed_valid_mask,
                         bool clip_boxes,
                         const float score_shift) {
#ifdef PADDLE_WITH_HIP
  PADDLE_ENFORCE_GPU_SUCCESS(
      hipMemsetAsync(num_detections, 0, num_images * sizeof(int), stream));
  const int BS = 256;
#else
  PADDLE_ENFORCE_GPU_SUCCESS(
      cudaMemsetAsync(num_detections, 0, num_images * sizeof(int), stream));
  const int BS = 32;
#endif
  const int GS = 32;
  GatherNMSOutputsKernel<T_BBOX, T_SCORE, BS>
      <<<GS, BS, 0, stream>>>(share_location,
                              num_images,
                              num_preds_per_class,
                              num_classes,
                              top_k,
                              keep_top_k,
                              reinterpret_cast<const int*>(indices),
                              reinterpret_cast<const T_SCORE*>(scores),
                              reinterpret_cast<const T_BBOX*>(bbox_data),
                              reinterpret_cast<int*>(num_detections),
                              reinterpret_cast<T_BBOX*>(nmsed_boxes),
                              reinterpret_cast<T_BBOX*>(nmsed_scores),
                              reinterpret_cast<T_BBOX*>(nmsed_classes),
                              reinterpret_cast<int*>(nmsed_indices),
                              reinterpret_cast<int*>(nmsed_valid_mask),
                              clip_boxes,
                              T_SCORE(score_shift));
#ifdef PADDLE_WITH_HIP
  PADDLE_ENFORCE_GPU_SUCCESS(hipGetLastError());
#else
  PADDLE_ENFORCE_GPU_SUCCESS(cudaGetLastError());
#endif
}

template <typename T_SCORE>
void SortScoresPerImageGPU(gpuStream_t stream,
                           const int num_images,
                           const int num_items_per_image,
                           void* unsorted_scores,
                           void* unsorted_bbox_indices,
                           void* sorted_scores,
                           void* sorted_bbox_indices,
                           void* workspace,
                           int score_bits) {
  void* d_offsets = workspace;
  void* cubWorkspace = GetNextWorkspacePtr(reinterpret_cast<int8_t*>(d_offsets),
                                           (num_images + 1) * sizeof(int));

  SetUniformOffsets(stream,
                    num_images,
                    num_items_per_image,
                    reinterpret_cast<int*>(d_offsets));

  const int array_len = num_images * num_items_per_image;
  size_t temp_storage_bytes =
      CalcCubSortPairsWorkspaceSize<T_SCORE, int>(array_len, num_images);
  size_t begin_bit = 0;
  size_t end_bit = sizeof(T_SCORE) * 8;
  if (sizeof(T_SCORE) == 2 && score_bits > 0 && score_bits <= 10) {
    end_bit = 10;
    begin_bit = end_bit - score_bits;
  }
  cub::DeviceSegmentedRadixSort::SortPairsDescending(
      cubWorkspace,
      temp_storage_bytes,
      reinterpret_cast<T_SCORE*>(unsorted_scores),
      reinterpret_cast<T_SCORE*>(sorted_scores),
      reinterpret_cast<int*>(unsorted_bbox_indices),
      reinterpret_cast<int*>(sorted_bbox_indices),
      array_len,
      num_images,
      reinterpret_cast<int*>(d_offsets),
      reinterpret_cast<int*>(d_offsets) + 1,
      begin_bit,
      end_bit,
      stream);
#ifdef PADDLE_WITH_HIP
  PADDLE_ENFORCE_GPU_SUCCESS(hipGetLastError());
#else
  PADDLE_ENFORCE_GPU_SUCCESS(cudaGetLastError());
#endif
}

template <typename T>
void InferNMS(gpuStream_t stream,
              const int N,
              const int per_batch_boxes_size,
              const int per_batch_scores_size,
              const bool share_location,
              const int background_label_id,
              const int num_preds_per_class,
              const int num_classes,
              const int top_k,
              const int keep_top_k,
              const float score_threshold,
              const float iou_threshold,
              const void* loc_data,
              const void* conf_data,
              void* keep_count,
              void* nmsed_boxes,
              void* nmsed_scores,
              void* nmsed_classes,
              void* nmsed_indices,
              void* nmsed_valid_mask,
              void* workspace,
              bool is_normalized,
              bool conf_sigmoid,
              bool clip_boxes,
              int score_bits,
              bool caffe_semantics) {
  PADDLE_ENFORCE_EQ(
      share_location,
      true,
      common::errors::Unimplemented("share_location=false is not supported."));

  // Prepare workspaces
  size_t bbox_data_size =
      CalcDetectionForwardBBoxDataSize<T>(N, per_batch_boxes_size);
  void* bbox_data_raw = workspace;
#ifdef PADDLE_WITH_HIP
  PADDLE_ENFORCE_GPU_SUCCESS(hipMemcpyAsync(bbox_data_raw,
                                            loc_data,
                                            bbox_data_size,
                                            hipMemcpyDeviceToDevice,
                                            stream));
#else
  PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpyAsync(bbox_data_raw,
                                             loc_data,
                                             bbox_data_size,
                                             cudaMemcpyDeviceToDevice,
                                             stream));
#endif
  void* bbox_data = bbox_data_raw;

  const int num_scores = N * per_batch_scores_size;
  size_t total_scores_size =
      CalcDetectionForwardPreNMSSize<T>(N, per_batch_scores_size);
  void* scores =
      GetNextWorkspacePtr(reinterpret_cast<int8_t*>(bbox_data), bbox_data_size);
#ifdef PADDLE_WITH_HIP
  PADDLE_ENFORCE_GPU_SUCCESS(hipMemcpyAsync(
      scores, conf_data, total_scores_size, hipMemcpyDeviceToDevice, stream));
#else
  PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpyAsync(
      scores, conf_data, total_scores_size, cudaMemcpyDeviceToDevice, stream));
#endif

  size_t indices_size =
      CalcDetectionForwardPreNMSSize<int>(N, per_batch_scores_size);
  void* indices =
      GetNextWorkspacePtr(reinterpret_cast<int8_t*>(scores), total_scores_size);

  size_t post_nms_scores_size =
      CalcDetectionForwardPostNMSSize<T>(N, num_classes, top_k);
  size_t post_nms_indices_size = CalcDetectionForwardPostNMSSize<int>(
      N, num_classes, top_k);  // indices are full int32
  void* post_nms_scores =
      GetNextWorkspacePtr(reinterpret_cast<int8_t*>(indices), indices_size);
  void* post_nms_indices = GetNextWorkspacePtr(
      reinterpret_cast<int8_t*>(post_nms_scores), post_nms_scores_size);

  void* sorting_workspace = GetNextWorkspacePtr(
      reinterpret_cast<int8_t*>(post_nms_indices), post_nms_indices_size);
  // Sort the scores so that the following NMS could be applied.
  float score_shift = 0.f;
  SortScoresPerClassGPU<T>(stream,
                           N,
                           num_classes,
                           num_preds_per_class,
                           background_label_id,
                           score_threshold,
                           scores,
                           indices,
                           sorting_workspace,
                           score_bits,
                           score_shift);

  // This is set to true as the input bounding boxes are of the format [ymin,
  // xmin, ymax, xmax]. The default implementation assumes [xmin, ymin, xmax,
  // ymax]
  bool flip_xy = true;
  // NMS
  AllClassNMSGPU<T, T>(stream,
                       N,
                       num_classes,
                       num_preds_per_class,
                       top_k,
                       iou_threshold,
                       share_location,
                       is_normalized,
                       bbox_data,
                       scores,
                       indices,
                       post_nms_scores,
                       post_nms_indices,
                       flip_xy,
                       score_shift,
                       caffe_semantics);

  // Sort the bounding boxes after NMS using scores
  SortScoresPerImageGPU<T>(stream,
                           N,
                           num_classes * top_k,
                           post_nms_scores,
                           post_nms_indices,
                           scores,
                           indices,
                           sorting_workspace,
                           score_bits);

  // Gather data from the sorted bounding boxes after NMS
  GatherNMSOutputsGPU<T, T>(stream,
                            share_location,
                            N,
                            num_preds_per_class,
                            num_classes,
                            top_k,
                            keep_top_k,
                            indices,
                            scores,
                            bbox_data,
                            keep_count,
                            nmsed_boxes,
                            nmsed_scores,
                            nmsed_classes,
                            nmsed_indices,
                            nmsed_valid_mask,
                            clip_boxes,
                            score_shift);
}

template <typename T, typename Context>
void MultiClassNMSGPUKernel(const Context& dev_ctx,
                            const DenseTensor& bboxes,
                            const DenseTensor& scores,
                            const optional<DenseTensor>& rois_num,
                            float score_threshold,
                            int nms_top_k,
                            int keep_top_k,
                            float nms_threshold,
                            bool normalized,
                            float nms_eta,
                            int background_label,
                            DenseTensor* out,
                            DenseTensor* index,
                            DenseTensor* nms_rois_num) {
  bool return_index = index != nullptr;
  bool has_roisnum = rois_num.get_ptr() != nullptr;
  auto score_dims = scores.dims();
  auto score_size = score_dims.size();

  bool is_supported = (score_size == 3) && (nms_top_k >= 0) &&
                      (nms_top_k <= 4096) && (keep_top_k >= 0) &&
                      (nms_eta == 1.0) && !has_roisnum;
  if (!is_supported) {
    VLOG(6)
        << "This configuration is not supported by GPU kernel. Falling back to "
           "CPU kernel. "
           "Expect (score_size == 3) && (nms_top_k >= 0) && (nms_top_k <= 4096)"
           "(keep_top_k >= 0) && (nms_eta == 1.0) && !has_roisnum, "
           "got score_size="
        << score_size << ", nms_top_k=" << nms_top_k
        << ", keep_top_k=" << keep_top_k << ", nms_eta=" << nms_eta
        << ", has_roisnum=" << has_roisnum;

    DenseTensor bboxes_cpu, scores_cpu, rois_num_cpu_tenor;
    DenseTensor out_cpu, index_cpu, nms_rois_num_cpu;
    optional<DenseTensor> rois_num_cpu(paddle::none);
    auto cpu_place = CPUPlace();
    auto gpu_place = dev_ctx.GetPlace();

    // copy from GPU to CPU
    Copy(dev_ctx, bboxes, cpu_place, false, &bboxes_cpu);
    Copy(dev_ctx, scores, cpu_place, false, &scores_cpu);
    if (has_roisnum) {
      Copy(dev_ctx, *rois_num.get_ptr(), cpu_place, false, &rois_num_cpu_tenor);
      rois_num_cpu = optional<DenseTensor>(rois_num_cpu_tenor);
    }
    dev_ctx.Wait();
    DeviceContextPool& pool = DeviceContextPool::Instance();
    auto* cpu_ctx = static_cast<CPUContext*>(pool.Get(cpu_place));
    MultiClassNMSKernel<T, CPUContext>(*cpu_ctx,
                                       bboxes_cpu,
                                       scores_cpu,
                                       rois_num_cpu,
                                       score_threshold,
                                       nms_top_k,
                                       keep_top_k,
                                       nms_threshold,
                                       normalized,
                                       nms_eta,
                                       background_label,
                                       &out_cpu,
                                       &index_cpu,
                                       &nms_rois_num_cpu);
    // copy back
    Copy(dev_ctx, out_cpu, gpu_place, false, out);
    Copy(dev_ctx, index_cpu, gpu_place, false, index);
    Copy(dev_ctx, nms_rois_num_cpu, gpu_place, false, nms_rois_num);
    return;
  }

  // Calculate input shapes
  int64_t batch_size = score_dims[0];
  const int64_t per_batch_boxes_size =
      bboxes.dims()[1] * bboxes.dims()[2];  // M * 4
  const int64_t per_batch_scores_size =
      scores.dims()[1] * scores.dims()[2];       // C * M
  const int64_t num_priors = bboxes.dims()[1];   // M
  const int64_t num_classes = scores.dims()[1];  // C
  const bool share_location = true;
  auto stream = reinterpret_cast<const Context&>(dev_ctx).stream();
  // Sanity check
  PADDLE_ENFORCE_LE(
      nms_top_k,
      num_priors,
      common::errors::InvalidArgument("Expect nms_top_k (%d)"
                                      " <= num of boxes per batch (%d).",
                                      nms_top_k,
                                      num_priors));
  PADDLE_ENFORCE_LE(keep_top_k,
                    nms_top_k,
                    common::errors::InvalidArgument("Expect keep_top_k (%d)"
                                                    " <= nms_top_k (%d).",
                                                    keep_top_k,
                                                    nms_top_k));

  // Transform the layout of bboxes and scores
  // bboxes: [N,M,4] -> [N,1,M,4]
  DenseTensor transformed_bboxes(bboxes.type());
  transformed_bboxes.ShareDataWith(bboxes).Resize(
      {bboxes.dims()[0], 1, bboxes.dims()[1], bboxes.dims()[2]});
  // scores: [N, C, M] => [N, C, M, 1]
  DenseTensor transformed_scores(scores.type());
  transformed_scores.ShareDataWith(scores).Resize(
      {scores.dims()[0], scores.dims()[1], scores.dims()[2], 1});

  // Prepare intermediate outputs for NMS kernels
  DenseTensor keep_count(DataType::INT32);
  keep_count.Resize({batch_size});
  if (nms_rois_num != nullptr) {
    nms_rois_num->Resize({batch_size});
    dev_ctx.template Alloc<int>(nms_rois_num);
    keep_count.ShareDataWith(*nms_rois_num);
  } else {
    dev_ctx.template Alloc<int>(&keep_count);
  }

  DenseTensor nmsed_indices(DataType::INT32);
  nmsed_indices.Resize({batch_size * keep_top_k, 1});
  dev_ctx.template Alloc<int>(&nmsed_indices);

  DenseTensor nmsed_valid_mask(DataType::INT32);
  nmsed_valid_mask.Resize({batch_size * keep_top_k});
  dev_ctx.template Alloc<int>(&nmsed_valid_mask);

  DenseTensor nmsed_boxes(bboxes.dtype());
  DenseTensor nmsed_scores(scores.dtype());
  DenseTensor nmsed_classes(scores.dtype());
  nmsed_boxes.Resize({batch_size * keep_top_k, 4});
  nmsed_scores.Resize({batch_size * keep_top_k, 1});
  nmsed_classes.Resize({batch_size * keep_top_k, 1});
  dev_ctx.template Alloc<T>(&nmsed_boxes);
  dev_ctx.template Alloc<T>(&nmsed_scores);
  dev_ctx.template Alloc<T>(&nmsed_classes);

  auto workspace_size =
      CalcDetectionInferenceWorkspaceSize<T>(share_location,
                                             batch_size,
                                             per_batch_boxes_size,
                                             per_batch_scores_size,
                                             num_classes,
                                             num_priors,
                                             nms_top_k);

  DenseTensor workspace = DenseTensor();
  workspace.Resize({static_cast<int64_t>(workspace_size)});
  T* workspace_ptr = dev_ctx.template Alloc<T>(&workspace);

  // Launch the NMS kernel
  InferNMS<T>(stream,
              batch_size,
              per_batch_boxes_size,
              per_batch_scores_size,
              share_location,
              background_label,
              num_priors,
              num_classes,
              nms_top_k,
              keep_top_k,
              score_threshold,
              nms_threshold,
              transformed_bboxes.data<T>(),
              transformed_scores.data<T>(),
              keep_count.data<int>(),
              nmsed_boxes.data<T>(),
              nmsed_scores.data<T>(),
              nmsed_classes.data<T>(),
              nmsed_indices.data<int>(),
              nmsed_valid_mask.data<int>(),
              workspace_ptr,
              normalized,
              false,
              false,
              0,
              true);

  // Post-processing to get the final outputs
  // Concat the individual class, score and boxes outputs
  // into a [N * M, 6] tensor.
  DenseTensor raw_out;
  raw_out.Resize({batch_size * keep_top_k, 6});
  dev_ctx.template Alloc<T>(&raw_out);
  funcs::ConcatFunctor<Context, T> concat;
  concat(dev_ctx, {nmsed_classes, nmsed_scores, nmsed_boxes}, 1, &raw_out);

  // Output of NMS kernel may include invalid entries, which is
  // marked by nmsed_valid_mask. Eliminate the invalid entries
  // by gathering the valid ones.

  // 1. Get valid indices
  DenseTensor valid_indices;
  NonZeroKernel<int, Context>(dev_ctx, nmsed_valid_mask, &valid_indices);
  // 2. Perform gathering
  const int64_t valid_samples = valid_indices.dims()[0];
  out->Resize({valid_samples, 6});
  dev_ctx.template Alloc<T>(out);
  funcs::GPUGatherNd<T, int64_t>(dev_ctx, raw_out, valid_indices, out);
  index->Resize({valid_samples, 1});
  dev_ctx.template Alloc<int>(index);
  funcs::GPUGatherNd<int, int64_t>(
      dev_ctx, nmsed_indices, valid_indices, index);
}

}  // namespace phi

PD_REGISTER_KERNEL(
    multiclass_nms3, GPU, ALL_LAYOUT, phi::MultiClassNMSGPUKernel, float) {
  kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
  kernel->OutputAt(2).SetDataType(phi::DataType::INT32);
}