/*************************************************************************
 * 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 <memory>
#include <string>
#include <vector>
#include "comm_overlap.h"
#include "kernels/moe/combine_result.mluh"
#include "torch_api/torch_ops_api.h"
#include "torch_api/utils.h"

namespace tmo {
namespace torch_api {

using namespace torch::indexing;

struct GroupGemmCombineResultSplitN {
  const at::Tensor &a_;
  const at::Tensor &b_;
  const at::Tensor &m_list_;
  const at::Tensor &combine_idx_;
  const at::Tensor &combine_weight_;
  const c10::optional<at::Tensor> &c_;
  const c10::optional<at::Tensor> &alpha_;
  const c10::optional<at::Tensor> &beta_;
  const c10::optional<at::Tensor> &a_scale_;
  const c10::optional<at::Tensor> &b_scale_;
  const c10::optional<std::string> &data_type_;
  int64_t num_token_;
  int64_t topk_;
  int64_t block_n_;
  int64_t expert_num_;
  int64_t n_;
  bool has_c_;
  bool has_alpha_;
  bool has_beta_;
  bool has_a_scale_;
  bool has_b_scale_;
  bool has_data_type_;
  int64_t n_per_blk_;
  cnrtQueue_t queue_;
  at::Tensor output_buff_;
  at::Tensor b_scale_trans_;
  c10::optional<at::Tensor> b_offset_ = c10::nullopt;
  cnnlDataType_t cnnl_dtype_;

  GroupGemmCombineResultSplitN(const at::Tensor &a_tensor,
                               const at::Tensor &b_tensor,
                               const at::Tensor &m_list,
                               const at::Tensor &combine_idx,
                               const at::Tensor &combine_weight,
                               const c10::optional<at::Tensor> &c_tensor,
                               const c10::optional<at::Tensor> &alpha,
                               const c10::optional<at::Tensor> &beta,
                               const c10::optional<at::Tensor> &a_scale,
                               const c10::optional<at::Tensor> &b_scale,
                               const c10::optional<std::string> &data_type,
                               const int64_t num_token,
                               const int64_t topk,
                               const int64_t block_n)
      : a_(a_tensor),
        b_(b_tensor),
        m_list_(m_list),
        combine_idx_(combine_idx),
        combine_weight_(combine_weight),
        c_(c_tensor),
        alpha_(alpha),
        beta_(beta),
        a_scale_(a_scale),
        b_scale_(b_scale),
        data_type_(data_type),
        num_token_(num_token),
        topk_(topk),
        block_n_(block_n) {
    has_c_ = c_.has_value();
    has_alpha_ = alpha_.has_value();
    has_beta_ = beta_.has_value();
    has_a_scale_ = a_scale_.has_value();
    has_b_scale_ = b_scale_.has_value();
    has_data_type_ = data_type_.has_value();
    auto b_shape = b_.sizes();
    expert_num_ = b_.dim() == 3 ? b_shape[0] : b_shape[0] * b_shape[2];
    n_ = b_shape[1];
    auto k = b_.size(-1);
    auto m = a_.size(0);
    set_block(m);

    queue_ = torch_mlu::getCurMLUStream();
    auto a_options = a_.options();
    if (has_data_type_) {
      auto dtype = data_type_.value();
      TORCH_CHECK(dtype == "float" || dtype == "half" || dtype == "bfloat16",
                  "data type must be 'float', 'half' or 'bfloat16'");
      auto torch_dtype = str2TorchDtype(dtype);
      cnnl_dtype_ = str2CnnlDtype(dtype);
      output_buff_ = at::empty({block_n_, num_token_, n_per_blk_}, a_options.dtype(torch_dtype));
    } else {
      output_buff_ = at::empty({block_n_, num_token_, n_per_blk_}, a_options);
      cnnl_dtype_ = getCnnlDataType(a_.scalar_type());
    }
    if (has_b_scale_) {
      b_scale_trans_ = b_scale_.value().view({expert_num_, block_n_, -1});
      b_scale_trans_ = b_scale_trans_.transpose(0, 1).contiguous();
    }

    auto b_offset = at::empty({expert_num_}, a_options.dtype(at::kLong).device(at::kCPU));
    auto element_size = b_.element_size();
    if (block_n_ > 1 && b_.dim() == 3) {
      for (int64_t i = 0; i < expert_num_; i++) {
        b_offset[i] = n_ * k * i * element_size;
      }
      b_offset_ = b_offset;
    } else if (b_.dim() == 4) {
      for (int64_t i = 0; i < expert_num_; i++) {
        b_offset[i] = (i / b_shape[2] * b_.stride(0) + i % b_shape[2] * b_shape[3]) * element_size;
      }
      b_offset_ = b_offset;
    }
  }

  auto split(const at::Tensor &a, int64_t dim, int64_t start, int64_t block) {
    return a.narrow(dim, start, block);
  }

  at::Tensor forward(const int64_t block_id) {
    auto gg_o =
        group_gemm(a_, split(b_, 1, block_id * n_per_blk_, n_per_blk_), m_list_, c10::nullopt,
                   (has_c_) ? split(c_.value(), 1, block_id * n_per_blk_, n_per_blk_) : c_, alpha_,
                   beta_, a_scale_, has_b_scale_ ? b_scale_trans_[block_id] : b_scale_,
                   c10::nullopt, data_type_, c10::nullopt, b_offset_, num_token_);

    tmo::invokeMoeCombineResultKernel(
        queue_, getAtTensorPtr(output_buff_[block_id]), getAtTensorPtr(gg_o), nullptr, nullptr,
        (float *)getAtTensorPtr(combine_weight_), nullptr, (int *)getAtTensorPtr(combine_idx_),
        num_token_, topk_, expert_num_, n_per_blk_, 0, expert_num_, cnnl_dtype_);
    return output_buff_[block_id];
  }

  void set_block(int64_t m) {
    if (block_n_ < 1) {
      if (m >= 4096 && m < 8192 && n_ % 2048 == 0) {
        n_per_blk_ = 2048;
        block_n_ = n_ / n_per_blk_;
      } else if (m >= 8192 && n_ % 1024 == 0) {
        n_per_blk_ = 1024;
        block_n_ = n_ / n_per_blk_;
      } else {
        n_per_blk_ = n_;
        block_n_ = 1;
      }
    } else {
      TORCH_CHECK(n_ % block_n_ == 0, "n must be divisible by block_n");
      n_per_blk_ = n_ / block_n_;
    }
  }

  int64_t getLoopNum() const { return block_n_; }
  at::Tensor getOutput() const { return output_buff_.transpose(0, 1).reshape({num_token_, n_}); }
};

at::Tensor group_gemm_combine_result_allreduce(int64_t cncl_comm,
                                               const at::Tensor &a_tensor,
                                               const at::Tensor &b_tensor,
                                               const at::Tensor &m_list,
                                               const at::Tensor &combine_idx,
                                               const at::Tensor &combine_weight,
                                               const c10::optional<at::Tensor> &c_tensor,
                                               const c10::optional<at::Tensor> &alpha,
                                               const c10::optional<at::Tensor> &beta,
                                               const c10::optional<at::Tensor> &a_scale,
                                               const c10::optional<at::Tensor> &b_scale,
                                               const c10::optional<std::string> &data_type,
                                               const int64_t num_token,
                                               const int64_t topk,
                                               const int64_t block_n) {
  GroupGemmCombineResultSplitN gg(a_tensor, b_tensor, m_list, combine_idx, combine_weight, c_tensor,
                                  alpha, beta, a_scale, b_scale, data_type, num_token, topk,
                                  block_n);
  ParallelAllReduce<GroupGemmCombineResultSplitN> parallel_rs(cncl_comm);
  return parallel_rs(gg);
}
}  //  namespace torch_api
}  //  namespace tmo
