//   Copyright (c) 2019 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 <string>

#include "paddle/fluid/framework/mixed_vector.h"
#include "paddle/fluid/operators/fused/fused_seqpool_cvm_op.h"
#include "paddle/fluid/platform/device/gpu/gpu_info.h"
#include "paddle/fluid/platform/device/gpu/gpu_launch_config.h"

namespace paddle {
namespace operators {

template <typename T>
using Vector = framework::Vector<T>;

#define CUDA_KERNEL_LOOP(i, n)                                  \
  for (auto i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); \
       i += blockDim.x * gridDim.x)

// normal
template <typename T>
__global__ void FusedSeqpoolKernelNormal(const size_t N,
                                         T **input_values,
                                         T **seqpool_output_values,
                                         size_t **lods_values,
                                         const int batch_size,
                                         const int embedding_size,
                                         const float pad_value) {
  CUDA_KERNEL_LOOP(i, N) {
    int key = i / embedding_size;
    int offset = i % embedding_size;
    int x = key / batch_size;  // slot id
    int y = key % batch_size;  // ins id
    auto &start = *(lods_values[x] + y);
    auto &end = *(lods_values[x] + y + 1);

    T val = static_cast<T>(pad_value);
    for (auto k = start; k < end; ++k) {
      val += *(input_values[x] + k * embedding_size + offset);
    }
    *(seqpool_output_values[x] + y * embedding_size + offset) = val;
  }
}

// join need show click input
template <typename T>
__global__ void FusedCVMKernelWithCVM(const size_t N,
                                      T **output_values,
                                      T **seqpool_output_values,
                                      const int batch_size,
                                      const int embedding_size,
                                      const int cvm_offset) {
  CUDA_KERNEL_LOOP(i, N) {
    int key = i / embedding_size;
    int offset = i % embedding_size;
    int x = key / batch_size;  // slot id
    int y = key % batch_size;  // ins id
    if (offset == 0) {         // show
      *(output_values[x] + y * embedding_size) =
          log(*(seqpool_output_values[x] + y * embedding_size) + 1);
    } else if (offset == 1) {  // click
      *(output_values[x] + y * embedding_size + offset) =
          log(*(seqpool_output_values[x] + y * embedding_size + 1) + 1) -
          log(*(seqpool_output_values[x] + y * embedding_size) + 1);
    } else {
      *(output_values[x] + y * embedding_size + offset) =
          *(seqpool_output_values[x] + y * embedding_size + offset);
    }
  }
}

// update not need show click input
template <typename T>
__global__ void FusedCVMKernelNoCVM(const size_t N,
                                    T **output_values,
                                    T **seqpool_output_values,
                                    const int batch_size,
                                    const int no_cvm_embedding_size,
                                    const int cvm_offset) {
  CUDA_KERNEL_LOOP(i, N) {
    int key = i / no_cvm_embedding_size;
    int offset = i % no_cvm_embedding_size;
    int x = key / batch_size;  // slot id
    int y = key % batch_size;  // ins id
    // no cvm
    *(output_values[x] + y * no_cvm_embedding_size + offset) =
        *(seqpool_output_values[x] + y * (no_cvm_embedding_size + cvm_offset) +
          offset + cvm_offset);
  }
}

template <typename T>
void FusedSeqpoolCVM(const framework::ExecutionContext
                         &ctx,  // const paddle::platform::Place &place,
                     const std::vector<const T *> &input_data,
                     const std::vector<T *> &output_data,
                     const std::vector<T *> &seqpool_output_data,
                     std::vector<const size_t *> lods,
                     const int batch_size,
                     const int slot_num,
                     const int embedding_size,
                     const float padding_value,
                     const bool use_cvm,
                     const int cvm_offset) {
  auto stream =
      ctx.template device_context<platform::CUDADeviceContext>().stream();
  auto &dev_ctx = ctx.template device_context<platform::CUDADeviceContext>();
  size_t total_ptr_len = input_data.size() + output_data.size() +
                         seqpool_output_data.size() + lods.size();
  auto temp_ptr =
      memory::AllocShared(ctx.GetPlace(), total_ptr_len * sizeof(void *));
  void *ptr = temp_ptr->ptr();

#ifdef PADDLE_WITH_HIP
  T **gpu_input_values = reinterpret_cast<T **>(temp_ptr->ptr());
  platform::GpuMemcpyAsync(gpu_input_values,
                           input_data.data(),
                           input_data.size() * sizeof(T *),
                           hipMemcpyHostToDevice,
                           stream);
  T **gpu_output_values =
      reinterpret_cast<T **>(&gpu_input_values[input_data.size()]);
  platform::GpuMemcpyAsync(gpu_output_values,
                           output_data.data(),
                           output_data.size() * sizeof(T *),
                           hipMemcpyHostToDevice,
                           stream);
  T **gpu_seqpool_output_values =
      reinterpret_cast<T **>(&gpu_output_values[output_data.size()]);
  platform::GpuMemcpyAsync(gpu_seqpool_output_values,
                           seqpool_output_data.data(),
                           seqpool_output_data.size() * sizeof(T *),
                           hipMemcpyHostToDevice,
                           stream);
  size_t **lods_values = reinterpret_cast<size_t **>(
      &gpu_seqpool_output_values[seqpool_output_data.size()]);
  platform::GpuMemcpyAsync(lods_values,
                           lods.data(),
                           lods.size() * sizeof(size_t *),
                           hipMemcpyHostToDevice,
                           stream);
#else
  T **gpu_input_values = reinterpret_cast<T **>(temp_ptr->ptr());
  platform::GpuMemcpyAsync(gpu_input_values,
                           input_data.data(),
                           input_data.size() * sizeof(T *),
                           cudaMemcpyHostToDevice,
                           stream);
  T **gpu_output_values =
      reinterpret_cast<T **>(&gpu_input_values[input_data.size()]);
  platform::GpuMemcpyAsync(gpu_output_values,
                           output_data.data(),
                           output_data.size() * sizeof(T *),
                           cudaMemcpyHostToDevice,
                           stream);
  T **gpu_seqpool_output_values =
      reinterpret_cast<T **>(&gpu_output_values[output_data.size()]);
  platform::GpuMemcpyAsync(gpu_seqpool_output_values,
                           seqpool_output_data.data(),
                           seqpool_output_data.size() * sizeof(T *),
                           cudaMemcpyHostToDevice,
                           stream);
  size_t **lods_values = reinterpret_cast<size_t **>(
      &gpu_seqpool_output_values[seqpool_output_data.size()]);
  platform::GpuMemcpyAsync(lods_values,
                           lods.data(),
                           lods.size() * sizeof(size_t *),
                           cudaMemcpyHostToDevice,
                           stream);
#endif

  size_t N = static_cast<size_t>(batch_size * slot_num * embedding_size);
  platform::GpuLaunchConfig config = platform::GetGpuLaunchConfig1D(dev_ctx, N);
  // first sum pool
  FusedSeqpoolKernelNormal<<<config.block_per_grid.x,
                             config.thread_per_block.x,
                             0,
                             stream>>>(N,
                                       gpu_input_values,
                                       gpu_seqpool_output_values,
                                       lods_values,
                                       batch_size,
                                       embedding_size,
                                       padding_value);
  // second log
  if (use_cvm) {
    FusedCVMKernelWithCVM<<<config.block_per_grid.x,
                            config.thread_per_block.x,
                            0,
                            stream>>>(N,
                                      gpu_output_values,
                                      gpu_seqpool_output_values,
                                      batch_size,
                                      embedding_size,
                                      cvm_offset);
  } else {
    // not need show click input
    N = static_cast<size_t>(batch_size * slot_num *
                            (embedding_size - cvm_offset));
    platform::GpuLaunchConfig config =
        platform::GetGpuLaunchConfig1D(dev_ctx, N);
    FusedCVMKernelNoCVM<<<config.block_per_grid.x,
                          config.thread_per_block.x,
                          0,
                          stream>>>(N,
                                    gpu_output_values,
                                    gpu_seqpool_output_values,
                                    batch_size,
                                    (embedding_size - cvm_offset),
                                    cvm_offset);
  }
}

// join grad
template <typename T>
__global__ void FusedSeqpoolCVMGradKernelWithCVM(const size_t N,
                                                 T **out_grads_values,
                                                 T **in_grads_values,
                                                 T **cvm_values,
                                                 size_t **lods_values,
                                                 const int batch_size,
                                                 const int embedding_size,
                                                 const int cvm_offset) {
  CUDA_KERNEL_LOOP(i, N) {
    int key = i / embedding_size;
    int offset = i % embedding_size;  // embedx offset
    int x = key / batch_size;         // slot id
    int y = key % batch_size;         // ins id

    T &val = (offset < cvm_offset)
                 ? *(cvm_values[x] + y * cvm_offset + offset)
                 : *(out_grads_values[x] + y * embedding_size + offset);

    auto &start = *(lods_values[x] + y);
    auto &end = *(lods_values[x] + y + 1);
    for (auto k = start; k < end; ++k) {
      *(in_grads_values[x] + k * embedding_size + offset) = val;
    }
  }
}

// join only show not has click
template <typename T>
__global__ void FusedSeqpoolCVMGradKernelWithShow(const size_t N,
                                                  T **out_grads_values,
                                                  T **in_grads_values,
                                                  T **cvm_values,
                                                  size_t **lods_values,
                                                  const int batch_size,
                                                  const int embedding_size,
                                                  const int cvm_offset) {
  CUDA_KERNEL_LOOP(i, N) {
    int key = i / embedding_size;
    int offset = i % embedding_size;  // embedx offset
    int x = key / batch_size;         // slot id
    int y = key % batch_size;         // ins id

    T &val =
        (offset < cvm_offset)
            ? *(cvm_values[x] + y * cvm_offset + offset)
            : *(out_grads_values[x] + y * (embedding_size - 1) + offset - 1);

    auto &start = *(lods_values[x] + y);
    auto &end = *(lods_values[x] + y + 1);
    for (auto k = start; k < end; ++k) {
      *(in_grads_values[x] + k * embedding_size + offset) = val;
    }
  }
}

// update grad
template <typename T>
__global__ void FusedSeqpoolCVMGradKernelNoCVM(const size_t N,
                                               T **out_grads_values,
                                               T **in_grads_values,
                                               T **cvm_values,
                                               size_t **lods_values,
                                               const int batch_size,
                                               const int embedding_size,
                                               const int cvm_offset) {
  CUDA_KERNEL_LOOP(i, N) {
    int key = i / embedding_size;
    int offset = i % embedding_size;  // embedx offset
    int x = key / batch_size;         // slot id
    int y = key % batch_size;         // ins id

    T &val = (offset < cvm_offset)
                 ? *(cvm_values[x] + y * cvm_offset + offset)
                 : *(out_grads_values[x] + y * (embedding_size - cvm_offset) +
                     offset - cvm_offset);

    auto &start = *(lods_values[x] + y);
    auto &end = *(lods_values[x] + y + 1);
    for (auto k = start; k < end; ++k) {
      *(in_grads_values[x] + k * embedding_size + offset) = val;
    }
  }
}

template <typename T>
void FusedSeqpoolCVMGrad(const framework::ExecutionContext &ctx,
                         const std::vector<const T *> &out_grads_data,
                         const std::vector<T *> &in_grads_data,
                         const std::vector<const T *> &cvm_data,
                         const std::vector<const size_t *> &lods,
                         const int batch_size,
                         const int slot_num,
                         const int embedding_size,
                         const bool use_cvm,
                         const int cvm_offset) {
  auto stream =
      ctx.template device_context<platform::CUDADeviceContext>().stream();
  auto &dev_ctx = ctx.template device_context<platform::CUDADeviceContext>();
  size_t total_ptr_len = out_grads_data.size() + in_grads_data.size() +
                         cvm_data.size() + lods.size();
  auto temp_ptr =
      memory::AllocShared(ctx.GetPlace(), total_ptr_len * sizeof(void *));
#ifdef PADDLE_WITH_HIP
  T **gpu_out_grads_values = reinterpret_cast<T **>(temp_ptr->ptr());
  platform::GpuMemcpyAsync(gpu_out_grads_values,
                           out_grads_data.data(),
                           out_grads_data.size() * sizeof(T *),
                           hipMemcpyHostToDevice,
                           stream);

  T **gpu_in_grads_values =
      reinterpret_cast<T **>(&gpu_out_grads_values[out_grads_data.size()]);
  platform::GpuMemcpyAsync(gpu_in_grads_values,
                           in_grads_data.data(),
                           in_grads_data.size() * sizeof(T *),
                           hipMemcpyHostToDevice,
                           stream);

  T **gpu_cvm_values =
      reinterpret_cast<T **>(&gpu_in_grads_values[in_grads_data.size()]);
  platform::GpuMemcpyAsync(gpu_cvm_values,
                           cvm_data.data(),
                           cvm_data.size() * sizeof(T *),
                           hipMemcpyHostToDevice,
                           stream);

  size_t **lods_values =
      reinterpret_cast<size_t **>(&gpu_cvm_values[cvm_data.size()]);
  platform::GpuMemcpyAsync(lods_values,
                           lods.data(),
                           lods.size() * sizeof(size_t *),
                           hipMemcpyHostToDevice,
                           stream);
#else
  T **gpu_out_grads_values = reinterpret_cast<T **>(temp_ptr->ptr());
  platform::GpuMemcpyAsync(gpu_out_grads_values,
                           out_grads_data.data(),
                           out_grads_data.size() * sizeof(T *),
                           cudaMemcpyHostToDevice,
                           stream);

  T **gpu_in_grads_values =
      reinterpret_cast<T **>(&gpu_out_grads_values[out_grads_data.size()]);
  platform::GpuMemcpyAsync(gpu_in_grads_values,
                           in_grads_data.data(),
                           in_grads_data.size() * sizeof(T *),
                           cudaMemcpyHostToDevice,
                           stream);

  T **gpu_cvm_values =
      reinterpret_cast<T **>(&gpu_in_grads_values[in_grads_data.size()]);
  platform::GpuMemcpyAsync(gpu_cvm_values,
                           cvm_data.data(),
                           cvm_data.size() * sizeof(T *),
                           cudaMemcpyHostToDevice,
                           stream);

  size_t **lods_values =
      reinterpret_cast<size_t **>(&gpu_cvm_values[cvm_data.size()]);
  platform::GpuMemcpyAsync(lods_values,
                           lods.data(),
                           lods.size() * sizeof(size_t *),
                           cudaMemcpyHostToDevice,
                           stream);
#endif

  size_t N = static_cast<size_t>(batch_size * slot_num * embedding_size);
  auto config = platform::GetGpuLaunchConfig1D(dev_ctx, N);
  if (use_cvm) {
    // join grad
    FusedSeqpoolCVMGradKernelWithCVM<<<config.block_per_grid.x,
                                       config.thread_per_block.x,
                                       0,
                                       stream>>>(N,
                                                 gpu_out_grads_values,
                                                 gpu_in_grads_values,
                                                 gpu_cvm_values,
                                                 lods_values,
                                                 batch_size,
                                                 embedding_size,
                                                 cvm_offset);
  } else {
    // update grad
    FusedSeqpoolCVMGradKernelNoCVM<<<config.block_per_grid.x,
                                     config.thread_per_block.x,
                                     0,
                                     stream>>>(N,
                                               gpu_out_grads_values,
                                               gpu_in_grads_values,
                                               gpu_cvm_values,
                                               lods_values,
                                               batch_size,
                                               embedding_size,
                                               cvm_offset);
  }
}

template <typename T>
class FusedSeqpoolCVMCUDAKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext &ctx) const override {
    auto inputs = ctx.MultiInput<LoDTensor>("X");
    auto outputs = ctx.MultiOutput<framework::Tensor>("Out");

    const auto slot_size = inputs.size();
    std::vector<const float *> input_data(slot_size);
    std::vector<const size_t *> lods_data(slot_size);
    std::vector<T *> output_data(slot_size);

    std::vector<LoDTensor> seqpool_outputs(slot_size);
    std::vector<T *> seqpool_output_data(slot_size);

    auto padding_value = ctx.Attr<float>("pad_value");
    auto use_cvm = ctx.Attr<bool>("use_cvm");
    const int cvm_offset = ctx.Attr<int>("cvm_offset");

    int embedding_size = inputs[0]->numel() / inputs[0]->dims()[0];
    int batch_size = -1;
    std::vector<paddle::framework::MixVector<size_t> *> mix_lods_v(slot_size);

    for (size_t i = 0; i < slot_size; ++i) {
      const auto *input = inputs[i];

      Vector<size_t> lods;
      if (input->lod().size() != 0) {
        auto lod = input->lod();
        lods = lod[0];
      } else {
        lods.push_back(0);
        for (int i = 0; i < input->dims()[0]; i++) {
          lods.push_back(i + 1);
        }
      }
      int cur_batch_size =
          input->lod().size() ? input->lod()[0].size() - 1 : input->dims()[0];
      if (batch_size == -1) {
        batch_size = cur_batch_size;
      } else {
        PADDLE_ENFORCE_EQ(batch_size,
                          cur_batch_size,
                          platform::errors::PreconditionNotMet(
                              "The batch size of all input should be same, "
                              "please cheack, last batchsize is %d, current "
                              "batchsize is %d",
                              batch_size,
                              cur_batch_size));
      }
      input_data[i] = reinterpret_cast<const T *>(input->data<T>());

      auto *output = outputs[i];
      if (use_cvm) {
        output->Resize({batch_size, embedding_size});
      } else {
        output->Resize({batch_size, embedding_size - cvm_offset});
      }
      output_data[i] =
          reinterpret_cast<T *>(output->mutable_data<T>(ctx.GetPlace()));
      mix_lods_v[i] = new paddle::framework::MixVector<size_t>(&lods);
      lods_data[i] = mix_lods_v[i]->CUDAData(ctx.GetPlace());
      seqpool_output_data[i] =
          reinterpret_cast<T *>(seqpool_outputs[i].mutable_data<T>(
              {batch_size, embedding_size}, ctx.GetPlace()));
    }

    FusedSeqpoolCVM(ctx,
                    input_data,
                    output_data,
                    seqpool_output_data,
                    lods_data,
                    batch_size,
                    slot_size,
                    embedding_size,
                    padding_value,
                    use_cvm,
                    cvm_offset);

    for (int i = 0; i < slot_size; i++) {
      delete mix_lods_v[i];
    }
  }
};

template <typename T>
class FusedSeqpoolCVMGradCUDAKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext &ctx) const override {
    auto out_grads = ctx.MultiInput<LoDTensor>(framework::GradVarName("Out"));
    auto in_grads = ctx.MultiOutput<LoDTensor>(framework::GradVarName("X"));
    auto *cvm = ctx.Input<LoDTensor>("CVM");

    std::string pooltype = ctx.Attr<std::string>("pooltype");
    auto use_cvm = ctx.Attr<bool>("use_cvm");
    const int cvm_offset = ctx.Attr<int>("cvm_offset");

    const auto slot_size = in_grads.size();
    std::vector<const T *> out_grads_data(slot_size);
    std::vector<T *> in_grads_data(slot_size);
    std::vector<const T *> cvm_data(slot_size);
    std::vector<const size_t *> lods_data(slot_size);

    int embedding_size = in_grads[0]->numel() / in_grads[0]->dims()[0];
    int batch_size = -1;
    std::vector<paddle::framework::MixVector<size_t> *> mix_lods_v(slot_size);

    for (size_t i = 0; i < slot_size; ++i) {
      auto *in_grad = in_grads[i];

      Vector<size_t> lods;
      if (in_grad->lod().size() != 0) {
        auto lod = in_grad->lod();
        lods = lod[0];
      } else {
        lods.push_back(0);
        for (int i = 0; i < in_grad->dims()[0]; i++) {
          lods.push_back(i + 1);
        }
      }

      int cur_batch_size = in_grad->lod().size() ? in_grad->lod()[0].size() - 1
                                                 : in_grad->dims()[0];
      if (batch_size == -1) {
        batch_size = cur_batch_size;
      } else {
        PADDLE_ENFORCE_EQ(batch_size,
                          cur_batch_size,
                          platform::errors::PreconditionNotMet(
                              "The batch size of all input should be same, "
                              "please cheack, last batchsize is %d, current "
                              "batchsize is %d",
                              batch_size,
                              cur_batch_size));
      }

      auto *out_grad = out_grads[i];
      out_grads_data[i] = reinterpret_cast<const T *>(out_grad->data<T>());

      in_grads_data[i] =
          reinterpret_cast<T *>(in_grad->mutable_data<T>(ctx.GetPlace()));
      mix_lods_v[i] = new paddle::framework::MixVector<size_t>(&lods);
      lods_data[i] = mix_lods_v[i]->CUDAData(ctx.GetPlace());
      cvm_data[i] = reinterpret_cast<const T *>(cvm->data<T>());
    }
    FusedSeqpoolCVMGrad(ctx,
                        out_grads_data,
                        in_grads_data,
                        cvm_data,
                        lods_data,
                        batch_size,
                        slot_size,
                        embedding_size,
                        use_cvm,
                        cvm_offset);

    for (int i = 0; i < slot_size; i++) {
      delete mix_lods_v[i];
    }
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
REGISTER_OP_CUDA_KERNEL(fused_seqpool_cvm,
                        ops::FusedSeqpoolCVMCUDAKernel<float>);

REGISTER_OP_CUDA_KERNEL(fused_seqpool_cvm_grad,
                        ops::FusedSeqpoolCVMGradCUDAKernel<float>);