/*************************************************************************
 * 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.
 *************************************************************************/
#ifndef CSRC_TORCH_API_TORCH_OPS_API_H_
#define CSRC_TORCH_API_TORCH_OPS_API_H_

#include <cstdint>
#include <map>
#include <optional>
#include <vector>
#include "op_theory.h"
#include "ops/kernel_api.h"
#include "torch/extension.h"
#include "utils.h"

namespace tmo {
namespace torch_api {

void fused_layernorm(const at::Tensor &input,
                     const at::Tensor &out,
                     const c10::optional<at::Tensor> &residual,
                     const c10::optional<at::Tensor> &gamma,
                     const c10::optional<at::Tensor> &beta,
                     const c10::optional<at::Tensor> &bias,
                     const c10::optional<at::Tensor> &quant_scale,
                     const c10::optional<at::Tensor> &residual_out,
                     const c10::optional<at::Tensor> &smooth_quant_scale,
                     const std::string &norm_mode,
                     double eps,
                     bool store_output_before_norm,
                     bool dynamic_quant);

std::vector<at::Tensor> attention_project(const at::Tensor &input,
                                          const at::Tensor &q_weight,
                                          const c10::optional<at::Tensor> &q_bias,
                                          const c10::optional<at::Tensor> &k_weight,
                                          const c10::optional<at::Tensor> &k_bias,
                                          const c10::optional<at::Tensor> &v_weight,
                                          const c10::optional<at::Tensor> &v_bias,
                                          const c10::optional<at::Tensor> &norm_weight,
                                          const c10::optional<at::Tensor> &norm_bias,
                                          const c10::optional<at::Tensor> &residual,
                                          const std::string &out_layout,
                                          int64_t head_size,
                                          double eps,
                                          double alpha,
                                          double beta,
                                          bool norm_out);

at::Tensor ffn(const at::Tensor &input,
               const at::Tensor &up_fc_weight,
               const c10::optional<at::Tensor> &up_fc_bias,
               const at::Tensor &down_proj_weight,
               const c10::optional<at::Tensor> &down_proj_bias,
               const c10::optional<at::Tensor> &gate_up_proj_weight,
               const c10::optional<at::Tensor> &gate_up_proj_bias,
               const c10::optional<at::Tensor> &layernorm_weight,
               const c10::optional<at::Tensor> &layernorm_bias,
               const std::string &act_mode,
               const std::string &residual_is,
               double eps,
               double alpha,
               double beta);

void flash_attention(const at::Tensor &q,
                     const at::Tensor &k,
                     const at::Tensor &v,
                     const at::Tensor &out,
                     const c10::optional<at::Tensor> &output_lse,
                     const c10::optional<at::Tensor> &cu_seq_lens_q,
                     const c10::optional<at::Tensor> &cu_seq_lens_kv,
                     const c10::optional<at::Tensor> &alibi_slope,
                     const c10::optional<at::Tensor> &attn_bias,
                     const c10::optional<at::Tensor> &k_cache_quant_scale,
                     const c10::optional<at::Tensor> &v_cache_quant_scale,
                     const c10::optional<at::Tensor> &block_tables,
                     const int64_t max_seq_len_q,
                     const int64_t max_seq_len_kv,
                     const double softmax_scale,
                     const bool is_causal,
                     const int64_t window_size_left,
                     const int64_t window_size_right,
                     const std::string &compute_dtype,
                     bool return_lse);

void single_query_cached_kv_attn(
    const torch::Tensor &q_ori,
    const torch::Tensor &k_cache,
    const torch::Tensor &v_cache,
    const torch::Tensor &output,
    const torch::Tensor &block_tables,
    const torch::Tensor &context_lens,  // [batch]
    const c10::optional<torch::Tensor> &output_lse,
    const c10::optional<torch::Tensor> &k_cache_quant_scale,
    const c10::optional<torch::Tensor> &v_cache_quant_scale,
    const c10::optional<torch::Tensor> &alibi_slopes,  // [bs, head_num] or [head_num]
    int64_t max_context_len,
    int64_t windows_size_left,
    int64_t windows_size_right,
    double softmax_scale,
    bool return_lse,
    int64_t kv_cache_quant_bit_size);

void reshape_linear_cache(const at::Tensor &key,
                          const c10::optional<at::Tensor> &value,
                          at::Tensor &key_cache,
                          const c10::optional<at::Tensor> &value_cache,
                          const at::Tensor &context_lengths,
                          const int64_t max_context_len,
                          bool packed,
                          const c10::optional<at::Tensor> &context_seq_offset,
                          const c10::optional<at::Tensor> &cache_bs_id,
                          const c10::optional<at::Tensor> &cache_seqlen_offset);

void reshape_paged_cache(const torch::Tensor &k,
                         const c10::optional<torch::Tensor> &v,
                         torch::Tensor &k_cache,
                         const c10::optional<torch::Tensor> &v_cache,
                         const torch::Tensor &slot_mapping);

void quant_to_paged_cache(const torch::Tensor &k,
                          const c10::optional<torch::Tensor> &v,
                          torch::Tensor &k_cache,
                          const c10::optional<torch::Tensor> &v_cache,
                          torch::Tensor &k_cache_scale,
                          const c10::optional<torch::Tensor> &v_cache_scale,
                          const torch::Tensor &slot_mapping);

void offline_quant_to_paged_cache(const torch::Tensor &k,
                                  const c10::optional<torch::Tensor> &v,
                                  const torch::Tensor &k_cache_scale,
                                  const c10::optional<torch::Tensor> &v_cache_scale,
                                  const torch::Tensor &slot_mapping,
                                  torch::Tensor &k_cache,
                                  const c10::optional<torch::Tensor> &v_cache);

void quant_to_linear_cache(const at::Tensor &key,
                           const c10::optional<at::Tensor> &value,
                           at::Tensor &key_cache,
                           const c10::optional<at::Tensor> &value_cache,
                           at::Tensor &key_cache_scale,
                           const c10::optional<at::Tensor> &value_cache_scale,
                           const at::Tensor &context_lengths,
                           const int64_t max_context_len,
                           bool packed,
                           const c10::optional<at::Tensor> &context_seq_offset,
                           const c10::optional<at::Tensor> &cache_bs_id,
                           const c10::optional<at::Tensor> &cache_seqlen_offset,
                           const int64_t quant_bit);

void offline_quant_to_linear_cache(const at::Tensor &key,
                                   const c10::optional<at::Tensor> &value,
                                   at::Tensor &key_cache,
                                   const c10::optional<at::Tensor> &value_cache,
                                   const at::Tensor &key_cache_scale,
                                   const c10::optional<at::Tensor> &value_cache_scale,
                                   const at::Tensor &context_lengths,
                                   const int64_t max_context_len,
                                   const int64_t quant_mode,
                                   const bool packed,
                                   const c10::optional<at::Tensor> &context_seq_offset,
                                   const c10::optional<at::Tensor> &cache_bs_id,
                                   const c10::optional<at::Tensor> &cache_seqlen_offset);

void apply_rotary(const torch::Tensor &input,
                  const torch::Tensor &sin_cache,
                  const torch::Tensor &cos_cache,
                  const c10::optional<torch::Tensor> &position_ids,
                  const c10::optional<torch::Tensor> &cu_seqlens,
                  bool interleaved,
                  bool discrete,
                  bool dynamic_ntk,
                  int64_t max_seqlen);

void swap_blocks(torch::Tensor &dst,
                 const torch::Tensor &src,
                 const c10::Dict<int64_t, int64_t> &block_mapping);

void copy_blocks(const std::vector<torch::Tensor> &k_caches,
                 const std::vector<torch::Tensor> &v_caches,
                 const c10::Dict<int64_t, c10::List<int64_t>> &block_mapping);
std::tuple<std::vector<at::Tensor>, std::vector<at::Tensor>> copy_blocks_out_of_place(
    const std::vector<torch::Tensor> &k_caches,
    const std::vector<torch::Tensor> &v_caches,
    const c10::Dict<int64_t, c10::List<int64_t>> &block_mapping);

at::Tensor fused_moe(const at::Tensor &hidden_states,
                     const at::Tensor &gating_output,
                     const at::Tensor &w1,
                     const at::Tensor &w2,
                     const c10::optional<at::Tensor> &bias1,
                     const c10::optional<at::Tensor> &bias2,
                     const c10::optional<at::Tensor> &residual,
                     const c10::optional<at::Tensor> &input_smooth,
                     const c10::optional<at::Tensor> &act_smooth,
                     const c10::optional<at::Tensor> &w1_scale,
                     const c10::optional<at::Tensor> &w2_scale,
                     const c10::optional<at::List<int64_t>> &w1_quant_flag,
                     const c10::optional<at::List<int64_t>> &w2_quant_flag,
                     const int64_t topk,
                     const bool renormalize,
                     const bool gated,
                     const std::string &act_mode,
                     const int64_t start_expert_id,
                     const int64_t block_n = 0,
                     const int64_t cncl_comm = 0);

// d = (a * b + bias) * alpha + c * beta
// d = (((a * a_scale * b * b_scale) * gemm_output_scale + bias) * alpha + c * beta) * d_scale
at::Tensor quant_matmul(
    const at::Tensor &a_tensor,                          // input
    const c10::optional<at::Tensor> &a_scale,            // input scale, smooth quant
    const c10::optional<at::Tensor> &a_zero,             // input zero
    const at::Tensor &b_tensor,                          // weight
    const c10::optional<at::Tensor> &b_scale,            // weight scale, smooth quant
    const c10::optional<at::Tensor> &b_zero,             // weight zero
    const c10::optional<at::Tensor> &bias,               // bias
    const c10::optional<at::Tensor> &c_tensor,           // residual
    const c10::optional<at::Tensor> &c_scale,            // residual scale
    const c10::optional<at::Tensor> &c_zero,             // residual zero
    const c10::optional<at::Tensor> &gemm_output_scale,  // for dequant, weight only
    const c10::optional<at::Tensor> &gemm_output_zero,   // for dequant, preserve
    const c10::optional<std::string> &data_type,         // output data type
    const c10::optional<at::Tensor> &d,                  // output
    const std::string &quant_algo,                       // quant type
    const std::string &a_quant_layout,                   // input quant mode
    const std::string &b_quant_layout,                   // weight quant mode
    int64_t quant_bit_size = 8,                          // weight quant bit size
    const std::string &act_mode = "none",                // act mode
    bool use_hp_active = false,  // for active, whether uses high precision mode
    double act_coef = 1.f,       // active_coef
    double alpha = 1.f,          // for quant matmul
    double beta = 1.f,           // for residual
    bool trans_a = false,        // whether transpose input
    bool trans_b = true);        // whether transpose weight

at::Tensor quant_matmul_allreduce(const int64_t cncl_comm,
                                  const at::Tensor &a_tensor,
                                  const c10::optional<at::Tensor> &a_scale,
                                  const c10::optional<at::Tensor> &a_zero,
                                  const at::Tensor &b_tensor,
                                  const c10::optional<at::Tensor> &b_scale,
                                  const c10::optional<at::Tensor> &b_zero,
                                  const c10::optional<at::Tensor> &bias,
                                  const c10::optional<at::Tensor> &c_tensor,
                                  const c10::optional<at::Tensor> &c_scale,
                                  const c10::optional<at::Tensor> &c_zero,
                                  const c10::optional<at::Tensor> &gemm_output_scale,
                                  const c10::optional<at::Tensor> &gemm_output_zero,
                                  const c10::optional<std::string> &data_type,
                                  const c10::optional<at::Tensor> &d,
                                  const std::string &quant_algo,
                                  const std::string &a_quant_layout,
                                  const std::string &b_quant_layout,
                                  int64_t quant_bit_size,
                                  double alpha,
                                  double beta,
                                  bool trans_a,
                                  bool trans_b,
                                  const int64_t block_m);

void active(const torch::Tensor &input,
            const torch::Tensor &output,
            const c10::optional<torch::Tensor> &bias,
            const c10::optional<torch::Tensor> &cusum_token_count,
            const std::string &act_mode,
            bool is_gated,
            int64_t start_expert_id,
            int64_t expert_size,
            double active_coef = 1.0);

void smooth_quant(const at::Tensor &input,
                  const at::Tensor &input_scale,
                  const at::Tensor &output,
                  const at::Tensor &output_scale,
                  const c10::optional<at::Tensor> &input_zero,
                  const c10::optional<at::Tensor> &token_count,
                  const c10::optional<at::Tensor> &gather_index,
                  const c10::optional<at::Tensor> &gather_index_start_position,
                  const std::string &quant_mode,
                  const bool &dynamic_quant);

at::Tensor matmul(const at::Tensor &a,
                  const at::Tensor &b,
                  const c10::optional<at::Tensor> &d,
                  const c10::optional<at::Tensor> &bias,
                  const c10::optional<at::Tensor> &c,
                  const c10::optional<std::string> &dtype,
                  const std::string &act_mode,
                  double alpha,
                  double beta,
                  bool fast_act,
                  bool approximate,
                  double a_scale,
                  double b_scale,
                  bool trans_a,
                  bool trans_b);

at::Tensor matmul_allreduce(const int64_t cncl_comm,
                            const at::Tensor &a,
                            const at::Tensor &b,
                            const c10::optional<at::Tensor> &bias,
                            const c10::optional<at::Tensor> &c,
                            const c10::optional<at::Tensor> &d,
                            const double alpha,
                            const double beta,
                            const int64_t block_m);

at::Tensor group_gemm(const at::Tensor &a_tensor,
                      const at::Tensor &b_tensor,
                      const at::Tensor &m_list,
                      const c10::optional<at::Tensor> &gather_idx,
                      const c10::optional<at::Tensor> &c_tensor,
                      const c10::optional<at::Tensor> &alpha,
                      const c10::optional<at::Tensor> &beta,
                      const c10::optional<at::Tensor> &a_scale,
                      const c10::optional<at::Tensor> &b_scale,
                      const c10::optional<at::Tensor> &bias,
                      const c10::optional<std::string> &data_type,
                      const c10::optional<at::List<int64_t>> &quant_flag,
                      const c10::optional<at::Tensor> &b_offset,
                      const int64_t max_m /*max m in m_list*/);

void preload(const torch::Tensor &weight, const int64_t size);

at::Tensor group_gemm_combine_result_allreduce(const int64_t cncl_comm,
                                               const at::Tensor &a_tensor,
                                               const at::Tensor &b_tensor,
                                               const at::Tensor &m_list,
                                               const at::Tensor &combine_idx,
                                               const at::Tensor &combine_weight,
                                               const c10::optional<at::Tensor> &c_tensor,
                                               const c10::optional<at::Tensor> &alpha,
                                               const c10::optional<at::Tensor> &beta,
                                               const c10::optional<at::Tensor> &a_scale,
                                               const c10::optional<at::Tensor> &b_scale,
                                               const c10::optional<std::string> &data_type,
                                               const int64_t num_token,
                                               const int64_t topk,
                                               const int64_t block_n);

at::Tensor flash_attn_sq_mm_allreduce(const int64_t cncl_comm,
                                      const at::Tensor &q,
                                      const at::Tensor &k,
                                      const at::Tensor &v,
                                      const c10::optional<at::Tensor> &cu_seq_lens_q,
                                      const c10::optional<at::Tensor> &cu_seq_lens_kv,
                                      const c10::optional<at::Tensor> &alibi_slope,
                                      const c10::optional<at::Tensor> &attn_bias,
                                      const at::Tensor &smooth,
                                      const at::Tensor &weight,
                                      const at::Tensor &weight_scale,
                                      const c10::optional<at::Tensor> &bias,
                                      const int64_t max_seq_len_q,
                                      const int64_t max_seq_len_kv,
                                      const double softmax_scale,
                                      const bool is_causal,
                                      const int64_t window_size_left,
                                      const int64_t window_size_right,
                                      const std::string &compute_dtype,
                                      const int64_t block_seq);

std::vector<at::Tensor> moe_softmax_topk(const at::Tensor &input,
                                         int64_t topk,
                                         int64_t num_expert_group,
                                         int64_t topk_group,
                                         bool normalize,
                                         const c10::optional<at::Tensor> &mask,
                                         const std::string &normed_by);

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);

std::vector<at::Tensor> moe_gen_idx(const torch::Tensor &expert_id, int64_t expert_num);

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);

void fused_rope(at::Tensor &qkv,
                at::Tensor &key_cache_hp,
                at::Tensor &value_cache_hp,
                const c10::optional<at::Tensor> &key_cache_lp,
                const c10::optional<at::Tensor> &value_cache_lp,
                const at::Tensor &sin_table,
                const at::Tensor &cos_table,
                const at::Tensor &position_ids,
                const at::Tensor &gamma,
                const at::Tensor &beta,
                const c10::optional<at::Tensor> &key_scale_hp,
                const c10::optional<at::Tensor> &value_scale_hp,
                const c10::optional<at::Tensor> &key_scale_lp,
                const c10::optional<at::Tensor> &value_scale_lp,
                const c10::optional<at::Tensor> &cache_bs_id_hp,
                const c10::optional<at::Tensor> &cache_seq_offsets_hp,
                const c10::optional<at::Tensor> &cache_bs_id_lp,
                const c10::optional<at::Tensor> &cache_seq_offsets_lp,
                const c10::optional<at::Tensor> &slot_mapping_hp,
                const c10::optional<at::Tensor> &slot_mapping_lp,
                const double eps);

void batch_matmul(const at::Tensor &a,
                  const at::Tensor &b,
                  const at::Tensor &c,
                  double alpha,
                  double beta,
                  double a_scale,
                  double b_scale,
                  bool trans_a,
                  bool trans_b);

at::Tensor moe_cast_gating(const at::Tensor &input, const at::Tensor &weight);

void update_out_and_lse(at::Tensor &out,
                        at::Tensor &lse,
                        const at::Tensor &block_out,
                        const at::Tensor &block_lse,
                        const c10::optional<at::Tensor> &seq_offsets,
                        const c10::optional<at::Tensor> &cu_seqs,
                        const c10::optional<at::Tensor> &block_cu_seqs);

void cnpxPush(const OpTheory &op);

void cnpxPop(void);

void dequant_from_linear_cache(at::Tensor &key,
                               const c10::optional<at::Tensor> &value,
                               const at::Tensor &key_cache,
                               const c10::optional<at::Tensor> &value_cache,
                               const at::Tensor &key_quant_scale,
                               const c10::optional<at::Tensor> &value_quant_scale,
                               const at::Tensor &context_lengths,
                               const int64_t max_context_len,
                               const c10::optional<at::Tensor> &context_seq_offset,
                               const c10::optional<at::Tensor> &cache_bs_id,
                               const c10::optional<at::Tensor> &cache_seq_offset,
                               const int64_t quant_mode,
                               const int64_t quant_bit);

void dequant_from_paged_cache(at::Tensor &key,
                              const c10::optional<at::Tensor> &value,
                              const at::Tensor &key_cache,
                              const c10::optional<at::Tensor> &value_cache,
                              const at::Tensor &key_cache_quant_scale,
                              const c10::optional<at::Tensor> &value_cache_quant_scale,
                              const at::Tensor &context_lengths,
                              int64_t max_context_len,
                              const c10::optional<at::Tensor> &context_seq_offset,
                              const at::Tensor &block_tables,
                              int64_t quant_mode,
                              int64_t quant_bit);

}  // namespace torch_api
}  // namespace tmo

#endif  // CSRC_TORCH_API_TORCH_OPS_API_H_
