// Usage: choreo -gs -t cute -arch=sm_90a v1_manual_baseline.co -o bench.sh
//        bash bench.sh --execute

// Sequential load + compute + store with SIMT-level memory copy (dma.copy)
// D=128, bf16, causal

#define BLOCK_M 128
#define BLOCK_N 128
#define WG_M 64
#define WG_K 16

#define DIM 128

__co__ void flash_atten(
  global bf16[B, Q_SEQ, H, DIM] Q, global bf16[B, KV_SEQ, H, DIM] K,
  global bf16[B, KV_SEQ, H, DIM] V,
  global bf16[B, Q_SEQ, H, DIM] O
) {
  float scale = 0.12751743f;
  int past_len = KV_SEQ - Q_SEQ;
  assert(BLOCK_M % WG_M == 0, "BLOCK_M must be a multiple of WG_M");
  assert(BLOCK_N % WG_K == 0, "BLOCK_N must be a multiple of WG_K");

  parallel.async {bm, head, batch} by [cdiv(Q_SEQ, BLOCK_M), H, B] : block {
    Q_head = Q.subspan(1, Q_SEQ, 1, DIM).at(batch, 0, head, 0).sqz();
    K_head = K.subspan(1, KV_SEQ, 1, DIM).at(batch, 0, head, 0).sqz();
    V_head = V.subspan(1, KV_SEQ, 1, DIM).at(batch, 0, head, 0).sqz();
    O_head = O.subspan(1, Q_SEQ, 1, DIM).at(batch, 0, head, 0).sqz();

    q_s = dma.copy.swiz<128> Q_head.subspan(BLOCK_M, DIM).at(bm, 0) => shared;
    parallel wgm by [BLOCK_M / WG_M]: group-4 {
      frag bf16[WG_M, BLOCK_N] acc_s_cast;
      frag f32[WG_M] scores_max {-inf};
      frag f32[WG_M] scores_max_prev;
      frag f32[WG_M] scores_scale;
      frag f32[WG_M] scores_sum;
      frag f32[WG_M] logsum {0.0f};
      acc_o = mma.fill.f32 0.0f;

      kv_bound = __min(cdiv((bm + 1) * BLOCK_M + past_len, BLOCK_N), cdiv(KV_SEQ, BLOCK_N));

      foreach {bn} in [kv_bound] {
        acc_s = mma.fill.f32 0.0f;
        k_s = dma.copy.swiz<128> K_head.subspan(BLOCK_N, DIM).at(bn, 0) => shared;
        mma.row.row acc_s, q_s.chunkat(wgm, _), k_s;
        apply {i, j} in acc_s.span {
          if ((bn * BLOCK_N + j) > (bm * BLOCK_M + wgm * WG_M + i + past_len))
            acc_s.at(i, j) = -inf;
        }

        copy(scores_max_prev, scores_max);
        reduce_max(scores_max, acc_s, 1);
        apply {i} in scores_max.span {
          scores_max.at(i) = __max(scores_max.at(i), scores_max_prev.at(i));
          scores_scale.at(i) = __exp2f(scores_max_prev.at(i) * scale - scores_max.at(i) * scale);
        }
        apply {i, j} in acc_s.span
          acc_s.at(i, j) = __exp2f(acc_s.at(i, j) * scale - scores_max.at(i) * scale);
        reduce_sum(scores_sum, acc_s, 1);
        apply {i} in logsum.span
          logsum.at(i) = logsum.at(i) * scores_scale.at(i) + scores_sum.at(i);
        apply {i, j} in acc_o.span
          acc_o.at(i, j) = acc_o.at(i, j) * scores_scale.at(i);
        apply {i, j} in acc_s_cast.span
          acc_s_cast.at(i, j) = __to<bf16>(acc_s.at(i, j));

        v_s = dma.copy.swiz<128> V_head.subspan(BLOCK_N, DIM).at(bn, 0) => shared;
        mma.row.col acc_o, acc_s_cast, v_s;
      }

      apply {i, j} in acc_o.span
        acc_o.at(i, j) = acc_o.at(i, j) / logsum.at(i);
      mma.store acc_o, O_head.subspan(WG_M, DIM).at(bm # wgm, 0);
    }
  }
}

#define MHA_DIM DIM
#define MHA_DTYPE choreo::bf16
#include "../fa_helper.hpp"

int main() {
  #include "../bench_configs.inc"
  return mha_helper::RunBenchmarks(
      "Choreo Flash Attention v1 Manual Baseline (D=128 bf16 causal)",
      kBenchConfigs, kBenchConfigCount,
      [&](const mha_helper::BenchConfig& cfg,
          const mha_helper::TensorViews& views) {
        flash_atten(views.Q, views.K, views.V, views.O);
      });
}