/*************************************************************************
 * 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 <torch/library.h>
#include "glue_ops.h"
#include "torch_ops_api.h"

TORCH_LIBRARY_FRAGMENT(torch_mlu_ops, m) {
  // torch 2.1.0 does not support impl_abstract_pystub, enable it when torch 2.3.0 is released
#if TORCH_VERSION_MAJOR >= 2 && TORCH_VERSION_MINOR >= 3
  m.impl_abstract_pystub("torch_mlu_ops.abstract");
#endif
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::attention_project(Tensor input, Tensor q_weight, Tensor? q_bias, Tensor? "
      "k_weight, "
      "Tensor? k_bias, Tensor? v_weight, Tensor? v_bias, Tensor? norm_weight, Tensor? norm_bias, "
      "Tensor? residual, str out_layout, int head_size, float eps, float alpha, float beta, bool "
      "norm_out) -> (Tensor[])"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::ffn(Tensor input, Tensor up_fc_weight, Tensor? up_fc_bias, Tensor "
      "down_proj_weight, "
      "Tensor? down_proj_bias, Tensor? gate_up_proj_weight, Tensor? gate_up_proj_bias, Tensor? "
      "layernorm_weight, Tensor? layernorm_bias, str act_mode, str residual_is, float eps, float "
      "alpha, float beta) -> Tensor"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::flash_attention(Tensor q, Tensor k, Tensor v, Tensor(a!) out, Tensor(b!)? "
      "out_lse, Tensor? cu_seq_lens_q, Tensor? cu_seq_lens_kv, "
      "Tensor? alibi_slope, Tensor? attn_bias, Tensor? k_quant_scale, Tensor? v_quant_scale, "
      "Tensor? block_tables, int max_seq_len_q, int "
      "max_seq_len_kv, float softmax_scale, bool is_causal, int window_size_left, int "
      "window_size_right, str compute_dtype, bool return_lse) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::single_query_cached_kv_attn(Tensor q_ori, Tensor k_cache, Tensor "
      "v_cache, Tensor(a!) output, Tensor block_tables, Tensor context_lens, "
      "Tensor(b!)? out_lse, Tensor? k_cache_quant_scale, Tensor? v_cache_quant_scale, "
      "Tensor? alibi_slopes, int max_contxt_len, int windows_size_left, int "
      "windows_size_right, float softmax_scale, bool return_lse, int kv_cache_quant_bit_size) -> "
      "()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::apply_rotary(Tensor(a!) input, Tensor sin_cache, Tensor cos_cache, "
      "Tensor? position_ids, Tensor? cu_seqlens, bool interleaved, bool "
      "discrete, bool dynamic_ntk, int max_seqlen) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::reshape_linear_cache(Tensor key, Tensor? value, Tensor(a!) key_cache, "
      "Tensor(b!)? "
      "value_cache, Tensor context_lengths, int max_context_len, bool packed, Tensor? "
      "context_seq_offset, Tensor? cache_bs_id, Tensor? cache_seqlen_offset) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::reshape_paged_cache(Tensor k, Tensor? v, Tensor(a!) k_cache, "
      "Tensor(b!)? v_cache, Tensor slot_mapping) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::quant_to_paged_cache(Tensor k, Tensor? v, Tensor(a!) k_cache, Tensor(b!)? "
      "v_cache, "
      "Tensor(c!) k_cache_scale, Tensor(d!)? v_cache_scale, Tensor slot_mapping) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::offline_quant_to_paged_cache(Tensor k, Tensor? v, Tensor k_cache_scale,"
      "Tensor? v_cache_scale, Tensor slot_mapping, Tensor(a!) k_cache, Tensor(b!)? v_cache)-> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::quant_to_linear_cache(Tensor key, Tensor? value, Tensor(a!) key_cache, "
      "Tensor(b!)? "
      "value_cache, Tensor(c!) key_cache_scale, Tensor(d!)? value_cache_scale, Tensor "
      "context_lengths, "
      "int max_context_len, bool packed, Tensor? context_seq_offset, Tensor? cache_bs_id, Tensor? "
      "cache_seqlen_offset, int quant_bit=8) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::offline_quant_to_linear_cache(Tensor key, Tensor? value, Tensor(a!) "
      "key_cache,"
      " Tensor(b!)? value_cache, Tensor(c!) key_cache_scale, Tensor(d!)? value_cache_scale, "
      "Tensor context_lengths, int max_context_len, int quant_mode, bool packed, "
      "Tensor? context_seq_offset, Tensor? cache_bs_id, Tensor? "
      "cache_seqlen_offset) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::quant_matmul(Tensor a_tensor, Tensor? a_scale, Tensor? a_zero, Tensor "
      "b_tensor, "
      "Tensor? b_scale, Tensor? b_zero, Tensor? bias, Tensor? c_tensor, Tensor? c_scale, Tensor? "
      "c_zero, Tensor? gemm_output_scale, Tensor? gemm_output_zero, str? data_type, Tensor? d, "
      "str "
      "quant_algo, str a_quant_layout, str b_quant_layout, int quant_bit_size=8, str "
      "act_mode='none', bool use_hp_active=False, float act_coef=1.0, float alpha=1.0, float "
      "beta=1.0, bool trans_a=False, bool trans_b=True) -> Tensor"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::quant_matmul_allreduce(int cncl_comm, Tensor a_tensor, Tensor? a_scale, "
      "Tensor? a_zero, Tensor b_tensor, Tensor? b_scale, Tensor? b_zero, "
      "Tensor? bias, Tensor? c_tensor, Tensor? c_scale, Tensor? "
      "c_zero, Tensor? gemm_output_scale, Tensor? gemm_output_zero, str? data_type, Tensor? d, "
      "str quant_algo, str a_quant_layout, str b_quant_layout, int quant_bit_size=8, "
      "float alpha=1.0, float beta=1.0, bool trans_a=False, bool trans_b=True, int block_m=0) -> "
      "Tensor"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::active(Tensor input, Tensor(a!) output, Tensor? bias, "
      "Tensor? cusum_token_count, str act_mode, bool is_gated, int start_expert_id, "
      "int expert_size, float active_coef=1.0) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::smooth_quant(Tensor input, Tensor input_scale, Tensor(a!) output, "
      "Tensor(b!) output_scale, "
      "Tensor? input_zero, Tensor? token_count, Tensor? gather_index, "
      "Tensor? gather_index_start_position, str quant_mode, bool dynamic_quant) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::fused_layernorm(Tensor input, Tensor(a!) out, "
      "Tensor? residual, Tensor? gamma, Tensor? beta, Tensor? bias, Tensor? "
      "quant_scale, Tensor(b!)? residual_out, Tensor? smooth_quant_scale, str norm_mode, "
      "float eps, bool store_output_before_norm, bool dynamic_quant) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::fused_moe(Tensor hidden_states, Tensor gating_output, Tensor "
      "w1, Tensor w2, Tensor? bias1, Tensor? bias2, Tensor? residual, Tensor? "
      "input_smooth, Tensor? act_smooth, Tensor? w1_scale, Tensor? w2_scale, "
      "int[]? w1_quant_flag, int[]? w2_quant_flag, "
      "int topk, bool renormalize, bool gated, str act_mode, int start_expert_id, "
      "int block_n, int cncl_comm) -> Tensor"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::matmul(Tensor a, Tensor b, Tensor? d, Tensor? bias, Tensor? c, str? dtype,"
      "  str act_mode, float alpha, float beta, bool fast_act, bool approximate, float a_scale, "
      "float "
      "b_scale, bool trans_a, bool trans_b) -> Tensor"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::batch_matmul(Tensor a, Tensor b, Tensor(a!) c,"
      "float alpha, float beta, float a_scale, float b_scale, bool trans_a, bool trans_b) -> ()"));
  m.def(
      TORCH_SELECTIVE_SCHEMA("torch_mlu_ops::matmul_allreduce(int cncl_comm, Tensor a, Tensor b, "
                             "Tensor? bias, Tensor? c, Tensor? d, "
                             "float alpha, float beta, int block_m) -> Tensor"));
  m.def(
      TORCH_SELECTIVE_SCHEMA("torch_mlu_ops::swap_blocks(Tensor(a!) dst, Tensor src, Dict(int, "
                             "int) block_mapping) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::copy_blocks(Tensor(a!)[] k_caches, Tensor(b!)[] v_caches, "
      "Dict(int, int[]) block_mapping) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::copy_blocks_out_of_place(Tensor[] k_caches, Tensor[] v_caches, "
      "Dict(int, int[]) block_mapping) -> (Tensor[], Tensor[])"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::group_gemm(Tensor a, Tensor b, Tensor m_list, Tensor? idx, Tensor? c, "
      "Tensor? alpha, Tensor? beta, Tensor? a_scale, Tensor? b_scale, Tensor? bias,"
      "str? data_type, int[]? quant_flag, Tensor? b_offset, int max_m) -> Tensor"));
  m.def(TORCH_SELECTIVE_SCHEMA("torch_mlu_ops::preload(Tensor(a!) weight, int size) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::group_gemm_combine_result_allreduce(int cncl_comm, Tensor a, Tensor b, "
      "Tensor m_list, Tensor combine_idx, Tensor combine_weight, Tensor? c, "
      "Tensor? alpha, Tensor? beta, Tensor? a_scale, Tensor? b_scale, "
      "str? data_type, int num_token, int topk, int block_n) -> Tensor"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::flash_attn_sq_mm_allreduce(int cncl_comm, Tensor q, Tensor k, "
      "Tensor v, Tensor? cu_seq_lens_q, Tensor? cu_seq_lens_k, Tensor? alibi_slope, "
      "Tensor? attn_bias, Tensor smooth, Tensor weight, Tensor weight_scale, "
      "Tensor? bias, int max_seq_len_q, int max_seq_len_kv, float softmax_scale, bool is_causal, "
      "int window_size_left, int window_size_right, "
      "str compute_dtype, int block_seq) -> Tensor"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::moe_softmax_topk(Tensor input, int topk, int num_expert_group, "
      "int topk_group, bool normalize, Tensor? mask, str normed_by) -> Tensor[]"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::moe_expand_input(Tensor input, Tensor gather_idx, Tensor ? "
      "cusum_token_count, int start_expert_id, int expert_size) -> Tensor"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::moe_gen_idx(Tensor expert_id, int expert_num) -> Tensor[]"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::moe_combine_result(Tensor input, Tensor reduce_weight,"
      "Tensor gather_ids, Tensor? residual,"
      "Tensor? cusum_token_count, int start_expert_id, int expert_size, Tensor? bias) -> Tensor"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::fused_rope(Tensor(a!) input, Tensor(b!) k_cache_hp, Tensor(c!) v_cache_hp, "
      "Tensor(d!)? k_cache_lp, Tensor(e!)? v_cache_lp, Tensor sin_table, Tensor cos_table, "
      "Tensor position_ids, Tensor gamma, Tensor beta, Tensor? k_scale_hp, Tensor? v_scale_hp, "
      "Tensor(f!)? k_scale_lp, Tensor(g!)? v_scale_lp, Tensor? cache_bs_id_hp, "
      "Tensor? cache_seq_offsets_hp, Tensor? cache_bs_id_lp, Tensor? cache_seq_offsets_lp, "
      "Tensor? slot_mapping_hp, Tensor? slot_mapping_lp, float eps) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::moe_cast_gating(Tensor input, Tensor weight) -> Tensor"));
  m.def(
      TORCH_SELECTIVE_SCHEMA("torch_mlu_ops::update_out_and_lse(Tensor(a!) out, Tensor(b!) lse, "
                             "Tensor block_out, Tensor block_lse,"
                             "Tensor? seq_offsets, Tensor? cu_seqs, Tensor? block_cu_seqs) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::dequant_from_linear_cache(Tensor(a!) key, Tensor(b!)? value, "
      "Tensor key_cache, Tensor? value_cache, Tensor key_cache_quant_scale, "
      "Tensor? value_cache_quant_scale, Tensor context_lengths, int max_context_len, "
      "Tensor? context_seq_offset, Tensor? cache_bs_id, Tensor? cache_seq_offset, "
      "int quant_mode=0, int quant_bit=8) -> ()"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torch_mlu_ops::dequant_from_paged_cache(Tensor(a!) key, Tensor(b!)? value, "
      "Tensor key_cache, Tensor? value_cache, Tensor key_cache_quant_scale, "
      "Tensor? value_cache_quant_scale, Tensor context_lengths, int max_context_len, "
      "Tensor? context_seq_offset, Tensor block_tables, int quant_mode=0, int quant_bit=8) -> ()"));
}

TORCH_LIBRARY_IMPL(torch_mlu_ops, PrivateUse1, m) {
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::attention_project"),
         TORCH_FN(tmo::torch_api::attention_project));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::ffn"), TORCH_FN(tmo::torch_api::ffn));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::flash_attention"),
         TORCH_FN(tmo::torch_api::flash_attention));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::single_query_cached_kv_attn"),
         TORCH_FN(tmo::torch_api::single_query_cached_kv_attn));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::apply_rotary"),
         TORCH_FN(tmo::torch_api::apply_rotary));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::reshape_linear_cache"),
         TORCH_FN(tmo::torch_api::reshape_linear_cache));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::reshape_paged_cache"),
         TORCH_FN(tmo::torch_api::reshape_paged_cache));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::quant_to_paged_cache"),
         TORCH_FN(tmo::torch_api::quant_to_paged_cache));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::offline_quant_to_paged_cache"),
         TORCH_FN(tmo::torch_api::offline_quant_to_paged_cache));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::quant_to_linear_cache"),
         TORCH_FN(tmo::torch_api::quant_to_linear_cache));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::offline_quant_to_linear_cache"),
         TORCH_FN(tmo::torch_api::offline_quant_to_linear_cache));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::quant_matmul"),
         TORCH_FN(tmo::torch_api::quant_matmul));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::quant_matmul_allreduce"),
         TORCH_FN(tmo::torch_api::quant_matmul_allreduce));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::active"), TORCH_FN(tmo::torch_api::active));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::smooth_quant"),
         TORCH_FN(tmo::torch_api::smooth_quant));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::fused_layernorm"),
         TORCH_FN(tmo::torch_api::fused_layernorm));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::fused_moe"), TORCH_FN(tmo::torch_api::fused_moe));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::matmul"), TORCH_FN(tmo::torch_api::matmul));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::batch_matmul"),
         TORCH_FN(tmo::torch_api::batch_matmul));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::matmul_allreduce"),
         TORCH_FN(tmo::torch_api::matmul_allreduce));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::swap_blocks"), TORCH_FN(tmo::torch_api::swap_blocks));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::copy_blocks"), TORCH_FN(tmo::torch_api::copy_blocks));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::copy_blocks_out_of_place"),
         TORCH_FN(tmo::torch_api::copy_blocks_out_of_place));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::group_gemm"), TORCH_FN(tmo::torch_api::group_gemm));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::preload"), TORCH_FN(tmo::torch_api::preload));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::group_gemm_combine_result_allreduce"),
         TORCH_FN(tmo::torch_api::group_gemm_combine_result_allreduce));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::flash_attn_sq_mm_allreduce"),
         TORCH_FN(tmo::torch_api::flash_attn_sq_mm_allreduce));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::moe_softmax_topk"),
         TORCH_FN(tmo::torch_api::moe_softmax_topk));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::moe_expand_input"),
         TORCH_FN(tmo::torch_api::moe_expand_input));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::moe_gen_idx"), TORCH_FN(tmo::torch_api::moe_gen_idx));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::moe_combine_result"),
         TORCH_FN(tmo::torch_api::moe_combine_result));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::moe_cast_gating"),
         TORCH_FN(tmo::torch_api::moe_cast_gating));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::fused_rope"), TORCH_FN(tmo::torch_api::fused_rope));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::update_out_and_lse"),
         TORCH_FN(tmo::torch_api::update_out_and_lse));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::dequant_from_linear_cache"),
         TORCH_FN(tmo::torch_api::dequant_from_linear_cache));
  m.impl(TORCH_SELECTIVE_NAME("torch_mlu_ops::dequant_from_paged_cache"),
         TORCH_FN(tmo::torch_api::dequant_from_paged_cache));
}

// Function(copy_blocks__functionalization_glue) is no longer needed in the torch 2.5
TORCH_LIBRARY_IMPL(torch_mlu_ops, Functionalize, m) {
  m.impl("torch_mlu_ops::copy_blocks", copy_blocks__functionalization_glue);
}
