/* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#include <algorithm>
#include <string>

#include "paddle/common/macros.h"
#include "paddle/phi/backends/gpu/gpu_primitives.h"
#include "paddle/phi/core/mixed_vector.h"
#include "paddle/phi/kernels/funcs/math_function.h"
#include "paddle/phi/kernels/funcs/sequence_pooling.h"

namespace phi {
namespace funcs {

template <typename T>
struct MaxPoolFunctor {
  HOSTDEVICE void operator()(const T* input,
                             const T pad_value,
                             const size_t start,
                             const size_t end,
                             const size_t item_dim,
                             T* output,
                             int* index) {
    for (size_t tid = threadIdx.x; tid < item_dim; tid += blockDim.x) {
      T max_val = static_cast<T>(-FLT_MAX);
      int max_index = -1;
      if (start == end) {
        output[tid] = pad_value;
        index[tid] = -1;
      } else {
        for (size_t i = start; i < end; ++i) {
          if (max_val < input[item_dim * i + tid]) {
            max_val = input[item_dim * i + tid];
            max_index = i;
          }
        }
        output[tid] = max_val;
        index[tid] = max_index;
      }
    }
  }
};

template <typename T>
struct AvgPoolFunctor {
  HOSTDEVICE void operator()(const T* input,
                             const T pad_value,
                             const size_t start,
                             const size_t end,
                             const size_t item_dim,
                             T* output,
                             int* index) {
    for (size_t tid = threadIdx.x; tid < item_dim; tid += blockDim.x) {
      if (start == end) {
        output[tid] = pad_value;
      } else {
        T val = static_cast<T>(0);
        for (size_t i = start; i < end; ++i) {
          val += input[item_dim * i + tid];
        }
        // end, start is lod, so end - start != 0
        output[tid] = val / static_cast<T>(end - start);
      }
    }
  }
};

template <typename T>
struct SumPoolFunctor {
  HOSTDEVICE void operator()(const T* input,
                             const T pad_value,
                             const size_t start,
                             const size_t end,
                             const size_t item_dim,
                             T* output,
                             int* index) {
    for (size_t tid = threadIdx.x; tid < item_dim; tid += blockDim.x) {
      if (start == end) {
        output[tid] = pad_value;
      } else {
        T val = static_cast<T>(0);
        for (size_t i = start; i < end; ++i) {
          val += input[item_dim * i + tid];
        }
        output[tid] = val;
      }
    }
  }
};

template <typename T>
struct SqrtPoolFunctor {
  HOSTDEVICE void operator()(const T* input,
                             const T pad_value,
                             const size_t start,
                             const size_t end,
                             const size_t item_dim,
                             T* output,
                             int* index) {
    for (size_t tid = threadIdx.x; tid < item_dim; tid += blockDim.x) {
      if (start == end) {
        output[tid] = pad_value;
      } else {
        T val = static_cast<T>(0);
        for (size_t i = start; i < end; ++i) {
          val += input[item_dim * i + tid];
        }
        // end, start is lod, so end - start != 0
        output[tid] = val / sqrt(end - start);
      }
    }
  }
};

template <typename T>
struct LastPoolFunctor {
  HOSTDEVICE void operator()(const T* input,
                             const T pad_value,
                             const size_t start,
                             const size_t end,
                             const size_t item_dim,
                             T* output,
                             int* index) {
    for (size_t tid = threadIdx.x; tid < item_dim; tid += blockDim.x) {
      if (start == end) {
        output[tid] = pad_value;
      } else {
        output[tid] = input[item_dim * (end - 1) + tid];
      }
    }
  }
};

template <typename T>
struct FirstPoolFunctor {
  HOSTDEVICE void operator()(const T* input,
                             const T pad_value,
                             const size_t start,
                             const size_t end,
                             const size_t item_dim,
                             T* output,
                             int* index) {
    for (size_t tid = threadIdx.x; tid < item_dim; tid += blockDim.x) {
      if (start == end) {
        output[tid] = pad_value;
      } else {
        output[tid] = input[item_dim * start + tid];
      }
    }
  }
};

template <typename T, typename Range_OP>
__global__ void sequence_pool_kernel(Range_OP op,
                                     const T* input,
                                     const T pad_value,
                                     const size_t* lod,
                                     const size_t lod_size,
                                     const size_t item_dim,
                                     T* output,
                                     int* index) {
  int bid = blockIdx.x;
  if (bid >= lod_size - 1) return;
  size_t start = lod[bid];
  size_t end = lod[bid + 1];
  int* index_offset = nullptr;
  if (index != nullptr) {
    index_offset = &index[bid * item_dim];
  }
  op(input,
     pad_value,
     start,
     end,
     item_dim,
     &output[bid * item_dim],
     index_offset);
}

template <typename T>
class SequencePoolFunctor<GPUContext, T> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const std::string pooltype,
                  T pad_value,
                  const DenseTensor& input,
                  DenseTensor* output,
                  bool is_test,
                  DenseTensor* index = nullptr) {
    auto lod_level = input.lod().size();
    auto& lod = input.lod()[lod_level - 1];
    const size_t item_dim = output->numel() / output->dims()[0];
    dim3 threads(1024, 1);
    dim3 grid(std::max(static_cast<int>(lod.size()) - 1, 1), 1);
    phi::MixVector<size_t> mix_vector(&lod);
    if (pooltype == "MAX") {
      sequence_pool_kernel<T, MaxPoolFunctor<T>>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              MaxPoolFunctor<T>(),
              input.data<T>(),
              pad_value,
              mix_vector.CUDAData(dev_ctx.GetPlace()),
              lod.size(),
              item_dim,
              dev_ctx.template Alloc<T>(output),
              index->data<int>());
    } else if (pooltype == "AVERAGE") {
      sequence_pool_kernel<T, AvgPoolFunctor<T>>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              AvgPoolFunctor<T>(),
              input.data<T>(),
              pad_value,
              mix_vector.CUDAData(dev_ctx.GetPlace()),
              lod.size(),
              item_dim,
              dev_ctx.template Alloc<T>(output),
              nullptr);
    } else if (pooltype == "SUM") {
      sequence_pool_kernel<T, SumPoolFunctor<T>>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              SumPoolFunctor<T>(),
              input.data<T>(),
              pad_value,
              mix_vector.CUDAData(dev_ctx.GetPlace()),
              lod.size(),
              item_dim,
              dev_ctx.template Alloc<T>(output),
              nullptr);
    } else if (pooltype == "SQRT") {
      sequence_pool_kernel<T, SqrtPoolFunctor<T>>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              SqrtPoolFunctor<T>(),
              input.data<T>(),
              pad_value,
              mix_vector.CUDAData(dev_ctx.GetPlace()),
              lod.size(),
              item_dim,
              dev_ctx.template Alloc<T>(output),
              nullptr);
    } else if (pooltype == "LAST") {
      sequence_pool_kernel<T, LastPoolFunctor<T>>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              LastPoolFunctor<T>(),
              input.data<T>(),
              pad_value,
              mix_vector.CUDAData(dev_ctx.GetPlace()),
              lod.size(),
              item_dim,
              dev_ctx.template Alloc<T>(output),
              nullptr);
    } else if (pooltype == "FIRST") {
      sequence_pool_kernel<T, FirstPoolFunctor<T>>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              FirstPoolFunctor<T>(),
              input.data<T>(),
              pad_value,
              mix_vector.CUDAData(dev_ctx.GetPlace()),
              lod.size(),
              item_dim,
              dev_ctx.template Alloc<T>(output),
              nullptr);
    } else {
      PADDLE_THROW(errors::InvalidArgument(
          "unsupported pooling pooltype: %s. Only support \"MAX\", "
          "\"AVERAGE\", \"SUM\", \"SQRT\", \"LAST\" and \"FIRST\"",
          pooltype));
    }
  }
};

template <typename T>
struct MaxPoolGradFunctor {
  HOSTDEVICE void operator()(const T* out_grad,
                             const size_t start,
                             const size_t end,
                             const size_t item_dim,
                             T* in_grad,
                             const int* index) {
    for (size_t tid = threadIdx.x; tid < item_dim; tid += blockDim.x) {
      for (size_t i = start; i < end; ++i) {
        if (i == index[tid]) {
          in_grad[item_dim * i + tid] = out_grad[tid];
        } else {
          in_grad[item_dim * i + tid] = static_cast<T>(0);
        }
      }
    }
  }
};

template <typename T>
struct AvgPoolGradFunctor {
  HOSTDEVICE void operator()(const T* out_grad,
                             const size_t start,
                             const size_t end,
                             const size_t item_dim,
                             T* in_grad,
                             const int* index) {
    for (size_t tid = threadIdx.x; tid < item_dim; tid += blockDim.x) {
      for (size_t i = start; i < end; ++i) {
        in_grad[item_dim * i + tid] = out_grad[tid] / (end - start);
      }
    }
  }
};

template <typename T>
struct SumPoolGradFunctor {
  HOSTDEVICE void operator()(const T* out_grad,
                             const size_t start,
                             const size_t end,
                             const size_t item_dim,
                             T* in_grad,
                             const int* index) {
    for (size_t tid = threadIdx.x; tid < item_dim; tid += blockDim.x) {
      for (size_t i = start; i < end; ++i) {
        in_grad[item_dim * i + tid] = out_grad[tid];
      }
    }
  }
};

template <typename T>
struct SqrtPoolGradFunctor {
  HOSTDEVICE void operator()(const T* out_grad,
                             const size_t start,
                             const size_t end,
                             const size_t item_dim,
                             T* in_grad,
                             const int* index) {
    for (size_t tid = threadIdx.x; tid < item_dim; tid += blockDim.x) {
      for (size_t i = start; i < end; ++i) {
        in_grad[item_dim * i + tid] =
            out_grad[tid] / (sqrt(static_cast<T>(end - start)));
      }
    }
  }
};

template <typename T>
struct LastPoolGradFunctor {
  HOSTDEVICE void operator()(const T* out_grad,
                             const size_t start,
                             const size_t end,
                             const size_t item_dim,
                             T* in_grad,
                             const int* index) {
    for (size_t tid = threadIdx.x; tid < item_dim; tid += blockDim.x) {
      for (size_t i = start; i < end; ++i) {
        if (i == end - 1) {
          in_grad[item_dim * i + tid] = out_grad[tid];
        } else {
          in_grad[item_dim * i + tid] = static_cast<T>(0);
        }
      }
    }
  }
};

template <typename T>
struct FirstPoolGradFunctor {
  HOSTDEVICE void operator()(const T* out_grad,
                             const size_t start,
                             const size_t end,
                             const size_t item_dim,
                             T* in_grad,
                             const int* index) {
    for (size_t tid = threadIdx.x; tid < item_dim; tid += blockDim.x) {
      for (size_t i = start; i < end; ++i) {
        if (i == start) {
          in_grad[item_dim * i + tid] = out_grad[tid];
        } else {
          in_grad[item_dim * i + tid] = static_cast<T>(0);
        }
      }
    }
  }
};

template <typename T, typename Range_OP>
__global__ void sequence_pool_grad_kernel(Range_OP op,
                                          const T* out_grad,
                                          const size_t* lod,
                                          const size_t lod_size,
                                          const size_t item_dim,
                                          T* in_grad,
                                          const int* index) {
  int bid = blockIdx.x;
  if (bid >= lod_size - 1) return;
  size_t start = lod[bid];
  size_t end = lod[bid + 1];
  const int* index_offset = nullptr;
  if (index != nullptr) {
    index_offset = &index[bid * item_dim];
  }
  op(&out_grad[bid * item_dim], start, end, item_dim, in_grad, index_offset);
}

template <typename T>
class SequencePoolGradFunctor<GPUContext, T> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const std::string pooltype,
                  const DenseTensor& out_grad,
                  DenseTensor* in_grad,
                  /* max pool has index */
                  const DenseTensor* index = nullptr) {
    auto lod_level = in_grad->lod().size();
    auto& lod = in_grad->lod()[lod_level - 1];
    const size_t item_dim = in_grad->numel() / in_grad->dims()[0];
    dim3 threads(1024, 1);
    dim3 grid(std::max(static_cast<int>(lod.size()) - 1, 1), 1);
    phi::MixVector<size_t> mix_vector(&lod);
    if (pooltype == "MAX") {
      sequence_pool_grad_kernel<T, MaxPoolGradFunctor<T>>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              MaxPoolGradFunctor<T>(),
              out_grad.data<T>(),
              mix_vector.CUDAData(dev_ctx.GetPlace()),
              lod.size(),
              item_dim,
              dev_ctx.template Alloc<T>(in_grad),
              index->data<int>());
    } else if (pooltype == "AVERAGE") {
      sequence_pool_grad_kernel<T, AvgPoolGradFunctor<T>>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              AvgPoolGradFunctor<T>(),
              out_grad.data<T>(),
              mix_vector.CUDAData(dev_ctx.GetPlace()),
              lod.size(),
              item_dim,
              dev_ctx.template Alloc<T>(in_grad),
              nullptr);
    } else if (pooltype == "SUM") {
      sequence_pool_grad_kernel<T, SumPoolGradFunctor<T>>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              SumPoolGradFunctor<T>(),
              out_grad.data<T>(),
              mix_vector.CUDAData(dev_ctx.GetPlace()),
              lod.size(),
              item_dim,
              dev_ctx.template Alloc<T>(in_grad),
              nullptr);
    } else if (pooltype == "SQRT") {
      sequence_pool_grad_kernel<T, SqrtPoolGradFunctor<T>>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              SqrtPoolGradFunctor<T>(),
              out_grad.data<T>(),
              mix_vector.CUDAData(dev_ctx.GetPlace()),
              lod.size(),
              item_dim,
              dev_ctx.template Alloc<T>(in_grad),
              nullptr);
    } else if (pooltype == "LAST") {
      sequence_pool_grad_kernel<T, LastPoolGradFunctor<T>>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              LastPoolGradFunctor<T>(),
              out_grad.data<T>(),
              mix_vector.CUDAData(dev_ctx.GetPlace()),
              lod.size(),
              item_dim,
              dev_ctx.template Alloc<T>(in_grad),
              nullptr);
    } else if (pooltype == "FIRST") {
      sequence_pool_grad_kernel<T, FirstPoolGradFunctor<T>>
          <<<grid, threads, 0, dev_ctx.stream()>>>(
              FirstPoolGradFunctor<T>(),
              out_grad.data<T>(),
              mix_vector.CUDAData(dev_ctx.GetPlace()),
              lod.size(),
              item_dim,
              dev_ctx.template Alloc<T>(in_grad),
              nullptr);

    } else {
      PADDLE_THROW(errors::InvalidArgument(
          "unsupported pooling pooltype: %s. Only support \"MAX\", "
          "\"AVERAGE\", \"SUM\", \"SQRT\", \"LAST\" and \"FIRST\"",
          pooltype));
    }
  }
};

// sequence pooling
template class SequencePoolFunctor<GPUContext, float>;
template class SequencePoolFunctor<GPUContext, double>;
template class PADDLE_API SequencePoolGradFunctor<GPUContext, float>;
template class SequencePoolGradFunctor<GPUContext, double>;

}  // namespace funcs
}  // namespace phi