// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hattn compute, round 6 (see reader.cpp): row 0 of every tile is the query token.
//   s_j  = sum_d q_d @ K_{j,d}^T            (MT score tiles, fp32)
//   m    = rowmax(s);  e_j = exp(s_j - m)   (bf16, the PV operand)
//   r    = 1 / rowsum(e)
//   o_d  = (sum_j e_j @ V_{j,d}) * r        -> 2 bf16 tiles
// Round 6 status: correct (PCC 0.999997 vs fp64) but NOT wired -- 10.1 us kernel / ~14 us per call vs ~10.7 us for
// head split + library SDPA + head concat: the accurate exp of 6 score tiles costs 3.6 us even in row mode.
// compile-time: MT
#include <cstdint>
#include "api/compute/matmul.h"
#include "api/compute/reduce.h"
#include "api/compute/bcast.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/exp.h"
#include "api/compute/eltwise_unary/recip.h"
#include "api/compute/compute_kernel_hw_startup.h"
#ifdef HA_PROF
#include "tools/profiler/kernel_profiler.hpp"
#define PZ(n) DeviceZoneScopedN(n)
#else
#define PZ(n)
#endif

void kernel_main() {
    constexpr uint32_t MT = get_compile_time_arg_val(0);
    constexpr uint32_t c_q = 0, c_k = 1, c_v = 2, c_scal = 3, c_out = 16;
    constexpr uint32_t c_s = 24, c_m = 25, c_e = 26, c_r = 27, c_o = 28;
    compute_kernel_hw_startup(c_q, c_k, c_s);
    {PZ("WAIT");
    cb_wait_front(c_q, 2);
    cb_wait_front(c_k, 2 * MT);
    cb_wait_front(c_v, 2 * MT);
    cb_wait_front(c_scal, 1);}
    // scores
    {PZ("S");
    reconfig_data_format(c_k, c_q);  // srcA = in1 = K, srcB = in0 = q
    pack_reconfig_data_format(c_s);
    matmul_init(c_q, c_k, 1);
    cb_reserve_back(c_s, MT);
    for (uint32_t j0 = 0; j0 < MT; j0 += 4) {
        const uint32_t n = MT - j0 < 4 ? MT - j0 : 4;
        tile_regs_acquire();
        for (uint32_t j = 0; j < n; ++j) {
            for (uint32_t d = 0; d < 2; ++d) {
                matmul_tiles(c_q, c_k, d, (j0 + j) * 2 + d, j);
            }
        }
        tile_regs_commit();
        tile_regs_wait();
        for (uint32_t j = 0; j < n; ++j) {
            pack_tile(j, c_s);
        }
        tile_regs_release();
    }
    cb_push_back(c_s, MT);
    cb_wait_front(c_s, MT);}
    // row max
    {PZ("MAX");
    reconfig_data_format(c_s, c_scal);
    pack_reconfig_data_format(c_m);
    reduce_init<PoolType::MAX, ReduceDim::REDUCE_ROW>(c_s, c_scal, c_m);
    cb_reserve_back(c_m, 1);
    tile_regs_acquire();
    for (uint32_t j = 0; j < MT; ++j) {
        reduce_tile<PoolType::MAX, ReduceDim::REDUCE_ROW>(c_s, c_scal, j, 0, 0);
    }
    reduce_uninit();
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_m);
    tile_regs_release();
    cb_push_back(c_m, 1);
    cb_wait_front(c_m, 1);}
    // e = exp(s - m)
    {PZ("EXP");
    reconfig_data_format(c_s, c_m);
    pack_reconfig_data_format(c_e);
    cb_reserve_back(c_e, MT);
    for (uint32_t j0 = 0; j0 < MT; j0 += 4) {
        const uint32_t n = MT - j0 < 4 ? MT - j0 : 4;
        sub_bcast_cols_init(c_s, c_m);
        tile_regs_acquire();
        for (uint32_t j = 0; j < n; ++j) {
            sub_tiles_bcast_cols(c_s, c_m, j0 + j, 0, j);
        }
        exp_tile_init<false>();
        for (uint32_t j = 0; j < n; ++j) {
            exp_tile<false>(j, VectorMode::R);  // only row 0 (the query token) matters
        }
        tile_regs_commit();
        tile_regs_wait();
        for (uint32_t j = 0; j < n; ++j) {
            pack_tile(j, c_e);
        }
        tile_regs_release();
    }
    cb_push_back(c_e, MT);
    cb_pop_front(c_m, 1);
    cb_pop_front(c_s, MT);
    cb_wait_front(c_e, MT);}
    // r = 1 / rowsum(e)
    {PZ("SUM");  // REDUCE_ROW SUM: scaler -> SrcA, data -> SrcB
    reconfig_data_format(c_scal, c_e);
    pack_reconfig_data_format(c_r);
    reduce_init<PoolType::SUM, ReduceDim::REDUCE_ROW>(c_e, c_scal, c_r);
    cb_reserve_back(c_r, 1);
    tile_regs_acquire();
    for (uint32_t j = 0; j < MT; ++j) {
        reduce_tile<PoolType::SUM, ReduceDim::REDUCE_ROW>(c_e, c_scal, j, 0, 0);
    }
    reduce_uninit();
    recip_tile_init();
    recip_tile(0);
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_r);
    tile_regs_release();
    cb_push_back(c_r, 1);}
    // o = e @ V
    {PZ("PV");
    reconfig_data_format(c_v, c_e);
    pack_reconfig_data_format(c_o);
    matmul_init(c_e, c_v, 0);
    cb_reserve_back(c_o, 2);
    tile_regs_acquire();
    for (uint32_t d = 0; d < 2; ++d) {
        for (uint32_t j = 0; j < MT; ++j) {
            matmul_tiles(c_e, c_v, j, j * 2 + d, d);
        }
    }
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_o);
    pack_tile(1, c_o);
    tile_regs_release();
    cb_push_back(c_o, 2);}
    // * r
    {PZ("NORM");
    cb_wait_front(c_o, 2);
    cb_wait_front(c_r, 1);
    reconfig_data_format(c_o, c_r);
    pack_reconfig_data_format(c_out);
    mul_bcast_cols_init(c_o, c_r);
    cb_reserve_back(c_out, 2);
    tile_regs_acquire();
    mul_tiles_bcast_cols(c_o, c_r, 0, 0, 0);
    mul_tiles_bcast_cols(c_o, c_r, 1, 0, 1);
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_out);
    pack_tile(1, c_out);
    tile_regs_release();
    cb_push_back(c_out, 2);}
    cb_pop_front(c_o, 2);
    cb_pop_front(c_r, 1);
    cb_pop_front(c_e, MT);
    cb_pop_front(c_q, 2);
    cb_pop_front(c_k, 2 * MT);
    cb_pop_front(c_v, 2 * MT);
    cb_pop_front(c_scal, 1);
}