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

namespace tmo {
namespace torch_api {

at::Tensor moe_expand_input(const torch::Tensor &input,
                            const torch::Tensor &gather_idx,
                            const c10::optional<torch::Tensor> &cusum_token_count,
                            int64_t start_expert_id,
                            int64_t expert_size) {
  TORCH_CHECK(input.dim() == 2, "input dim must be equal to 2.")
  TORCH_CHECK(gather_idx.dim() == 1, "gather_idx dim must be equal to 1.")
  // check device and dtype
  TORCH_CHECK(isMlu(input), "input tensor must on mlu.")
  TORCH_CHECK(isMlu(gather_idx), "gather_idx must on mlu.")
  TORCH_CHECK(gather_idx.dtype() == torch::kInt32, "data type of gather_idx must be int32.")
  int64_t expand_token_num = gather_idx.size(0);
  int64_t token_num = input.size(0);
  int64_t hidden_size = input.size(1);
  TORCH_CHECK(expand_token_num % token_num == 0, "expand_token_num % token_num == 0.")
  int64_t topk = expand_token_num / token_num;
  int64_t expert_num = 0;
  if (cusum_token_count.has_value()) {
    TORCH_CHECK(isMlu(cusum_token_count.value()), "cusum_token_count must on mlu.")
    TORCH_CHECK(cusum_token_count.value().dtype() == torch::kInt32,
                "data type of cusum_token_count must be int32.")
    expert_num = cusum_token_count.value().size(0) - 1;
    TORCH_CHECK(start_expert_id >= 0 && start_expert_id < expert_num,
                "start_expert_id >=0 && start_expert_id < expert_num.")
    TORCH_CHECK((start_expert_id + expert_size) <= expert_num,
                "start_expert_id + expert_size <= expert_num.")
  }
  const torch_mlu::mlu::MLUGuard device_guard(input.device());
  std::vector<int64_t> output_shape = {expand_token_num, hidden_size};
  auto output = torch::empty(output_shape, input.options());

  auto dtype = getCnnlDataType(input.scalar_type());
  auto queue = torch_mlu::getCurMLUStream();

  tmo::invokeMoeExpandInputKernel(queue, getAtTensorPtr(output), getAtTensorPtr(input),
                                  (int *)getAtTensorPtr(gather_idx),
                                  (int *)getAtTensorPtr(cusum_token_count), token_num, hidden_size,
                                  topk, dtype, expert_num, start_expert_id, expert_size);
  return output;
}
}  // namespace torch_api
}  // namespace tmo
