// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hgemv compute, round 6 (see reader.cpp).  One output tile per core.
//   HG_LN (K = 1024: the whole vector is ONE dense tile): mean = sum(x) / K, d = x - mean, rstd = rsqrt(sum(d^2) / K
//          + eps), h = d * rstd (scalar broadcasts; gamma / beta are folded into W / b on the host) -> c_hd (bf16),
//          which the writer spreads into the row-form operand c_hrow.
//   y = sum_k hrow_k @ W_k (fp32 DST)  (+ bias)  [GELU, the model's erf form]  [+ residual]  -> c_out
// compile-time: KT, inv_k bits, eps bits
#include <cstdint>
#include "api/compute/matmul.h"
#include "api/compute/reduce.h"
#include "api/compute/bcast.h"
#include "api/compute/eltwise_binary.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/rsqrt.h"
#include "api/compute/eltwise_unary/binop_with_scalar.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/reconfig_data_format.h"
#ifdef HG_PROF
#include "tools/profiler/kernel_profiler.hpp"
#define PZ(n) DeviceZoneScopedN(n)
#else
#define PZ(n)
#endif
#ifdef HG_GELU
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#ifdef TRISC_MATH
#include "sfpi.h"
namespace ckernel::sfpu {
// the erf-form GELU of kernels/eltwise/gelu_compute.cpp (degree-8 CDF polynomial, fp32 SFPU)
inline void hg_gelu_face() {
    for (size_t i = 0; i < 8; i += 2) {
        sfpi::vFloat x0 = sfpi::dst_reg[i];
        sfpi::vFloat x1 = sfpi::dst_reg[i + 1];
        sfpi::vFloat u0 = sfpi::min(x0 * x0, 16.0f);
        sfpi::vFloat u1 = sfpi::min(x1 * x1, 16.0f);
        sfpi::vFloat a0, a1;
        { const sfpi::vFloat c8 = 8.040229593e-11f; const sfpi::vFloat c7 = -6.988785106e-09f;
        a0 = u0 * c8 + c7;
        a1 = u1 * c8 + c7; }
        { const sfpi::vFloat c = 2.712364733e-07f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = -6.289864814e-06f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = 9.886907355e-05f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = -1.133757411e-03f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = 9.877118282e-03f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = -6.641025841e-02f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = 3.989226222e-01f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        sfpi::vFloat p0 = sfpi::clamp(x0 * a0 + 0.5f, 0.0f, 1.0f);
        sfpi::vFloat p1 = sfpi::clamp(x1 * a1 + 0.5f, 0.0f, 1.0f);
        sfpi::dst_reg[i] = x0 * p0;
        sfpi::dst_reg[i + 1] = x1 * p1;
    }
}
}  // namespace ckernel::sfpu
#endif
inline void hg_gelu_tile(uint32_t idst) {
    MATH(SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, hg_gelu_face, idst, VectorMode::RC));
}
#endif

#ifdef HG_LN
// scalar statistic of the dense tile in c_src: sum * inv_k [+ eps -> rsqrt], packed (only [0, 0] is used)
template <bool RSTD>
FORCE_INLINE void scalar_stat(uint32_t c_src, uint32_t c_scal, uint32_t c_dst, uint32_t inv_k, uint32_t eps) {
    reconfig_data_format(c_src, c_scal);  // REDUCE_SCALAR: data -> SrcA, scaler -> SrcB (not swapped)
    pack_reconfig_data_format(c_dst);
    reduce_init<PoolType::SUM, ReduceDim::REDUCE_SCALAR>(c_src, c_scal, c_dst);
    cb_reserve_back(c_dst, 1);
    tile_regs_acquire();
    reduce_tile<PoolType::SUM, ReduceDim::REDUCE_SCALAR>(c_src, c_scal, 0, 0, 0);
    reduce_uninit();
    binop_with_scalar_tile_init();
    mul_unary_tile(0, inv_k);
    if constexpr (RSTD) {
        add_unary_tile(0, eps);
        rsqrt_tile_init();
        rsqrt_tile(0);
    }
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_dst);
    tile_regs_release();
    cb_push_back(c_dst, 1);
}
#endif

void kernel_main() {
    constexpr uint32_t KT = get_compile_time_arg_val(0);
    constexpr uint32_t INV_K_BITS = get_compile_time_arg_val(1);
    constexpr uint32_t EPS_BITS = get_compile_time_arg_val(2);
    constexpr uint32_t c_xd = 0, c_w = 1, c_bias = 2, c_scal = 3, c_res = 4, c_hrow = 5, c_hd = 6, c_out = 16;
    constexpr uint32_t c_mean = 24, c_d = 25, c_sq = 26, c_rstd = 27, c_f32 = 28;  // c_f32: fp32 format reference
#ifdef HG_LN
    compute_kernel_hw_startup(c_xd, c_scal, c_mean);
    {PZ("WX");
    cb_wait_front(c_xd, 1);}
    {PZ("LN");
    cb_wait_front(c_scal, 1);
    scalar_stat<false>(c_xd, c_scal, c_mean, INV_K_BITS, EPS_BITS);
    // d = x - mean
    cb_wait_front(c_mean, 1);
    reconfig_data_format(c_xd, c_mean);
    pack_reconfig_data_format(c_d);
    sub_bcast_scalar_init(c_xd, c_mean);
    cb_reserve_back(c_d, 1);
    tile_regs_acquire();
    sub_tiles_bcast_scalar(c_xd, c_mean, 0, 0, 0);
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_d);
    tile_regs_release();
    cb_push_back(c_d, 1);
    cb_pop_front(c_mean, 1);
    cb_pop_front(c_xd, 1);
    // d^2
    cb_wait_front(c_d, 1);
    reconfig_data_format(c_d, c_d);
    pack_reconfig_data_format(c_sq);
    mul_init(c_d, c_d);
    cb_reserve_back(c_sq, 1);
    tile_regs_acquire();
    mul_tiles(c_d, c_d, 0, 0, 0);
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_sq);
    tile_regs_release();
    cb_push_back(c_sq, 1);
    cb_wait_front(c_sq, 1);
    scalar_stat<true>(c_sq, c_scal, c_rstd, INV_K_BITS, EPS_BITS);
    cb_pop_front(c_sq, 1);
    // h = d * rstd -> bf16 dense tile for the writer
    cb_wait_front(c_rstd, 1);
    reconfig_data_format(c_d, c_rstd);
    pack_reconfig_data_format(c_hd);
    mul_bcast_scalar_init(c_d, c_rstd);
    cb_reserve_back(c_hd, 1);
    tile_regs_acquire();
    mul_tiles_bcast_scalar(c_d, c_rstd, 0, 0, 0);
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_hd);
    tile_regs_release();
    cb_push_back(c_hd, 1);
#ifdef HG_TINY
    // round 10: the row-major dense LN output IS the matmul operand: c_hrow = KT [1, 32] tiny tiles over c_hd's memory
    // (no spread by the writer)
    cb_reserve_back(c_hrow, KT);
    cb_push_back(c_hrow, KT);
#endif
    cb_pop_front(c_rstd, 1);
    cb_pop_front(c_d, 1);
    cb_pop_front(c_scal, 1);
    }
#else
    compute_kernel_hw_startup(c_hrow, c_w, c_out);
#endif
    // y = sum_k hrow_k @ W_k
    {PZ("WH");
    cb_wait_front(c_hrow, KT);}
    {PZ("WW");
    cb_wait_front(c_w, KT);}
    PZ("MM");
#ifdef HG_TINY
    // round 10: [1, 32] tiny in0 / output tiles -- reprogram the tile / face geometry, not only the formats
    reconfig_full_operand(c_w, c_hrow);  // srcA = in1 = W, srcB = in0 = h
    pack_reconfig_data_format<true>(c_out);
#else
    reconfig_data_format(c_w, c_hrow);  // srcA = in1 = W, srcB = in0 = h
    pack_reconfig_data_format(c_out);
#endif
    matmul_init(c_hrow, c_w);
    tile_regs_acquire();
    for (uint32_t k = 0; k < KT; ++k) {
        matmul_tiles(c_hrow, c_w, k, k, 0);
    }
    // + bias [GELU] [+ residual]: SFPU fp32 adds in DST (operands copied in with unpack-to-dest fp32, so the matmul
    // result never goes through a 19-bit source register)
    cb_wait_front(c_bias, 1);
#ifdef HG_TINY
    reconfig_full_operand_srca(c_bias);
    copy_tile_to_dst_init_short(c_bias);
#else
    copy_tile_to_dst_init_short_with_dt(c_w, c_bias);
#endif
    copy_tile(c_bias, 0, 1);
    add_binary_tile_init();
    add_binary_tile(0, 1, 0);
#ifdef HG_GELU
    hg_gelu_tile(0);
#endif
#ifdef HG_RES
    cb_wait_front(c_res, 1);
#ifdef HG_TINY
    reconfig_full_operand_srca(c_res);
    copy_tile_to_dst_init_short(c_res);
#else
    copy_tile_to_dst_init_short_with_dt(c_bias, c_res);
#endif
    copy_tile(c_res, 0, 1);
    add_binary_tile_init();
    add_binary_tile(0, 1, 0);
#endif
    tile_regs_commit();
    cb_reserve_back(c_out, 1);
    tile_regs_wait();
    pack_tile(0, c_out);
    tile_regs_release();
    cb_push_back(c_out, 1);
    cb_pop_front(c_hrow, KT);
    cb_pop_front(c_w, KT);
    cb_pop_front(c_bias, 1);
#ifdef HG_RES
    cb_pop_front(c_res, 1);
#endif
}