// 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.

#pragma once

#include "paddle/phi/api/include/tensor.h"
#include "paddle/phi/backends/device_guard.h"
#include "paddle/phi/backends/device_manager.h"
#include "paddle/phi/core/device_context.h"
#include "paddle/phi/kernels/funcs/concat_and_split_functor.h"

namespace paddle {
namespace pybind {

template <typename Context, typename T>
struct ConcatDenseTensor {
  void operator()(const Context &dev_ctx,
                  const std::vector<DenseTensor> &in,
                  DenseTensor *out,
                  int axis = 0) {
    phi::funcs::ConcatFunctor<Context, T> concat_functor;
    concat_functor(dev_ctx, in, axis, out);
  }
};

template <typename Context, typename T>
struct SplitDenseTensor {
  void operator()(const Context &dev_ctx,
                  const DenseTensor &in,
                  std::vector<DenseTensor *> *out,
                  int axis = 0) {
    std::vector<const DenseTensor *> shape_refer;
    shape_refer.reserve(out->size());
    for (auto *p_tensor : *out) {
      shape_refer.emplace_back(p_tensor);
    }
    phi::funcs::SplitFunctor<Context, T> split_functor;
    split_functor(dev_ctx, in, shape_refer, axis, out);
  }
};

#ifdef PADDLE_WITH_CUSTOM_DEVICE
template <typename T>
struct ConcatDenseTensor<phi::CustomContext, T> {
  void operator()(const phi::CustomContext &dev_ctx,
                  const std::vector<DenseTensor> &in,
                  DenseTensor *out,
                  int axis UNUSED = 0) {
    VLOG(10) << "ConcatDenseTensor: " << in.size();
    auto kernel_result =
        phi::KernelFactory::Instance().SelectKernelOrThrowError(
            "concat",
            phi::KernelKey(phi::TransToPhiBackend(dev_ctx.GetPlace()),
                           phi::DataLayout::ALL_LAYOUT,
                           phi::CppTypeToDataType<T>::Type()));
    const auto &kernel = kernel_result.kernel;
    using kernel_signature = void (*)(const phi::DeviceContext &,
                                      const std::vector<const DenseTensor *> &,
                                      const phi::Scalar &,
                                      DenseTensor *);
    auto *kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
    std::vector<const DenseTensor *> inputs;
    (*kernel_fn)(dev_ctx, inputs, phi::Scalar(0), out);
  }
};

template <typename T>
struct SplitDenseTensor<phi::CustomContext, T> {
  void operator()(const phi::CustomContext &dev_ctx,
                  const DenseTensor &in,
                  std::vector<DenseTensor *> *out,
                  int axis UNUSED = 0) {
    VLOG(10) << "SplitDenseTensor: " << out->size();
    auto kernel_result =
        phi::KernelFactory::Instance().SelectKernelOrThrowError(
            "split_with_num",
            phi::KernelKey(phi::TransToPhiBackend(dev_ctx.GetPlace()),
                           phi::DataLayout::ALL_LAYOUT,
                           phi::CppTypeToDataType<T>::Type()));
    const auto &kernel = kernel_result.kernel;
    using kernel_signature = void (*)(const phi::DeviceContext &,
                                      const DenseTensor &,
                                      int,
                                      const phi::Scalar &,
                                      std::vector<DenseTensor *>);
    auto *kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();

    auto in_dims = common::vectorize(in.dims());
    auto origin_out_dims = common::vectorize(out->at(0)->dims());
    for (auto *tensor : *out) {
      if (origin_out_dims.size() != in_dims.size()) {
        std::vector<int> new_dims({1});
        new_dims.insert(
            new_dims.end(), origin_out_dims.begin(), origin_out_dims.end());
        tensor->Resize(common::make_ddim(new_dims));
      }
    }
    (*kernel_fn)(dev_ctx, in, out->size(), phi::Scalar(0), *out);
    for (auto *tensor : *out) {
      auto tensor_dims = common::vectorize(tensor->dims());
      if (tensor_dims.size() != origin_out_dims.size()) {
        tensor->Resize(common::make_ddim(origin_out_dims));
      }
    }
  }
};
#endif

template <typename Context>
void ConcatDenseTensorWithType(const Context &dev_ctx,
                               const std::vector<DenseTensor> &t_list,
                               DenseTensor *p_out,
                               DataType type) {
  switch (type) {
    case DataType::BOOL:
      ConcatDenseTensor<Context, bool>()(dev_ctx, t_list, p_out);
      break;
    case DataType::UINT8:
      ConcatDenseTensor<Context, uint8_t>()(dev_ctx, t_list, p_out);
      break;
    case DataType::INT8:
      ConcatDenseTensor<Context, int8_t>()(dev_ctx, t_list, p_out);
      break;
    case DataType::INT32:
      ConcatDenseTensor<Context, int32_t>()(dev_ctx, t_list, p_out);
      break;
    case DataType::INT64:
      ConcatDenseTensor<Context, int64_t>()(dev_ctx, t_list, p_out);
      break;
    case DataType::FLOAT16:
      ConcatDenseTensor<Context, phi::float16>()(dev_ctx, t_list, p_out);
      break;
    case DataType::BFLOAT16:
      ConcatDenseTensor<Context, phi::bfloat16>()(dev_ctx, t_list, p_out);
      break;
    case DataType::FLOAT32:
      ConcatDenseTensor<Context, float>()(dev_ctx, t_list, p_out);
      break;
    case DataType::FLOAT64:
      ConcatDenseTensor<Context, double>()(dev_ctx, t_list, p_out);
      break;
    default:
      PADDLE_THROW(common::errors::Unimplemented(
          "Data type (%s) is not supported when it concats tensors.", type));
  }
}

#ifdef PADDLE_WITH_XPU
template <>
void ConcatDenseTensorWithType(const phi::XPUContext &dev_ctx,
                               const std::vector<DenseTensor> &t_list,
                               DenseTensor *p_out,
                               DataType type) {
  switch (type) {
    case DataType::FLOAT16:
      ConcatDenseTensor<phi::XPUContext, phi::float16>()(
          dev_ctx, t_list, p_out);
      break;
    case DataType::BFLOAT16:
      ConcatDenseTensor<phi::XPUContext, phi::bfloat16>()(
          dev_ctx, t_list, p_out);
      break;
    case DataType::FLOAT32:
      ConcatDenseTensor<phi::XPUContext, float>()(dev_ctx, t_list, p_out);
      break;
    case DataType::INT32:
      ConcatDenseTensor<phi::XPUContext, int32_t>()(dev_ctx, t_list, p_out);
      break;
    case DataType::INT64:
      ConcatDenseTensor<phi::XPUContext, int64_t>()(dev_ctx, t_list, p_out);
      break;
    case DataType::UINT8:
      ConcatDenseTensor<phi::XPUContext, uint8_t>()(dev_ctx, t_list, p_out);
      break;
    default:
      PADDLE_THROW(common::errors::Unimplemented(
          "Data type (%s) is not supported when it concats tensors.", type));
  }
}
#endif

template <typename Context>
void SplitDenseTensorWithType(const Context &dev_ctx,
                              const DenseTensor &t_in,
                              std::vector<DenseTensor *> *p_list,
                              DataType type) {
  switch (type) {
    case DataType::BOOL:
      SplitDenseTensor<Context, bool>()(dev_ctx, t_in, p_list);
      break;
    case DataType::UINT8:
      SplitDenseTensor<Context, uint8_t>()(dev_ctx, t_in, p_list);
      break;
    case DataType::INT8:
      SplitDenseTensor<Context, int8_t>()(dev_ctx, t_in, p_list);
      break;
    case DataType::FLOAT8_E4M3FN:
      SplitDenseTensor<Context, phi::float8_e4m3fn>()(dev_ctx, t_in, p_list);
      break;
    case DataType::FLOAT8_E5M2:
      SplitDenseTensor<Context, phi::float8_e5m2>()(dev_ctx, t_in, p_list);
      break;
    case DataType::INT32:
      SplitDenseTensor<Context, int32_t>()(dev_ctx, t_in, p_list);
      break;
    case DataType::INT64:
      SplitDenseTensor<Context, int64_t>()(dev_ctx, t_in, p_list);
      break;
    case DataType::FLOAT16:
      SplitDenseTensor<Context, phi::float16>()(dev_ctx, t_in, p_list);
      break;
    case DataType::BFLOAT16:
      SplitDenseTensor<Context, phi::bfloat16>()(dev_ctx, t_in, p_list);
      break;
    case DataType::FLOAT32:
      SplitDenseTensor<Context, float>()(dev_ctx, t_in, p_list);
      break;
    case DataType::FLOAT64:
      SplitDenseTensor<Context, double>()(dev_ctx, t_in, p_list);
      break;
    default:
      PADDLE_THROW(common::errors::Unimplemented(
          "Data type (%s) is not supported when it splits tensors.", type));
  }
}

#ifdef PADDLE_WITH_XPU
template <>
void SplitDenseTensorWithType(const phi::XPUContext &dev_ctx,
                              const DenseTensor &t_in,
                              std::vector<DenseTensor *> *p_list,
                              DataType type) {
  switch (type) {
    case DataType::FLOAT16:
      SplitDenseTensor<phi::XPUContext, phi::float16>()(dev_ctx, t_in, p_list);
      break;
    case DataType::BFLOAT16:
      SplitDenseTensor<phi::XPUContext, phi::bfloat16>()(dev_ctx, t_in, p_list);
      break;
    case DataType::FLOAT32:
      SplitDenseTensor<phi::XPUContext, float>()(dev_ctx, t_in, p_list);
      break;
    case DataType::INT32:
      SplitDenseTensor<phi::XPUContext, int32_t>()(dev_ctx, t_in, p_list);
      break;
    case DataType::INT64:
      SplitDenseTensor<phi::XPUContext, int64_t>()(dev_ctx, t_in, p_list);
      break;
    case DataType::UINT8:
      SplitDenseTensor<phi::XPUContext, uint8_t>()(dev_ctx, t_in, p_list);
      break;
    default:
      PADDLE_THROW(common::errors::Unimplemented(
          "Data type (%s) is not supported when it splits tensors.", type));
  }
}
#endif

void ConcatTensor(const phi::DeviceContext &dev_ctx,
                  const std::vector<DenseTensor> &tensor_list,
                  const Tensor *tensor) {
  auto *dense_tensor =
      std::dynamic_pointer_cast<DenseTensor>(tensor->impl()).get();

  const auto &place = dev_ctx.GetPlace();
  if (phi::is_gpu_place(place)) {
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
    ConcatDenseTensorWithType(static_cast<const phi::GPUContext &>(dev_ctx),
                              tensor_list,
                              dense_tensor,
                              tensor->dtype());
#else
    PADDLE_THROW(common::errors::PermissionDenied(
        "Paddle can't concat tensor since it's not support GPU, please "
        "recompile or reinstall Paddle with GPU support."));
#endif
  } else if (phi::is_xpu_place(place)) {
#ifdef PADDLE_WITH_XPU
    ConcatDenseTensorWithType(static_cast<const phi::XPUContext &>(dev_ctx),
                              tensor_list,
                              dense_tensor,
                              tensor->dtype());
#else
    PADDLE_THROW(common::errors::PermissionDenied(
        "Paddle can't concat tensor since it's not support XPU, please "
        "recompile or reinstall Paddle with XPU support."));
#endif
  } else if (phi::is_custom_place(place)) {
#ifdef PADDLE_WITH_CUSTOM_DEVICE
    ConcatDenseTensorWithType(static_cast<const phi::CustomContext &>(dev_ctx),
                              tensor_list,
                              dense_tensor,
                              tensor->dtype());
#else
    PADDLE_THROW(common::errors::PermissionDenied(
        "Paddle can't concat tensor since it's not compiled with "
        "CUSTOM_DEVICE, please recompile or reinstall Paddle with "
        "CUSTOM_DEVICE support."));
#endif
  } else if (phi::is_cpu_place(place)) {
    ConcatDenseTensorWithType(static_cast<const phi::CPUContext &>(dev_ctx),
                              tensor_list,
                              dense_tensor,
                              tensor->dtype());
  } else {
    PADDLE_THROW(common::errors::Unimplemented(
        "Concat tensor not supported on place (%s)", place));
  }
}

void SplitTensor(const phi::DeviceContext &dev_ctx,
                 const DenseTensor &tensor,
                 const std::vector<Tensor> *tensor_list) {
  std::vector<DenseTensor *> dense_list;
  for (auto &tensor : *tensor_list) {
    auto *p_tensor =
        std::dynamic_pointer_cast<DenseTensor>(tensor.impl()).get();
    dense_list.emplace_back(p_tensor);
  }

  const auto &place = dev_ctx.GetPlace();
  if (phi::is_gpu_place(place)) {
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
    SplitDenseTensorWithType(static_cast<const phi::GPUContext &>(dev_ctx),
                             tensor,
                             &dense_list,
                             tensor.dtype());
#else
    PADDLE_THROW(common::errors::PermissionDenied(
        "Paddle can't split tensor since it's not support GPU, please "
        "recompile or reinstall Paddle with GPU support."));
#endif
  } else if (phi::is_xpu_place(place)) {
#ifdef PADDLE_WITH_XPU
    SplitDenseTensorWithType(static_cast<const phi::XPUContext &>(dev_ctx),
                             tensor,
                             &dense_list,
                             tensor.dtype());
#else
    PADDLE_THROW(common::errors::PermissionDenied(
        "Paddle can't split tensor since it's not compiled with XPU, "
        "please recompile or reinstall Paddle with XPU support."));
#endif
  } else if (phi::is_custom_place(place)) {
#ifdef PADDLE_WITH_CUSTOM_DEVICE
    SplitDenseTensorWithType(static_cast<const phi::CustomContext &>(dev_ctx),
                             tensor,
                             &dense_list,
                             tensor.dtype());
#else
    PADDLE_THROW(common::errors::PermissionDenied(
        "Paddle can't split tensor since it's not compiled with CUSTOM_DEVICE, "
        "please recompile or reinstall Paddle with CUSTOM_DEVICE support."));
#endif
  } else if (phi::is_cpu_place(place)) {
    SplitDenseTensorWithType(static_cast<const phi::CPUContext &>(dev_ctx),
                             tensor,
                             &dense_list,
                             tensor.dtype());
  } else {
    PADDLE_THROW(common::errors::Unimplemented(
        "Split tensor not supported on place (%s)", place));
  }
}

inline std::vector<int64_t> GetDefaultSplitSizes(const DenseTensor &tensor,
                                                 int world_size) {
  return std::vector<int64_t>(world_size, tensor.dims()[0] / world_size);
}

inline std::vector<DenseTensor> ToDenseTensors(
    const std::vector<Tensor> &tensors) {
  std::vector<DenseTensor> ret;
  for (auto &t : tensors) {
    ret.emplace_back(*std::dynamic_pointer_cast<DenseTensor>(t.impl()));
  }
  return ret;
}

}  //  namespace pybind
}  //  namespace paddle