/*************************************************************************
 * 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 flash_attention(const at::Tensor &q,
                     const at::Tensor &k,
                     const at::Tensor &v,
                     const at::Tensor &output,
                     const c10::optional<at::Tensor> &output_lse,
                     const c10::optional<at::Tensor> &cu_seq_lens_q,
                     const c10::optional<at::Tensor> &cu_seq_lens_kv,
                     const c10::optional<at::Tensor> &alibi_slope,
                     const c10::optional<at::Tensor> &attn_bias,
                     const c10::optional<at::Tensor> &k_cache_quant_scale,
                     const c10::optional<at::Tensor> &v_cache_quant_scale,
                     const c10::optional<at::Tensor> &block_tables,
                     const int64_t max_seq_len_q,
                     const int64_t max_seq_len_kv,
                     const double softmax_scale,
                     const bool is_causal,
                     const int64_t window_size_left,
                     const int64_t window_size_right,
                     const std::string &compute_dtype,
                     bool return_lse) {
  TORCH_CHECK(compute_dtype == "float" || compute_dtype == "half" || compute_dtype == "bfloat16",
              "compute_dtype must be 'float', 'half' or 'bfloat16'.");
  TORCH_CHECK(k_cache_quant_scale.has_value() == false, "k_cache_scale for reserve.");
  TORCH_CHECK(v_cache_quant_scale.has_value() == false, "v_cache_scale for reserve.");
  bool has_block_table = block_tables.has_value();
  bool is_pack = cu_seq_lens_q.has_value();
  int64_t batch = is_pack ? cu_seq_lens_q.value().size(0) - 1 : q.size(0);
  int qk_head_size = q.size(-1);
  int v_head_size = v.size(-1);

  // Check tensor type and tensor device.
  checkTensorSameAttr<TensorAttr::ALL>(q, k, v, output);
  // 3d for packed
  TORCH_CHECK(q.dim() == 3 || q.dim() == 4, "query must be 3d or 4d.");
  if (has_block_table) {
    TORCH_CHECK(block_tables.value().dim() == 2, "block_tables must be 2d.");
    TORCH_CHECK(k.dim() == 4, "with block table, key_cache must be 4d.");
    TORCH_CHECK(v.dim() == 4, "with block table, value_cache must be 4d.");
    int max_num_blocks_per_seq = block_tables.value().size(1);
    int num_blocks = k.size(0);
    int block_size = k.size(2);
    int k_head_num = k.size(1);
    CHECK_SHAPE(k, num_blocks, k_head_num, block_size, qk_head_size);
    CHECK_SHAPE(v, num_blocks, k_head_num, block_size, v_head_size);
    if (max_num_blocks_per_seq > 1) {  // paged
      CHECK_SHAPE(cu_seq_lens_kv.value(), batch + 1);
    }
  } else {
    // 3d for packed
    TORCH_CHECK(k.dim() == 3 || k.dim() == 4, "key_cache must be 3d or 4d.");
    TORCH_CHECK(v.dim() == 3 || v.dim() == 4, "value_cache must be 3d or 4d.");
    if (k.dim() == 3) {  // packed_kv
      CHECK_SHAPE(cu_seq_lens_kv.value(), batch + 1);
    }
  }
  // Convert torch tensor to tensor descs
  auto descs = createTensorDescs(
      {q, k, v, cu_seq_lens_q.value_or(at::Tensor()), cu_seq_lens_kv.value_or(at::Tensor()),
       alibi_slope.value_or(at::Tensor()), attn_bias.value_or(at::Tensor()),
       block_tables.value_or(at::Tensor()), output, output_lse.value_or(at::Tensor())});

  // Get current handle.
  const torch_mlu::mlu::MLUGuard device_guard(q.device());
  cnnlHandle_t handle = torch_mlu::getCurrentHandle();
  // Get workspace size and malloc workspace.
  cnnlDataType_t cnnl_compute_dtype = compute_dtype == "float"  ? CNNL_DTYPE_FLOAT
                                      : compute_dtype == "half" ? CNNL_DTYPE_HALF
                                                                : CNNL_DTYPE_BFLOAT16;
  size_t workspace_size = 0;
  CNNL_CHECK_FATAL(cnnlGetScaledDotProductAttnWorkspaceSize_v2(
      handle, nullptr /*op_desc*/, nullptr /*quant_desc*/, descs[0].get(), descs[1].get(),
      descs[2].get(), descs[3].get(), descs[4].get(), descs[6].get(), descs[5].get(),
      descs[7].get(), max_seq_len_q, max_seq_len_kv, is_causal, window_size_left, window_size_right,
      return_lse, CNNL_ACTIVATION_FAST, cnnl_compute_dtype, &workspace_size));
  auto workspace = at::empty({static_cast<int64_t>(workspace_size)}, q.options().dtype(at::kByte));

  int64_t total_q = is_pack ? q.size(0) : q.size(0) * q.size(1);
  int64_t total_k =
      has_block_table ? k.size(0) * k.size(2) : (k.dim() == 3 ? k.size(0) : k.size(0) * k.size(1));
  int64_t head_q = q.size(-2);
  int64_t head_k = has_block_table ? k.size(1) : k.size(-2);
  cnnlDataType_t data_dtype = getCnnlDataType(q.scalar_type());
  FlashAttnTheory obj(batch, total_q, total_k, head_q, head_k, qk_head_size, v_head_size, is_causal,
                      data_dtype);
  cnpxPush(obj);
  // call cnnl extra op.
  CNNL_CHECK_FATAL(cnnlScaledDotProductAttn_v3(
      handle, nullptr, nullptr, descs[0].get(), getAtTensorPtr(q), descs[1].get(),
      getAtTensorPtr(k), descs[2].get(), getAtTensorPtr(v), descs[3].get(),
      getAtTensorPtr(cu_seq_lens_q), descs[4].get(), getAtTensorPtr(cu_seq_lens_kv), nullptr,
      nullptr, descs[6].get(), getAtTensorPtr(attn_bias), descs[5].get(),
      getAtTensorPtr(alibi_slope), descs[7].get(), getAtTensorPtr(block_tables), max_seq_len_q,
      max_seq_len_kv, is_causal, window_size_left, window_size_right, softmax_scale,
      CNNL_ACTIVATION_FAST, cnnl_compute_dtype, getAtTensorPtr(workspace), workspace_size,
      return_lse, descs[9].get(), getAtTensorPtr(output_lse), descs[8].get(),
      getAtTensorPtr(output)));
  cnpxPop();
}

}  // namespace torch_api
}  // namespace tmo
