#pragma once
#include "mha_components.cuh"
#include "mha_stdheaders.cuh"
#include "tma.h"
#include "utils.cuh"

template <uint32_t tileM, uint32_t totalNbRows>
__device__ inline ThrdRegRowMaxT<tileM> loadShmRowMax(Vec<float, totalNbRows> const& shm,
                                                      uint32_t tileBaseRow,
                                                      uint32_t lane = laneId()) {
  ThrdRegRowMaxT<tileM> result{};
#pragma unroll
  for (uint32_t i = 0; i < result.size; i++) {
    result[i] = shm[tileBaseRow + i * warp_size + lane];
  }
  return result;
}

template <uint32_t tileM, uint32_t totalNbRows>
__device__ inline void storeRowMax(Vec<float, totalNbRows>& shm, ThrdRegRowMaxT<tileM> const& src,
                                   uint32_t tileBaseRow, uint32_t lane = laneId()) {
#pragma unroll
  for (uint32_t i = 0; i < src.size; i++) {
    shm[tileBaseRow + i * warp_size + lane] = src[i];
  }
}

template <uint32_t tileM, uint32_t totalNbRows>
__device__ inline void storeRowMaxAsync(CgaBarrier& bar, Vec<float, totalNbRows>& shm,
                                        ThrdRegRowMaxT<tileM> const& src, uint32_t tileBaseRow,
                                        uint32_t lane = laneId()) {
#pragma unroll
  for (uint32_t i = 0; i < src.size; i++) {
    tma::storeAsync(&shm[tileBaseRow + i * warp_size + lane], src[i], bar);
  }
}

template <uint32_t tileM, uint32_t tileN>
__device__ inline QuadRegRowMaxT<tileM> computeRowMax(WarpAccT<tileM, tileN> const& acc) {
  QuadRegRowMaxT<tileM> rowMaxLog2e{};
// compute per-thread row max
#pragma unroll
  for (uint32_t n = 0; n < acc.cols; n++) {
#pragma unroll
    for (uint32_t j = 0; j < InstAcc::cols; j++) {
#pragma unroll
      for (uint32_t m = 0; m < acc.rows; m++) {
#pragma unroll
        for (uint32_t i = 0; i < InstAcc::rows; i++) {
          float& dst = rowMaxLog2e[m * InstAcc::rows + i];
          dst = ((n == 0 && j == 0) ? acc(m, n)(i, j) : fmaxf(dst, acc(m, n)(i, j)));
        }
      }
    }
  }
// compute warp row max
#pragma unroll
  for (uint32_t xorMask = 2; xorMask != 0; xorMask /= 2) {
#pragma unroll
    for (uint32_t i = 0; i < rowMaxLog2e.size; i++) {
      rowMaxLog2e[i] = fmaxf(rowMaxLog2e[i], __shfl_xor_sync(~0U, rowMaxLog2e[i], xorMask));
    }
  }
  return rowMaxLog2e;
}

template <typename T, uint32_t n>
__device__ inline uint32_t hashRegData(Vec<T, n> const& data) {
  static_assert(sizeof(T) == 4);
  uint32_t result = 0;
#pragma unroll
  for (uint32_t i = 0; i < n; i++) {
    result ^= reinterpret_cast<uint32_t const&>(data[i]);
  }
  return result;
}