/*************************************************************************
 * 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 {

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) {
  bool use_beta = beta == 0 ? false : true;
  int batch = a.size(0);
  int m = trans_a ? a.size(-1) : a.size(-2);
  int n = trans_b ? b.size(-2) : b.size(-1);
  checkTensorSameAttr<TensorAttr::DEVICE>(a, b, c);
  // check contiguous
  TORCH_CHECK(a.is_contiguous(), "a must be contiguous.")
  TORCH_CHECK(b.is_contiguous(), "b must be contiguous.")
  TORCH_CHECK(c.is_contiguous(), "c must be contiguous.")
  CHECK_SHAPE(c, batch, m, n);

  // get cnnl data type and init output
  auto a_dtype = getCnnlDataType(a.scalar_type());
  auto b_dtype = getCnnlDataType(b.scalar_type());
  TORCH_CHECK(a_dtype == b_dtype, "a, b must be same dtype.");
  auto c_dtype = getCnnlDataType(c.scalar_type());
  if (a_dtype == CNNL_DTYPE_BFLOAT16) c_dtype = CNNL_DTYPE_FLOAT;
  // create tensor desc
  auto descs = createTensorDescs({a, b, c});
  CNNL_CHECK_FATAL(cnnlSetTensorDescriptorOnchipDataType(descs[0].get(), a_dtype));
  CNNL_CHECK_FATAL(cnnlSetTensorDescriptorOnchipDataType(descs[1].get(), b_dtype));
  CNNL_CHECK_FATAL(cnnlSetTensorDescriptorOnchipDataType(descs[2].get(), c_dtype));

  if (a_dtype == CNNL_DTYPE_INT8) {
    float int_max = 127.0;
    int quant_bit = 8;
    float max_a = int_max / a_scale;
    float max_b = int_max / b_scale;
    int pos_a = std::floor(std::log2(max_a) - (quant_bit - 2));
    int pos_b = std::floor(std::log2(max_b) - (quant_bit - 2));
    float new_a_scale = std::pow(2.0f, pos_a) * a_scale;
    float new_b_scale = std::pow(2.0f, pos_b) * b_scale;
    CNNL_CHECK_FATAL(cnnlSetTensorDescriptorPositionAndScale(descs[0].get(), pos_a, new_a_scale));
    CNNL_CHECK_FATAL(cnnlSetTensorDescriptorPositionAndScale(descs[1].get(), pos_b, new_b_scale));
    TORCH_CHECK(c_dtype != CNNL_DTYPE_BFLOAT16,
                "output dtype cannot be bfloat16 when a/b is fixed-point")
  }

  // get && set Op desc
  cnnlMatMulHeuristicResult_t heuristic_result;
  cnnlMatMulAlgo_t algo;
  auto bmm_ex_desc =
      tmo::op_desc::BatchMatMulDesc(heuristic_result, algo, use_beta, trans_a, trans_b);

  const torch_mlu::mlu::MLUGuard device_guard(a.device());
  auto handle = torch_mlu::getCurrentHandle();

  size_t workspace_size = 0;
  int requested_algo_count = 1;
  int returned_algo_count = 0;
  CNNL_CHECK_FATAL(cnnlGetBatchMatMulExAlgoHeuristic(
      handle, bmm_ex_desc, descs[0].get(), descs[1].get(), descs[2].get(), nullptr,
      requested_algo_count, &heuristic_result, &returned_algo_count));
  CNNL_CHECK_FATAL(cnnlGetBatchMatMulExHeuristicResult(heuristic_result, algo, &workspace_size));
  auto workspace = at::empty({static_cast<int64_t>(workspace_size)}, a.options().dtype(at::kByte));

  // run forward
  float alpha_f = alpha;
  float beta_f = beta;
  CNNL_CHECK_FATAL(cnnlBatchMatMulEx(handle, bmm_ex_desc, algo, &alpha_f, descs[0].get(),
                                     getAtTensorPtr(a), descs[1].get(), getAtTensorPtr(b), &beta_f,
                                     descs[2].get(), getAtTensorPtr(c), getAtTensorPtr(workspace),
                                     workspace_size));
}
}  //  namespace torch_api
}  //  namespace tmo
