/*************************************************************************
 * 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 "comm_overlap.h"
#include "torch_api/torch_ops_api.h"
namespace tmo {
namespace torch_api {

using namespace torch::indexing;

struct FlashAttnSmoothQuantMatmul {
  at::Tensor q_;
  at::Tensor k_;
  at::Tensor v_;
  const c10::optional<at::Tensor> &cu_seq_lens_q_;
  const c10::optional<at::Tensor> &cu_seq_lens_kv_;
  const at::Tensor &smooth_;
  const at::Tensor &weight_;
  const at::Tensor &weight_scale_;
  const c10::optional<at::Tensor> &bias_;
  const std::string &compute_dtype_;
  std::string input_dtype_str_;
  int64_t max_seq_len_q_;
  int64_t max_seq_len_kv_;
  double softmax_scale_;
  bool is_causal_;
  int64_t block_seq_;
  std::vector<int> cumsum_seq_;
  at::Tensor output_;

  FlashAttnSmoothQuantMatmul(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 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 std::string &compute_dtype,
                             const int64_t block_seq)
      : q_(q),
        k_(k),
        v_(v),
        cu_seq_lens_q_(cu_seq_lens_q),
        cu_seq_lens_kv_(cu_seq_lens_kv),
        smooth_(smooth),
        weight_(weight),
        weight_scale_(weight_scale),
        bias_(bias),
        compute_dtype_(compute_dtype),
        max_seq_len_q_(max_seq_len_q),
        max_seq_len_kv_(max_seq_len_kv),
        softmax_scale_(softmax_scale),
        is_causal_(is_causal) {
    input_dtype_str_ = q.scalar_type() == at::kBFloat16 ? "bfloat16"
                       : q.scalar_type() == at::kHalf   ? "half"
                                                        : "float";
    bool is_pack = cu_seq_lens_q.has_value();
    if (is_pack) {
      TORCH_CHECK(cu_seq_lens_q.value().size(0) == 2 && cu_seq_lens_kv.value().size(0) == 2,
                  "only support 1 batch.")
      TORCH_CHECK(q.dim() == 3 && k.dim() == 3 && v.dim() == 3, "q,k,v must be 3-d in pack mode.")
      // 1 batch pack to pad
      q_ = q_.unsqueeze(0);
      k_ = k_.unsqueeze(0);
      v_ = v_.unsqueeze(0);
    } else {
      TORCH_CHECK(q.size(0) == 1, "only support 1 batch.")
      TORCH_CHECK(q.dim() == 4, "q must be 4-d in pad mode.")
    }
    int total_seq_q = q_.size(0) * q_.size(1);
    output_ = at::empty({total_seq_q, weight.size(0)}, q.options());
    if (total_seq_q >= 4096) {
      block_seq_ = 4;
      split_4(total_seq_q);
    } else {
      block_seq_ = 1;
      split_1(total_seq_q);
    }
  }

  void split_1(int seq) {
    cumsum_seq_.resize(2);
    cumsum_seq_[0] = 0;
    cumsum_seq_[1] = seq;
  }

  void split_4(int seq) {
    auto pad_up = [](int x, int y) -> int { return (x + y - 1) / y * y; };
    int seq_4 = pad_up(seq / 8, 256);
    int seq_3 = pad_up(seq / 8, 256);
    int seq_2 = pad_up(seq / 4, 256);
    int seq_1 = seq - seq_2 - seq_3 - seq_4;
    cumsum_seq_.resize(5);
    cumsum_seq_[0] = 0;
    cumsum_seq_[1] = cumsum_seq_[0] + seq_1;
    cumsum_seq_[2] = cumsum_seq_[1] + seq_2;
    cumsum_seq_[3] = cumsum_seq_[2] + seq_3;
    cumsum_seq_[4] = cumsum_seq_[3] + seq_4;
  }

  auto split_tensor(const at::Tensor &a, int64_t dim, int64_t start, int64_t end) {
    return a.narrow(dim, start, end - start);
  }

  at::Tensor forward(const int64_t block_id) {
    // flash attn
    int64_t end_kv = is_causal_ ? cumsum_seq_[block_id + 1] : cumsum_seq_[block_seq_];
    auto q_i = split_tensor(q_, 1, cumsum_seq_[block_id], cumsum_seq_[block_id + 1]);
    auto k_i = split_tensor(k_, 1, 0, end_kv);
    auto v_i = split_tensor(v_, 1, 0, end_kv);
    auto attn_out = at::empty_like(q_i);
    flash_attention(q_i, k_i, v_i, attn_out, c10::nullopt, c10::nullopt, c10::nullopt, c10::nullopt,
                    c10::nullopt, c10::nullopt, c10::nullopt, c10::nullopt, q_i.size(1),
                    k_i.size(1), softmax_scale_, is_causal_, -1, -1, compute_dtype_, false);
    // smooth quant
    auto smooth_quant_input =
        attn_out.flatten(-2, -1).flatten(0, 1);  // (b, s, hn, hs) -> (b*s, hn*hs)
    auto quant_out =
        at::empty(smooth_quant_input.sizes(),
                  torch::TensorOptions().dtype(torch::kInt8).device(smooth_quant_input.device()));
    auto quant_out_scale = at::empty({smooth_quant_input.size(0)}, smooth_.options());
    smooth_quant(smooth_quant_input, smooth_, quant_out, quant_out_scale, c10::nullopt,
                 c10::nullopt, c10::nullopt, c10::nullopt, "per_token", true);
    // quant matmul
    auto d_i = quant_matmul(
        quant_out, quant_out_scale, c10::nullopt, weight_, weight_scale_, c10::nullopt, bias_,
        c10::nullopt, c10::nullopt, c10::nullopt, weight_scale_, c10::nullopt, input_dtype_str_,
        split_tensor(output_, 0, cumsum_seq_[block_id], cumsum_seq_[block_id + 1]), "smooth_quant",
        "quantize_per_token", "quantize_per_channel", 8, "none", false, 1.0, 1.0, 1.0, false, true);
    return d_i;
  }

  int64_t getLoopNum() const { return block_seq_; }
  at::Tensor getOutput() const { return output_; }
};

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) {
  FlashAttnSmoothQuantMatmul mm(q, k, v, cu_seq_lens_q, cu_seq_lens_kv, smooth, weight,
                                weight_scale, bias, max_seq_len_q, max_seq_len_kv, softmax_scale,
                                is_causal, compute_dtype, block_seq);
  ParallelAllReduce<FlashAttnSmoothQuantMatmul> parallel_rs(cncl_comm);
  return parallel_rs(mm);
}

}  //  namespace torch_api
}  //  namespace tmo
