#include <iostream>
#include <cuda_fp16.h>
#include "cutlass/cutlass.h"
#include "cutlass/gemm/device/default_gemm_configuration.h"

#include "mem_eff_attention/gemm_kernel_utils.h"
#include "mem_eff_attention/kernel_forward.h"


using namespace gemm_kernel_utils;


#ifndef CUDA_CHECK_ME_ATTN
#define CUDA_CHECK_ME_ATTN(expr, msg)                                          \
  do {                                                                         \
    cudaError_t status = (expr);                                               \
    if (status != cudaSuccess) {                                               \
      std::cerr << msg << " at " << __FILE__ << ": " << __LINE__ << std::endl; \
      throw std::runtime_error(cudaGetErrorString(status));                    \
    }                                                                          \
  } while (0)
#endif // CUDA_CHECK_ME_ATTN



void mem_eff_attention_418(void* output,
                   void* query,
                   void* key,
                   void* value,
                   int64_t* batch_size,
                   int64_t* seq_len_kv,
                   int64_t* seq_len_q,
                   int num_heads,
                   int head_size,
                   int head_size_v,
                   float p_dropout,
                   float softmax_scale,
                   bool is_causal,
                   bool fixed_seq_length_kv,
                   int32_t* lengths_kv,
                   bool fixed_seq_length_q,
                   int32_t* lengths_q,
                   void* workspace,
                   cudaStream_t stream)
    
{

    /*
    The code is based on fused_multihead_attention_fixed_seqlen.cu example in CUTLASS repo:
    https://github.com/NVIDIA/cutlass/blob/209faf7b94ce4ba573d27389fb643962e75d0581/examples/41_fused_multi_head_attention/fused_multihead_attention_fixed_seqlen.cu

    problem_sizes0 [b, m, n, k]
    [head_number * batch_size, m, mkv, k0]
    [head_number * batch_size, seq_length_q, seq_length_kv, head_size]

    problem_sizes1
    [head_number * batch_size, m, k1, mkv]
    [head_number * batch_size, seq_length_q, head_size_v, seq_length_kv]

    m = seq_len_q
    n = seq_len_kv
    k = head_size

    Q: B, M, K
    K: B, N, K
    P: B, M, N
    V: B, N, K
    O: B, M, K
    output: bs, seq_len_q, num_head, head_size
    */


    using ArchTag = cutlass::arch::Sm80;
    constexpr bool kIs64x64 = false;
    constexpr bool kSingleValueIteration = true;

    // Set grid size
    constexpr int64_t kQueriesPerBlock = kIs64x64 ? 64 : 32;
    constexpr int64_t kKeysPerBlock = kIs64x64 ? 64 : 128;
    if (kIs64x64 && head_size_v > kKeysPerBlock) {
        std::cerr << "WARNING: you will get better performance with `kIs64x64=false`";
    }
    if (kSingleValueIteration && head_size_v > kKeysPerBlock) {
        std::cerr << "ERROR  : Use kSingleValueIteration to keep output in RF. "         "This requires to have `head_size <= kKeysPerBlock` "         "but head_size_v=" << head_size_v << " and kKeysPerBlock=" << kKeysPerBlock << "";
        return;
    }
    if (!kSingleValueIteration && head_size_v <= kKeysPerBlock) {
        std::cerr << "WARNING: you will get better performance with `kSingleValueIteration=true` (keeps the output in RF rather than GMEM)";
    }

    using GemmType = DefaultGemmType<ArchTag, cutlass::half_t>;
    using OpClass = typename GemmType::OpClass;
    using DefaultConfig =
        typename cutlass::gemm::device::DefaultGemmConfiguration<
            OpClass,
            ArchTag,
            cutlass::half_t,
            cutlass::half_t,
            cutlass::half_t, // ElementC
            float // ElementAccumulator
            >;

    // If the head_size already meets the alignment requirement, then
    // it's safe to mark mem_align to be true to maximize the alignment
    // benefit. Otherwise, assign false to it to use the minimal alignment.
    constexpr const bool mem_align =
        (72 % DefaultConfig::kAlignmentA == 0) &&
        (72 % DefaultConfig::kAlignmentB == 0);
    using Attention = AttentionKernel<
        cutlass::half_t, // scalar_t
        ArchTag,
        mem_align, // memory is aligned
        kQueriesPerBlock,
        kKeysPerBlock,
        kSingleValueIteration
    >;

    typename Attention::Params p;
    {
        // set parameters
        p.query_ptr = static_cast<cutlass::half_t*>(query);
        p.key_ptr = static_cast<cutlass::half_t*>(key);
        p.value_ptr = static_cast<cutlass::half_t*>(value);
        p.logsumexp_ptr = nullptr; // Only needed for bw
        p.output_accum_ptr = nullptr;

        if (!fixed_seq_length_q) {
            p.seqlens_q_ptr = lengths_q;
        }
        if (!fixed_seq_length_kv) {
            p.seqlens_k_ptr = lengths_kv;
        }

        if (Attention::kNeedsOutputAccumulatorBuffer) {
          p.output_accum_ptr = static_cast<float*>(workspace);
        }
        p.output_ptr = static_cast<cutlass::half_t*>(output);

        p.num_heads = num_heads;
        p.num_batches = *batch_size;
        p.head_dim = head_size;
        p.head_dim_value = head_size_v;
        p.num_queries = *seq_len_q;
        p.num_keys = *seq_len_kv;
        p.causal = is_causal;


        p.q_strideM = head_size;
        p.k_strideM = head_size;
        p.v_strideM = head_size_v;

        p.q_strideH = p.q_strideM * (*seq_len_q);
        p.k_strideH = p.k_strideM * (*seq_len_kv);
        p.v_strideH = p.v_strideM * (*seq_len_kv);
        p.o_strideH = head_size_v;
        p.q_strideB = p.q_strideH * num_heads;
        p.k_strideB = p.k_strideH * num_heads;
        p.v_strideB = p.v_strideH * num_heads;
        p.o_strideB = head_size_v * (*seq_len_q) * num_heads;
    }

    // launch kernel
    constexpr auto kernel_fn = attention_kernel_batched_impl<Attention>;
    int smem_bytes = sizeof(typename Attention::SharedStorage);
    if (smem_bytes > 0xc000) {
      cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
    }
    if (!Attention::check_supported(p)) {
      std::string error_msg = std::string("Got error: kernel does not support these inputs") +
           " at " + __FILE__ + ": " + std::to_string(__LINE__);
      throw std::runtime_error(error_msg);
    }
    kernel_fn<<<p.getBlocksGrid(), p.getThreadsGrid(), smem_bytes, stream>>>(p);
}