/*************************************************************************
 * 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 <atomic>
#include <vector>
#include "comm_overlap.h"
#include "torch_api/torch_ops_api.h"

namespace tmo {
namespace torch_api {

using namespace torch::indexing;

struct QuantMatmulSplitM {
  const at::Tensor &a_;
  const c10::optional<at::Tensor> &a_scale_;
  const c10::optional<at::Tensor> &a_zero_;
  const at::Tensor &b_;
  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_;
  const c10::optional<at::Tensor> &c_scale_;
  const c10::optional<at::Tensor> &c_zero_;
  const c10::optional<at::Tensor> &output_scale_;
  const c10::optional<at::Tensor> &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_;
  int64_t block_m_;
  bool has_a_scale_;
  bool has_a_zero_;
  bool has_b_scale_;
  bool has_b_zero_;
  bool has_bias_;
  bool has_c_;
  bool has_c_scale_;
  bool has_c_zero_;
  bool has_output_scale_;
  bool has_output_zero_;
  bool has_output_;
  bool has_dtype_;
  at::Tensor output_;

  QuantMatmulSplitM(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)
      : a_(a_tensor),
        a_scale_(a_scale),
        a_zero_(a_zero),
        b_(b_tensor),
        b_scale_(b_scale),
        b_zero_(b_zero),
        bias_(bias),
        c_(c_tensor),
        c_scale_(c_scale),
        c_zero_(c_zero),
        output_scale_(gemm_output_scale),
        output_zero_(gemm_output_zero),
        data_type_(data_type),
        d_(d),
        quant_algo_(quant_algo),
        a_quant_layout_(a_quant_layout),
        b_quant_layout_(b_quant_layout),
        quant_bit_size_(quant_bit_size),
        alpha_(alpha),
        beta_(beta),
        trans_a_(trans_a),
        trans_b_(trans_b),
        block_m_(block_m) {
    auto m = a_.size(0);
    if (block_m_ > m) {
      block_m_ = 1;
    } else if (block_m_ < 1) {
      block_m_ = m > 8192 ? 8 : (m < 2048 ? 1 : 4);
    }
    has_a_scale_ = a_scale_.has_value();
    has_a_zero_ = a_zero_.has_value();
    has_b_scale_ = b_scale_.has_value();
    has_b_zero_ = b_zero_.has_value();
    has_bias_ = bias_.has_value();
    has_c_ = c_.has_value();
    has_c_scale_ = c_scale_.has_value();
    has_c_zero_ = c_zero_.has_value();
    has_output_scale_ = output_scale_.has_value();
    has_output_zero_ = output_zero_.has_value();
    has_output_ = d_.has_value();
    has_dtype_ = data_type_.has_value();
    auto a_options = a_.options();
    if (has_output_) {
      output_ = d_.value();
    } else if (has_dtype_) {
      auto dtype = data_type_.value();
      TORCH_CHECK(dtype == "float" || dtype == "half" || dtype == "bfloat16",
                  "data type must be 'float', 'half' or 'bfloat16'");
      auto torch_dtype = str2TorchDtype(dtype);
      output_ = at::empty({a_.size(0), b_.size(0)}, a_options.dtype(torch_dtype));
    } else {
      output_ = at::empty({a_.size(0), b_.size(0)}, a_options);
    }
  }

  auto split(const at::Tensor &a, const int64_t block_id) {
    auto m = a.size(0);
    auto m_per_blk = m / block_m_;
    auto remain = m % block_m_;
    auto start = block_id * m_per_blk + std::min(block_id, remain);
    auto end = (block_id + 1) * m_per_blk + std::min(block_id + 1, remain);
    return a.narrow(0, start, end - start);
  }

  at::Tensor forward(const int64_t block_id) {
    auto Di = quant_matmul(
        split(a_, block_id), has_a_scale_ ? split(a_scale_.value(), block_id) : a_scale_,
        has_a_zero_ ? split(a_zero_.value(), block_id) : a_zero_, b_, b_scale_, b_zero_, bias_,
        has_c_ ? split(c_.value(), block_id) : c_,
        has_c_scale_ ? split(c_scale_.value(), block_id) : c_scale_,
        has_c_zero_ ? split(c_zero_.value(), block_id) : c_zero_, output_scale_, output_zero_,
        data_type_, split(output_, block_id), quant_algo_, a_quant_layout_, b_quant_layout_,
        quant_bit_size_, "none", false, 1.0, alpha_, beta_, trans_a_, trans_b_);
    return Di;
  }

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

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) {
  TORCH_CHECK(!trans_a && trans_b, "trans_a must be false and trans_b must be true");
  QuantMatmulSplitM quant_mm(a_tensor, a_scale, a_zero, b_tensor, b_scale, b_zero, bias, c_tensor,
                             c_scale, c_zero, gemm_output_scale, gemm_output_zero, data_type, d,
                             quant_algo, a_quant_layout, b_quant_layout, quant_bit_size, alpha,
                             beta, trans_a, trans_b, block_m);
  ParallelAllReduce<QuantMatmulSplitM> parallel_rs(cncl_comm);
  return parallel_rs(quant_mm);
}

}  //  namespace torch_api
}  //  namespace tmo
