// Warp-specialized 1p2c with 2-stage pipeline, TMA + 128B swizzle
// D=128, bf16, causal prefill
//
// Compared to v1_manual_baseline (sequential dma.copy):
//   - Async TMA loads for Q, K, V (single-thread producer)
//   - 1p2c warp specialization: 1 producer warpgroup + 2 consumer warpgroups
//   - 2-stage double buffering with shared events for K and V
//   - Producer/consumer overlap hides memory latency
//   - 128B swizzle on TMA and WGMMA to avoid bank conflicts
//
// The compiler auto-splits the 128-element inner dim into [2, 64] groups
// so that each TMA box fits the 128B swizzle constraint (64 * 2B = 128B).

#define BLOCK_M 128
#define BLOCK_N 128
#define WG_M 64
#define WG_K 16
#define SWIZ 128
#define STAGES 2

#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;

  [[launch_bounds(_, 1)]]
  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();

    shared bf16[BLOCK_M, DIM] q_shared;
    shared bf16[BLOCK_N, DIM] k_buf[STAGES];
    shared bf16[BLOCK_N, DIM] v_buf[STAGES];

    shared event qf, kf[STAGES], ke[STAGES], vf[STAGES], ve[STAGES];

    kv_bound = cdiv((bm + 1) * BLOCK_M, BLOCK_N);
    parallel p by [BLOCK_M/WG_M + 1] : group-4, t by 128 : thread {
      inthreads.async (p == 0) {
        call croq::cuda::setreg_dec(24);
        q_load = tma.copy.async.swiz<SWIZ>
          Q_head.subspan(BLOCK_M, DIM).at(bm, 0) => q_shared;
        trigger qf after q_load;

        foreach {bn} in [kv_bound] {
          stage = bn % STAGES;
          wait ke.at(bn);
          k_load = tma.copy.async.swiz<SWIZ>
            K_head.subspan(BLOCK_N, DIM).at(bn, 0) => k_buf[stage];
          trigger kf.at(bn) after k_load;

          wait ve.at(bn);
          v_load = tma.copy.async.swiz<SWIZ>
            V_head.subspan(BLOCK_N, DIM).at(bn, 0) => v_buf[stage];
          trigger vf.at(bn) after v_load;
        }
      }

      inthreads.async (p > 0) {
        call croq::cuda::setreg_inc(240);
        cid = p - 1;
        foreach {s} in [STAGES] {
          trigger ke.at(s);
          trigger ve.at(s);
        }

        wait qf;

        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;

        foreach {bn} in [kv_bound] {
          stage = bn % STAGES;
          acc_s = mma.fill.f32 0.0f;

          wait kf.at(bn);
          mma.row.row acc_s, q_shared.subspan(WG_M, DIM).at(cid, 0), k_buf[stage];
          trigger ke.at(bn);

          int causal_row_base = bm * BLOCK_M + cid * WG_M;
          if ((bn + 1) * BLOCK_N > causal_row_base + 1) {
            apply {i, j} in acc_s.span {
              if ((bn * BLOCK_N + j) > (bm * BLOCK_M + cid * WG_M + i))
                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);

          frag bf16[WG_M, BLOCK_N] acc_s_cast;
          apply {i, j} in acc_s_cast.span
            acc_s_cast.at(i, j) = __to<bf16>(acc_s.at(i, j));

          wait vf.at(bn);
          mma.row.col acc_o, acc_s_cast, v_buf[stage];
          trigger ve.at(bn);
        }

        apply {i, j} in acc_o.span
          acc_o.at(i, j) = acc_o.at(i, j) / logsum.at(i);
        shared bf16[BLOCK_M, DIM] o_buf;
        mma.store acc_o, o_buf.subspan(WG_M, DIM).at(cid, 0);
        // Publish the complete tile after both consumers have stored their rows.
        tma.copy o_buf => O_head.subspan(BLOCK_M, DIM).at(bm, 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 v2 Manual S2 1p2c TMA (D=128 bf16 causal)",
      kBenchConfigs, kBenchConfigCount,
      [&](const mha_helper::BenchConfig& cfg,
          const mha_helper::TensorViews& views) {
        (void)cfg;
        flash_atten(views.Q, views.K, views.V, views.O);
      });
}