/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/

#include <cstdint>
#include <vector>
#include "kernels/moe/moe.mluh"
#include "torch_ops_api.h"
#include "utils.h"

namespace tmo {
namespace torch_api {

const std::string arch_370 = "MLU370";
using GroupGemmDesc = tmo::op_desc::GroupGemmDesc;
using QuantMode = tmo::op_desc::GroupGemmDesc::QuantMode;

at::Tensor fused_moe(const at::Tensor &hidden_states,
                     const at::Tensor &gating_output,
                     const at::Tensor &w1,
                     const at::Tensor &w2,
                     const c10::optional<at::Tensor> &bias1,
                     const c10::optional<at::Tensor> &bias2,
                     const c10::optional<at::Tensor> &residual,
                     const c10::optional<at::Tensor> &input_smooth,
                     const c10::optional<at::Tensor> &act_smooth,
                     const c10::optional<at::Tensor> &w1_scale,
                     const c10::optional<at::Tensor> &w2_scale,
                     const c10::optional<at::List<int64_t>> &w1_quant_flag,
                     const c10::optional<at::List<int64_t>> &w2_quant_flag,
                     const int64_t topk,
                     const bool renormalize,
                     const bool gated,
                     const std::string &act_mode,
                     const int64_t start_expert_id,
                     const int64_t block_n,
                     const int64_t cncl_comm) {
  auto sizes_1 = hidden_states.sizes();
  auto sizes_2 = gating_output.sizes();
  auto w2_shape = w2.sizes();
  TORCH_CHECK(sizes_1.size() == sizes_2.size(),
              "hidden_states and gating_output must have the same rank.")
  TORCH_CHECK(sizes_1.size() == 2 || sizes_1.size() == 3, "hidden_states must be 2-D or 3-D.")
  TORCH_CHECK(sizes_1[0] == sizes_2[0],
              "hidden_states and gating_output must have the same batch.");
  if (sizes_1.size() == 3) {
    TORCH_CHECK(sizes_1[1] == sizes_2[1],
                "hidden_states and gating_output must have the same seq.");
  }
  if (residual.has_value()) {
    TORCH_CHECK(residual.value().sizes() == sizes_1,
                "hidden_states and residual must have the same shape.");
  }
  TORCH_CHECK(!bias1.has_value() && !bias2.has_value(), "Currently not support bias1 and bias2.")
  TORCH_CHECK(w1.dim() == 3 || w1.dim() == 1, "w1 should be 1-D or 3-D.")
  TORCH_CHECK(w2.dim() == 3 || w2.dim() == 4 || w2.dim() == 1, "w2 should be 1-D or 3-D or 4-D.")
  TORCH_CHECK((w1.dim() == 1 && w2.dim() == 1) || (w1.dim() != 1 && w2.dim() != 1),
              "w1 and w2 should be both 1-D or not both 1-D.")
  if (w1.dim() == 1) {
    TORCH_CHECK(w1_quant_flag.has_value() && w2_quant_flag.has_value(),
                "w1_quant_flag and w2_quant_flag need to exist simultaneously.");
  }
  // check contiguous
  CHECK_TENSOR_CONTIGUOUS(hidden_states)
  CHECK_TENSOR_CONTIGUOUS(gating_output)
  CHECK_TENSOR_CONTIGUOUS(w1)
  CHECK_TENSOR_CONTIGUOUS(w2)
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(act_smooth)
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(w1_scale)
  const int64_t hidden_size = hidden_states.size(-1);
  const int64_t num_expert = gating_output.size(-1);
  const int64_t expert_size = w1_quant_flag.has_value() ? w1_scale.value().size(1) : w1.size(0);
  auto hidden_states_ = hidden_states.view({-1, hidden_size});
  auto gating_output_ = gating_output.view({-1, num_expert});
  const int64_t num_token = hidden_states_.size(0);
  const int64_t num_expand_token = num_token * topk;
  const int64_t gemm1_co = w1_quant_flag.has_value() ? w1_scale.value().size(2) : w1.size(1);
  const int64_t inner_size = gemm1_co / (1 + gated);
  bool has_input_smooth = input_smooth.has_value();
  bool has_act_smooth = act_smooth.has_value();
  bool has_w1_scale = w1_scale.has_value();
  bool has_w2_scale = w2_scale.has_value();
  int opt_num = has_input_smooth + has_act_smooth + has_w1_scale + has_w2_scale;
  QuantMode quant_mode = QuantMode::noQuant;
  bool quant_grouped = false;

  bool per_token_sq = opt_num == 4;
  TORCH_CHECK(opt_num == 0 || opt_num == 4,
              "input_smooth, act_smooth, w1_scale and w2_scale must be present and absent at the "
              "same time.")
  if (per_token_sq) {
    TORCH_CHECK(input_smooth.value().dtype() == torch::kFloat32,
                "the data type of input_smooth must be float.");
    checkTensorSameAttr<TensorAttr::DTYPE>(input_smooth, act_smooth, w1_scale, w2_scale);
    CHECK_SHAPE(input_smooth.value(), expert_size, hidden_size);
    CHECK_SHAPE(act_smooth.value(), expert_size, inner_size);

    quant_grouped = w1_scale.value().dim() == 3 ? true : false;
    if (quant_grouped) {
      CHECK_SHAPE(w1_scale.value(), w1_scale.value().size(0), expert_size, gemm1_co);
      CHECK_SHAPE(w2_scale.value(), w2_scale.value().size(0), expert_size, hidden_size);
    } else {
      CHECK_SHAPE(w1_scale.value(), expert_size, gemm1_co);
      CHECK_SHAPE(w2_scale.value(), expert_size, hidden_size);
    }

    if (w1_quant_flag.has_value()) {
      quant_mode = QuantMode::W4W8;
      quant_grouped = true;
    } else {
      TORCH_CHECK(hidden_size == w1.size(-1) || hidden_size == w1.size(-1) * 2,
                  "hidden_size == w1.size(-1) || hidden_size == w1.size(-1) * 2.");
      quant_mode = hidden_size == w1.size(-1) ? QuantMode::W8 : QuantMode::W4;
    }
  }

  if (quant_mode == QuantMode::W8 || quant_mode == QuantMode::noQuant) {
    CHECK_SHAPE(w1, expert_size, gemm1_co, hidden_size);
    if (w2.dim() == 3) {
      CHECK_SHAPE(w2, expert_size, hidden_size, inner_size);
    } else {
      TORCH_CHECK(w2_shape[0] * w2_shape[2] == expert_size,
                  "w2_shape[0] * w2_shape[2] == expert_size");
      CHECK_SHAPE(w2, w2_shape[0], hidden_size, w2_shape[2], inner_size);
    }
  } else if (quant_mode == QuantMode::W4) {
    CHECK_SHAPE(w1, expert_size, gemm1_co, hidden_size / 2);
    CHECK_SHAPE(w2, expert_size, hidden_size, inner_size / 2);
  }

  if (bias1.has_value()) {
    CHECK_SHAPE(bias1.value(), num_expert, gemm1_co);
  }
  if (bias2.has_value()) {
    CHECK_SHAPE(bias2.value(), num_expert, hidden_size);
  }
  TORCH_CHECK(topk <= num_expert, "topk <= num_expert.")
  TORCH_CHECK(act_mode == "silu" || act_mode == "gelu",
              "act_mode must be 'silu' or 'gelu', but got ", act_mode)

  checkTensorSameAttr<TensorAttr::DTYPE>(hidden_states, bias1, bias2);
  checkTensorSameAttr<TensorAttr::DEVICE>(hidden_states, gating_output, w1, w2, residual,
                                          input_smooth, act_smooth, w1_scale, w2_scale);

  const torch_mlu::mlu::MLUGuard device_guard(hidden_states_.device());
  torch_mlu::DeviceProp *dev_prop = torch_mlu::getDeviceProperties(hidden_states.get_device());
  std::string dev_name = dev_prop->name;
  bool is_mlu370 = arch_370.compare(3, 3, dev_name, 3, 3) >= 0 ? true : false;
  auto handle = torch_mlu::getCurrentHandle();
  auto gating_output_dtype = getCnnlDataType(gating_output.scalar_type());
  auto queue = torch_mlu::getCurMLUStream();
  auto tensor_options = hidden_states_.options();
  auto data_dtype = getCnnlDataType(hidden_states.scalar_type());
  auto weight_dtype = getCnnlDataType(w1.scalar_type());
  auto input_dtype = per_token_sq ? CNNL_DTYPE_INT8 : weight_dtype;
  if (quant_mode == QuantMode::W4) {
    weight_dtype = CNNL_DTYPE_INT4X2;
  }

  // create tensors
  auto reduce_weight = at::empty({num_token, topk}, tensor_options.dtype(torch::kFloat));
  auto expert_id = at::empty({num_token, topk}, tensor_options.dtype(torch::kInt32));
  auto int32_idx = at::empty({2, num_token, topk}, tensor_options.dtype(torch::kInt32));
  auto gather_expand_idx = int32_idx[0];
  auto gather_combine_idx = int32_idx[1];
  auto token_count = at::empty({num_expert}, tensor_options.dtype(torch::kInt32));
  auto cumsum_token_count = at::empty({num_expert + 1}, tensor_options.dtype(torch::kInt32));
  auto gen_idx_workspace =
      at::empty({num_expert + 1 + num_expand_token}, tensor_options.dtype(torch::kInt32));
  int64_t mid_dim = (1 + gated) * (1 + per_token_sq);
  auto gemm1_output = at::empty({num_expand_token, mid_dim, inner_size}, w1.options());
  auto gemm2_output = at::empty({num_expand_token, hidden_size}, tensor_options);
  auto quant_input = at::Tensor();
  auto input_scale = at::Tensor();
  auto act_scale = at::Tensor();
  auto expand_hidden_state = at::Tensor();
  if (is_mlu370 || !per_token_sq) {
    expand_hidden_state = at::empty({num_expand_token, hidden_size}, tensor_options);
  }
  if (per_token_sq) {
    quant_input = at::empty({num_expand_token, hidden_size}, tensor_options.dtype(torch::kInt8));
    input_scale = at::empty({num_expand_token}, tensor_options.dtype(torch::kFloat));
    act_scale = at::empty({num_expand_token}, tensor_options.dtype(torch::kFloat));
  }

  //=========================================1:topk_softmax=========================================
  int normalize_mode = renormalize ? 1 : 0;
  tmo::invokeMoeSoftmaxTopkKernel(queue, (float *)getAtTensorPtr(reduce_weight),
                                  (int *)getAtTensorPtr(expert_id), getAtTensorPtr(gating_output_),
                                  nullptr, num_token, num_expert, 1, topk, -1, 0,
                                  gating_output_dtype, normalize_mode);
  //=========================================2:generate_idx=========================================
  tmo::invokeMoeGenIdxKernel(
      queue, (int *)getAtTensorPtr(gather_expand_idx), (int *)getAtTensorPtr(gather_combine_idx),
      (int *)getAtTensorPtr(token_count), (int *)getAtTensorPtr(cumsum_token_count),
      getAtTensorPtr(gen_idx_workspace), getAtTensorPtr(expert_id), num_token, num_expert, topk);
  if (per_token_sq) {
    //========================================3:pertoken_sq(optinal)==================================
    if (is_mlu370) {
      tmo::invokeMoeExpandInputKernel(
          queue, getAtTensorPtr(expand_hidden_state), getAtTensorPtr(hidden_states_),
          (int *)getAtTensorPtr(gather_expand_idx), (int *)getAtTensorPtr(cumsum_token_count),
          num_token, hidden_size, topk, data_dtype, num_expert, start_expert_id, expert_size);
      SmoothQuantTheory sm_obj1(num_expand_token, hidden_size, data_dtype, "per_token");
      cnpxPush(sm_obj1);
      tmo::ops::SmoothQuant(
          handle, getAtTensorPtr(expand_hidden_state), getAtTensorPtr(input_smooth),
          getAtTensorPtr(token_count[start_expert_id]), nullptr, nullptr,
          getAtTensorPtr(quant_input), getAtTensorPtr(input_scale), num_expand_token, hidden_size,
          expert_size, hidden_size, hidden_size, 1, data_dtype);
      cnpxPop();
    } else {
      SmoothQuantTheory sm_obj1(num_expand_token, hidden_size, data_dtype, "per_token");
      cnpxPush(sm_obj1);
      tmo::ops::SmoothQuant(handle, getAtTensorPtr(hidden_states_), getAtTensorPtr(input_smooth),
                            getAtTensorPtr(token_count[start_expert_id]),
                            getAtTensorPtr(gather_expand_idx),
                            getAtTensorPtr(cumsum_token_count[start_expert_id]),
                            getAtTensorPtr(quant_input), getAtTensorPtr(input_scale), num_token,
                            hidden_size, expert_size, hidden_size, hidden_size, topk, data_dtype);
      cnpxPop();
    }
  } else {
    //========================================3:expand_input========================================
    tmo::invokeMoeExpandInputKernel(
        queue, getAtTensorPtr(expand_hidden_state), getAtTensorPtr(hidden_states_),
        (int *)getAtTensorPtr(gather_expand_idx), (int *)getAtTensorPtr(cumsum_token_count),
        num_token, hidden_size, topk, data_dtype, num_expert, start_expert_id, expert_size);
  }
  //========================================4:group_gemm1===========================================
  GroupGemmDesc group_gemm_desc1(
      expert_size, num_token /*the maximum number of token that may be processed by each expert*/,
      gemm1_co, hidden_size, data_dtype, false, quant_mode);
  group_gemm_desc1.setInputOutputTensor(
      input_dtype, weight_dtype, data_dtype, CNNL_DTYPE_INT32,
      per_token_sq ? getAtTensorPtr(quant_input) : getAtTensorPtr(expand_hidden_state),
      getAtTensorPtr(w1), nullptr, getAtTensorPtr(gemm1_output), nullptr, hidden_size, hidden_size,
      num_expand_token, false);
  std::vector<int> w1_flag_vec;
  if (per_token_sq) {
    if (w1_quant_flag.has_value()) {
      auto vec = w1_quant_flag.value().vec();
      w1_flag_vec.resize(vec.size());
      std::copy(vec.begin(), vec.end(), w1_flag_vec.begin());
    }
    group_gemm_desc1.setPerRowColScaleBiasAct(
        getAtTensorPtr(input_scale), getAtTensorPtr(w1_scale), w1_flag_vec.data(), nullptr,
        data_dtype, quant_grouped ? hidden_size : 0,
        quant_grouped ? hidden_size / w1_scale.value().size(0) : 0, 0);
  }
  size_t group_gemm1_wsize =
      tmo::ops::getGroupGemmWorkspaceSize(handle, group_gemm_desc1, expert_size);
  auto group_gemm1_workspace = at::empty({static_cast<int64_t>(group_gemm1_wsize)},
                                         hidden_states.options().dtype(at::kByte));
  std::vector<int> ldb_array;
  if (quant_mode != QuantMode::W4W8) ldb_array.assign(expert_size, hidden_size);
  GroupGemmTheory gg_obj1(num_expand_token, expert_size, hidden_size, gemm1_co, false /*has_res*/,
                          input_dtype, data_dtype);
  cnpxPush(gg_obj1);
  tmo::ops::GroupGemm(handle, group_gemm_desc1, getAtTensorPtr(token_count[start_expert_id]),
                      nullptr, nullptr, getAtTensorPtr(group_gemm1_workspace), group_gemm1_wsize,
                      expert_size, hidden_size, /*k*/
                      gemm1_co,                 /*n*/
                      hidden_size /*lda*/, ldb_array /*ldb*/);
  cnpxPop();
  //========================================5:activation============================================
  cnnlActivationMode_t act_type = act_mode == "silu" ? CNNL_ACTIVATION_SWISH : CNNL_ACTIVATION_GELU;
  GroupAddBiasActiveTheory add_bias_obj(expert_size, num_expand_token, inner_size, gated,
                                        bias1.has_value(), data_dtype, act_mode);
  cnpxPush(add_bias_obj);
  tmo::invokeGroupAddBiasActivationKernel(
      queue, getAtTensorPtr(gemm1_output), getAtTensorPtr(gemm1_output),
      nullptr /*getAtTensorPtr(bias1)*/, (int *)getAtTensorPtr(cumsum_token_count), num_expert,
      num_expand_token, inner_size, gemm1_co, data_dtype, gated, act_type, start_expert_id,
      expert_size, 1.0f);
  cnpxPop();
  //========================================6:smooth_quant==========================================
  if (per_token_sq) {
    SmoothQuantTheory sm_obj2(num_expand_token, inner_size * expert_size, data_dtype, "per_token");
    cnpxPush(sm_obj2);
    tmo::ops::SmoothQuant(handle, getAtTensorPtr(gemm1_output), getAtTensorPtr(act_smooth),
                          getAtTensorPtr(token_count[start_expert_id]), nullptr /*gather_idx*/,
                          nullptr, getAtTensorPtr(gemm1_output), getAtTensorPtr(act_scale),
                          num_expand_token, inner_size, expert_size, gemm1_co /*input_stride*/,
                          2 * gemm1_co /*output_stride*/, 1 /*topk*/, data_dtype);
    cnpxPop();
  }
  if (cncl_comm > 0) {
    TORCH_CHECK(num_expert == expert_size, "expert_size must be num_expert when cncl_comm > 0");
    c10::optional<at::Tensor> act_scale_opt(act_scale);
    c10::optional<at::Tensor> w2_scale_opt(w2_scale);
    std::string dtype_s = torchDtype2Str(hidden_states.scalar_type());
    auto gg_input =
        gemm1_output.as_strided({num_expand_token, inner_size}, {mid_dim * inner_size, 1});
    auto output = group_gemm_combine_result_allreduce(
        cncl_comm, gg_input, w2, token_count, gather_combine_idx, reduce_weight, c10::nullopt,
        c10::nullopt, c10::nullopt, per_token_sq ? act_scale_opt : c10::nullopt,
        per_token_sq ? w2_scale_opt : c10::nullopt, dtype_s, num_token, topk, block_n);
    if (residual.has_value()) {
      output.view(sizes_1) += residual.value();
    }
    return output.view(sizes_1);
  } else {
    //========================================8:group_gemm2===========================================
    GroupGemmDesc group_gemm_desc2(expert_size, num_token, hidden_size, inner_size, data_dtype,
                                   false, quant_mode);
    std::vector<int64_t> b_offset(expert_size);
    int w2_ldb = inner_size;
    if (w2.dim() == 4) {
      w2_ldb = w2_shape[2] * inner_size;
      auto elem_size = w2.element_size();
      for (int64_t i = 0; i < expert_size; i++) {
        b_offset[i] = (i / w2_shape[2] * w2.stride(0) + i % w2_shape[2] * w2_shape[3]) * elem_size;
      }
    }
    group_gemm_desc2.setInputOutputTensor(input_dtype, weight_dtype, data_dtype, CNNL_DTYPE_INT32,
                                          getAtTensorPtr(gemm1_output), getAtTensorPtr(w2), nullptr,
                                          getAtTensorPtr(gemm2_output), nullptr, inner_size,
                                          gemm1_co * (per_token_sq + 1), num_expand_token, false,
                                          w2.dim() == 4 ? b_offset.data() : nullptr);
    std::vector<int> w2_flag_vec;
    if (per_token_sq) {
      if (w2_quant_flag.has_value()) {
        auto vec = w2_quant_flag.value().vec();
        w2_flag_vec.resize(vec.size());
        std::copy(vec.begin(), vec.end(), w2_flag_vec.begin());
      }
      group_gemm_desc2.setPerRowColScaleBiasAct(
          getAtTensorPtr(act_scale), getAtTensorPtr(w2_scale), w2_flag_vec.data(), nullptr,
          data_dtype, quant_grouped ? inner_size : 0,
          quant_grouped ? inner_size / w2_scale.value().size(0) : 0, 0);
    }

    size_t group_gemm2_wsize =
        tmo::ops::getGroupGemmWorkspaceSize(handle, group_gemm_desc2, expert_size);
    auto group_gemm2_workspace = at::empty({static_cast<int64_t>(group_gemm2_wsize)},
                                           hidden_states.options().dtype(at::kByte));
    std::vector<int> w2_ldb_array;
    if (quant_mode != QuantMode::W4W8) w2_ldb_array.assign(expert_size, w2_ldb);
    GroupGemmTheory gg_obj2(num_expand_token, expert_size, inner_size, hidden_size,
                            false /*has_res*/, input_dtype, data_dtype);
    cnpxPush(gg_obj2);
    tmo::ops::GroupGemm(handle, group_gemm_desc2, getAtTensorPtr(token_count[start_expert_id]),
                        nullptr, nullptr, getAtTensorPtr(group_gemm2_workspace), group_gemm2_wsize,
                        expert_size, inner_size, /*k*/
                        hidden_size,             /*n*/
                        gemm1_co * (per_token_sq + 1) /*lda*/, w2_ldb_array /*ldb*/);
    cnpxPop();
    //========================================9:combine_result=======================================
    auto output = at::empty(hidden_states_.sizes(), hidden_states_.options());
    MoeCombineResultTheory cr_obj(num_token, topk, hidden_size, expert_size, bias2.has_value(),
                                  residual.has_value(), data_dtype);
    cnpxPush(cr_obj);
    tmo::invokeMoeCombineResultKernel(
        queue, getAtTensorPtr(output), getAtTensorPtr(gemm2_output), nullptr,
        getAtTensorPtr(residual), (float *)getAtTensorPtr(reduce_weight),
        (int *)getAtTensorPtr(cumsum_token_count), (int *)getAtTensorPtr(gather_combine_idx),
        num_token, topk, num_expert, hidden_size, start_expert_id, expert_size, data_dtype);
    cnpxPop();
    return output.view(sizes_1);
  }
}

}  //  namespace torch_api
}  //  namespace tmo
