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

namespace tmo {
namespace torch_api {
at::Tensor moe_combine_result(const at::Tensor &input,
                              const at::Tensor &reduce_weight,
                              const at::Tensor &gather_ids,
                              const c10::optional<at::Tensor> &residual,
                              const c10::optional<at::Tensor> &cusum_token_count,
                              int64_t start_expert_id,
                              int64_t expert_size,
                              const c10::optional<at::Tensor> &bias) {
  // check device and dtype
  checkTensorSameAttr<TensorAttr::DEVICE>(input, reduce_weight, residual, bias, cusum_token_count,
                                          gather_ids);

  // check contiguous
  CHECK_TENSOR_CONTIGUOUS(input)
  CHECK_TENSOR_CONTIGUOUS(reduce_weight)
  CHECK_TENSOR_CONTIGUOUS(gather_ids)
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(residual)
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(bias)
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(cusum_token_count)

  // do not support bias
  TORCH_CHECK(!bias.has_value(), "bias is not supported.");

  // check input
  TORCH_CHECK(input.dim() == 2, "input should be with shape [num_tokens, hidden_size].");
  TORCH_CHECK(input.dtype() == torch::kFloat32 || input.dtype() == torch::kFloat16 ||
                  input.dtype() == torch::kBFloat16,
              "input dtype should be float32/float16/bfloat16.");

  int num_tokens = input.size(0);
  int hidden_size = input.size(1);

  // check reduce weight
  TORCH_CHECK(reduce_weight.dim() == 2, "reduce_weight should be with shape [num_token, topk].");
  TORCH_CHECK(reduce_weight.dtype() == torch::kFloat32,
              "reduce_weight should be with dtype float32.");
  int topk = reduce_weight.size(1);
  TORCH_CHECK(reduce_weight.numel() == num_tokens,
              "reduce_weight.numel() should equal to input.size(0).");
  int num_token = num_tokens / topk;

  // check gather_ids
  TORCH_CHECK(gather_ids.dim() == 1, "gather_ids should be with shape [num_tokens].");
  TORCH_CHECK(gather_ids.dtype() == torch::kInt32, "gather_ids should be with dtype int32.");
  TORCH_CHECK(gather_ids.size(0) == num_tokens,
              "gather_ids.numel() should equal to input.size(0).");

  if (residual.has_value()) {
    TORCH_CHECK(residual.value().dim() == 2,
                "residual should be with shape [num_token, hidden_size].");
    TORCH_CHECK(residual.value().dtype() == input.dtype(),
                "residual should have same dtype with input.");
    TORCH_CHECK(residual.value().size(0) == num_token,
                "residual.size(0) should equal to input.size(0).");
    TORCH_CHECK(residual.value().size(1) == hidden_size,
                "residual.size(1) should equal to input.size(1).");
  }

  // check bias and cusum_token_count
  bool has_bias = bias.has_value();
  bool has_cusum_token_count = cusum_token_count.has_value();
  TORCH_CHECK((!has_bias || (has_bias && has_cusum_token_count)),
              "if bias is not None, cusum_token_count should not be None.");

  int num_expert = -1;
  if (has_bias) {
    TORCH_CHECK(bias.value().dim() == 2, "bias should be with shape [num_expert, hidden_size].");
    TORCH_CHECK(bias.value().dtype() == input.dtype(), "bias should have same dtype with input.");
    TORCH_CHECK(bias.value().size(1) == hidden_size, "bias.size(1) should equal to input.size(1).");
    num_expert = bias.value().size(0);
  }

  if (has_cusum_token_count) {
    TORCH_CHECK(cusum_token_count.value().dim() == 1,
                "cusum_token_count should be with shape [num_expert + 1].");
    TORCH_CHECK(cusum_token_count.value().dtype() == torch::kInt32,
                "cusum_token_count should be with dtype int32.");
    if (num_expert > 0) {
      TORCH_CHECK(cusum_token_count.value().dim() == 1,
                  "cusum_token_count should be with shape [num_expert + 1].");
    } else {
      num_expert = cusum_token_count.value().size(0) - 1;
    }
  }

  // check expert
  TORCH_CHECK(start_expert_id >= 0, "start_expert_id shoule be larger or equal to 0.");
  TORCH_CHECK(num_expert == -1 || num_expert >= (start_expert_id + expert_size),
              "num_expert shape shoule be larger or equal to start_expert_id + expert_size.");
  if (num_expert == -1) {
    num_expert = expert_size;
  }

  const torch_mlu::mlu::MLUGuard device_guard(input.device());
  auto queue = torch_mlu::getCurMLUStream();
  auto output_shape = std::vector<int64_t>({num_token, hidden_size});
  auto output = at::empty(output_shape, input.options());
  auto dtype = getCnnlDataType(input.scalar_type());

  invokeMoeCombineResultKernel(queue, output.data_ptr(), input.data_ptr(), getAtTensorPtr(bias),
                               getAtTensorPtr(residual), (float *)reduce_weight.data_ptr(),
                               (int *)getAtTensorPtr(cusum_token_count),
                               (int *)getAtTensorPtr(gather_ids), num_token, topk, num_expert,
                               hidden_size, start_expert_id, expert_size, dtype);
  return output;
}
}  //  namespace torch_api
}  //  namespace tmo
