/*************************************************************************
 * 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_ops_api.h"

namespace tmo {
namespace torch_api {

at::Tensor quant_matmul(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,
                        const std::string &act_mode,
                        bool use_hp_active,
                        double act_coef,
                        double alpha,
                        double beta,
                        bool trans_a,
                        bool trans_b) {
  // check device and dtype
  // only support weight only or smooth quant
  TORCH_CHECK(quant_algo == "weight_only" || quant_algo == "smooth_quant", "illegal quant mode");
  TORCH_CHECK(quant_bit_size == 4 || quant_bit_size == 8, "illegal quant bit size");
  TORCH_CHECK(trans_a == false && trans_b == true, "illegal trans mode");
  // create output tensor
  auto tensor_options = a_tensor.options();
  std::vector<int64_t> output_shape = {a_tensor.sizes()[0], b_tensor.sizes()[0]};
  at::Tensor output_tensor;
  if (d.has_value()) {
    output_tensor = d.value();
  } else if (data_type.has_value()) {
    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_tensor = at::empty(output_shape, tensor_options.dtype(torch_dtype));
  } else {
    output_tensor = at::empty(output_shape, tensor_options);
  }

  // create tensor desc
  auto descs = createTensorDescs(
      {a_tensor, a_scale.value_or(torch::Tensor()), a_zero.value_or(torch::Tensor()), b_tensor,
       b_scale.value_or(torch::Tensor()), b_zero.value_or(torch::Tensor()),
       bias.value_or(torch::Tensor()), c_tensor.value_or(torch::Tensor()),
       c_scale.value_or(torch::Tensor()), c_zero.value_or(torch::Tensor()), output_tensor,
       gemm_output_scale.value_or(torch::Tensor()), gemm_output_zero.value_or(torch::Tensor())});

  // get && set Op desc
  auto compute_dtype = getCnnlDataType(output_tensor.scalar_type());
  auto op_desc = tmo::op_desc::QuantMatmulDesc(quant_algo, a_quant_layout, b_quant_layout,
                                               quant_bit_size, compute_dtype, act_mode,
                                               use_hp_active, act_coef, trans_a, trans_b);

  auto handle = torch_mlu::getCurrentHandle();
  // get workspace size
  size_t workspace_size = 0;
  CNNL_CHECK_FATAL(cnnlGetLLMQuantMatmulWorkspaceSize(handle, op_desc, descs[0].get(),
                                                      descs[3].get(), descs[7].get(),
                                                      descs[10].get(), &workspace_size));
  auto workspace =
      at::empty({static_cast<int64_t>(workspace_size)}, a_tensor.options().dtype(at::kByte));

  // run forward
  const cnnlTensorDescriptor_t a_descs[] = {descs[0].get(), descs[1].get(), descs[2].get()};
  const cnnlTensorDescriptor_t b_descs[] = {descs[3].get(), descs[4].get(), descs[5].get()};
  const cnnlTensorDescriptor_t c_descs[] = {descs[7].get(), descs[8].get(), descs[9].get()};
  const cnnlTensorDescriptor_t d_descs[] = {descs[10].get(), nullptr, nullptr};
  const void *a_tensors[] = {getAtTensorPtr(a_tensor), getAtTensorPtr(a_scale),
                             getAtTensorPtr(a_zero)};
  const void *b_tensors[] = {getAtTensorPtr(b_tensor), getAtTensorPtr(b_scale),
                             getAtTensorPtr(b_zero)};
  const void *c_tensors[] = {getAtTensorPtr(c_tensor), getAtTensorPtr(c_scale),
                             getAtTensorPtr(c_zero)};
  void *d_tensors[] = {getAtTensorPtr(output_tensor), nullptr, nullptr};
  float alpha_arr[1] = {(float)alpha};
  float beta_arr[1] = {(float)beta};
  bool has_res = c_tensor.has_value();
  QuantMatmulTheory obj(a_tensor.size(0), a_tensor.size(1), b_tensor.size(0), quant_bit_size,
                        has_res, compute_dtype, quant_algo);
  cnpxPush(obj);
  CNNL_CHECK_FATAL(cnnlLLMQuantMatmul(
      handle, op_desc, nullptr, alpha_arr, a_descs, a_tensors, b_descs, b_tensors, nullptr,
      beta_arr, c_descs, c_tensors, descs[6].get(), getAtTensorPtr(bias), descs[11].get(),
      getAtTensorPtr(gemm_output_scale), descs[12].get(), getAtTensorPtr(gemm_output_zero),
      getAtTensorPtr(workspace), workspace_size, d_descs, d_tensors));
  cnpxPop();
  return output_tensor;
}

}  //  namespace torch_api
}  //  namespace tmo
