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

namespace tmo {
namespace torch_api {

std::vector<at::Tensor> moe_gen_idx(const torch::Tensor &expert_id, int64_t expert_num) {
  TORCH_CHECK(expert_id.dim() == 2, "expert_id dim must be equal to 2.")
  int64_t token_num = expert_id.size(0);
  int64_t topk = expert_id.size(1);

  // check device and dtype
  TORCH_CHECK(isMlu(expert_id), "expert_id must on mlu.")
  TORCH_CHECK(expert_id.dtype() == torch::kInt32, "data type of expert_id must be int32.")

  const torch_mlu::mlu::MLUGuard device_guard(expert_id.device());
  auto expand_idx = at::empty({token_num * topk}, expert_id.options());
  auto combine_idx = at::empty({token_num * topk}, expert_id.options());
  auto token_count = at::empty({expert_num}, expert_id.options());
  auto cusum_token_count = at::empty({expert_num + 1}, expert_id.options());
  auto gen_idx_workspace = at::empty({expert_num + 1 + token_num * topk}, expert_id.options());

  auto queue = torch_mlu::getCurMLUStream();
  tmo::invokeMoeGenIdxKernel(
      queue, (int *)getAtTensorPtr(expand_idx), (int *)getAtTensorPtr(combine_idx),
      (int *)getAtTensorPtr(token_count), (int *)getAtTensorPtr(cusum_token_count),
      getAtTensorPtr(gen_idx_workspace), getAtTensorPtr(expert_id), token_num, expert_num, topk);
  std::vector<at::Tensor> output = {expand_idx, combine_idx, token_count, cusum_token_count};
  return output;
}
}  // namespace torch_api
}  // namespace tmo
