/*************************************************************************
 * 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 <cnnl.h>
#include <torch/extension.h>
#include "kernels/quantize.mluh"
#include "register_api.h"
#include "utils.h"

namespace tmo {
namespace kernel_test_api {
void per_head_quantize_test(at::Tensor &dst,
                            at::Tensor &scale,
                            const at::Tensor &src,
                            int bs,
                            int seq_len,
                            int head_num,
                            int head_size,
                            int dst_bs_stride,
                            int dst_seq_stride,
                            int dst_head_stride,
                            int src_bs_stride,
                            int src_seq_stride,
                            int src_head_stride) {
  MLU_TENSOR_CHECK_FATAL(dst);
  MLU_TENSOR_CHECK_FATAL(scale);
  MLU_TENSOR_CHECK_FATAL(src);

  cnnlHandle_t handle = torch_mlu::getCurrentHandle();
  cnrtQueue_t queue;
  cnnlGetQueue(handle, &queue);

  cnnlDataType_t dst_dtype = torch2cnnlDataType(torch::typeMetaToScalarType(dst.dtype()));
  cnnlDataType_t scale_dtype = torch2cnnlDataType(torch::typeMetaToScalarType(scale.dtype()));
  cnnlDataType_t src_dtype = torch2cnnlDataType(torch::typeMetaToScalarType(src.dtype()));

  size_t io_bytes =
      bs * seq_len * head_num * head_size * (dtype_size_map[src_dtype] + dtype_size_map[dst_dtype]);

  cnrtNotifier_t time_begin;
  cnrtNotifier_t time_end;

  CNRT_CHECK(cnrtNotifierCreate(&time_begin));
  CNRT_CHECK(cnrtNotifierCreate(&time_end));

  CNRT_CHECK(cnrtPlaceNotifier(time_begin, queue));
  TMO_KERNEL_CHECK_FATAL(invokeMluQuantizePerHead(
      queue, (void *)dst.data_ptr(), (void *)scale.data_ptr(), (void *)src.data_ptr(), dst_dtype,
      scale_dtype, src_dtype, bs, seq_len, head_num, head_size, dst_bs_stride, dst_seq_stride,
      dst_head_stride, src_bs_stride, src_seq_stride, src_head_stride));
  CNRT_CHECK(cnrtPlaceNotifier(time_end, queue));
  cnrtQueueSync(queue);
  float usec = 0.0f;
  CNRT_CHECK(cnrtNotifierDuration(time_begin, time_end, &usec));
  print_info(usec, io_bytes);
}

}  // namespace kernel_test_api
}  // namespace tmo
