// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hattn2 compute, round 15 (D5): 1-query attention of one head with a DENSE score layout (round 6's hattn exp'd six
// row-form score tiles; here the 192 scores are the first 192 elements = face 0, rows 0..11 of ONE fp32 tile):
//   s_j = q @ K_j^T        six [1, 32] tiny fp32 tiles (q: two [1, 32] tiny tiles, K tiles transposed in the matmul),
//                          packed back to back into c_st = the start of the dense tile c_sd (the reader pre-filled the
//                          rest of c_sd with a large negative pad, so the whole-tile max sees only the scores)
//   m = max(c_sd)          (REDUCE_SCALAR)
//   e = exp((s - m) * scale) on face 0 only (VectorMode::None), packed fp32 -> c_e; its first 192 elements are six
//                          [1, 32] tiny fp32 tiles (alias c_et); the other faces are never read
//   o_d = sum_j e_j @ V_{j,d},  z = sum_j e_j @ ONES   (one DST acquire; z = sum(e) in every column)
//   o_d *= 1 / z           (SFPU in DST: recip of the tiny z tile, then two tiny multiplies) -> two [1, 32] bf16 -> c_out
// compile-time: scale bits (fp32; exact as bf16 for 1/8)
#include <cstdint>
#include "api/compute/matmul.h"
#include "api/compute/reduce.h"
#include "api/compute/bcast.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/eltwise_binary_sfpu.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"
#include "api/compute/reconfig_data_format.h"
#ifdef HA2_PROF
#include "tools/profiler/kernel_profiler.hpp"
#define PZ(n) DeviceZoneScopedN(n)
#else
#define PZ(n)
#endif

void kernel_main() {
    constexpr uint32_t SCALE_BITS = get_compile_time_arg_val(0);
    constexpr uint16_t SCALE_BF16 = (uint16_t)(SCALE_BITS >> 16);
    static_assert((SCALE_BITS & 0xFFFFu) == 0, "the exp scale must be exact in bf16");
    constexpr uint32_t MT = 6;
    constexpr uint32_t c_qt = 0, c_k = 1, c_v = 2, c_scal = 3, c_st = 4, c_sd = 5, c_one = 6, c_et = 7, c_out = 16;
    constexpr uint32_t c_m = 24, c_e = 25;
    compute_kernel_hw_startup(c_qt, c_k, c_st);
    {PZ("WAIT");
    cb_wait_front(c_qt, 2);
    cb_wait_front(c_k, 2 * MT);}  // (V arrives separately, waited for before the PV phase)
    // scores: six tiny tiles
    {PZ("S");
    reconfig_full_operand(c_k, c_qt);  // srcA = in1 = K, srcB = in0 = q
    pack_reconfig_data_format<true>(c_st);
    matmul_init(c_qt, c_k, 1);
    cb_reserve_back(c_st, MT);
    for (uint32_t j = 0; j < MT; ++j) {
        tile_regs_acquire();
        matmul_tiles(c_qt, c_k, 0, j * 2, 0);
        matmul_tiles(c_qt, c_k, 1, j * 2 + 1, 0);
        tile_regs_commit();
        tile_regs_wait();
        pack_tile(0, c_st);
        tile_regs_release();
    }
    cb_push_back(c_st, MT);
    // the same memory as one dense fp32 tile
    cb_reserve_back(c_sd, 1);
    cb_push_back(c_sd, 1);
    }
    cb_wait_front(c_scal, 1);
    // m = max over the dense tile
    {PZ("MAX");
    reconfig_full_operand(c_sd, c_scal);  // REDUCE_SCALAR: data -> SrcA, scaler -> SrcB
    pack_reconfig_data_format<true>(c_m);
    reduce_init<PoolType::MAX, ReduceDim::REDUCE_SCALAR>(c_sd, c_scal, c_m);
    cb_wait_front(c_sd, 1);
    cb_reserve_back(c_m, 1);
    tile_regs_acquire();
    reduce_tile<PoolType::MAX, ReduceDim::REDUCE_SCALAR>(c_sd, c_scal, 0, 0, 0);
    reduce_uninit();
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_m);
    tile_regs_release();
    cb_push_back(c_m, 1);}
    // e = exp((s - m) * scale), face 0 only
    {PZ("EXP");
    cb_wait_front(c_m, 1);
    reconfig_data_format(c_sd, c_m);
    pack_reconfig_data_format(c_e);
    sub_bcast_scalar_init(c_sd, c_m);
    cb_reserve_back(c_e, 1);
    tile_regs_acquire();
    sub_tiles_bcast_scalar(c_sd, c_m, 0, 0, 0);
    exp_tile_init<false>();
    exp_tile<false, true>(0, VectorMode::None, SCALE_BF16);
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_e);
    tile_regs_release();
    cb_push_back(c_e, 1);
    cb_pop_front(c_m, 1);
    cb_reserve_back(c_et, MT);
    cb_push_back(c_et, MT);}
    // o = (e @ V) / (e @ ONES)
    {PZ("PV");
    cb_wait_front(c_v, 2 * MT);
    cb_wait_front(c_one, 1);
    cb_wait_front(c_et, MT);
    reconfig_full_operand(c_v, c_et);  // srcA = in1 = V / ONES, srcB = in0 = e
    pack_reconfig_data_format<true>(c_out);
    matmul_init(c_et, c_v, 0);
    cb_reserve_back(c_out, 2);
    tile_regs_acquire();
    for (uint32_t j = 0; j < MT; ++j) {
        matmul_tiles(c_et, c_v, j, j * 2, 0);
        matmul_tiles(c_et, c_v, j, j * 2 + 1, 1);
        matmul_tiles(c_et, c_one, j, 0, 2);
    }
    // the tiny [1, 32] tiles span DST faces 0 and 1 only: VectorMode::R
    recip_tile_init();
    recip_tile(2, VectorMode::R);
    mul_binary_tile_init();
    MATH((SFPU_BINARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, calculate_sfpu_binary_mul,
                           (APPROX, ckernel::BinaryOp::MUL, 8, DST_ACCUM_MODE), 0, 2, 0, VectorMode::R)));
    MATH((SFPU_BINARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, calculate_sfpu_binary_mul,
                           (APPROX, ckernel::BinaryOp::MUL, 8, DST_ACCUM_MODE), 1, 2, 1, VectorMode::R)));
    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_et, MT);
    cb_pop_front(c_e, 1);
    cb_pop_front(c_sd, 1);
    cb_pop_front(c_one, 1);
    cb_pop_front(c_qt, 2);
    cb_pop_front(c_k, 2 * MT);
    cb_pop_front(c_v, 2 * MT);
    cb_pop_front(c_scal, 1);
}