/*************************************************************************
 * 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_OP_THEORY_H_
#define CSRC_TORCH_API_OP_THEORY_H_

#include "utils.h"

namespace tmo {
namespace torch_api {

// base class
struct OpTheory {
  virtual size_t getTheoryCalc(void) const = 0;
  virtual size_t getTheoryIO(void) const = 0;
  virtual cnnlDataType_t getCalcDtype(void) const = 0;
  virtual const std::string &getOpName(void) const = 0;
};

struct MatmulTheory : public OpTheory {
  int m_, k_, n_;
  bool has_res_;
  cnnlDataType_t calc_type_, data_type_;
  std::string name_;
  MatmulTheory(int m,
               int k,
               int n,
               bool has_res,
               cnnlDataType_t calc_type,
               cnnlDataType_t data_type)
      : OpTheory(),
        m_(m),
        k_(k),
        n_(n),
        has_res_(has_res),
        calc_type_(calc_type),
        data_type_(data_type),
        name_("matmul") {}
  // lt: m * n * k, ct: add_bias + act
  size_t getTheoryCalc(void) const override {
    size_t size_calc = (size_t)m_ * n_ * k_;
    return size_calc;
  }
  size_t getTheoryIO(void) const override {
    size_t calc_dw, data_dw;
    cnnlGetSizeOfDataType(calc_type_, &calc_dw);
    cnnlGetSizeOfDataType(data_type_, &data_dw);
    return calc_dw * (m_ * k_ + n_ * k_) + data_dw * m_ * n_ * (1 + has_res_);
  }
  cnnlDataType_t getCalcDtype(void) const override { return calc_type_; }
  const std::string &getOpName(void) const override { return name_; }
};

struct QuantMatmulTheory : public OpTheory {
  int m_, k_, n_, quant_bit_size_;
  bool has_res_;
  cnnlDataType_t data_type_;
  std::string quant_algo_, name_;
  QuantMatmulTheory(int m,
                    int k,
                    int n,
                    int quant_bit_size,
                    bool has_res,
                    cnnlDataType_t data_type,
                    std::string quant_algo)
      : OpTheory(),
        m_(m),
        k_(k),
        n_(n),
        quant_bit_size_(quant_bit_size),
        has_res_(has_res),
        data_type_(data_type),
        quant_algo_(quant_algo),
        name_("quant_matmul") {}
  size_t getTheoryCalc(void) const override {
    size_t size_calc = (size_t)m_ * n_ * k_;
    return size_calc;
  }
  size_t getTheoryIO(void) const override {
    size_t io_size = 0, data_dw = 0;
    cnnlGetSizeOfDataType(data_type_, &data_dw);
    if (quant_algo_ == "weight_only") {
      io_size = (data_dw * m_ * k_ + m_ * n_ * (1 + has_res_)) + n_ * k_ * (quant_bit_size_ / 8.0);
    } else {
      io_size = data_dw * m_ * n_ * (1 + has_res_) + (m_ * k_ + n_ * k_) * (quant_bit_size_ / 8.0);
    }
    return io_size;
  }
  cnnlDataType_t getCalcDtype(void) const override {
    if (quant_algo_ == "weight_only")
      return data_type_;
    else
      return CNNL_DTYPE_INT8;
  }
  const std::string &getOpName(void) const override { return name_; }
};

struct GroupGemmTheory : public OpTheory {
  int total_m_, experts_num_, k_, n_, quant_bit_size_;
  bool has_res_;
  cnnlDataType_t calc_type_, data_type_;
  std::string name_;
  GroupGemmTheory(int total_m,
                  int experts_num,  // The max number of processed experts.
                  int k,
                  int n,
                  bool has_res,
                  cnnlDataType_t calc_type,
                  cnnlDataType_t data_type)
      : OpTheory(),
        total_m_(total_m),
        experts_num_(experts_num),
        k_(k),
        n_(n),
        has_res_(has_res),
        calc_type_(calc_type),
        data_type_(data_type),
        name_("group_gemm") {}  // smooth_quant or float-point
  size_t getTheoryCalc(void) const override {
    size_t size_calc = (size_t)total_m_ * n_ * k_;
    return size_calc;
  }
  size_t getTheoryIO(void) const override {
    size_t calc_dw, data_dw;
    cnnlGetSizeOfDataType(calc_type_, &calc_dw);
    cnnlGetSizeOfDataType(data_type_, &data_dw);
    // assume load max_num weights
    size_t size_io = calc_dw * ((size_t)total_m_ * k_ + (size_t)experts_num_ * n_ * k_) +
                     data_dw * total_m_ * n_ * (1 + has_res_);
    return size_io;
  }
  cnnlDataType_t getCalcDtype(void) const override { return calc_type_; }
  const std::string &getOpName(void) const override { return name_; }
};

struct FlashAttnTheory : public OpTheory {
  int batch_, total_q_, total_k_, head_num_q_, head_num_kv_, head_size_qk_, head_size_v_;
  bool is_causal_;
  cnnlDataType_t data_type_;
  std::string name_;
  FlashAttnTheory(int batch,
                  int total_q,
                  int total_k,
                  int head_num_q,
                  int head_num_kv,
                  int head_size_qk,
                  int head_size_v,
                  bool is_causal,
                  cnnlDataType_t data_type)
      : OpTheory(),
        batch_(batch),
        total_q_(total_q),
        total_k_(total_k),
        head_num_q_(head_num_q),
        head_num_kv_(head_num_kv),
        head_size_qk_(head_size_qk),
        head_size_v_(head_size_v),
        is_causal_(is_causal),
        data_type_(data_type),
        name_("flash_attention") {}
  size_t getTheoryCalc(void) const override {
    // q * k + qk * v
    size_t seq_q = total_q_ / batch_;
    size_t seq_kv = total_k_ / batch_;
    if (is_causal_) seq_kv = seq_kv - 0.5 * seq_q;
    size_t size_lt = (size_t)batch_ * head_num_q_ * seq_q * seq_kv * (head_size_qk_ + head_size_v_);
    return size_lt;
  }
  size_t getTheoryIO(void) const override {
    size_t data_dw;
    cnnlGetSizeOfDataType(data_type_, &data_dw);
    size_t size_io = data_dw * (head_num_q_ * total_q_ + head_num_kv_ * total_k_) *
                     (head_size_qk_ + head_size_v_);
    return size_io;
  }
  cnnlDataType_t getCalcDtype(void) const override { return data_type_; }
  const std::string &getOpName(void) const override { return name_; }
};

struct SingleQueryCachedKvAttnTheory : public OpTheory {
  int batch_, seq_q_, head_num_q_, head_num_kv_, num_block_, block_size_, head_size_qk_,
      head_size_v_, kv_cache_quant_bit_size_;
  cnnlQuantizeLayout_t cache_quant_layout_;
  cnnlDataType_t data_type_;
  std::string name_;
  SingleQueryCachedKvAttnTheory(int batch,
                                int seq_q,
                                int head_num_q,
                                int head_num_kv,
                                int num_block,
                                int block_size,
                                int head_size_qk,
                                int head_size_v,
                                int kv_cache_quant_bit_size,
                                cnnlQuantizeLayout_t cache_quant_layout,
                                cnnlDataType_t data_type)
      : OpTheory(),
        batch_(batch),
        seq_q_(seq_q),
        head_num_q_(head_num_q),
        head_num_kv_(head_num_kv),
        num_block_(num_block),
        block_size_(block_size),
        head_size_qk_(head_size_qk),
        head_size_v_(head_size_v),
        kv_cache_quant_bit_size_(kv_cache_quant_bit_size),
        cache_quant_layout_(cache_quant_layout),
        data_type_(data_type),
        name_("single_query_cached_kv_attn") {}
  size_t getTheoryCalc(void) const override {
    // q * k + qk * v
    size_t size_lt =
        (size_t)seq_q_ * head_num_q_ * num_block_ * block_size_ * (head_size_qk_ + head_size_v_);
    return size_lt;
  }
  size_t getTheoryIO(void) const override {
    size_t data_dw = 0, dw_float = 0, size_scale = 0;
    cnnlGetSizeOfDataType(data_type_, &data_dw);
    cnnlGetSizeOfDataType(CNNL_DTYPE_FLOAT, &dw_float);
    size_t cache_dw = kv_cache_quant_bit_size_;
    if (kv_cache_quant_bit_size_ == -1) {
      cache_dw = data_dw * 8;
    } else {
      if (cache_quant_layout_ == CNNL_QUANTIZE_PER_CHANNEL) {
        size_scale = dw_float * head_num_kv_ * (head_size_qk_ + head_size_v_);
      } else {
        size_scale = dw_float * num_block_ * head_num_kv_ * block_size_ * 2;
      }
    }
    size_t size_io =
        size_scale + data_dw * batch_ * seq_q_ * head_num_q_ * (head_size_qk_ + head_size_v_) +
        (cache_dw / 8) * head_num_kv_ * num_block_ * block_size_ * (head_size_qk_ + head_size_v_);
    return size_io;
  }
  cnnlDataType_t getCalcDtype(void) const override { return data_type_; }
  const std::string &getOpName(void) const override { return name_; }
};

struct FusedNormTheory : public OpTheory {
  int num_token_, hidden_size_;
  bool has_res_, has_bias_, quant_out_, dynamic_quant_, store_output_before_norm_;
  cnnlDataType_t data_type_;
  std::string norm_mode_, name_;
  FusedNormTheory(int num_token,
                  int hidden_size,
                  bool has_res,
                  bool has_bias,
                  bool quant_out,
                  bool dynamic_quant,
                  bool store_output_before_norm,
                  cnnlDataType_t data_type,
                  std::string norm_mode)
      : OpTheory(),
        num_token_(num_token),
        hidden_size_(hidden_size),
        has_res_(has_res),
        has_bias_(has_bias),
        quant_out_(quant_out),
        dynamic_quant_(dynamic_quant),
        store_output_before_norm_(store_output_before_norm),
        data_type_(data_type),
        norm_mode_(norm_mode),
        name_("fused_norm") {}
  size_t getTheoryCalc(void) const override {
    size_t size_ct = 0;
    if (norm_mode_ == "layernorm") {
      size_ct = 7;  // avg + sub + square + sum + cycle_mul + mul gamma + add beta,
    } else if (norm_mode_ == "rmsnorm") {
      size_ct = 4;  // square + avg + cycle_mul + mul gamma
    }
    size_ct += (has_res_ + has_bias_);
    if (dynamic_quant_)
      size_ct += 4;
    else if (quant_out_)
      size_ct += 1;
    return size_ct * num_token_ * hidden_size_;
  }
  size_t getTheoryIO(void) const override {
    size_t data_dw, dw_int8, dw_float;
    cnnlGetSizeOfDataType(CNNL_DTYPE_FLOAT, &dw_float);
    cnnlGetSizeOfDataType(data_type_, &data_dw);
    cnnlGetSizeOfDataType(CNNL_DTYPE_INT8, &dw_int8);
    size_t size_io = data_dw * num_token_ * hidden_size_;
    if (has_res_) size_io += data_dw * num_token_ * hidden_size_;
    if (has_bias_) size_io += data_dw * hidden_size_;
    size_io += hidden_size_ * data_dw * (1 + (norm_mode_ == "layernorm"));  // gamma&beta
    if (quant_out_) {
      size_io += dw_int8 * num_token_ * hidden_size_;                          // output
      size_io += dw_float * (hidden_size_ + dynamic_quant_ ? num_token_ : 0);  // scale
    } else {
      size_io += data_dw * num_token_ * hidden_size_;
    }
    if (store_output_before_norm_) size_io += data_dw * num_token_ * hidden_size_;
    return size_io;
  }
  cnnlDataType_t getCalcDtype(void) const override { return data_type_; }
  const std::string &getOpName(void) const override { return name_; }
};

struct SmoothQuantTheory : public OpTheory {
  int m_, ci_;
  cnnlDataType_t data_type_;
  bool dynamic_quant_;
  std::string name_;
  SmoothQuantTheory(int m, int ci, cnnlDataType_t data_type, bool dynamic_quant)
      : OpTheory(),
        m_(m),
        ci_(ci),
        data_type_(data_type),
        dynamic_quant_(dynamic_quant),
        name_("smooth_quant_online") {}
  size_t getTheoryCalc(void) const override {
    if (!dynamic_quant_) {
      return 3 * (size_t)m_ * ci_;  // convert + mul scale + convert
    }
    size_t calc_scale = m_;  // 128 / maxvalue
    size_t size_calc = (2 /*max + min*/ + 2 /*mul scale*/ + 2 /*convert*/) * (size_t)m_ * ci_;
    return calc_scale + size_calc;
  }
  size_t getTheoryIO(void) const override {
    size_t data_dw, dw_int8, dw_float;
    cnnlGetSizeOfDataType(CNNL_DTYPE_FLOAT, &dw_float);
    cnnlGetSizeOfDataType(CNNL_DTYPE_INT8, &dw_int8);
    cnnlGetSizeOfDataType(data_type_, &data_dw);
    size_t scale_size = dw_float * (ci_ + dynamic_quant_ ? m_ : 0);
    return scale_size + (data_dw + dw_int8) * (m_ * ci_);
  }
  cnnlDataType_t getCalcDtype(void) const override { return data_type_; }
  const std::string &getOpName(void) const override { return name_; }
};

struct GroupAddBiasActiveTheory : public OpTheory {
  int expert_size_, num_expand_token_, ci_, inner_size_;
  bool has_gate_, is_addbias_;
  cnnlDataType_t data_type_;
  std::string act_mode_, name_;
  GroupAddBiasActiveTheory(int expert_size,
                           int num_expand_token,
                           int inner_size,
                           bool has_gate,
                           bool is_addbias,
                           cnnlDataType_t data_type,
                           std::string act_mode)
      : OpTheory(),
        expert_size_(expert_size),
        num_expand_token_(num_expand_token),
        inner_size_(inner_size),
        has_gate_(has_gate),
        is_addbias_(is_addbias),
        data_type_(data_type),
        act_mode_(act_mode),
        name_("group_add_bias_active") {}
  size_t getTheoryCalc(void) const override {
    size_t size_vec = 0, act_size = 0;
    if (is_addbias_) size_vec += (size_t)num_expand_token_ * inner_size_ * (1 + has_gate_);
    if (act_mode_ == "silu") {  // silu = sigmoid + mul
      act_size += 2 * (size_t)num_expand_token_ * inner_size_;
    } else {  // gelu = 1 vector_compution
      act_size += (size_t)num_expand_token_ * inner_size_;
    }
    if (has_gate_) size_vec += (size_t)num_expand_token_ * inner_size_;
    return size_vec + act_size;
  }
  size_t getTheoryIO(void) const override {
    size_t data_dw = 0;
    cnnlGetSizeOfDataType(data_type_, &data_dw);
    size_t bias_io = is_addbias_ ? expert_size_ * ci_ * (1 + has_gate_) : 0;
    return data_dw * (bias_io + num_expand_token_ * inner_size_ * (2 + has_gate_));
  }
  cnnlDataType_t getCalcDtype(void) const override { return data_type_; }
  const std::string &getOpName(void) const override { return name_; }
};

struct MoeCombineResultTheory : public OpTheory {
  int num_token_, topk_, ci_, expert_size_;
  bool has_bias_, has_res_;
  cnnlDataType_t data_type_;
  std::string name_;
  MoeCombineResultTheory(int num_token,
                         int topk,
                         int ci,
                         int expert_size,
                         bool has_bias,
                         bool has_res,
                         cnnlDataType_t data_type)
      : OpTheory(),
        num_token_(num_token),
        topk_(topk),
        ci_(ci),
        expert_size_(expert_size),
        has_bias_(has_bias),
        has_res_(has_res),
        data_type_(data_type),
        name_("moe_combine_result") {}
  size_t getTheoryCalc(void) const override {
    size_t size_calc =
        (size_t)((has_bias_ + 3 /*cvt+fma+cvt*/) * topk_ + has_res_) * num_token_ * ci_;
    return size_calc;
  }
  size_t getTheoryIO(void) const override {
    size_t data_dw = 0;
    cnnlGetSizeOfDataType(data_type_, &data_dw);
    size_t bias_io = has_bias_ ? expert_size_ * ci_ : 0;
    size_t res_io = has_res_ ? num_token_ * ci_ : 0;
    return data_dw * (bias_io + res_io + num_token_ * topk_ * ci_ /*input*/ +
                      num_token_ * ci_ /*output*/ + num_token_ * topk_ /*reduce_weight*/);
  }
  cnnlDataType_t getCalcDtype(void) const override { return data_type_; }
  const std::string &getOpName(void) const override { return name_; }
};

}  // namespace torch_api
}  // namespace tmo

#endif  // CSRC_TORCH_API_OP_THEORY_H_
