/*************************************************************************
 * 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 <numeric>
#include "kernels/moe/cast_gating.mluh"
#include "torch_ops_api.h"
namespace tmo {
namespace torch_api {

at::Tensor moe_cast_gating(const at::Tensor &input, const at::Tensor &weight) {
  // check
  checkTensorSameAttr<TensorAttr::DEVICE>(input, weight);
  TORCH_CHECK(input.is_contiguous(), "input must be contiguous")
  TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous")
  TORCH_CHECK(weight.dim() == 2, "weight.dim() == 2")
  TORCH_CHECK(input.size(-1) == weight.size(-1), "input.size(-1) == weight.size(-1)")
  TORCH_CHECK(input.scalar_type() == torch::kFloat16 || input.scalar_type() == torch::kBFloat16,
              "input type need be torch::kFloat16 or torch::kBFloat16")
  TORCH_CHECK(weight.scalar_type() == torch::kFloat32, "weight type need be torch::kFloat32")

  auto hidden_size = input.size(-1);
  auto expert_num = weight.size(0);
  TORCH_CHECK(hidden_size > 0 && hidden_size <= 16384, "hidden_size > 0 && hidden_size <= 16384")
  TORCH_CHECK(expert_num > 0 && expert_num <= 128, "expert_num > 0 && expert_num <= 128")
  auto input_shape = input.sizes();  // [..., hidden_size]
  auto input_row =
      std::accumulate(input_shape.begin(), input_shape.end() - 1, 1, std::multiplies<int>());
  std::vector<int64_t> output_shape(input_shape.begin(), input_shape.end() - 1);
  output_shape.push_back(expert_num);

  const torch_mlu::mlu::MLUGuard device_guard(input.device());
  auto output = at::empty(output_shape, weight.options());
  const int64_t workspace_size = 16 * 1024 * 1024;
  auto workspace = at::empty({workspace_size}, weight.options().dtype(at::kByte));
  auto queue = torch_mlu::getCurMLUStream();
  auto input_dtype = getCnnlDataType(input.scalar_type());
  TMO_KERNEL_CHECK_FATAL(tmo::invokeCastGating(
      queue, input.data_ptr(), weight.data_ptr(), output.data_ptr(), input_row, expert_num,
      hidden_size, input_dtype, workspace.data_ptr(), workspace_size));
  return output;
}
}  //  namespace torch_api
}  //  namespace tmo