#include "flash_fwd_mla_kernel.h"

static constexpr int MaxBatchSize = 4096;

__global__ void __launch_bounds__(256, 1, 1) get_mla_metadata_kernel(
    __grid_constant__ const Mla_metadata_params params) {
  int* seqlens_k_ptr = params.seqlens_k_ptr;
  int* tile_scheduler_metadata_ptr = params.tile_scheduler_metadata_ptr;
  int* num_splits_ptr = params.num_splits_ptr;
  int batch_size = params.batch_size;
  int block_size_n = params.block_size_n;
  int fixed_overhead_num_blocks = params.fixed_overhead_num_blocks;
  int num_sm_parts = params.num_sm_parts;

  __shared__ int num_blocks_shared[MaxBatchSize];
  __shared__ int num_splits_shared[MaxBatchSize];

  int total_num_blocks = 0;
  for (int i = threadIdx.x; i < batch_size; i += 32) {
    int num_blocks = cutlass::ceil_div(seqlens_k_ptr[i], block_size_n);
    total_num_blocks += num_blocks + fixed_overhead_num_blocks;
    num_blocks_shared[i] = num_blocks;
  }
  for (int offset = 16; offset >= 1; offset /= 2) {
    total_num_blocks += __shfl_xor_sync(uint32_t(-1), total_num_blocks, offset);
  }
  __syncwarp();

  if (threadIdx.x == 0) {
    int payload = cutlass::ceil_div(total_num_blocks, num_sm_parts) +
                  fixed_overhead_num_blocks;

    int now_idx = 0, now_block = 0, now_n_split_idx = 0, cum_num_splits = 0;
    num_splits_shared[0] = 0;
    for (int i = 0; i < num_sm_parts; ++i) {
      int tile_scheduler_metadata0[4], tile_scheduler_metadata1;
      tile_scheduler_metadata0[0] = now_idx;
      tile_scheduler_metadata0[1] = now_block * block_size_n;
      tile_scheduler_metadata1 = now_n_split_idx;
      int remain_payload = payload;
      while (now_idx < batch_size) {
        int num_blocks = num_blocks_shared[now_idx];
        int now_remain_blocks = num_blocks - now_block;
        if (remain_payload >= now_remain_blocks + fixed_overhead_num_blocks) {
          cum_num_splits += now_n_split_idx + 1;
          num_splits_shared[now_idx + 1] = cum_num_splits;
          remain_payload -= now_remain_blocks + fixed_overhead_num_blocks;
          ++now_idx;
          now_block = 0;
          now_n_split_idx = 0;
        } else {
          if (remain_payload - fixed_overhead_num_blocks > 0) {
            now_block += remain_payload - fixed_overhead_num_blocks;
            ++now_n_split_idx;
            remain_payload = 0;
          }
          break;
        }
      }
      tile_scheduler_metadata0[2] = now_block > 0 ? now_idx : now_idx - 1;
      tile_scheduler_metadata0[3] =
          now_block > 0 ? now_block * block_size_n : seqlens_k_ptr[now_idx - 1];
      *reinterpret_cast<int4*>(tile_scheduler_metadata_ptr +
                               i * TileSchedulerMetaDataSize) =
          *reinterpret_cast<int4*>(tile_scheduler_metadata0);
      tile_scheduler_metadata_ptr[i * TileSchedulerMetaDataSize + 4] =
          tile_scheduler_metadata1;
    }
    FLASH_DEVICE_ASSERT(now_idx == batch_size && now_block == 0 &&
                        now_n_split_idx == 0);
  }
  __syncwarp();

  for (int i = threadIdx.x; i <= batch_size; i += 32) {
    num_splits_ptr[i] = num_splits_shared[i];
  }
}

void get_mla_metadata_func(Mla_metadata_params& params, cudaStream_t stream) {
  FLASH_ASSERT(params.batch_size < MaxBatchSize);
  get_mla_metadata_kernel<<<1, 32, 0, stream>>>(params);
  CHECK_CUDA_KERNEL_LAUNCH();
}