/* 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 <algorithm>
#include <vector>
#include "paddle/phi/kernels/funcs/index_elementwise.cu.h"
#ifdef __NVCC__
#include <curand_kernel.h>
#endif
#ifdef __HIPCC__
#include <hiprand_kernel.h>
#endif

#include "paddle/common/flags.h"
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
#include "paddle/phi/backends/gpu/gpu_primitives.h"
#include "paddle/phi/kernels/funcs/distribution_helper.h"
#include "paddle/phi/kernels/funcs/fast_divmod.h"
#include "paddle/phi/kernels/funcs/pooling.h"
#include "paddle/phi/kernels/funcs/random.cuh"
#include "paddle/phi/kernels/funcs/reduce_function.h"

COMMON_DECLARE_bool(use_accuracy_compatible_kernel);

namespace phi {
namespace funcs {

template <typename IndexT>
struct FastDivModForPooling {
 public:
  FastDivMod<IndexT> channel;
  FastDivMod<IndexT> width;
  FastDivMod<IndexT> height;

  explicit HOSTDEVICE FastDivModForPooling(const IndexT channels,
                                           const IndexT output_width,
                                           const IndexT output_height)
      : channel(channels), width(output_width), height(output_height) {}
};

template <typename IndexT>
struct FastDivModForPooling3D {
 public:
  FastDivMod<IndexT> channel;
  FastDivMod<IndexT> width;
  FastDivMod<IndexT> height;
  FastDivMod<IndexT> depth;

  explicit HOSTDEVICE FastDivModForPooling3D(const IndexT channels,
                                             const IndexT output_width,
                                             const IndexT output_height,
                                             const IndexT output_depth)
      : channel(channels),
        width(output_width),
        height(output_height),
        depth(output_depth) {}
};

template <typename IndexT>
struct FastDivModForPoolingWithMoreStaff {
 public:
  FastDivMod<IndexT> channel;
  FastDivMod<IndexT> width;
  FastDivMod<IndexT> height;
  FastDivMod<IndexT> ksize_w;
  FastDivMod<IndexT> ksize_h;
  FastDivMod<IndexT> stride_w;
  FastDivMod<IndexT> stride_h;

  explicit HOSTDEVICE FastDivModForPoolingWithMoreStaff(
      const IndexT channels,
      const IndexT input_width,
      const IndexT input_height,
      const IndexT ksize_width,
      const IndexT ksize_height,
      const IndexT stride_width,
      const IndexT stride_height)
      : channel(channels),
        width(input_width),
        height(input_height),
        ksize_w(ksize_width),
        ksize_h(ksize_height),
        stride_w(stride_width),
        stride_h(stride_height) {}
};

static __device__ inline int p_start(int size,
                                     int pad,
                                     int kernel,
                                     int stride) {
  return (size + pad < kernel) ? 0 : (size + pad - kernel) / stride + 1;
}

static __device__ inline int p_end(int size,
                                   int pad,
                                   int pooled_size,
                                   int stride) {
  return std::min((size + pad) / stride + 1, pooled_size);
}

template <typename FastDivModForPooling, typename IndexT>
__device__ void OffsetPreparationFor4Dimension(IndexT index,
                                               bool channel_last,
                                               FastDivModForPooling divmods,
                                               const IndexT pad_width,
                                               const IndexT pad_height,
                                               const IndexT aux_width,
                                               const IndexT aux_height,
                                               IndexT* w_offset,
                                               IndexT* h_offset,
                                               IndexT* c_offset,
                                               IndexT* stride) {
  if (!channel_last) { /* NCHW */
    auto input_width_divmod = divmods.width.Divmod(index);
    auto input_height_divmod = divmods.height.Divmod(input_width_divmod.val[0]);
    auto channel_divmod = divmods.channel.Divmod(input_height_divmod.val[0]);
    *w_offset = input_width_divmod.val[1] + pad_width;
    *h_offset = input_height_divmod.val[1] + pad_height;
    *c_offset = channel_divmod.val[1];
    *stride = (channel_divmod.val[0] * divmods.channel.divisor + *c_offset) *
              aux_height * aux_width;
  } else { /* NHWC */
    auto c_divmod = divmods.channel.Divmod(index);
    auto input_width_divmod = divmods.width.Divmod(c_divmod.val[0]);
    auto input_height_divmod = divmods.height.Divmod(input_width_divmod.val[0]);
    *c_offset = c_divmod.val[1];
    *w_offset = input_width_divmod.val[1] + pad_width;
    *h_offset = input_height_divmod.val[1] + pad_height;
    *stride = input_height_divmod.val[0] * aux_height * aux_width *
              divmods.channel.divisor;
  }
}

template <typename IndexT>
__device__ void PreparationPoolSize(IndexT index,
                                    IndexT input_size,
                                    IndexT output_size,
                                    FastDivMod<IndexT> divmods,
                                    IndexT* tmp_size

) {
  IndexT left = (index == 0) ? 0 : divmods.Div(index * input_size);
  IndexT right = (index == output_size - 1)
                     ? input_size
                     : divmods.DivCeil((index + 1) * input_size);
  *tmp_size = right - left;
}

template <typename PoolProcess, typename T, typename IndexT>
__global__ void KernelPool2D(const IndexT nthreads,
                             const T* input_data,
                             const IndexT channels,
                             const IndexT input_height,
                             const IndexT input_width,
                             const IndexT output_height,
                             const IndexT output_width,
                             const IndexT ksize_height,
                             const IndexT ksize_width,
                             const IndexT stride_height,
                             const IndexT stride_width,
                             const IndexT padding_height,
                             const IndexT padding_width,
                             FastDivModForPooling<IndexT> divmods,
                             PoolProcess pool_process,
                             bool exclusive,
                             T* output_data,
                             bool channel_last = false) {
  const IndexT start_index =
      static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
  const IndexT step = static_cast<IndexT>(blockDim.x) * gridDim.x;
  for (IndexT index = start_index; index < nthreads; index += step) {
    IndexT hstart, hend, wstart, wend;
    IndexT w_offset, h_offset, c_offset, input_offset;
    OffsetPreparationFor4Dimension<FastDivModForPooling<IndexT>, IndexT>(
        index,
        channel_last,
        divmods,
        0,
        0,
        input_width,
        input_height,
        &w_offset,
        &h_offset,
        &c_offset,
        &input_offset);
    input_data += input_offset;

    hstart = h_offset * stride_height - padding_height;
    hend = min(hstart + ksize_height, input_height);
    hstart = max(hstart, static_cast<IndexT>(0));
    wstart = w_offset * stride_width - padding_width;
    wend = min(wstart + ksize_width, input_width);
    wstart = max(wstart, static_cast<IndexT>(0));

    T ele = pool_process.initial();
    for (IndexT h = hstart; h < hend; ++h) {
      for (IndexT w = wstart; w < wend; ++w) {
        auto input_idx = channel_last
                             ? (h * input_width + w) * channels + c_offset
                             : h * input_width + w;
        pool_process.compute(input_data[input_idx], &ele);
      }
    }
    IndexT pool_size = exclusive ? (hend - hstart) * (wend - wstart)
                                 : ksize_height * ksize_width;
    pool_process.finalize(static_cast<T>(pool_size), &ele);
    output_data[index] = ele;
  }
}

template <typename PoolProcess, typename T, typename IndexT>
__global__ void AdaptiveKernelPool2D(const IndexT nthreads,
                                     const T* input_data,
                                     const IndexT channels,
                                     const IndexT input_height,
                                     const IndexT input_width,
                                     const IndexT output_height,
                                     const IndexT output_width,
                                     const IndexT ksize_height,
                                     const IndexT ksize_width,
                                     const IndexT stride_height,
                                     const IndexT stride_width,
                                     const IndexT padding_height,
                                     const IndexT padding_width,
                                     const IndexT batch_size,
                                     FastDivModForPooling<IndexT> divmods,
                                     PoolProcess pool_process,
                                     bool exclusive,
                                     T* output_data,
                                     bool channel_last = false) {
  for (IndexT n_offset = blockIdx.y; n_offset < batch_size;
       n_offset += gridDim.y) {
    const IndexT c_offset =
        static_cast<IndexT>(blockIdx.x) * blockDim.y + threadIdx.y;
    if (c_offset >= channels) {
      return;
    }
    IndexT hstart, hend, wstart, wend;
    IndexT input_offset =
        channel_last
            ? n_offset * input_height * input_width * channels
            : (n_offset * channels + c_offset) * input_height * input_width;
    IndexT output_offset =
        channel_last
            ? n_offset * output_height * output_width * channels
            : (n_offset * channels + c_offset) * output_height * output_width;
    for (IndexT hw_offset = threadIdx.x;
         hw_offset < output_height * output_width;
         hw_offset += blockDim.x) {
      IndexT w_offset = hw_offset % output_width;
      IndexT h_offset = hw_offset / output_width;
      hstart = AdaptStartIndex(h_offset, input_height, output_height);
      hend = AdaptEndIndex(h_offset, input_height, output_height);
      wstart = AdaptStartIndex(w_offset, input_width, output_width);
      wend = AdaptEndIndex(w_offset, input_width, output_width);

      T ele = pool_process.initial();
      for (IndexT h = hstart; h < hend; ++h) {
        for (IndexT w = wstart; w < wend; ++w) {
          auto input_idx = channel_last
                               ? (h * input_width + w) * channels + c_offset
                               : h * input_width + w;
          pool_process.compute(input_data[input_offset + input_idx], &ele);
        }
      }
      IndexT pool_size = (hend - hstart) * (wend - wstart);
      pool_process.finalize(static_cast<T>(pool_size), &ele);
      IndexT output_idx =
          channel_last
              ? (h_offset * output_width + w_offset) * channels + c_offset
              : h_offset * output_width + w_offset;
      output_data[output_offset + output_idx] = ele;
    }
  }
}

template <typename T, typename PoolProcess, typename IndexT>
__global__ void KernelPool2DGrad(
    const IndexT nthreads,
    const T* __restrict__ input_data,
    const T* __restrict__ output_data,
    const T* __restrict__ output_grad,
    const IndexT output_width,
    const IndexT output_height,
    const IndexT input_width,
    const IndexT input_height,
    const IndexT ksize_width,
    const IndexT ksize_height,
    const IndexT stride_width,
    const IndexT stride_height,
    const IndexT padding_width,
    const IndexT padding_height,
    FastDivModForPoolingWithMoreStaff<IndexT> divmods,
    PoolProcess pool_process,
    bool exclusive,
    bool adaptive,
    T* __restrict__ input_grad,
    bool channel_last = false) {
  const IndexT start_index =
      static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;

  for (IndexT index =
           static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
       index < nthreads;
       index += static_cast<IndexT>(blockDim.x) * gridDim.x) {
    T input = static_cast<T>(0);
    T input_grad_data = static_cast<T>(0);
    IndexT phstart, phend, pwstart, pwend;
    IndexT w_offset, h_offset, c_offset, output_offset;
    OffsetPreparationFor4Dimension<FastDivModForPoolingWithMoreStaff<IndexT>,
                                   IndexT>(index,
                                           channel_last,
                                           divmods,
                                           padding_width,
                                           padding_height,
                                           output_width,
                                           output_height,
                                           &w_offset,
                                           &h_offset,
                                           &c_offset,
                                           &output_offset);
    if (pool_process.use_x) {
      input = input_data[index];
      output_data += output_offset;
    }
    output_grad += output_offset;

    if (adaptive) {
      auto tmp_phstart = divmods.height.Divmod(h_offset * output_height);
      auto tmp_pwstart = divmods.width.Divmod(w_offset * output_width);
      auto tmp_phend = divmods.height.Divmod((h_offset + 1) * output_height);
      auto tmp_pwend = divmods.width.Divmod((w_offset + 1) * output_width);
      phstart = tmp_phstart.val[0];
      pwstart = tmp_pwstart.val[0];
      phend = tmp_phend.val[1] > 0 ? tmp_phend.val[0] + 1 : tmp_phend.val[0];
      pwend = tmp_pwend.val[1] > 0 ? tmp_pwend.val[0] + 1 : tmp_pwend.val[0];

      IndexT pool_height, pool_width;
      for (IndexT ph = phstart; ph < phend; ++ph) {
        PreparationPoolSize(
            ph, input_height, output_height, divmods.ksize_h, &pool_height);

        for (IndexT pw = pwstart; pw < pwend; ++pw) {
          PreparationPoolSize(
              pw, input_width, output_width, divmods.ksize_w, &pool_width);
          IndexT pool_size = pool_height * pool_width;
          IndexT tmp_idx = ph * output_width + pw;
          IndexT output_sub_idx =
              channel_last ? tmp_idx * divmods.channel.divisor + c_offset
                           : tmp_idx;
          T output_value = pool_process.use_x ? output_data[output_sub_idx]
                                              : static_cast<T>(0);
          pool_process.compute(input,
                               output_value,
                               output_grad[output_sub_idx],
                               static_cast<T>(1.0 / pool_size),
                               &input_grad_data);
        }
      }
    } else {
      auto stride_height_div = divmods.stride_h.Div(h_offset - ksize_height);
      auto stride_width_div = divmods.stride_w.Div(w_offset - ksize_width);
      phstart = (h_offset < ksize_height) ? 0 : stride_height_div + 1;
      pwstart = (w_offset < ksize_width) ? 0 : stride_width_div + 1;
      phend = min(divmods.stride_h.Div(h_offset) + 1, output_height);
      pwend = min(divmods.stride_w.Div(w_offset) + 1, output_width);

      if (exclusive) {
        for (IndexT ph = phstart; ph < phend; ++ph) {
          for (IndexT pw = pwstart; pw < pwend; ++pw) {
            IndexT hstart = ph * stride_height - padding_height;
            IndexT wstart = pw * stride_width - padding_width;
            IndexT hend = min(hstart + ksize_height, input_height);
            IndexT wend = min(wstart + ksize_width, input_width);
            hstart = max(hstart, static_cast<IndexT>(0));
            wstart = max(wstart, static_cast<IndexT>(0));
            IndexT pool_size = (hend - hstart) * (wend - wstart);
            IndexT tmp_idx = ph * output_width + pw;
            IndexT output_sub_idx =
                channel_last ? tmp_idx * divmods.channel.divisor + c_offset
                             : tmp_idx;
            T output_value = pool_process.use_x ? output_data[output_sub_idx]
                                                : static_cast<T>(0);
            pool_process.compute(input,
                                 output_value,
                                 output_grad[output_sub_idx],
                                 static_cast<T>(1.0 / pool_size),
                                 &input_grad_data);
          }
        }
      } else {
        for (IndexT ph = phstart; ph < phend; ++ph) {
          for (IndexT pw = pwstart; pw < pwend; ++pw) {
            IndexT pool_size = ksize_height * ksize_width;
            IndexT tmp_idx = ph * output_width + pw;
            IndexT output_sub_idx =
                channel_last ? tmp_idx * divmods.channel.divisor + c_offset
                             : tmp_idx;
            T output_value = pool_process.use_x ? output_data[output_sub_idx]
                                                : static_cast<T>(0);
            pool_process.compute(input,
                                 output_value,
                                 output_grad[output_sub_idx],
                                 static_cast<T>(1.0 / pool_size),
                                 &input_grad_data);
          }
        }
      }
    }
    input_grad[index] = input_grad_data;
  }
}

template <typename T, typename IndexT>
__global__ void KernelMaxPool2DGrad(const IndexT nthreads,
                                    const T* input_data,
                                    const T* output_data,
                                    const T* output_grad,
                                    const IndexT channels,
                                    const IndexT input_height,
                                    const IndexT input_width,
                                    const IndexT output_height,
                                    const IndexT output_width,
                                    const IndexT ksize_height,
                                    const IndexT ksize_width,
                                    const IndexT stride_height,
                                    const IndexT stride_width,
                                    const IndexT padding_height,
                                    const IndexT padding_width,
                                    T* input_grad,
                                    FastDivModForPooling<IndexT> divmods,
                                    bool channel_last = false) {
  const IndexT start_index =
      static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
  const IndexT step = static_cast<IndexT>(blockDim.x) * gridDim.x;
  for (IndexT index = start_index; index < nthreads; index += step) {
    IndexT w_offset, h_offset, c_offset, input_offset;
    OffsetPreparationFor4Dimension<FastDivModForPooling<IndexT>, IndexT>(
        index,
        channel_last,
        divmods,
        0,
        0,
        input_width,
        input_height,
        &w_offset,
        &h_offset,
        &c_offset,
        &input_offset);
    input_data += input_offset;
    input_grad += input_offset;

    IndexT hstart = h_offset * stride_height - padding_height;
    IndexT hend = min(hstart + ksize_height, input_height);
    hstart = max(hstart, static_cast<IndexT>(0));

    IndexT wstart = w_offset * stride_width - padding_width;
    IndexT wend = min(wstart + ksize_width, input_width);
    wstart = max(wstart, static_cast<IndexT>(0));

    T ele = output_data[index];
    IndexT maxIndex = -1;
    bool stop = false;
    for (IndexT h = hstart; h < hend && !stop; ++h) {
      for (IndexT w = wstart; w < wend && !stop; ++w) {
        IndexT input_data_idx =
            channel_last ? (h * input_width + w) * channels + c_offset
                         : h * input_width + w;
        if (ele == input_data[input_data_idx]) {
          maxIndex = input_data_idx;
          stop = true;
        }
      }
    }

    if (maxIndex != -1) {
      // atomic add
      CudaAtomicAdd(input_grad + maxIndex, output_grad[index]);
    }
  }
}

template <typename T, typename IndexT>
__global__ void KernelMaxPool2DGradCompatible(
    const T* input_data,
    const T* output_data,
    const T* output_grad,
    const IndexT batch_size,
    const IndexT channels,
    const IndexT input_height,
    const IndexT input_width,
    const IndexT output_height,
    const IndexT output_width,
    const IndexT ksize_height,
    const IndexT ksize_width,
    const IndexT stride_height,
    const IndexT stride_width,
    const IndexT padding_height,
    const IndexT padding_width,
    T* input_grad,
    FastDivModForPooling<IndexT> divmods,
    bool channel_last = false) {
  using MT = typename MPTypeTrait<T>::Type;

  CUDA_KERNEL_LOOP(index, input_height * input_width) {
    IndexT h = index / input_width;
    IndexT w = index - h * input_width;
    IndexT phstart = p_start(h, padding_height, ksize_height, stride_height);
    IndexT phend = p_end(h, padding_height, output_height, stride_height);
    IndexT pwstart = p_start(w, padding_width, ksize_width, stride_width);
    IndexT pwend = p_end(w, padding_width, output_width, stride_width);
    T input_data_value = input_data[h * input_width + w];
    for (IndexT n = blockIdx.y; n < batch_size; n += gridDim.y) {
      for (IndexT c = blockIdx.z; c < channels; c += gridDim.z) {
        MT gradient = static_cast<MT>(0.0f);
        IndexT offset = (n * channels + c) * output_height * output_width;
        for (int ph = phstart; ph < phend; ++ph) {
          for (int pw = pwstart; pw < pwend; ++pw) {
            T output_data_value = output_data[ph * output_width + pw + offset];
            if (output_data_value == input_data_value) {
              gradient +=
                  static_cast<MT>(output_grad[ph * output_width + pw + offset]);
            }
          }
        }
        input_grad[(n * channels + c) * input_height * input_width + index] =
            static_cast<MT>(gradient);
      }
    }
  }
}

template <typename PoolProcess, typename T>
void Pool2dDirectCUDAFunctor<PoolProcess, T>::operator()(
    const T* input,
    const std::vector<int>& input_shape,
    const std::vector<int>& output_shape,
    const std::vector<int>& ksize,
    const std::vector<int>& strides,
    const std::vector<int>& paddings,
    bool exclusive,
    bool adaptive,
    T* output,
    gpuStream_t stream,
    PoolProcess pool_compute) {
  const int batch_size = input_shape[0];
  const int input_channels = input_shape[1];
  const int input_height = input_shape[2];
  const int input_width = input_shape[3];
  const int output_channels = output_shape[1];
  const int output_height = output_shape[2];
  const int output_width = output_shape[3];
  const int ksize_height = ksize[0];
  const int ksize_width = ksize[1];
  const int stride_height = strides[0];
  const int stride_width = strides[1];
  const int padding_height = paddings[0];
  const int padding_width = paddings[1];
  int64_t nthreads = static_cast<int64_t>(batch_size) * output_channels *
                     output_height * output_width;
  auto pool_divmods =
      FastDivModForPooling<int>(input_channels, output_width, output_height);
  if (adaptive) {
    int64_t max_threads = 512;
    int64_t thread_num = std::min(
        funcs::details::GetLastPow2(output_height * output_width), max_threads);
    int64_t blocks = std::min(max_threads / thread_num,
                              static_cast<int64_t>(output_channels));
    auto max_grid_dim = backends::gpu::GetGpuMaxGridDimSize(
        phi::backends::gpu::GetCurrentDeviceId());
    dim3 threads(thread_num, blocks, 1);
    dim3 grid(std::max((output_channels + blocks - 1) / blocks,
                       static_cast<int64_t>(1)),
              std::min(batch_size, static_cast<int>(max_grid_dim[1])),
              1);
    AdaptiveKernelPool2D<PoolProcess, T, int>
        <<<grid, threads, 0, stream>>>(nthreads,
                                       input,
                                       input_channels,
                                       input_height,
                                       input_width,
                                       output_height,
                                       output_width,
                                       ksize_height,
                                       ksize_width,
                                       stride_height,
                                       stride_width,
                                       padding_height,
                                       padding_width,
                                       batch_size,
                                       pool_divmods,
                                       pool_compute,
                                       exclusive,
                                       output);

  } else {
    int thread_num = 1024;
#ifdef WITH_NV_JETSON
    // backends::gpu::ChangeThreadNum(dev_ctx, &thread_num);
    thread_num = 512;
#endif
    int blocks = (nthreads + thread_num - 1) / thread_num;
    dim3 threads(thread_num, 1);
    dim3 grid(blocks, 1);
    KernelPool2D<PoolProcess, T, int>
        <<<grid, threads, 0, stream>>>(nthreads,
                                       input,
                                       input_channels,
                                       input_height,
                                       input_width,
                                       output_height,
                                       output_width,
                                       ksize_height,
                                       ksize_width,
                                       stride_height,
                                       stride_width,
                                       padding_height,
                                       padding_width,
                                       pool_divmods,
                                       pool_compute,
                                       exclusive,
                                       output);
  }
}

/*
 * Tensors are in NCHW or NHWC format.
 * Ksize, strides are two elements. These two elements represent height
 * and width, respectively.
 * Paddings are four elements. These four elements represent height_up,
 * height_down, width_left and width_right, respectively.
 */
template <typename PoolProcess, typename T>
class Pool2dFunctor<GPUContext, PoolProcess, T> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& input,
                  const std::vector<int64_t>& ksize,
                  const std::vector<int64_t>& strides,
                  const std::vector<int64_t>& paddings,
                  const std::string data_format,
                  bool exclusive,
                  bool adaptive,
                  DenseTensor* output,
                  PoolProcess pool_process) {
    bool channel_last = (data_format == "NHWC");
    const int64_t batch_size = input.dims()[0];

    const int64_t input_channels =
        channel_last ? input.dims()[3] : input.dims()[1];
    const int64_t input_height =
        channel_last ? input.dims()[1] : input.dims()[2];
    const int64_t input_width =
        channel_last ? input.dims()[2] : input.dims()[3];

    const int64_t output_channels =
        channel_last ? output->dims()[3] : output->dims()[1];
    const int64_t output_height =
        channel_last ? output->dims()[1] : output->dims()[2];
    const int64_t output_width =
        channel_last ? output->dims()[2] : output->dims()[3];

    const int64_t ksize_height = ksize[0];
    const int64_t ksize_width = ksize[1];

    const int64_t stride_height = strides[0];
    const int64_t stride_width = strides[1];

    const int64_t padding_height = paddings[0];
    const int64_t padding_width = paddings[1];

    const T* input_data = input.data<T>();
    T* output_data = dev_ctx.template Alloc<T>(output);

    std::array<unsigned int, 3> max_grid_dim = dev_ctx.GetCUDAMaxGridDimSize();
    int64_t nthreads =
        batch_size * output_channels * output_height * output_width;
    if (adaptive) {
      int64_t max_threads = 512;
      int64_t thread_num =
          std::min(funcs::details::GetLastPow2(output_height * output_width),
                   max_threads);
      int64_t blocks = std::min(max_threads / thread_num,
                                static_cast<int64_t>(output_channels));
      dim3 threads(thread_num, blocks, 1);
      dim3 grid(std::max((output_channels + blocks - 1) / blocks,
                         static_cast<int64_t>(1)),
                std::min(batch_size, static_cast<int64_t>(max_grid_dim[1])),
                1);
      if (input.numel() <= std::numeric_limits<int>::max()) {
        auto pool_divmods = FastDivModForPooling<int>(
            input_channels, output_width, output_height);
        AdaptiveKernelPool2D<PoolProcess, T, int>
            <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                     input_data,
                                                     input_channels,
                                                     input_height,
                                                     input_width,
                                                     output_height,
                                                     output_width,
                                                     ksize_height,
                                                     ksize_width,
                                                     stride_height,
                                                     stride_width,
                                                     padding_height,
                                                     padding_width,
                                                     batch_size,
                                                     pool_divmods,
                                                     pool_process,
                                                     exclusive,
                                                     output_data,
                                                     channel_last);
      } else {
        auto pool_divmods = FastDivModForPooling<int64_t>(
            input_channels, output_width, output_height);
        AdaptiveKernelPool2D<PoolProcess, T, int64_t>
            <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                     input_data,
                                                     input_channels,
                                                     input_height,
                                                     input_width,
                                                     output_height,
                                                     output_width,
                                                     ksize_height,
                                                     ksize_width,
                                                     stride_height,
                                                     stride_width,
                                                     padding_height,
                                                     padding_width,
                                                     batch_size,
                                                     pool_divmods,
                                                     pool_process,
                                                     exclusive,
                                                     output_data,
                                                     channel_last);
      }
    } else {
      int thread_num = 1024;
#ifdef WITH_NV_JETSON
      backends::gpu::ChangeThreadNum(dev_ctx, &thread_num);
#endif
      int64_t blocks = (nthreads + thread_num - 1) / thread_num;
      dim3 threads(thread_num, 1);
      dim3 grid(blocks, 1);
      if (input.numel() <= std::numeric_limits<int>::max()) {
        auto pool_divmods = FastDivModForPooling<int>(
            input_channels, output_width, output_height);
        KernelPool2D<PoolProcess, T, int>
            <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                     input_data,
                                                     input_channels,
                                                     input_height,
                                                     input_width,
                                                     output_height,
                                                     output_width,
                                                     ksize_height,
                                                     ksize_width,
                                                     stride_height,
                                                     stride_width,
                                                     padding_height,
                                                     padding_width,
                                                     pool_divmods,
                                                     pool_process,
                                                     exclusive,
                                                     output_data,
                                                     channel_last);
      } else {
        auto pool_divmods = FastDivModForPooling<int64_t>(
            input_channels, output_width, output_height);
        KernelPool2D<PoolProcess, T, int64_t>
            <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                     input_data,
                                                     input_channels,
                                                     input_height,
                                                     input_width,
                                                     output_height,
                                                     output_width,
                                                     ksize_height,
                                                     ksize_width,
                                                     stride_height,
                                                     stride_width,
                                                     padding_height,
                                                     padding_width,
                                                     pool_divmods,
                                                     pool_process,
                                                     exclusive,
                                                     output_data,
                                                     channel_last);
      }
    }
  }
};
/*
 * Tensors are in NCHW or NHWC format.
 * Ksize, strides are two elements. These two elements represent height
 * and width, respectively.
 * Paddings are four elements. These four elements represent height_up,
 * height_down, width_left and width_right, respectively.
 */
template <typename PoolProcess, typename T>
class Pool2dGradFunctor<GPUContext, PoolProcess, T> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& input,
                  const DenseTensor& output,
                  const DenseTensor& output_grad,
                  const std::vector<int64_t>& ksize,
                  const std::vector<int64_t>& strides,
                  const std::vector<int64_t>& paddings,
                  const std::string data_format,
                  bool exclusive,
                  bool adaptive,
                  DenseTensor* input_grad,
                  PoolProcess pool_process) {
    bool channel_last = (data_format == "NHWC");

    const int64_t batch_size = input.dims()[0];
    const int64_t input_channels =
        channel_last ? input.dims()[3] : input.dims()[1];
    const int64_t input_height =
        channel_last ? input.dims()[1] : input.dims()[2];
    const int64_t input_width =
        channel_last ? input.dims()[2] : input.dims()[3];

    const int64_t output_channels =
        channel_last ? output.dims()[3] : output.dims()[1];
    const int64_t output_height =
        channel_last ? output.dims()[1] : output.dims()[2];
    const int64_t output_width =
        channel_last ? output.dims()[2] : output.dims()[3];

    const int64_t ksize_height = ksize[0];
    const int64_t ksize_width = ksize[1];

    const int64_t stride_height = strides[0];
    const int64_t stride_width = strides[1];

    const int64_t padding_height = paddings[0];
    const int64_t padding_width = paddings[1];

    const T* input_data = input.data<T>();
    const T* output_data = output.data<T>();
    const T* output_grad_data = output_grad.data<T>();
    T* input_grad_data = dev_ctx.template Alloc<T>(input_grad);

    int64_t nthreads = batch_size * input_channels * input_height * input_width;
    auto config = phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, nthreads);
    if (input.numel() <= std::numeric_limits<int>::max() &&
        output.numel() <= std::numeric_limits<int>::max()) {
      auto pool_divmods = FastDivModForPoolingWithMoreStaff<int>(input_channels,
                                                                 input_width,
                                                                 input_height,
                                                                 ksize_width,
                                                                 ksize_height,
                                                                 stride_width,
                                                                 stride_height);
      KernelPool2DGrad<T, PoolProcess, int>
          <<<config.block_per_grid,
             config.thread_per_block,
             0,
             dev_ctx.stream()>>>(nthreads,
                                 input_data,
                                 output_data,
                                 output_grad_data,
                                 output_width,
                                 output_height,
                                 input_width,
                                 input_height,
                                 ksize_width,
                                 ksize_height,
                                 stride_width,
                                 stride_height,
                                 padding_width,
                                 padding_height,
                                 pool_divmods,
                                 pool_process,
                                 exclusive,
                                 adaptive,
                                 input_grad_data,
                                 channel_last);
    } else {
      auto pool_divmods =
          FastDivModForPoolingWithMoreStaff<int64_t>(input_channels,
                                                     input_width,
                                                     input_height,
                                                     ksize_width,
                                                     ksize_height,
                                                     stride_width,
                                                     stride_height);
      KernelPool2DGrad<T, PoolProcess, int64_t>
          <<<config.block_per_grid,
             config.thread_per_block,
             0,
             dev_ctx.stream()>>>(nthreads,
                                 input_data,
                                 output_data,
                                 output_grad_data,
                                 output_width,
                                 output_height,
                                 input_width,
                                 input_height,
                                 ksize_width,
                                 ksize_height,
                                 stride_width,
                                 stride_height,
                                 padding_width,
                                 padding_height,
                                 pool_divmods,
                                 pool_process,
                                 exclusive,
                                 adaptive,
                                 input_grad_data,
                                 channel_last);
    }
  }
};

/*
 * Tensors are in NCHW or NHWC format.
 * Ksize, strides are two elements. These two elements represent height
 * and width, respectively.
 * Paddings are four elements. These four elements represent height_up,
 * height_down, width_left and width_right, respectively.
 */
template <typename T>
class MaxPool2dGradFunctor<GPUContext, T> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& input,
                  const DenseTensor& output,
                  const DenseTensor& output_grad,
                  const std::vector<int64_t>& ksize,
                  const std::vector<int64_t>& strides,
                  const std::vector<int64_t>& paddings,
                  const std::string data_format,
                  DenseTensor* input_grad) {
    static const int kBlockThreads = 1024;

    bool channel_last = (data_format == "NHWC");

    const int64_t batch_size = input.dims()[0];

    const int64_t input_channels =
        channel_last ? input.dims()[3] : input.dims()[1];
    const int64_t input_height =
        channel_last ? input.dims()[1] : input.dims()[2];
    const int64_t input_width =
        channel_last ? input.dims()[2] : input.dims()[3];

    const int64_t output_channels =
        channel_last ? output.dims()[3] : output.dims()[1];
    const int64_t output_height =
        channel_last ? output.dims()[1] : output.dims()[2];
    const int64_t output_width =
        channel_last ? output.dims()[2] : output.dims()[3];

    const int64_t ksize_height = ksize[0];
    const int64_t ksize_width = ksize[1];

    const int64_t stride_height = strides[0];
    const int64_t stride_width = strides[1];

    const int64_t padding_height = paddings[0];
    const int64_t padding_width = paddings[1];

    const T* input_data = input.data<T>();
    const T* output_data = output.data<T>();
    const T* output_grad_data = output_grad.data<T>();
    T* input_grad_data = dev_ctx.template Alloc<T>(input_grad);

    int64_t nthreads =
        batch_size * output_channels * output_height * output_width;
    dim3 threads(kBlockThreads, 1);

    if (input.numel() <= std::numeric_limits<int>::max() &&
        output.numel() <= std::numeric_limits<int>::max()) {
      auto pool_divmods = FastDivModForPooling<int>(
          input_channels, output_width, output_height);
      if (FLAGS_use_accuracy_compatible_kernel) {
        int64_t blocks =
            (input_width * input_height + kBlockThreads - 1) / kBlockThreads;
        dim3 grid(blocks, batch_size, input_channels);
        // NOTE: input.numel() <= std::numeric_limits<int>::max() &&
        // output.numel() <= std::numeric_limits<int>::max()
        KernelMaxPool2DGradCompatible<T, int>
            <<<grid, threads, 0, dev_ctx.stream()>>>(input_data,
                                                     output_data,
                                                     output_grad_data,
                                                     batch_size,
                                                     input_channels,
                                                     input_height,
                                                     input_width,
                                                     output_height,
                                                     output_width,
                                                     ksize_height,
                                                     ksize_width,
                                                     stride_height,
                                                     stride_width,
                                                     padding_height,
                                                     padding_width,
                                                     input_grad_data,
                                                     pool_divmods,
                                                     channel_last);
      } else {
        int64_t blocks = (nthreads + kBlockThreads - 1) / kBlockThreads;
        dim3 grid(blocks, 1);
        // NOTE: input.numel() <= std::numeric_limits<int>::max() &&
        // output.numel() <= std::numeric_limits<int>::max()
        KernelMaxPool2DGrad<T, int>
            <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                     input_data,
                                                     output_data,
                                                     output_grad_data,
                                                     input_channels,
                                                     input_height,
                                                     input_width,
                                                     output_height,
                                                     output_width,
                                                     ksize_height,
                                                     ksize_width,
                                                     stride_height,
                                                     stride_width,
                                                     padding_height,
                                                     padding_width,
                                                     input_grad_data,
                                                     pool_divmods,
                                                     channel_last);
      }

    } else {
      auto pool_divmods = FastDivModForPooling<int64_t>(
          input_channels, output_width, output_height);
      if (FLAGS_use_accuracy_compatible_kernel) {
        int64_t blocks =
            (input_width * input_height + kBlockThreads - 1) / kBlockThreads;
        dim3 grid(blocks, batch_size, input_channels);
        KernelMaxPool2DGradCompatible<T, int64_t>
            <<<grid, threads, 0, dev_ctx.stream()>>>(input_data,
                                                     output_data,
                                                     output_grad_data,
                                                     batch_size,
                                                     input_channels,
                                                     input_height,
                                                     input_width,
                                                     output_height,
                                                     output_width,
                                                     ksize_height,
                                                     ksize_width,
                                                     stride_height,
                                                     stride_width,
                                                     padding_height,
                                                     padding_width,
                                                     input_grad_data,
                                                     pool_divmods,
                                                     channel_last);
      } else {
        int64_t blocks = (nthreads + kBlockThreads - 1) / kBlockThreads;
        dim3 grid(blocks, 1);
        KernelMaxPool2DGrad<T, int64_t>
            <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                     input_data,
                                                     output_data,
                                                     output_grad_data,
                                                     input_channels,
                                                     input_height,
                                                     input_width,
                                                     output_height,
                                                     output_width,
                                                     ksize_height,
                                                     ksize_width,
                                                     stride_height,
                                                     stride_width,
                                                     padding_height,
                                                     padding_width,
                                                     input_grad_data,
                                                     pool_divmods,
                                                     channel_last);
      }
    }
  }
};

template class PADDLE_API Pool2dDirectCUDAFunctor<MaxPool<float>, float>;
template class PADDLE_API Pool2dDirectCUDAFunctor<AvgPool<float>, float>;

template class MaxPool2dGradFunctor<GPUContext, float>;
template class MaxPool2dGradFunctor<GPUContext, double>;
template class MaxPool2dGradFunctor<GPUContext, dtype::float16>;
template class MaxPool2dGradFunctor<GPUContext, dtype::bfloat16>;

template class Pool2dFunctor<GPUContext, MaxPool<float>, float>;
template class Pool2dFunctor<GPUContext, AvgPool<float>, float>;
template class Pool2dFunctor<GPUContext, LPPool<float>, float>;
template class Pool2dGradFunctor<GPUContext, MaxPoolGrad<float>, float>;
template class Pool2dGradFunctor<GPUContext, AvgPoolGrad<float>, float>;
template class Pool2dGradFunctor<GPUContext, LPPoolGrad<float>, float>;
template class Pool2dFunctor<GPUContext, MaxPool<double>, double>;
template class Pool2dFunctor<GPUContext, AvgPool<double>, double>;
template class Pool2dFunctor<GPUContext, LPPool<double>, double>;
template class Pool2dGradFunctor<GPUContext, MaxPoolGrad<double>, double>;
template class Pool2dGradFunctor<GPUContext, AvgPoolGrad<double>, double>;
template class Pool2dGradFunctor<GPUContext, LPPoolGrad<double>, double>;

template class Pool2dFunctor<GPUContext,
                             MaxPool<dtype::float16>,
                             dtype::float16>;
template class Pool2dFunctor<GPUContext,
                             AvgPool<dtype::float16>,
                             dtype::float16>;
template class Pool2dFunctor<GPUContext,
                             LPPool<dtype::float16>,
                             dtype::float16>;
template class Pool2dGradFunctor<GPUContext,
                                 MaxPoolGrad<dtype::float16>,
                                 dtype::float16>;
template class Pool2dGradFunctor<GPUContext,
                                 AvgPoolGrad<dtype::float16>,
                                 dtype::float16>;
template class Pool2dGradFunctor<GPUContext,
                                 LPPoolGrad<dtype::float16>,
                                 dtype::float16>;
template class Pool2dFunctor<GPUContext,
                             MaxPool<dtype::bfloat16>,
                             dtype::bfloat16>;
template class Pool2dFunctor<GPUContext,
                             AvgPool<dtype::bfloat16>,
                             dtype::bfloat16>;
template class Pool2dFunctor<GPUContext,
                             LPPool<dtype::bfloat16>,
                             dtype::bfloat16>;
template class Pool2dGradFunctor<GPUContext,
                                 MaxPoolGrad<dtype::bfloat16>,
                                 dtype::bfloat16>;
template class Pool2dGradFunctor<GPUContext,
                                 AvgPoolGrad<dtype::bfloat16>,
                                 dtype::bfloat16>;
template class Pool2dGradFunctor<GPUContext,
                                 LPPoolGrad<dtype::bfloat16>,
                                 dtype::bfloat16>;

template <typename PoolProcess, typename T, typename IndexT>
__global__ void KernelPool3D(const IndexT nthreads,
                             const T* input_data,
                             const IndexT channels,
                             const IndexT input_depth,
                             const IndexT input_height,
                             const IndexT input_width,
                             const IndexT output_depth,
                             const IndexT output_height,
                             const IndexT output_width,
                             const IndexT ksize_depth,
                             const IndexT ksize_height,
                             const IndexT ksize_width,
                             const IndexT stride_depth,
                             const IndexT stride_height,
                             const IndexT stride_width,
                             const IndexT padding_depth,
                             const IndexT padding_height,
                             const IndexT padding_width,
                             PoolProcess pool_process,
                             bool exclusive,
                             bool adaptive,
                             T* output_data,
                             bool channel_last = false) {
  const IndexT start_index =
      static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
  const IndexT step = static_cast<IndexT>(blockDim.x) * gridDim.x;
  for (IndexT index = start_index; index < nthreads; index += step) {
    IndexT pw, ph, pd, c, batch_idx;
    if (!channel_last) {
      pw = index % output_width;
      ph = (index / output_width) % output_height;
      pd = (index / output_width / output_height) % output_depth;
      c = (index / output_width / output_height / output_depth) % channels;
      batch_idx =
          index / output_width / output_height / output_depth / channels;
    } else {
      c = index % channels;
      pw = (index / channels) % output_width;
      ph = (index / channels / output_width) % output_height;
      pd = (index / channels / output_width / output_height) % output_depth;
      batch_idx =
          index / channels / output_width / output_height / output_depth;
    }

    IndexT dstart, dend;
    IndexT hstart, hend;
    IndexT wstart, wend;
    if (adaptive) {
      dstart = AdaptStartIndex(pd, input_depth, output_depth);
      dend = AdaptEndIndex(pd, input_depth, output_depth);

      hstart = AdaptStartIndex(ph, input_height, output_height);
      hend = AdaptEndIndex(ph, input_height, output_height);

      wstart = AdaptStartIndex(pw, input_width, output_width);
      wend = AdaptEndIndex(pw, input_width, output_width);
    } else {
      dstart = pd * stride_depth - padding_depth;
      hstart = ph * stride_height - padding_height;
      wstart = pw * stride_width - padding_width;
      dend = min(dstart + ksize_depth, input_depth);
      hend = min(hstart + ksize_height, input_height);
      wend = min(wstart + ksize_width, input_width);
      dstart = max(dstart, static_cast<IndexT>(0));
      hstart = max(hstart, static_cast<IndexT>(0));
      wstart = max(wstart, static_cast<IndexT>(0));
    }

    IndexT input_data_stride;
    if (!channel_last) { /* NCDHW */
      input_data_stride =
          (batch_idx * channels + c) * input_depth * input_height * input_width;
    } else { /* NDHWC */
      input_data_stride =
          batch_idx * input_depth * input_height * input_width * channels;
    }
    input_data += input_data_stride;

    T ele = pool_process.initial();
    for (IndexT d = dstart; d < dend; ++d) {
      for (IndexT h = hstart; h < hend; ++h) {
        for (IndexT w = wstart; w < wend; ++w) {
          auto input_data_idx =
              channel_last
                  ? ((d * input_height + h) * input_width + w) * channels + c
                  : (d * input_height + h) * input_width + w;
          pool_process.compute(input_data[input_data_idx], &ele);
        }
      }
    }
    IndexT pool_size = (exclusive || adaptive)
                           ? (dend - dstart) * (hend - hstart) * (wend - wstart)
                           : ksize_depth * ksize_height * ksize_width;
    pool_process.finalize(static_cast<T>(pool_size), &ele);
    output_data[index] = ele;
  }
}

template <typename T, typename PoolProcess, typename IndexT>
__global__ void KernelPool3DGrad(const IndexT nthreads,
                                 const T* __restrict__ input_data,
                                 const T* __restrict__ output_data,
                                 const T* __restrict__ output_grad,
                                 const IndexT channels,
                                 const IndexT input_depth,
                                 const IndexT input_height,
                                 const IndexT input_width,
                                 const IndexT output_depth,
                                 const IndexT output_height,
                                 const IndexT output_width,
                                 const IndexT ksize_depth,
                                 const IndexT ksize_height,
                                 const IndexT ksize_width,
                                 const IndexT stride_depth,
                                 const IndexT stride_height,
                                 const IndexT stride_width,
                                 const IndexT padding_depth,
                                 const IndexT padding_height,
                                 const IndexT padding_width,
                                 PoolProcess pool_process,
                                 bool exclusive,
                                 bool adaptive,
                                 T* input_grad,
                                 bool channel_last = false) {
  const IndexT start_index =
      static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
  const IndexT step = static_cast<IndexT>(blockDim.x) * gridDim.x;
  for (IndexT index = start_index; index < nthreads; index += step) {
    IndexT w_offset, h_offset, d_offset, c_offset, batch_idx, output_stride;
    T input = static_cast<T>(0);
    if (!channel_last) { /* "NCDHW" */
      w_offset = index % input_width + padding_width;
      h_offset = (index / input_width) % input_height + padding_height;
      d_offset =
          (index / input_width / input_height) % input_depth + padding_depth;
      c_offset = (index / input_width / input_height / input_depth) % channels;
      batch_idx = index / input_width / input_height / input_depth / channels;
      output_stride = (batch_idx * channels + c_offset) * output_depth *
                      output_height * output_width;
    } else { /* "NDHWC" */
      c_offset = index % channels;
      w_offset = (index / channels) % input_width + padding_width;
      h_offset =
          (index / channels / input_width) % input_height + padding_height;
      d_offset = (index / channels / input_width / input_height) % input_depth +
                 padding_depth;
      batch_idx = index / channels / input_width / input_height / input_depth;
      output_stride =
          batch_idx * output_depth * output_height * output_width * channels;
    }

    IndexT pdstart, pdend;
    IndexT phstart, phend;
    IndexT pwstart, pwend;

    IndexT pool_depth, pool_height, pool_width;

    if (adaptive) {
      pdstart = AdaptStartIndex(d_offset, output_depth, input_depth);
      pdend = AdaptEndIndex(d_offset, output_depth, input_depth);

      phstart = AdaptStartIndex(h_offset, output_height, input_height);
      phend = AdaptEndIndex(h_offset, output_height, input_height);

      pwstart = AdaptStartIndex(w_offset, output_width, input_width);
      pwend = AdaptEndIndex(w_offset, output_width, input_width);
    } else {
      pdstart = (d_offset < ksize_depth)
                    ? 0
                    : (d_offset - ksize_depth) / stride_depth + 1;
      phstart = (h_offset < ksize_height)
                    ? 0
                    : (h_offset - ksize_height) / stride_height + 1;
      pwstart = (w_offset < ksize_width)
                    ? 0
                    : (w_offset - ksize_width) / stride_width + 1;
      pdend = min((d_offset) / stride_depth + 1, output_depth);
      phend = min((h_offset) / stride_height + 1, output_height);
      pwend = min((w_offset) / stride_width + 1, output_width);
    }
    if (pool_process.use_x) {
      input = input_data[index];
      output_data += output_stride;
    }
    output_grad += output_stride;
    T input_grad_data = static_cast<T>(0.0);

    IndexT pool_size;
    for (IndexT pd = pdstart; pd < pdend; ++pd) {
      for (IndexT ph = phstart; ph < phend; ++ph) {
        for (IndexT pw = pwstart; pw < pwend; ++pw) {
          // figure out the pooling size
          if (adaptive) {
            PreparationPoolSize(pd,
                                input_depth,
                                output_depth,
                                FastDivMod<IndexT>(output_depth),
                                &pool_depth);
            PreparationPoolSize(pw,
                                input_width,
                                output_width,
                                FastDivMod<IndexT>(output_width),
                                &pool_width);
            PreparationPoolSize(ph,
                                input_height,
                                output_height,
                                FastDivMod<IndexT>(output_height),
                                &pool_height);
            pool_size = pool_depth * pool_height * pool_width;
          } else {
            IndexT dstart = pd * stride_depth - padding_depth;
            IndexT hstart = ph * stride_height - padding_height;
            IndexT wstart = pw * stride_width - padding_width;
            IndexT dend = min(dstart + ksize_depth, input_depth);
            IndexT hend = min(hstart + ksize_height, input_height);
            IndexT wend = min(wstart + ksize_width, input_width);
            dstart = max(dstart, static_cast<IndexT>(0));
            hstart = max(hstart, static_cast<IndexT>(0));
            wstart = max(wstart, static_cast<IndexT>(0));
            pool_size =
                exclusive ? (dend - dstart) * (hend - hstart) * (wend - wstart)
                          : ksize_depth * ksize_height * ksize_width;
          }

          IndexT output_sub_idx =
              channel_last
                  ? ((pd * output_height + ph) * output_width + pw) * channels +
                        c_offset
                  : (pd * output_height + ph) * output_width + pw;
          T output_value = pool_process.use_x ? output_data[output_sub_idx]
                                              : static_cast<T>(0);
          pool_process.compute(input,
                               output_value,
                               output_grad[output_sub_idx],
                               static_cast<T>(1.0 / pool_size),
                               &input_grad_data);
        }
      }
    }
    input_grad[index] = input_grad_data;
  }
}

template <typename T, typename IndexT>
__global__ void KernelMaxPool3DGrad(const IndexT nthreads,
                                    const T* input_data,
                                    const T* output_data,
                                    const T* output_grad,
                                    const IndexT channels,
                                    const IndexT input_depth,
                                    const IndexT input_height,
                                    const IndexT input_width,
                                    const IndexT output_depth,
                                    const IndexT output_height,
                                    const IndexT output_width,
                                    const IndexT ksize_depth,
                                    const IndexT ksize_height,
                                    const IndexT ksize_width,
                                    const IndexT stride_depth,
                                    const IndexT stride_height,
                                    const IndexT stride_width,
                                    const IndexT padding_depth,
                                    const IndexT padding_height,
                                    const IndexT padding_width,
                                    T* input_grad,
                                    bool channel_last = false) {
  const IndexT start_index =
      static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
  const IndexT step = static_cast<IndexT>(blockDim.x) * gridDim.x;
  for (IndexT index = start_index; index < nthreads; index += step) {
    IndexT pw, ph, pd, c, batch_idx;

    if (!channel_last) { /*NCDHW*/
      pw = index % output_width;
      ph = (index / output_width) % output_height;
      pd = (index / output_width / output_height) % output_depth;
      c = (index / output_width / output_height / output_depth) % channels;
      batch_idx =
          index / output_width / output_height / output_depth / channels;
    } else { /*NDHWC*/
      c = index % channels;
      pw = (index / channels) % output_width;
      ph = (index / channels / output_width) % output_height;
      pd = (index / channels / output_width / output_height) % output_depth;
      batch_idx =
          index / channels / output_width / output_height / output_depth;
    }

    IndexT dstart = pd * stride_depth - padding_depth;
    IndexT hstart = ph * stride_height - padding_height;
    IndexT wstart = pw * stride_width - padding_width;

    IndexT dend = min(dstart + ksize_depth, input_depth);
    IndexT hend = min(hstart + ksize_height, input_height);
    IndexT wend = min(wstart + ksize_width, input_width);

    dstart = max(dstart, static_cast<IndexT>(0));
    hstart = max(hstart, static_cast<IndexT>(0));
    wstart = max(wstart, static_cast<IndexT>(0));

    T ele = output_data[index];
    bool stop = false;
    IndexT maxIdx = -1;

    IndexT input_stride;
    if (!channel_last) {
      input_stride =
          (batch_idx * channels + c) * input_depth * input_height * input_width;
    } else {
      input_stride =
          batch_idx * input_depth * input_height * input_width * channels;
    }
    input_data += input_stride;
    input_grad += input_stride;
    for (IndexT d = dstart; d < dend && !stop; ++d) {
      for (IndexT h = hstart; h < hend && !stop; ++h) {
        for (IndexT w = wstart; w < wend && !stop; ++w) {
          IndexT input_data_idx =
              channel_last
                  ? ((d * input_height + h) * input_width + w) * channels + c
                  : (d * input_height + h) * input_width + w;
          if (ele == input_data[input_data_idx]) {
            stop = true;
            maxIdx = input_data_idx;
          }
        }
      }
    }
    if (maxIdx != -1) {
      // atomic add
      CudaAtomicAdd(input_grad + maxIdx, output_grad[index]);
    }
  }
}

template <typename PoolProcess, typename T>
void Pool3dDirectCUDAFunctor<PoolProcess, T>::operator()(
    const T* input,
    const std::vector<int>& input_shape,
    const std::vector<int>& output_shape,
    const std::vector<int>& ksize,
    const std::vector<int>& strides,
    const std::vector<int>& paddings,
    bool exclusive,
    bool adaptive,
    T* output,
    gpuStream_t stream,
    PoolProcess pool_compute) {
  const int batch_size = input_shape[0];
  const int input_channels = input_shape[1];
  const int input_depth = input_shape[2];
  const int input_height = input_shape[3];
  const int input_width = input_shape[4];
  const int output_channels = output_shape[1];
  const int output_depth = output_shape[2];
  const int output_height = output_shape[3];
  const int output_width = output_shape[4];
  const int ksize_depth = ksize[0];
  const int ksize_height = ksize[1];
  const int ksize_width = ksize[2];
  const int stride_depth = strides[0];
  const int stride_height = strides[1];
  const int stride_width = strides[2];
  const int padding_depth = paddings[0];
  const int padding_height = paddings[1];
  const int padding_width = paddings[2];

  int nthreads = batch_size * output_channels * output_depth * output_height *
                 output_width;
  int thread_num = 1024;
#ifdef WITH_NV_JETSON
  thread_num = 512;
#endif
  int blocks = (nthreads + thread_num - 1) / thread_num;
  dim3 threads(thread_num, 1);
  dim3 grid(blocks, 1);

  KernelPool3D<PoolProcess, T, int>
      <<<grid, threads, 0, stream>>>(nthreads,
                                     input,
                                     input_channels,
                                     input_depth,
                                     input_height,
                                     input_width,
                                     output_depth,
                                     output_height,
                                     output_width,
                                     ksize_depth,
                                     ksize_height,
                                     ksize_width,
                                     stride_depth,
                                     stride_height,
                                     stride_width,
                                     padding_depth,
                                     padding_height,
                                     padding_width,
                                     pool_compute,
                                     exclusive,
                                     adaptive,
                                     output);
}

/*
 * Tensors are in NCDHW or NDHWC format.
 * Ksize, strides, paddings are three elements. These three elements represent
 * depth, height and width, respectively.
 * Paddings are six elements. These six elements represent depth_forth,
 * depth_back,
 * height_up, height_down, width_left and width_right, respectively.
 */
template <typename PoolProcess, class T>
class Pool3dFunctor<GPUContext, PoolProcess, T> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& input,
                  const std::vector<int64_t>& ksize,
                  const std::vector<int64_t>& strides,
                  const std::vector<int64_t>& paddings,
                  const std::string data_format,
                  bool exclusive,
                  bool adaptive,
                  DenseTensor* output,
                  PoolProcess pool_process) {
    bool channel_last = (data_format == "NDHWC");
    const int64_t batch_size = input.dims()[0];

    const int64_t input_channels =
        channel_last ? input.dims()[4] : input.dims()[1];
    const int64_t input_depth =
        channel_last ? input.dims()[1] : input.dims()[2];
    const int64_t input_height =
        channel_last ? input.dims()[2] : input.dims()[3];
    const int64_t input_width =
        channel_last ? input.dims()[3] : input.dims()[4];

    const int64_t output_channels =
        channel_last ? output->dims()[4] : output->dims()[1];
    const int64_t output_depth =
        channel_last ? output->dims()[1] : output->dims()[2];
    const int64_t output_height =
        channel_last ? output->dims()[2] : output->dims()[3];
    const int64_t output_width =
        channel_last ? output->dims()[3] : output->dims()[4];

    const int64_t ksize_depth = ksize[0];
    const int64_t ksize_height = ksize[1];
    const int64_t ksize_width = ksize[2];

    const int64_t stride_depth = strides[0];
    const int64_t stride_height = strides[1];
    const int64_t stride_width = strides[2];

    const int64_t padding_depth = paddings[0];
    const int64_t padding_height = paddings[1];
    const int64_t padding_width = paddings[2];

    const T* input_data = input.data<T>();
    T* output_data = dev_ctx.template Alloc<T>(output);

    int64_t nthreads = batch_size * output_channels * output_depth *
                       output_height * output_width;

    if (input.numel() <= std::numeric_limits<int>::max()) {
      int thread_num = 1024;
#ifdef WITH_NV_JETSON
      backends::gpu::ChangeThreadNum(dev_ctx, &thread_num);
#endif
      int64_t blocks = (nthreads + thread_num - 1) / thread_num;
      dim3 threads(thread_num, 1);
      dim3 grid(blocks, 1);
      KernelPool3D<PoolProcess, T, int>
          <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                   input_data,
                                                   input_channels,
                                                   input_depth,
                                                   input_height,
                                                   input_width,
                                                   output_depth,
                                                   output_height,
                                                   output_width,
                                                   ksize_depth,
                                                   ksize_height,
                                                   ksize_width,
                                                   stride_depth,
                                                   stride_height,
                                                   stride_width,
                                                   padding_depth,
                                                   padding_height,
                                                   padding_width,
                                                   pool_process,
                                                   exclusive,
                                                   adaptive,
                                                   output_data,
                                                   channel_last);
    } else {
      int thread_num = 512;
#ifdef WITH_NV_JETSON
      backends::gpu::ChangeThreadNum(dev_ctx, &thread_num);
#endif
      int64_t blocks = (nthreads + thread_num - 1) / thread_num;
      dim3 threads(thread_num, 1);
      dim3 grid(blocks, 1);
      KernelPool3D<PoolProcess, T, int64_t>
          <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                   input_data,
                                                   input_channels,
                                                   input_depth,
                                                   input_height,
                                                   input_width,
                                                   output_depth,
                                                   output_height,
                                                   output_width,
                                                   ksize_depth,
                                                   ksize_height,
                                                   ksize_width,
                                                   stride_depth,
                                                   stride_height,
                                                   stride_width,
                                                   padding_depth,
                                                   padding_height,
                                                   padding_width,
                                                   pool_process,
                                                   exclusive,
                                                   adaptive,
                                                   output_data,
                                                   channel_last);
    }
  }
};

/*
 * Tensors are in NCDHW or NDHWC format.
 * Ksize, strides, paddings are three elements. These three elements represent
 * depth, height and width, respectively.
 * Paddings are six elements. These six elements represent depth_forth,
 * depth_back,
 * height_up, height_down, width_left and width_right, respectively.
 */
template <typename PoolProcess, class T>
class Pool3dGradFunctor<GPUContext, PoolProcess, T> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& input,
                  const DenseTensor& output,
                  const DenseTensor& output_grad,
                  const std::vector<int64_t>& ksize,
                  const std::vector<int64_t>& strides,
                  const std::vector<int64_t>& paddings,
                  const std::string data_format,
                  bool exclusive,
                  bool adaptive,
                  DenseTensor* input_grad,
                  PoolProcess pool_process) {
    bool channel_last = (data_format == "NDHWC");

    const int64_t batch_size = input.dims()[0];
    const int64_t input_channels =
        channel_last ? input.dims()[4] : input.dims()[1];
    const int64_t input_depth =
        channel_last ? input.dims()[1] : input.dims()[2];
    const int64_t input_height =
        channel_last ? input.dims()[2] : input.dims()[3];
    const int64_t input_width =
        channel_last ? input.dims()[3] : input.dims()[4];

    const int64_t output_channels =
        channel_last ? output.dims()[4] : output.dims()[1];
    const int64_t output_depth =
        channel_last ? output.dims()[1] : output.dims()[2];
    const int64_t output_height =
        channel_last ? output.dims()[2] : output.dims()[3];
    const int64_t output_width =
        channel_last ? output.dims()[3] : output.dims()[4];

    const int64_t ksize_depth = ksize[0];
    const int64_t ksize_height = ksize[1];
    const int64_t ksize_width = ksize[2];

    const int64_t stride_depth = strides[0];
    const int64_t stride_height = strides[1];
    const int64_t stride_width = strides[2];

    const int64_t padding_depth = paddings[0];
    const int64_t padding_height = paddings[1];
    const int64_t padding_width = paddings[2];

    const T* input_data = input.data<T>();
    const T* output_data = output.data<T>();
    const T* output_grad_data = output_grad.data<T>();
    T* input_grad_data = dev_ctx.template Alloc<T>(input_grad);

    int64_t nthreads =
        batch_size * input_channels * input_depth * input_height * input_width;

    if (input.numel() <= std::numeric_limits<int>::max() &&
        output.numel() <= std::numeric_limits<int>::max()) {
      int thread_num = 1024;
      int64_t blocks = (nthreads + thread_num - 1) / thread_num;
      dim3 threads(thread_num, 1);
      dim3 grid(blocks, 1);
      KernelPool3DGrad<T, PoolProcess, int>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              nthreads,
              input_data,
              output_data,
              output_grad_data,
              input_channels,
              input_depth,
              input_height,
              input_width,
              output_depth,
              output_height,
              output_width,
              ksize_depth,
              ksize_height,
              ksize_width,
              stride_depth,
              stride_height,
              stride_width,
              padding_depth,
              padding_height,
              padding_width,
              pool_process,
              exclusive,
              adaptive,
              input_grad_data,
              channel_last);  // add channel_last
    } else {
      int thread_num = 512;
      int64_t blocks = (nthreads + thread_num - 1) / thread_num;
      dim3 threads(thread_num, 1);
      dim3 grid(blocks, 1);
      KernelPool3DGrad<T, PoolProcess, int64_t>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              nthreads,
              input_data,
              output_data,
              output_grad_data,
              input_channels,
              input_depth,
              input_height,
              input_width,
              output_depth,
              output_height,
              output_width,
              ksize_depth,
              ksize_height,
              ksize_width,
              stride_depth,
              stride_height,
              stride_width,
              padding_depth,
              padding_height,
              padding_width,
              pool_process,
              exclusive,
              adaptive,
              input_grad_data,
              channel_last);  // add channel_last
    }
  }
};

/*
 * tensors are in NCDHW or NDHWC format.
 * Ksize, strides, paddings are three elements. These three elements represent
 * depth, height and width, respectively.
 * Paddings are six elements. These six elements represent depth_forth,
 * depth_back,
 * height_up, height_down, width_left and width_right, respectively.
 */
template <class T>
class MaxPool3dGradFunctor<GPUContext, T> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& input,
                  const DenseTensor& output,
                  const DenseTensor& output_grad,
                  const std::vector<int64_t>& ksize,
                  const std::vector<int64_t>& strides,
                  const std::vector<int64_t>& paddings,
                  const std::string data_format,
                  DenseTensor* input_grad) {
    bool channel_last = (data_format == "NDHWC");
    const int64_t batch_size = input.dims()[0];

    const int64_t input_channels =
        channel_last ? input.dims()[4] : input.dims()[1];
    const int64_t input_depth =
        channel_last ? input.dims()[1] : input.dims()[2];
    const int64_t input_height =
        channel_last ? input.dims()[2] : input.dims()[3];
    const int64_t input_width =
        channel_last ? input.dims()[3] : input.dims()[4];

    const int64_t output_channels =
        channel_last ? output.dims()[4] : output.dims()[1];
    const int64_t output_depth =
        channel_last ? output.dims()[1] : output.dims()[2];
    const int64_t output_height =
        channel_last ? output.dims()[2] : output.dims()[3];
    const int64_t output_width =
        channel_last ? output.dims()[3] : output.dims()[4];

    const int64_t ksize_depth = ksize[0];
    const int64_t ksize_height = ksize[1];
    const int64_t ksize_width = ksize[2];

    const int64_t stride_depth = strides[0];
    const int64_t stride_height = strides[1];
    const int64_t stride_width = strides[2];

    const int64_t padding_depth = paddings[0];
    const int64_t padding_height = paddings[1];
    const int64_t padding_width = paddings[2];

    const T* input_data = input.data<T>();
    const T* output_data = output.data<T>();
    const T* output_grad_data = output_grad.data<T>();
    T* input_grad_data = dev_ctx.template Alloc<T>(input_grad);

    int64_t nthreads = batch_size * output_channels * output_depth *
                       output_height * output_width;
    int64_t blocks = (nthreads + 1024 - 1) / 1024;
    dim3 threads(1024, 1);
    dim3 grid(blocks, 1);
    if (input.numel() <= std::numeric_limits<int>::max() &&
        output.numel() <= std::numeric_limits<int>::max()) {
      KernelMaxPool3DGrad<T, int><<<grid, threads, 0, dev_ctx.stream()>>>(
          nthreads,
          input_data,
          output_data,
          output_grad_data,
          input_channels,
          input_depth,
          input_height,
          input_width,
          output_depth,
          output_height,
          output_width,
          ksize_depth,
          ksize_height,
          ksize_width,
          stride_depth,
          stride_height,
          stride_width,
          padding_depth,
          padding_height,
          padding_width,
          input_grad_data,
          channel_last);  // add channel_last
    } else {
      KernelMaxPool3DGrad<T, int64_t><<<grid, threads, 0, dev_ctx.stream()>>>(
          nthreads,
          input_data,
          output_data,
          output_grad_data,
          input_channels,
          input_depth,
          input_height,
          input_width,
          output_depth,
          output_height,
          output_width,
          ksize_depth,
          ksize_height,
          ksize_width,
          stride_depth,
          stride_height,
          stride_width,
          padding_depth,
          padding_height,
          padding_width,
          input_grad_data,
          channel_last);  // add channel_last
    }
  }
};

template class PADDLE_API Pool3dDirectCUDAFunctor<MaxPool<float>, float>;
template class PADDLE_API Pool3dDirectCUDAFunctor<AvgPool<float>, float>;

template class MaxPool3dGradFunctor<GPUContext, float>;
template class MaxPool3dGradFunctor<GPUContext, double>;
template class MaxPool3dGradFunctor<GPUContext, dtype::float16>;
template class MaxPool3dGradFunctor<GPUContext, dtype::bfloat16>;

template class Pool3dFunctor<GPUContext, MaxPool<float>, float>;
template class Pool3dFunctor<GPUContext, AvgPool<float>, float>;
template class Pool3dGradFunctor<GPUContext, MaxPoolGrad<float>, float>;
template class Pool3dGradFunctor<GPUContext, AvgPoolGrad<float>, float>;
template class Pool3dFunctor<GPUContext, MaxPool<double>, double>;
template class Pool3dFunctor<GPUContext, AvgPool<double>, double>;
template class Pool3dGradFunctor<GPUContext, MaxPoolGrad<double>, double>;
template class Pool3dGradFunctor<GPUContext, AvgPoolGrad<double>, double>;

template class Pool3dFunctor<GPUContext,
                             MaxPool<dtype::float16>,
                             dtype::float16>;
template class Pool3dFunctor<GPUContext,
                             AvgPool<dtype::float16>,
                             dtype::float16>;
template class Pool3dFunctor<GPUContext,
                             MaxPool<dtype::bfloat16>,
                             dtype::bfloat16>;
template class Pool3dFunctor<GPUContext,
                             AvgPool<dtype::bfloat16>,
                             dtype::bfloat16>;
template class Pool3dGradFunctor<GPUContext,
                                 MaxPoolGrad<dtype::float16>,
                                 dtype::float16>;
template class Pool3dGradFunctor<GPUContext,
                                 AvgPoolGrad<dtype::float16>,
                                 dtype::float16>;
template class Pool3dGradFunctor<GPUContext,
                                 MaxPoolGrad<dtype::bfloat16>,
                                 dtype::bfloat16>;
template class Pool3dGradFunctor<GPUContext,
                                 AvgPoolGrad<dtype::bfloat16>,
                                 dtype::bfloat16>;

template <typename T1, typename T2, typename IndexT>
__global__ void KernelMaxPool2dWithIdx(const IndexT nthreads,
                                       const T1* input_data,
                                       const IndexT channels,
                                       const IndexT input_height,
                                       const IndexT input_width,
                                       const IndexT output_height,
                                       const IndexT output_width,
                                       const IndexT ksize_height,
                                       const IndexT ksize_width,
                                       const IndexT stride_height,
                                       const IndexT stride_width,
                                       const IndexT padding_height,
                                       const IndexT padding_width,
                                       const IndexT dilation_height,
                                       const IndexT dilation_width,
                                       bool adaptive,
                                       T1* output_data,
                                       T2* mask_data,
                                       FastDivModForPooling<IndexT> divmods) {
  const IndexT start_index =
      static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
  const IndexT step = static_cast<IndexT>(blockDim.x) * gridDim.x;
  for (IndexT index = start_index; index < nthreads; index += step) {
    IndexT hstart, hend, wstart, wend;
    IndexT w_offset, h_offset, c_offset, input_offset;
    OffsetPreparationFor4Dimension<FastDivModForPooling<IndexT>, IndexT>(
        index,
        false,
        divmods,
        0,
        0,
        input_width,
        input_height,
        &w_offset,
        &h_offset,
        &c_offset,
        &input_offset);
    input_data += input_offset;

    const bool dilation = (dilation_height != 1 || dilation_width != 1);
    if (adaptive) {
      hstart = AdaptStartIndex(h_offset, input_height, output_height);
      hend = AdaptEndIndex(h_offset, input_height, output_height);

      wstart = AdaptStartIndex(w_offset, input_width, output_width);
      wend = AdaptEndIndex(w_offset, input_width, output_width);
    } else if (dilation) {
      hstart = h_offset * stride_height - padding_height;
      hend =
          min(hstart + dilation_height * (ksize_height - 1) + 1, input_height);
      if (hstart < 0) {
        hstart =
            ((-hstart - 1) / dilation_height + 1) * dilation_height + hstart;
      }

      wstart = w_offset * stride_width - padding_width;
      wend = min(wstart + dilation_width * (ksize_width - 1) + 1, input_width);
      if (wstart < 0) {
        wstart = ((-wstart - 1) / dilation_width + 1) * dilation_width + wstart;
      }
    } else {
      hstart = h_offset * stride_height - padding_height;
      hend = min(hstart + ksize_height, input_height);
      hstart = max(hstart, static_cast<IndexT>(0));

      wstart = w_offset * stride_width - padding_width;
      wend = min(wstart + ksize_width, input_width);
      wstart = max(wstart, static_cast<IndexT>(0));
    }

    T1 ele = static_cast<T1>(-FLT_MAX);
    IndexT max_index = -1;
    for (IndexT h = hstart; h < hend; h += dilation_height) {
      for (IndexT w = wstart; w < wend; w += dilation_width) {
        IndexT input_index = h * input_width + w;
        if (ele < input_data[input_index]) {
          max_index = input_index;
          ele = input_data[input_index];
        }
      }
    }
    output_data[index] = ele;
    mask_data[index] = max_index;
  }
}

template <typename T1, typename T2, typename IndexT>
__global__ void AdaptiveKernelMaxPool2dWithIdx(
    const IndexT nthreads,
    const T1* input_data,
    const IndexT channels,
    const IndexT input_height,
    const IndexT input_width,
    const IndexT output_height,
    const IndexT output_width,
    const IndexT ksize_height,
    const IndexT ksize_width,
    const IndexT stride_height,
    const IndexT stride_width,
    const IndexT padding_height,
    const IndexT padding_width,
    const IndexT batch_size,
    T1* output_data,
    T2* mask_data,
    FastDivModForPooling<IndexT> divmods) {
  for (IndexT n_offset = blockIdx.y; n_offset < batch_size;
       n_offset += gridDim.y) {
    const IndexT c_offset =
        static_cast<IndexT>(blockIdx.x) * blockDim.y + threadIdx.y;
    if (c_offset >= channels) {
      return;
    }
    IndexT hstart, hend, wstart, wend;
    IndexT input_offset =
        (n_offset * channels + c_offset) * input_height * input_width;
    IndexT output_offset =
        (n_offset * channels + c_offset) * output_height * output_width;
    for (IndexT hw_offset = threadIdx.x;
         hw_offset < output_height * output_width;
         hw_offset += blockDim.x) {
      IndexT w_offset = hw_offset % output_width;
      IndexT h_offset = hw_offset / output_width;
      hstart = AdaptStartIndex(h_offset, input_height, output_height);
      hend = AdaptEndIndex(h_offset, input_height, output_height);
      wstart = AdaptStartIndex(w_offset, input_width, output_width);
      wend = AdaptEndIndex(w_offset, input_width, output_width);

      T1 ele = static_cast<T1>(-FLT_MAX);
      IndexT max_index = -1;
      for (IndexT h = hstart; h < hend; ++h) {
        for (IndexT w = wstart; w < wend; ++w) {
          IndexT input_index = h * input_width + w;
          if (ele < input_data[input_offset + input_index]) {
            max_index = input_index;
            ele = input_data[input_offset + input_index];
          }
        }
      }
      IndexT output_idx = output_offset + h_offset * output_width + w_offset;
      output_data[output_idx] = ele;
      mask_data[output_idx] = max_index;
    }
  }
}

template <typename T1, typename T2, typename IndexT>
__global__ void KernelMaxPool2DWithIdxGrad(
    const IndexT nthreads,
    const T1* output_grad,
    const T2* mask_data,
    const IndexT channels,
    const IndexT input_height,
    const IndexT input_width,
    const IndexT output_height,
    const IndexT output_width,
    const IndexT ksize_height,
    const IndexT ksize_width,
    const IndexT stride_height,
    const IndexT stride_width,
    const IndexT padding_height,
    const IndexT padding_width,
    const IndexT dilation_height,
    const IndexT dilation_width,
    bool adaptive,
    T1* input_grad,
    FastDivModForPooling<IndexT> divmods) {
  const IndexT start_index =
      static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
  const IndexT step = static_cast<IndexT>(blockDim.x) * gridDim.x;
  for (IndexT index = start_index; index < nthreads; index += step) {
    IndexT phstart, phend, pwstart, pwend;
    IndexT w_offset, h_offset, c_offset, output_offset;
    OffsetPreparationFor4Dimension<FastDivModForPooling<IndexT>, IndexT>(
        index,
        false,
        divmods,
        0,
        0,
        output_width,
        output_height,
        &w_offset,
        &h_offset,
        &c_offset,
        &output_offset);
    mask_data += output_offset;
    output_grad += output_offset;
    const bool dilation = (dilation_height != 1 || dilation_width != 1);

    if (adaptive) {
      phstart = h_offset * output_height / input_height;
      phend =
          min((h_offset + 1) * output_height / input_height + 1, output_height);
      pwstart = w_offset * output_width / input_width;
      pwend =
          min((w_offset + 1) * output_width / input_width + 1, output_width);
    } else if (dilation) {
      const IndexT effective_ksize_height =
          (ksize_height - 1) * dilation_height + 1;
      const IndexT effective_ksize_width =
          (ksize_width - 1) * dilation_width + 1;
      phstart = (h_offset + padding_height < effective_ksize_height)
                    ? 0
                    : (h_offset + padding_height - effective_ksize_height) /
                              stride_height +
                          1;
      pwstart = (w_offset + padding_width < effective_ksize_width)
                    ? 0
                    : (w_offset + padding_width - effective_ksize_width) /
                              stride_width +
                          1;
      phend =
          min((h_offset + padding_height) / stride_height + 1, output_height);
      pwend = min((w_offset + padding_width) / stride_width + 1, output_width);
    } else {
      phstart =
          (h_offset + padding_height < ksize_height)
              ? 0
              : (h_offset + padding_height - ksize_height) / stride_height + 1;
      pwstart =
          (w_offset + padding_width < ksize_width)
              ? 0
              : (w_offset + padding_width - ksize_width) / stride_width + 1;
      phend =
          min((h_offset + padding_height) / stride_height + 1, output_height);
      pwend = min((w_offset + padding_width) / stride_width + 1, output_width);
    }

    T1 input_grad_data = static_cast<T1>(0);
    IndexT input_current_featuremap_idx = h_offset * input_width + w_offset;
    for (IndexT ph = phstart; ph < phend; ++ph) {
      for (IndexT pw = pwstart; pw < pwend; ++pw) {
        if (mask_data[ph * output_width + pw] == input_current_featuremap_idx)
          input_grad_data += output_grad[ph * output_width + pw];
      }
    }
    input_grad[index] = input_grad_data;
  }
}

/*
 * All tensors are in NCHW format.
 * Ksize, strides, paddings, dilations are two elements. These two elements
 * represent height and width, respectively.
 */
template <typename T1, typename T2>
class MaxPool2dWithIndexFunctor<GPUContext, T1, T2> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& input,
                  const std::vector<int64_t>& ksize,
                  const std::vector<int64_t>& strides,
                  const std::vector<int64_t>& paddings,
                  const std::vector<int64_t>& dilations,
                  bool adaptive,
                  DenseTensor* output,
                  DenseTensor* mask) {
    const int64_t batch_size = input.dims()[0];
    const int64_t input_channels = input.dims()[1];
    const int64_t input_height = input.dims()[2];
    const int64_t input_width = input.dims()[3];
    const int64_t output_channels = output->dims()[1];
    const int64_t output_height = output->dims()[2];
    const int64_t output_width = output->dims()[3];
    const int64_t ksize_height = ksize[0];
    const int64_t ksize_width = ksize[1];
    const int64_t stride_height = strides[0];
    const int64_t stride_width = strides[1];
    const int64_t padding_height = paddings[0];
    const int64_t padding_width = paddings[1];
    const int64_t dilation_height = dilations[0];
    const int64_t dilation_width = dilations[1];

    const T1* input_data = input.data<T1>();
    T1* output_data = dev_ctx.template Alloc<T1>(output);
    T2* mask_data = dev_ctx.template Alloc<T2>(mask);

    PADDLE_ENFORCE_GE(
        dilation_height,
        1,
        phi::errors::InvalidArgument(
            "dilation_height must be >= 1, but got [%d]", dilation_height));
    PADDLE_ENFORCE_GE(
        dilation_width,
        1,
        phi::errors::InvalidArgument(
            "dilation_width must be >= 1, but got [%d]", dilation_width));

    int64_t nthreads = static_cast<int64_t>(batch_size) * output_channels *
                       output_height * output_width;
    if (adaptive && output_height > 1 && output_width > 1) {
      int64_t max_threads = 512;
      int64_t thread_num =
          std::min(funcs::details::GetLastPow2(output_height * output_width),
                   max_threads);
      int64_t blocks = std::min(max_threads / thread_num,
                                static_cast<int64_t>(output_channels));
      dim3 threads(thread_num, blocks, 1);
      std::array<unsigned int, 3> max_grid_dim =
          dev_ctx.GetCUDAMaxGridDimSize();
      dim3 grid(std::max((output_channels + blocks - 1) / blocks,
                         static_cast<int64_t>(1)),
                std::min(batch_size, static_cast<int64_t>(max_grid_dim[1])),
                1);
      if (input.numel() <= std::numeric_limits<int>::max()) {
        auto pool_divmods = FastDivModForPooling<int>(
            input_channels, output_width, output_height);
        AdaptiveKernelMaxPool2dWithIdx<T1, T2, int>
            <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                     input_data,
                                                     input_channels,
                                                     input_height,
                                                     input_width,
                                                     output_height,
                                                     output_width,
                                                     ksize_height,
                                                     ksize_width,
                                                     stride_height,
                                                     stride_width,
                                                     padding_height,
                                                     padding_width,
                                                     batch_size,
                                                     output_data,
                                                     mask_data,
                                                     pool_divmods);
      } else {
        auto pool_divmods = FastDivModForPooling<int64_t>(
            input_channels, output_width, output_height);
        AdaptiveKernelMaxPool2dWithIdx<T1, T2, int64_t>
            <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                     input_data,
                                                     input_channels,
                                                     input_height,
                                                     input_width,
                                                     output_height,
                                                     output_width,
                                                     ksize_height,
                                                     ksize_width,
                                                     stride_height,
                                                     stride_width,
                                                     padding_height,
                                                     padding_width,
                                                     batch_size,
                                                     output_data,
                                                     mask_data,
                                                     pool_divmods);
      }
    } else {
      int thread_num = 1024;
#ifdef WITH_NV_JETSON
      backends::gpu::ChangeThreadNum(dev_ctx, &thread_num);
#endif
      int64_t blocks = (nthreads + thread_num - 1) / thread_num;
      dim3 threads(thread_num, 1);
      dim3 grid(blocks, 1);
      if (input.numel() <= std::numeric_limits<int>::max()) {
        auto pool_divmods = FastDivModForPooling<int>(
            input_channels, output_width, output_height);
        KernelMaxPool2dWithIdx<T1, T2, int>
            <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                     input_data,
                                                     input_channels,
                                                     input_height,
                                                     input_width,
                                                     output_height,
                                                     output_width,
                                                     ksize_height,
                                                     ksize_width,
                                                     stride_height,
                                                     stride_width,
                                                     padding_height,
                                                     padding_width,
                                                     dilation_height,
                                                     dilation_width,
                                                     adaptive,
                                                     output_data,
                                                     mask_data,
                                                     pool_divmods);
      } else {
        auto pool_divmods = FastDivModForPooling<int64_t>(
            input_channels, output_width, output_height);
        KernelMaxPool2dWithIdx<T1, T2, int64_t>
            <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                     input_data,
                                                     input_channels,
                                                     input_height,
                                                     input_width,
                                                     output_height,
                                                     output_width,
                                                     ksize_height,
                                                     ksize_width,
                                                     stride_height,
                                                     stride_width,
                                                     padding_height,
                                                     padding_width,
                                                     dilation_height,
                                                     dilation_width,
                                                     adaptive,
                                                     output_data,
                                                     mask_data,
                                                     pool_divmods);
      }
    }
  }
};

/*
 * All tensors are in NCHW format.
 * Ksize, strides, paddings, dilations are two elements. These two elements
 * represent height and width, respectively.
 */
template <typename T1, typename T2>
class MaxPool2dWithIndexGradFunctor<GPUContext, T1, T2> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& output_grad,
                  const DenseTensor& mask,
                  const std::vector<int64_t>& ksize,
                  const std::vector<int64_t>& strides,
                  const std::vector<int64_t>& paddings,
                  const std::vector<int64_t>& dilations,
                  bool adaptive,
                  DenseTensor* input_grad) {
    const int64_t batch_size = input_grad->dims()[0];
    const int64_t input_channels = input_grad->dims()[1];
    const int64_t input_height = input_grad->dims()[2];
    const int64_t input_width = input_grad->dims()[3];
    const int64_t output_height = output_grad.dims()[2];
    const int64_t output_width = output_grad.dims()[3];
    const int64_t ksize_height = ksize[0];
    const int64_t ksize_width = ksize[1];
    const int64_t stride_height = strides[0];
    const int64_t stride_width = strides[1];
    const int64_t padding_height = paddings[0];
    const int64_t padding_width = paddings[1];
    const int64_t dilation_height = dilations[0];
    const int64_t dilation_width = dilations[1];

    PADDLE_ENFORCE_GE(
        dilation_height,
        1,
        phi::errors::InvalidArgument(
            "dilation_height must be >= 1, but got [%d]", dilation_height));
    PADDLE_ENFORCE_GE(
        dilation_width,
        1,
        phi::errors::InvalidArgument(
            "dilation_width must be >= 1, but got [%d]", dilation_width));

    const T2* mask_data = mask.data<T2>();
    const T1* output_grad_data = output_grad.data<T1>();
    T1* input_grad_data = dev_ctx.template Alloc<T1>(input_grad);

    int64_t nthreads = static_cast<int64_t>(batch_size) * input_channels *
                       input_height * input_width;
    int64_t blocks = (nthreads + 1024 - 1) / 1024;
    dim3 threads(1024, 1);
    dim3 grid(blocks, 1);

    if (nthreads <= std::numeric_limits<int>::max()) {
      auto pool_divmods =
          FastDivModForPooling<int>(input_channels, input_width, input_height);
      KernelMaxPool2DWithIdxGrad<T1, T2, int>
          <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                   output_grad_data,
                                                   mask_data,
                                                   input_channels,
                                                   input_height,
                                                   input_width,
                                                   output_height,
                                                   output_width,
                                                   ksize_height,
                                                   ksize_width,
                                                   stride_height,
                                                   stride_width,
                                                   padding_height,
                                                   padding_width,
                                                   dilation_height,
                                                   dilation_width,
                                                   adaptive,
                                                   input_grad_data,
                                                   pool_divmods);
    } else {
      auto pool_divmods = FastDivModForPooling<int64_t>(
          input_channels, input_width, input_height);
      KernelMaxPool2DWithIdxGrad<T1, T2, int64_t>
          <<<grid, threads, 0, dev_ctx.stream()>>>(nthreads,
                                                   output_grad_data,
                                                   mask_data,
                                                   input_channels,
                                                   input_height,
                                                   input_width,
                                                   output_height,
                                                   output_width,
                                                   ksize_height,
                                                   ksize_width,
                                                   stride_height,
                                                   stride_width,
                                                   padding_height,
                                                   padding_width,
                                                   dilation_height,
                                                   dilation_width,
                                                   adaptive,
                                                   input_grad_data,
                                                   pool_divmods);
    }
  }
};

template class MaxPool2dWithIndexFunctor<GPUContext, float, int>;
template class MaxPool2dWithIndexGradFunctor<GPUContext, float, int>;
template class MaxPool2dWithIndexFunctor<GPUContext, double, int>;
template class MaxPool2dWithIndexGradFunctor<GPUContext, double, int>;
template class MaxPool2dWithIndexFunctor<GPUContext, dtype::float16, int>;
template class MaxPool2dWithIndexGradFunctor<GPUContext, dtype::float16, int>;
template class MaxPool2dWithIndexFunctor<GPUContext, dtype::bfloat16, int>;
template class MaxPool2dWithIndexGradFunctor<GPUContext, dtype::bfloat16, int>;

template <typename T1, typename T2, typename IndexT>
__global__ void KernelMaxPool3DWithIdx(
    const IndexT ncd,
    const T1* input_data,
    const IndexT channels,
    const IndexT input_depth,
    const IndexT input_height,
    const IndexT input_width,
    const IndexT output_depth,
    const IndexT output_height,
    const IndexT output_width,
    const IndexT ksize_depth,
    const IndexT ksize_height,
    const IndexT ksize_width,
    const IndexT stride_depth,
    const IndexT stride_height,
    const IndexT stride_width,
    const IndexT padding_depth,
    const IndexT padding_height,
    const IndexT padding_width,
    const IndexT dilation_depth,
    const IndexT dilation_height,
    const IndexT dilation_width,
    bool adaptive,
    T1* output_data,
    T2* mask_data,
    FastDivModForPooling3D<IndexT> divmods_output) {
  IndexT w_offset, h_offset, d_offset, nc_offset;
  IndexT dstart, dend, hstart, hend, wstart, wend;
  const T1* input_data_cur;

  w_offset = static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
  for (IndexT by = blockIdx.y; by < output_height; by += gridDim.y) {
    h_offset = by * blockDim.y + threadIdx.y;

    IndexT start_index =
        static_cast<IndexT>(blockIdx.z) * blockDim.z + threadIdx.z;
    IndexT step = static_cast<IndexT>(gridDim.z) * blockDim.z;
    if (w_offset < output_width && h_offset < output_height) {
      for (IndexT index_z = start_index; index_z < ncd; index_z += step) {
        auto output_depth_divmod = divmods_output.depth.Divmod(index_z);
        d_offset = output_depth_divmod.val[1];
        nc_offset = output_depth_divmod.val[0];
        IndexT output_index =
            nc_offset * output_depth * output_height * output_width +
            d_offset * output_height * output_width + h_offset * output_width +
            w_offset;
        IndexT input_offset =
            nc_offset * input_depth * input_height * input_width;
        input_data_cur = input_data + input_offset;

        const bool dilation = (dilation_depth != 1 || dilation_height != 1 ||
                               dilation_width != 1);
        if (adaptive) {
          dstart = AdaptStartIndex(d_offset, input_depth, output_depth);
          dend = AdaptEndIndex(d_offset, input_depth, output_depth);

          hstart = AdaptStartIndex(h_offset, input_height, output_height);
          hend = AdaptEndIndex(h_offset, input_height, output_height);

          wstart = AdaptStartIndex(w_offset, input_width, output_width);
          wend = AdaptEndIndex(w_offset, input_width, output_width);
        } else if (dilation) {
          dstart = d_offset * stride_depth - padding_depth;
          hstart = h_offset * stride_height - padding_height;
          wstart = w_offset * stride_width - padding_width;
          dend =
              min(dstart + dilation_depth * (ksize_depth - 1) + 1, input_depth);
          hend = min(hstart + dilation_height * (ksize_height - 1) + 1,
                     input_height);
          wend =
              min(wstart + dilation_width * (ksize_width - 1) + 1, input_width);
          if (dstart < 0) {
            dstart =
                ((-dstart - 1) / dilation_depth + 1) * dilation_depth + dstart;
          }
          if (hstart < 0) {
            hstart = ((-hstart - 1) / dilation_height + 1) * dilation_height +
                     hstart;
          }
          if (wstart < 0) {
            wstart =
                ((-wstart - 1) / dilation_width + 1) * dilation_width + wstart;
          }
        } else {
          dstart = d_offset * stride_depth - padding_depth;
          hstart = h_offset * stride_height - padding_height;
          wstart = w_offset * stride_width - padding_width;
          dend = min(dstart + ksize_depth, input_depth);
          hend = min(hstart + ksize_height, input_height);
          wend = min(wstart + ksize_width, input_width);
          dstart = max(dstart, static_cast<IndexT>(0));
          hstart = max(hstart, static_cast<IndexT>(0));
          wstart = max(wstart, static_cast<IndexT>(0));
        }

        T1 ele = static_cast<T1>(-FLT_MAX);
        IndexT max_index = -1;
        for (IndexT d = dstart; d < dend; d += dilation_depth) {
          for (IndexT h = hstart; h < hend; h += dilation_height) {
            for (IndexT w = wstart; w < wend; w += dilation_width) {
              if (ele <
                  input_data_cur[(d * input_height + h) * input_width + w]) {
                max_index = (d * input_height + h) * input_width + w;
                ele = input_data_cur[max_index];
              }
            }
          }
        }
        output_data[output_index] = ele;
        mask_data[output_index] = max_index;
      }
    }
  }
}

template <typename T1, typename T2, typename IndexT>
__global__ void KernelMaxPool3DWithIdxGrad(
    const IndexT ncd,
    const T1* output_grad,
    const T2* mask,
    const IndexT channels,
    const IndexT input_depth,
    const IndexT input_height,
    const IndexT input_width,
    const IndexT output_depth,
    const IndexT output_height,
    const IndexT output_width,
    const IndexT ksize_depth,
    const IndexT ksize_height,
    const IndexT ksize_width,
    const IndexT stride_depth,
    const IndexT stride_height,
    const IndexT stride_width,
    const IndexT padding_depth,
    const IndexT padding_height,
    const IndexT padding_width,
    bool adaptive,
    T1* input_grad,
    FastDivModForPooling3D<IndexT> divmods_output) {
  IndexT w_offset, h_offset, d_offset, nc_offset;

  w_offset = static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
  for (IndexT by = blockIdx.y; by < output_height; by += gridDim.y) {
    h_offset = by * blockDim.y + threadIdx.y;

    IndexT start_index =
        static_cast<IndexT>(blockIdx.z) * blockDim.z + threadIdx.z;
    IndexT step = static_cast<IndexT>(gridDim.z) * blockDim.z;
    if (w_offset < output_width && h_offset < output_height) {
      for (IndexT index_z = start_index; index_z < ncd; index_z += step) {
        auto output_depth_divmod = divmods_output.depth.Divmod(index_z);
        d_offset = output_depth_divmod.val[1];
        nc_offset = output_depth_divmod.val[0];
        IndexT output_index =
            nc_offset * output_depth * output_height * output_width +
            d_offset * output_height * output_width + h_offset * output_width +
            w_offset;
        IndexT max_index = mask[output_index];
        if (max_index != -1) {
          CudaAtomicAdd(
              &input_grad[nc_offset * input_depth * input_height * input_width +
                          max_index],
              output_grad[output_index]);
        }
      }
    }
  }
}

/*
 * All tensors are in NCDHW format.
 * Ksize, strides, paddings, dilations are three elements. These three elements
 * represent depth, height and width, respectively.
 */
template <typename T1, typename T2>
class MaxPool3dWithIndexFunctor<GPUContext, T1, T2> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& input,
                  const std::vector<int64_t>& ksize,
                  const std::vector<int64_t>& strides,
                  const std::vector<int64_t>& paddings,
                  const std::vector<int64_t>& dilations,
                  bool adaptive,
                  DenseTensor* output,
                  DenseTensor* mask) {
    const int64_t batch_size = input.dims()[0];
    const int64_t input_channels = input.dims()[1];
    const int64_t input_depth = input.dims()[2];
    const int64_t input_height = input.dims()[3];
    const int64_t input_width = input.dims()[4];
    const int64_t output_channels = output->dims()[1];
    const int64_t output_depth = output->dims()[2];
    const int64_t output_height = output->dims()[3];
    const int64_t output_width = output->dims()[4];
    const int64_t ksize_depth = ksize[0];
    const int64_t ksize_height = ksize[1];
    const int64_t ksize_width = ksize[2];
    const int64_t stride_depth = strides[0];
    const int64_t stride_height = strides[1];
    const int64_t stride_width = strides[2];
    const int64_t padding_depth = paddings[0];
    const int64_t padding_height = paddings[1];
    const int64_t padding_width = paddings[2];
    const int64_t dilation_depth = dilations[0];
    const int64_t dilation_height = dilations[1];
    const int64_t dilation_width = dilations[2];

    const T1* input_data = input.data<T1>();
    T1* output_data = dev_ctx.template Alloc<T1>(output);
    T2* mask_data = dev_ctx.template Alloc<T2>(mask);

    PADDLE_ENFORCE_GE(
        dilation_depth,
        1,
        phi::errors::InvalidArgument(
            "dilation_depth must be >= 1, but got [%d]", dilation_depth));
    PADDLE_ENFORCE_GE(
        dilation_height,
        1,
        phi::errors::InvalidArgument(
            "dilation_height must be >= 1, but got [%d]", dilation_height));
    PADDLE_ENFORCE_GE(
        dilation_width,
        1,
        phi::errors::InvalidArgument(
            "dilation_width must be >= 1, but got [%d]", dilation_width));

    int64_t ncd =
        static_cast<int64_t>(batch_size) * input_channels * output_depth;

    int64_t thread_x = 32;
    int64_t thread_y = 8;
    int64_t thread_z = 1;
    dim3 threads(thread_x, thread_y, thread_z);
    std::array<unsigned int, 3> max_grid_dim = dev_ctx.GetCUDAMaxGridDimSize();
    int64_t block_x = (output_width + threads.x - 1) / threads.x;
    int64_t block_y = (output_height > max_grid_dim[1] * threads.y)
                          ? max_grid_dim[1]
                          : (output_height + threads.y - 1) / threads.y;
    int64_t block_z = (ncd > max_grid_dim[2] * threads.z)
                          ? max_grid_dim[2]
                          : (ncd + threads.z - 1) / threads.z;
    dim3 grid(block_x, block_y, block_z);

    if (input.numel() <= std::numeric_limits<int>::max()) {
      auto pool_divmods_output = FastDivModForPooling3D<int>(
          input_channels, output_width, output_height, output_depth);
      KernelMaxPool3DWithIdx<T1, T2, int>
          <<<grid, threads, 0, dev_ctx.stream()>>>(ncd,
                                                   input_data,
                                                   input_channels,
                                                   input_depth,
                                                   input_height,
                                                   input_width,
                                                   output_depth,
                                                   output_height,
                                                   output_width,
                                                   ksize_depth,
                                                   ksize_height,
                                                   ksize_width,
                                                   stride_depth,
                                                   stride_height,
                                                   stride_width,
                                                   padding_depth,
                                                   padding_height,
                                                   padding_width,
                                                   dilation_depth,
                                                   dilation_height,
                                                   dilation_width,
                                                   adaptive,
                                                   output_data,
                                                   mask_data,
                                                   pool_divmods_output);
    } else {
      auto pool_divmods_output = FastDivModForPooling3D<int64_t>(
          input_channels, output_width, output_height, output_depth);
      KernelMaxPool3DWithIdx<T1, T2, int64_t>
          <<<grid, threads, 0, dev_ctx.stream()>>>(ncd,
                                                   input_data,
                                                   input_channels,
                                                   input_depth,
                                                   input_height,
                                                   input_width,
                                                   output_depth,
                                                   output_height,
                                                   output_width,
                                                   ksize_depth,
                                                   ksize_height,
                                                   ksize_width,
                                                   stride_depth,
                                                   stride_height,
                                                   stride_width,
                                                   padding_depth,
                                                   padding_height,
                                                   padding_width,
                                                   dilation_depth,
                                                   dilation_height,
                                                   dilation_width,
                                                   adaptive,
                                                   output_data,
                                                   mask_data,
                                                   pool_divmods_output);
    }
  }
};

/*
 * All tensors are in NCDHW format.
 * Ksize, strides, paddings, dilations are three elements. These three elements
 * represent depth, height and width, respectively.
 */
template <typename T1, typename T2>
class MaxPool3dWithIndexGradFunctor<GPUContext, T1, T2> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& output_grad,
                  const DenseTensor& mask,
                  const std::vector<int64_t>& ksize,
                  const std::vector<int64_t>& strides,
                  const std::vector<int64_t>& paddings,
                  const std::vector<int64_t>& dilations,
                  bool adaptive,
                  DenseTensor* input_grad) {
    const int64_t batch_size = input_grad->dims()[0];
    const int64_t input_channels = input_grad->dims()[1];
    const int64_t input_depth = input_grad->dims()[2];
    const int64_t input_height = input_grad->dims()[3];
    const int64_t input_width = input_grad->dims()[4];
    const int64_t output_depth = output_grad.dims()[2];
    const int64_t output_height = output_grad.dims()[3];
    const int64_t output_width = output_grad.dims()[4];
    const int64_t ksize_depth = ksize[0];
    const int64_t ksize_height = ksize[1];
    const int64_t ksize_width = ksize[2];
    const int64_t stride_depth = strides[0];
    const int64_t stride_height = strides[1];
    const int64_t stride_width = strides[2];
    const int64_t padding_depth = paddings[0];
    const int64_t padding_height = paddings[1];
    const int64_t padding_width = paddings[2];

    const T1* output_grad_data = output_grad.data<T1>();
    const T2* mask_data = mask.data<T2>();
    T1* input_grad_data = dev_ctx.template Alloc<T1>(input_grad);

    int64_t ncd = batch_size * input_channels * output_depth;

    int64_t thread_x = 32;
    int64_t thread_y = 8;
    int64_t thread_z = 1;
    dim3 threads(thread_x, thread_y, thread_z);
    std::array<unsigned int, 3> max_grid_dim = dev_ctx.GetCUDAMaxGridDimSize();
    int64_t block_x = (output_width + threads.x - 1) / threads.x;
    int64_t block_y = (output_height > max_grid_dim[1] * threads.y)
                          ? max_grid_dim[1]
                          : (output_height + threads.y - 1) / threads.y;
    int64_t block_z = (ncd > max_grid_dim[2] * threads.z)
                          ? max_grid_dim[2]
                          : (ncd + threads.z - 1) / threads.z;
    dim3 grid(block_x, block_y, block_z);

    if (input_grad->numel() <= std::numeric_limits<int>::max()) {
      auto pool_divmods_output = FastDivModForPooling3D<int>(
          input_channels, output_width, output_height, output_depth);
      KernelMaxPool3DWithIdxGrad<T1, T2, int>
          <<<grid, threads, 0, dev_ctx.stream()>>>(ncd,
                                                   output_grad_data,
                                                   mask_data,
                                                   input_channels,
                                                   input_depth,
                                                   input_height,
                                                   input_width,
                                                   output_depth,
                                                   output_height,
                                                   output_width,
                                                   ksize_depth,
                                                   ksize_height,
                                                   ksize_width,
                                                   stride_depth,
                                                   stride_height,
                                                   stride_width,
                                                   padding_depth,
                                                   padding_height,
                                                   padding_width,
                                                   adaptive,
                                                   input_grad_data,
                                                   pool_divmods_output);
    } else {
      auto pool_divmods_output = FastDivModForPooling3D<int64_t>(
          input_channels, output_width, output_height, output_depth);
      KernelMaxPool3DWithIdxGrad<T1, T2, int64_t>
          <<<grid, threads, 0, dev_ctx.stream()>>>(ncd,
                                                   output_grad_data,
                                                   mask_data,
                                                   input_channels,
                                                   input_depth,
                                                   input_height,
                                                   input_width,
                                                   output_depth,
                                                   output_height,
                                                   output_width,
                                                   ksize_depth,
                                                   ksize_height,
                                                   ksize_width,
                                                   stride_depth,
                                                   stride_height,
                                                   stride_width,
                                                   padding_depth,
                                                   padding_height,
                                                   padding_width,
                                                   adaptive,
                                                   input_grad_data,
                                                   pool_divmods_output);
    }
  }
};

template class MaxPool3dWithIndexFunctor<GPUContext, float, int>;
template class MaxPool3dWithIndexGradFunctor<GPUContext, float, int>;
template class MaxPool3dWithIndexFunctor<GPUContext, double, int>;
template class MaxPool3dWithIndexGradFunctor<GPUContext, double, int>;
template class MaxPool3dWithIndexFunctor<GPUContext, dtype::float16, int>;
template class MaxPool3dWithIndexGradFunctor<GPUContext, dtype::float16, int>;
template class MaxPool3dWithIndexFunctor<GPUContext, dtype::bfloat16, int>;
template class MaxPool3dWithIndexGradFunctor<GPUContext, dtype::bfloat16, int>;
// fractional max pool
template <typename T1, typename T2, typename IndexT>
__global__ void FractionalKernelMaxPool2d(
    const IndexT ncd,
    const T1* input_data,
    const IndexT channels,
    const IndexT input_height,
    const IndexT input_width,
    const IndexT output_height,
    const IndexT output_width,
    const IndexT pool_height,
    const IndexT pool_width,
    float random_u,
    uint64_t seed,
    uint64_t offset,
    T1* output_data,
    T2* mask_data,
    FastDivModForPooling<IndexT> divmods) {
  float alpha_height = 0, alpha_width = 0;
  float u_height = 0, u_width = 0;
  float u = 0;
  if (random_u == 0) {
    IndexT thread_idx =
        static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
#if defined(__NVCC__)
    curandStatePhilox4_32_10_t state;
    curand_init(seed, thread_idx, offset, &state);
#else
    hiprandStatePhilox4_32_10_t state;
    hiprand_init(seed, thread_idx, offset, &state);
#endif
    funcs::uniform_distribution<float> dist;
    float4 rand = dist(&state);
    u = (&rand.x)[0];
  } else {
    u = random_u;
  }

  alpha_height = static_cast<float>(input_height - pool_height) /
                 (output_height - (pool_height > 0 ? 1 : 0));
  alpha_width = static_cast<float>(input_width - pool_width) /
                (output_width - (pool_width > 0 ? 1 : 0));

  u_height = FractionalRationalU(
      u, alpha_height, input_height, output_height, pool_height);
  u_width = FractionalRationalU(
      u, alpha_width, input_width, output_width, pool_width);

  IndexT w_offset, h_offset, nc_offset;
  IndexT hstart, hend, wstart, wend;
  const T1* input_data_cur;

  w_offset = static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;

  IndexT start_index =
      static_cast<IndexT>(blockIdx.y) * blockDim.y + threadIdx.y;
  IndexT step = static_cast<IndexT>(gridDim.y) * blockDim.y;
  if (w_offset < output_width) {
    for (IndexT index_y = start_index; index_y < ncd; index_y += step) {
      auto output_height_divmod = divmods.height.Divmod(index_y);
      h_offset = output_height_divmod.val[1];
      nc_offset = output_height_divmod.val[0];

      IndexT output_index = nc_offset * output_height * output_width +
                            h_offset * output_width + w_offset;

      IndexT input_offset = nc_offset * input_height * input_width;
      input_data_cur = input_data + input_offset;

      hstart =
          FractionalStartIndex(h_offset, alpha_height, u_height, pool_height);
      hend = FractionalEndIndex(h_offset, alpha_height, u_height, pool_height);
      hstart = std::max(hstart, static_cast<IndexT>(0));
      hend = std::min(hend, input_height);

      wstart = FractionalStartIndex(w_offset, alpha_width, u_width, pool_width);
      wend = FractionalEndIndex(w_offset, alpha_width, u_width, pool_width);
      wstart = std::max(wstart, static_cast<IndexT>(0));
      wend = std::min(wend, input_width);

      T1 ele = static_cast<T1>(-FLT_MAX);
      IndexT max_index = -1;
      for (IndexT h = hstart; h < hend; ++h) {
        for (IndexT w = wstart; w < wend; ++w) {
          if (ele < input_data_cur[h * input_width + w]) {
            max_index = h * input_width + w;
            ele = input_data_cur[max_index];
          }
        }
      }

      output_data[output_index] = ele;
      mask_data[output_index] = max_index;
    }
  }
}

template <typename T1, typename T2, typename IndexT>
__global__ void FractionalKernelMaxPool2dGrad(
    const IndexT ncd,
    const T1* output_grad,
    const T2* mask_data,
    const IndexT channels,
    const IndexT input_height,
    const IndexT input_width,
    const IndexT output_height,
    const IndexT output_width,
    const IndexT pool_height,
    const IndexT pool_width,
    float random_u,
    uint64_t seed,
    uint64_t offset,
    T1* input_grad,
    FastDivModForPooling<IndexT> divmods) {
  IndexT w_offset, h_offset, nc_offset;

  w_offset = static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;

  IndexT start_index =
      static_cast<IndexT>(blockIdx.y) * blockDim.y + threadIdx.y;
  IndexT step = static_cast<IndexT>(gridDim.y) * blockDim.y;
  if (w_offset < output_width) {
    for (IndexT index_y = start_index; index_y < ncd; index_y += step) {
      auto output_height_divmod = divmods.height.Divmod(index_y);
      h_offset = output_height_divmod.val[1];
      nc_offset = output_height_divmod.val[0];

      IndexT output_index = nc_offset * output_height * output_width +
                            h_offset * output_width + w_offset;

      IndexT max_index = mask_data[output_index];
      if (max_index != -1) {
        CudaAtomicAdd(
            &input_grad[nc_offset * input_height * input_width + max_index],
            output_grad[output_index]);
      }
    }
  }
}

/*
 * All tensors are in NCHW format.
 */
template <typename T1, typename T2>
class FractionalMaxPool2dFunctor<GPUContext, T1, T2> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& input,
                  const std::vector<int64_t>& output_size,
                  const std::vector<int64_t>& kernel_size,
                  float random_u,
                  bool return_mask,
                  DenseTensor* output,
                  DenseTensor* mask) {
    const int64_t batch_size = input.dims()[0];
    const int64_t input_channels = input.dims()[1];
    const int64_t input_height = input.dims()[2];
    const int64_t input_width = input.dims()[3];
    const int64_t output_channels = output->dims()[1];
    const int64_t output_height = output->dims()[2];
    const int64_t output_width = output->dims()[3];
    const int64_t pool_height = kernel_size[0];
    const int64_t pool_width = kernel_size[1];

    PADDLE_ENFORCE_GE(
        input_height,
        output_height - 1 + pool_height,
        common::errors::InvalidArgument(
            "input_height [%d] is less than valid output_height [%d]",
            input_height,
            output_height - 1 + pool_height));
    PADDLE_ENFORCE_GE(
        input_width,
        output_width - 1 + pool_width,
        common::errors::InvalidArgument(
            "input_width [%d] is less than valid output_width [%d]",
            input_width,
            output_width - 1 + pool_width));

    const T1* input_data = input.data<T1>();
    T1* output_data = dev_ctx.template Alloc<T1>(output);
    T2* mask_data = dev_ctx.template Alloc<T2>(mask);

    int64_t ncd = batch_size * input_channels * output_height;

    int64_t thread_x = 32;
    int64_t thread_y = 1;
    int64_t thread_z = 1;
    dim3 threads(thread_x, thread_y, thread_z);
    std::array<unsigned int, 3> max_grid_dim = dev_ctx.GetCUDAMaxGridDimSize();
    int64_t block_x = (output_width + threads.x - 1) / threads.x;
    int64_t block_y = (ncd > max_grid_dim[1] * threads.y)
                          ? max_grid_dim[1]
                          : (ncd + threads.y - 1) / threads.y;
    int64_t block_z = 1;
    dim3 grid(block_x, block_y, block_z);

    uint64_t seed = 0;
    uint64_t offset = 0;
    // generate seed for fractional pool
    auto gen_cuda = dev_ctx.GetGenerator();
    constexpr uint64_t increment_offset = 1 * 4;  // one seed with multiple of 4
    auto seed_offset = gen_cuda->IncrementOffset(increment_offset);
    seed = seed_offset.first;
    offset = seed_offset.second;

    if (input.numel() <= std::numeric_limits<int>::max()) {
      auto pool_divmods = FastDivModForPooling<int>(
          input_channels, output_width, output_height);
      FractionalKernelMaxPool2d<T1, T2, int>
          <<<grid, threads, 0, dev_ctx.stream()>>>(ncd,
                                                   input_data,
                                                   input_channels,
                                                   input_height,
                                                   input_width,
                                                   output_height,
                                                   output_width,
                                                   pool_height,
                                                   pool_width,
                                                   random_u,
                                                   seed,
                                                   offset,
                                                   output_data,
                                                   mask_data,
                                                   pool_divmods);
    } else {
      auto pool_divmods = FastDivModForPooling<int64_t>(
          input_channels, output_width, output_height);
      FractionalKernelMaxPool2d<T1, T2, int64_t>
          <<<grid, threads, 0, dev_ctx.stream()>>>(ncd,
                                                   input_data,
                                                   input_channels,
                                                   input_height,
                                                   input_width,
                                                   output_height,
                                                   output_width,
                                                   pool_height,
                                                   pool_width,
                                                   random_u,
                                                   seed,
                                                   offset,
                                                   output_data,
                                                   mask_data,
                                                   pool_divmods);
    }
  }
};

/*
 * All tensors are in NCHW format.
 */
template <typename T1, typename T2>
class FractionalMaxPool2dGradFunctor<GPUContext, T1, T2> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& output_grad,
                  const DenseTensor& mask,
                  const std::vector<int64_t>& output_size,
                  const std::vector<int64_t>& kernel_size,
                  float random_u,
                  bool return_mask,
                  DenseTensor* input_grad) {
    const int64_t batch_size = input_grad->dims()[0];
    const int64_t input_channels = input_grad->dims()[1];
    const int64_t input_height = input_grad->dims()[2];
    const int64_t input_width = input_grad->dims()[3];
    const int64_t output_height = output_grad.dims()[2];
    const int64_t output_width = output_grad.dims()[3];
    const int64_t pool_height = kernel_size[0];
    const int64_t pool_width = kernel_size[1];

    const T2* mask_data = mask.data<T2>();
    const T1* output_grad_data = output_grad.data<T1>();
    T1* input_grad_data = dev_ctx.template Alloc<T1>(input_grad);

    int64_t ncd = batch_size * input_channels * output_height;

    int64_t thread_x = 32;
    int64_t thread_y = 1;
    int64_t thread_z = 1;
    dim3 threads(thread_x, thread_y, thread_z);
    std::array<unsigned int, 3> max_grid_dim = dev_ctx.GetCUDAMaxGridDimSize();
    int64_t block_x = (output_width + threads.x - 1) / threads.x;
    int64_t block_y = (ncd > max_grid_dim[1] * threads.y)
                          ? max_grid_dim[1]
                          : (ncd + threads.y - 1) / threads.y;
    int64_t block_z = 1;
    dim3 grid(block_x, block_y, block_z);

    uint64_t seed = 0;
    uint64_t offset = 0;
    // generate seed for fractional pool
    auto gen_cuda = dev_ctx.GetGenerator();
    constexpr uint64_t increment_offset = 1 * 4;  // one seed with multiple of 4
    auto seed_offset = gen_cuda->IncrementOffset(increment_offset);
    seed = seed_offset.first;
    offset = seed_offset.second;

    if (input_grad->numel() <= std::numeric_limits<int>::max()) {
      auto pool_divmods = FastDivModForPooling<int>(
          input_channels, output_width, output_height);
      FractionalKernelMaxPool2dGrad<T1, T2, int>
          <<<grid, threads, 0, dev_ctx.stream()>>>(ncd,
                                                   output_grad_data,
                                                   mask_data,
                                                   input_channels,
                                                   input_height,
                                                   input_width,
                                                   output_height,
                                                   output_width,
                                                   pool_height,
                                                   pool_width,
                                                   random_u,
                                                   seed,
                                                   offset,
                                                   input_grad_data,
                                                   pool_divmods);
    } else {
      auto pool_divmods = FastDivModForPooling<int64_t>(
          input_channels, output_width, output_height);
      FractionalKernelMaxPool2dGrad<T1, T2, int64_t>
          <<<grid, threads, 0, dev_ctx.stream()>>>(ncd,
                                                   output_grad_data,
                                                   mask_data,
                                                   input_channels,
                                                   input_height,
                                                   input_width,
                                                   output_height,
                                                   output_width,
                                                   pool_height,
                                                   pool_width,
                                                   random_u,
                                                   seed,
                                                   offset,
                                                   input_grad_data,
                                                   pool_divmods);
    }
  }
};

template class FractionalMaxPool2dFunctor<GPUContext, float, int>;
template class FractionalMaxPool2dGradFunctor<GPUContext, float, int>;
template class FractionalMaxPool2dFunctor<GPUContext, double, int>;
template class FractionalMaxPool2dGradFunctor<GPUContext, double, int>;
template class FractionalMaxPool2dFunctor<GPUContext, dtype::float16, int>;
template class FractionalMaxPool2dGradFunctor<GPUContext, dtype::float16, int>;
template class FractionalMaxPool2dFunctor<GPUContext, dtype::bfloat16, int>;
template class FractionalMaxPool2dGradFunctor<GPUContext, dtype::bfloat16, int>;

template <typename T1, typename T2, typename IndexT>
__global__ void FractionalKernelMaxPool3d(
    const IndexT ncd,
    const T1* input_data,
    const IndexT channels,
    const IndexT input_depth,
    const IndexT input_height,
    const IndexT input_width,
    const IndexT output_depth,
    const IndexT output_height,
    const IndexT output_width,
    const IndexT pool_depth,
    const IndexT pool_height,
    const IndexT pool_width,
    float random_u,
    uint64_t seed,
    uint64_t offset,
    T1* output_data,
    T2* mask_data,
    FastDivModForPooling3D<IndexT> divmods_output) {
  float alpha_height = 0, alpha_width = 0, alpha_depth = 0;
  float u_height = 0, u_width = 0, u_depth = 0;
  float u = 0;
  if (random_u == 0) {
    IndexT thread_idx =
        static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
#if defined(__NVCC__)
    curandStatePhilox4_32_10_t state;
    curand_init(seed, thread_idx, offset, &state);
#else
    hiprandStatePhilox4_32_10_t state;
    hiprand_init(seed, thread_idx, offset, &state);
#endif
    funcs::uniform_distribution<float> dist;
    float4 rand = dist(&state);
    u = (&rand.x)[0];
  } else {
    u = random_u;
  }

  alpha_depth = static_cast<float>(input_depth - pool_depth) /
                (output_depth - (pool_depth > 0 ? 1 : 0));
  alpha_height = static_cast<float>(input_height - pool_height) /
                 (output_height - (pool_height > 0 ? 1 : 0));
  alpha_width = static_cast<float>(input_width - pool_width) /
                (output_width - (pool_width > 0 ? 1 : 0));

  u_depth = FractionalRationalU(
      u, alpha_depth, input_depth, output_depth, pool_depth);
  u_height = FractionalRationalU(
      u, alpha_height, input_height, output_height, pool_height);
  u_width = FractionalRationalU(
      u, alpha_width, input_width, output_width, pool_width);

  IndexT w_offset, h_offset, d_offset, nc_offset;
  IndexT dstart, dend, hstart, hend, wstart, wend;
  const T1* input_data_cur;

  w_offset = static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
  for (IndexT by = blockIdx.y; by < output_height; by += gridDim.y) {
    h_offset = by * blockDim.y + threadIdx.y;

    IndexT start_index =
        static_cast<IndexT>(blockIdx.z) * blockDim.z + threadIdx.z;
    IndexT step = static_cast<IndexT>(gridDim.z) * blockDim.z;
    if (w_offset < output_width && h_offset < output_height) {
      for (IndexT index_z = start_index; index_z < ncd; index_z += step) {
        auto output_depth_divmod = divmods_output.depth.Divmod(index_z);
        d_offset = output_depth_divmod.val[1];
        nc_offset = output_depth_divmod.val[0];
        IndexT output_index =
            nc_offset * output_depth * output_height * output_width +
            d_offset * output_height * output_width + h_offset * output_width +
            w_offset;
        IndexT input_offset =
            nc_offset * input_depth * input_height * input_width;
        input_data_cur = input_data + input_offset;

        dstart =
            FractionalStartIndex(d_offset, alpha_depth, u_depth, pool_depth);
        dend = FractionalEndIndex(d_offset, alpha_depth, u_depth, pool_depth);
        dstart = std::max(dstart, static_cast<IndexT>(0));
        dend = std::min(dend, input_depth);

        hstart =
            FractionalStartIndex(h_offset, alpha_height, u_height, pool_height);
        hend =
            FractionalEndIndex(h_offset, alpha_height, u_height, pool_height);
        hstart = std::max(hstart, static_cast<IndexT>(0));
        hend = std::min(hend, input_height);

        wstart =
            FractionalStartIndex(w_offset, alpha_width, u_width, pool_width);
        wend = FractionalEndIndex(w_offset, alpha_width, u_width, pool_width);
        wstart = std::max(wstart, static_cast<IndexT>(0));
        wend = std::min(wend, input_width);

        T1 ele = static_cast<T1>(-FLT_MAX);
        IndexT max_index = -1;
        for (IndexT d = dstart; d < dend; ++d) {
          for (IndexT h = hstart; h < hend; ++h) {
            for (IndexT w = wstart; w < wend; ++w) {
              if (ele <
                  input_data_cur[(d * input_height + h) * input_width + w]) {
                max_index = (d * input_height + h) * input_width + w;
                ele = input_data_cur[max_index];
              }
            }
          }
        }
        output_data[output_index] = ele;
        mask_data[output_index] = max_index;
      }
    }
  }
}

template <typename T1, typename T2, typename IndexT>
__global__ void FractionalKernelMaxPool3dGrad(
    const IndexT ncd,
    const T1* output_grad,
    const T2* mask,
    const IndexT channels,
    const IndexT input_depth,
    const IndexT input_height,
    const IndexT input_width,
    const IndexT output_depth,
    const IndexT output_height,
    const IndexT output_width,
    const IndexT pool_depth,
    const IndexT pool_height,
    const IndexT pool_width,
    float random_u,
    T1* input_grad,
    FastDivModForPooling3D<IndexT> divmods_output) {
  IndexT w_offset, h_offset, d_offset, nc_offset;

  w_offset = static_cast<IndexT>(blockIdx.x) * blockDim.x + threadIdx.x;
  for (IndexT by = blockIdx.y; by < output_height; by += gridDim.y) {
    h_offset = by * blockDim.y + threadIdx.y;

    IndexT start_index =
        static_cast<IndexT>(blockIdx.z) * blockDim.z + threadIdx.z;
    IndexT step = static_cast<IndexT>(gridDim.z) * blockDim.z;
    if (w_offset < output_width && h_offset < output_height) {
      for (IndexT index_z = start_index; index_z < ncd; index_z += step) {
        auto output_depth_divmod = divmods_output.depth.Divmod(index_z);
        d_offset = output_depth_divmod.val[1];
        nc_offset = output_depth_divmod.val[0];
        IndexT output_index =
            nc_offset * output_depth * output_height * output_width +
            d_offset * output_height * output_width + h_offset * output_width +
            w_offset;
        IndexT max_index = mask[output_index];
        if (max_index != -1) {
          CudaAtomicAdd(
              &input_grad[nc_offset * input_depth * input_height * input_width +
                          max_index],
              output_grad[output_index]);
        }
      }
    }
  }
}

/*
 * All tensors are in NCDHW format.
 */
template <typename T1, typename T2>
class FractionalMaxPool3dFunctor<GPUContext, T1, T2> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& input,
                  const std::vector<int64_t>& output_size,
                  const std::vector<int64_t>& kernel_size,
                  float random_u,
                  bool return_mask,
                  DenseTensor* output,
                  DenseTensor* mask) {
    const int64_t batch_size = input.dims()[0];
    const int64_t input_channels = input.dims()[1];
    const int64_t input_depth = input.dims()[2];
    const int64_t input_height = input.dims()[3];
    const int64_t input_width = input.dims()[4];
    const int64_t output_channels = output->dims()[1];
    const int64_t output_depth = output->dims()[2];
    const int64_t output_height = output->dims()[3];
    const int64_t output_width = output->dims()[4];
    const int64_t pool_depth = kernel_size[0];
    const int64_t pool_height = kernel_size[1];
    const int64_t pool_width = kernel_size[2];

    PADDLE_ENFORCE_GE(
        input_depth,
        output_depth - 1 + pool_depth,
        common::errors::InvalidArgument(
            "input_depth [%d] is less than valid output_depth [%d]",
            input_depth,
            output_depth - 1 + pool_depth));
    PADDLE_ENFORCE_GE(
        input_height,
        output_height - 1 + pool_height,
        common::errors::InvalidArgument(
            "input_height [%d] is less than valid output_height [%d]",
            input_height,
            output_height - 1 + pool_height));
    PADDLE_ENFORCE_GE(
        input_width,
        output_width - 1 + pool_width,
        common::errors::InvalidArgument(
            "input_width [%d] is less than valid output_width [%d]",
            input_width,
            output_width - 1 + pool_width));

    const T1* input_data = input.data<T1>();
    T1* output_data = dev_ctx.template Alloc<T1>(output);
    T2* mask_data = dev_ctx.template Alloc<T2>(mask);

    int64_t ncd = batch_size * input_channels * output_depth;

    int64_t thread_x = 32;
    int64_t thread_y = 8;
    int64_t thread_z = 1;
    dim3 threads(thread_x, thread_y, thread_z);
    std::array<unsigned int, 3> max_grid_dim = dev_ctx.GetCUDAMaxGridDimSize();
    int64_t block_x = (output_width + threads.x - 1) / threads.x;
    int64_t block_y = (output_height > max_grid_dim[1] * threads.y)
                          ? max_grid_dim[1]
                          : (output_height + threads.y - 1) / threads.y;
    int64_t block_z = (ncd > max_grid_dim[2] * threads.z)
                          ? max_grid_dim[2]
                          : (ncd + threads.z - 1) / threads.z;
    dim3 grid(block_x, block_y, block_z);

    uint64_t seed = 0;
    uint64_t offset = 0;
    // generate seed for fractional pool
    auto gen_cuda = dev_ctx.GetGenerator();
    constexpr uint64_t increment_offset = 1 * 4;  // one seed with multiple of 4
    auto seed_offset = gen_cuda->IncrementOffset(increment_offset);
    seed = seed_offset.first;
    offset = seed_offset.second;

    if (input.numel() <= std::numeric_limits<int>::max()) {
      auto pool_divmods_output = FastDivModForPooling3D<int>(
          input_channels, output_width, output_height, output_depth);
      FractionalKernelMaxPool3d<T1, T2, int>
          <<<grid, threads, 0, dev_ctx.stream()>>>(ncd,
                                                   input_data,
                                                   input_channels,
                                                   input_depth,
                                                   input_height,
                                                   input_width,
                                                   output_depth,
                                                   output_height,
                                                   output_width,
                                                   pool_depth,
                                                   pool_height,
                                                   pool_width,
                                                   random_u,
                                                   seed,
                                                   offset,
                                                   output_data,
                                                   mask_data,
                                                   pool_divmods_output);
    } else {
      auto pool_divmods_output = FastDivModForPooling3D<int64_t>(
          input_channels, output_width, output_height, output_depth);
      FractionalKernelMaxPool3d<T1, T2, int64_t>
          <<<grid, threads, 0, dev_ctx.stream()>>>(ncd,
                                                   input_data,
                                                   input_channels,
                                                   input_depth,
                                                   input_height,
                                                   input_width,
                                                   output_depth,
                                                   output_height,
                                                   output_width,
                                                   pool_depth,
                                                   pool_height,
                                                   pool_width,
                                                   random_u,
                                                   seed,
                                                   offset,
                                                   output_data,
                                                   mask_data,
                                                   pool_divmods_output);
    }
  }
};

/*
 * All tensors are in NCDHW format.
 */
template <typename T1, typename T2>
class FractionalMaxPool3dGradFunctor<GPUContext, T1, T2> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& output_grad,
                  const DenseTensor& mask,
                  const std::vector<int64_t>& output_size,
                  const std::vector<int64_t>& kernel_size,
                  float random_u,
                  bool return_mask,
                  DenseTensor* input_grad) {
    const int64_t batch_size = input_grad->dims()[0];
    const int64_t input_channels = input_grad->dims()[1];
    const int64_t input_depth = input_grad->dims()[2];
    const int64_t input_height = input_grad->dims()[3];
    const int64_t input_width = input_grad->dims()[4];
    const int64_t output_depth = output_grad.dims()[2];
    const int64_t output_height = output_grad.dims()[3];
    const int64_t output_width = output_grad.dims()[4];
    const int64_t pool_depth = kernel_size[0];
    const int64_t pool_height = kernel_size[1];
    const int64_t pool_width = kernel_size[2];

    const T1* output_grad_data = output_grad.data<T1>();
    const T2* mask_data = mask.data<T2>();
    T1* input_grad_data = dev_ctx.template Alloc<T1>(input_grad);

    int64_t ncd = batch_size * input_channels * output_depth;

    int64_t thread_x = 32;
    int64_t thread_y = 8;
    int64_t thread_z = 1;
    dim3 threads(thread_x, thread_y, thread_z);
    std::array<unsigned int, 3> max_grid_dim = dev_ctx.GetCUDAMaxGridDimSize();
    int64_t block_x = (output_width + threads.x - 1) / threads.x;
    int64_t block_y = (output_height > max_grid_dim[1] * threads.y)
                          ? max_grid_dim[1]
                          : (output_height + threads.y - 1) / threads.y;
    int64_t block_z = (ncd > max_grid_dim[2] * threads.z)
                          ? max_grid_dim[2]
                          : (ncd + threads.z - 1) / threads.z;
    dim3 grid(block_x, block_y, block_z);

    if (input_grad->numel() <= std::numeric_limits<int>::max()) {
      auto pool_divmods_output = FastDivModForPooling3D<int>(
          input_channels, output_width, output_height, output_depth);

      FractionalKernelMaxPool3dGrad<T1, T2, int>
          <<<grid, threads, 0, dev_ctx.stream()>>>(ncd,
                                                   output_grad_data,
                                                   mask_data,
                                                   input_channels,
                                                   input_depth,
                                                   input_height,
                                                   input_width,
                                                   output_depth,
                                                   output_height,
                                                   output_width,
                                                   pool_depth,
                                                   pool_height,
                                                   pool_width,
                                                   random_u,
                                                   input_grad_data,
                                                   pool_divmods_output);
    } else {
      auto pool_divmods_output = FastDivModForPooling3D<int64_t>(
          input_channels, output_width, output_height, output_depth);

      FractionalKernelMaxPool3dGrad<T1, T2, int64_t>
          <<<grid, threads, 0, dev_ctx.stream()>>>(ncd,
                                                   output_grad_data,
                                                   mask_data,
                                                   input_channels,
                                                   input_depth,
                                                   input_height,
                                                   input_width,
                                                   output_depth,
                                                   output_height,
                                                   output_width,
                                                   pool_depth,
                                                   pool_height,
                                                   pool_width,
                                                   random_u,
                                                   input_grad_data,
                                                   pool_divmods_output);
    }
  }
};

template class FractionalMaxPool3dFunctor<GPUContext, float, int>;
template class FractionalMaxPool3dGradFunctor<GPUContext, float, int>;
template class FractionalMaxPool3dFunctor<GPUContext, double, int>;
template class FractionalMaxPool3dGradFunctor<GPUContext, double, int>;
template class FractionalMaxPool3dFunctor<GPUContext, dtype::float16, int>;
template class FractionalMaxPool3dGradFunctor<GPUContext, dtype::float16, int>;
template class FractionalMaxPool3dFunctor<GPUContext, dtype::bfloat16, int>;
template class FractionalMaxPool3dGradFunctor<GPUContext, dtype::bfloat16, int>;

}  // namespace funcs
}  // namespace phi