/*************************************************************************
 * 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/generate_alibi_slope.mluh"
#include "register_api.h"
#include "utils.h"

namespace tmo {
namespace kernel_test_api {
at::Tensor alibi_slope_test(const at::Tensor &true_seq_lens,
                            const int batch,
                            const int tp_head_num,
                            const int tp_num,
                            const bool use_dynamic,
                            const int max_sequence_length) {
  MLU_TENSOR_CHECK_FATAL(true_seq_lens);
  cnnlHandle_t handle = torch_mlu::getCurrentHandle();
  cnrtQueue_t queue;
  cnnlGetQueue(handle, &queue);

  int head_num = tp_head_num * tp_num;
  std::shared_ptr<torch::Device> dev = std::make_shared<torch::Device>(true_seq_lens.device());
  torch::Tensor output = torch::zeros(
      {batch, head_num}, torch::dtype(torch::kFloat32).device(*dev).requires_grad(false));

  for (int tp_id = 0; tp_id < tp_num; tp_id++) {
    torch::Tensor tp_slopes = torch::zeros(
        {batch, tp_head_num}, torch::dtype(torch::kFloat32).device(*dev).requires_grad(false));
    TMO_KERNEL_CHECK_FATAL(invokeGenerateAlibiSlope(
        queue, tp_slopes.data_ptr(), true_seq_lens.data_ptr(), batch, tp_id * tp_head_num,
        tp_head_num, head_num, max_sequence_length, use_dynamic));
    for (int batch_id = 0; batch_id < batch; batch_id++) {
      CNRT_CHECK(
          cnrtMemcpyAsync((float *)output.data_ptr() + batch_id * head_num + tp_id * tp_head_num,
                          (float *)tp_slopes.data_ptr() + batch_id * tp_head_num,
                          tp_head_num * sizeof(float), queue, cnrtMemcpyDevToDev));
    }
  }

  return output;
}

}  // namespace kernel_test_api
}  // namespace tmo
