/*************************************************************************
 * 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 "kernels/moe/softmax_topk.mluh"
#include "torch_ops_api.h"

namespace tmo {
namespace torch_api {

std::vector<at::Tensor> moe_softmax_topk(const at::Tensor &input,  // (num_token, num_expert)
                                         int64_t topk,
                                         int64_t num_expert_group,
                                         int64_t topk_group,
                                         bool normalize,
                                         const c10::optional<at::Tensor> &mask,
                                         const std::string &normed_by) {
  CHECK_TENSOR_CONTIGUOUS(input)
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(mask)
  int normalize_mode = 0;
  if (normalize) {
    TORCH_CHECK(normed_by == "topk_logit" || normed_by == "softmax_logit",
                "normed_by must be 'topk_logit' or 'softmax_logit'")
    if (normed_by == "topk_logit") {
      normalize_mode = 1;
    } else if (normed_by == "softmax_logit") {
      normalize_mode = 2;
    }
  }
  TORCH_CHECK(input.dim() >= 2, "input.dim() >= 2")
  int num_expert = input.size(-1);
  int num_token = (int)(input.numel() / num_expert);
  int num_mask = 1;
  auto input_shape = input.sizes();
  TORCH_CHECK(topk > 0 && topk <= num_expert, "topk > 0 && topk <= num_expert")
  bool has_mask = mask.has_value();
  if (has_mask) {
    TORCH_CHECK(mask.value().dim() == input.dim(), "the dim of mask should be the same as input")
    TORCH_CHECK(mask.value().size(-1) == num_expert,
                "the last dim of mask should be the same as the last dim of input")
    TORCH_CHECK(mask.value().size(-2) == input.size(-2),
                "the penultimate dim of mask should be the same as the penultimate dim of input")
    int high_dim = mask.value().numel() / (mask.value().size(-1) * mask.value().size(-2));
    TORCH_CHECK(high_dim == 1, "the product of all but the lower two dimensions of mask is 1")
    num_mask = mask.value().numel() / num_expert;
    TORCH_CHECK(mask.value().dtype() == input.dtype(),
                "the dtype of mask should be the same as input")
  }
  if (num_expert_group > 1) {
    TORCH_CHECK(has_mask == false, "if num_expert_group > 1, mask should be None")
    TORCH_CHECK(num_expert % num_expert_group == 0, "num_expert % num_expert_group == 0")
    TORCH_CHECK(topk_group > 0 && topk_group <= num_expert_group,
                "topk_group > 0 && topk_group <= num_expert_group")
    TORCH_CHECK(topk <= (num_expert / num_expert_group) * topk_group,
                "topk <= (num_expert / num_expert_group) * topk_group")
  }
  std::vector<int64_t> out_shape(input_shape.begin(), input_shape.end());
  out_shape.back() = topk;

  const torch_mlu::mlu::MLUGuard device_guard(input.device());
  auto tensor_options = input.options();
  auto reduce_weight = at::empty(out_shape, tensor_options.dtype(torch::kFloat));
  auto expert_id = at::empty(out_shape, tensor_options.dtype(torch::kInt32));

  auto queue = torch_mlu::getCurMLUStream();
  auto input_dtype = getCnnlDataType(input.scalar_type());
  TMO_KERNEL_CHECK_FATAL(invokeMoeSoftmaxTopkKernel(
      queue, (float *)reduce_weight.data_ptr(), (int *)expert_id.data_ptr(), input.data_ptr(),
      getAtTensorPtr(mask), num_token, num_expert, num_mask, topk, num_expert_group, topk_group,
      input_dtype, normalize_mode));
  return {reduce_weight, expert_id};
}

}  //  namespace torch_api
}  //  namespace tmo
