// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// Fused residual add + LayerNorm over a block-sharded [32*NY, 32*NT*NX] stream, one tile row per core
// (block_h = 1, NT <= 4 so every pass is ONE DST acquire with fp32 accumulation):
//   s    = x (+ y)                                   -> c_s (the new residual stream shard, HAS_Y only)
//   p1   = transpose(rowsum(s))                      -> c_p1: this core's per-token partial sums as a ROW vector
//          (the writer sends that 128-byte row to row `col` of a 1-tile gather buffer on each row peer)
//   mean = transpose(colsum(G1)) / W                 (column vector, per token)
//   d    = s - mean                                  -> c_xmm (fp32)
//   p2   = transpose(rowsum(d*d))                    -> c_p2, gathered the same way
//   rstd = rsqrt(transpose(colsum(G2)) / W + eps)
//   out  = (d * rstd) * gamma + beta                 -> c_out (gamma / beta are full 32-row tiles; dest reuse)
// Round 2's version of this kernel set hung chip 7: the writer (NOC_1) got an unswapped multicast
// rectangle.  This version has no multicast at all (two 64-byte unicast writes per peer and round).
// compile-time: NT, NX, inv_w bits, eps bits
#include <cstdint>
#include "api/compute/reduce.h"
#include "api/compute/bcast.h"
#include "api/compute/eltwise_binary.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/transpose_dest.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 "addln_ids.hpp"
#ifdef ADDLN_PROF
#include "tools/profiler/kernel_profiler.hpp"
#define ZONE(n) DeviceZoneScopedN(n)
#else
#define ZONE(n)
#endif

// round 9 (ADDLN_SFPU_C): the per-token statistics are a COLUMN vector after transpose_dest (column 0 = faces 0, 2), so
// their SFPU scalings / rsqrt only need VectorMode::C (half the tile); the other faces are never read
#ifdef ADDLN_SFPU_C
#define COL_MUL(i, p) MATH(SFPU_UNARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, calculate_binop_with_scalar, (APPROX, MUL_UNARY, 8), i, VectorMode::C, p))
#define COL_ADD(i, p) MATH(SFPU_UNARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, calculate_binop_with_scalar, (APPROX, ADD_UNARY, 8), i, VectorMode::C, p))
#define COL_RSQRT(i) MATH(SFPU_UNARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, calculate_rsqrt, (APPROX, 8, DST_ACCUM_MODE, false, false), i, VectorMode::C))
#else
#define COL_MUL(i, p) mul_unary_tile(i, p)
#define COL_ADD(i, p) add_unary_tile(i, p)
#define COL_RSQRT(i) rsqrt_tile(i)
#endif

template <uint32_t NT>
FORCE_INLINE void row_partial(uint32_t c_src, uint32_t c_scal, uint32_t c_dst) {
    // REDUCE_ROW SUM swaps operands (scaler -> SrcA, data -> SrcB): reconfig(scaler, data)
    reconfig_data_format(c_scal, c_src);
    pack_reconfig_data_format(c_dst);
    reduce_init<PoolType::SUM, ReduceDim::REDUCE_ROW>(c_src, c_scal, c_dst);
    cb_reserve_back(c_dst, 1);
    tile_regs_acquire();
    for (uint32_t i = 0; i < NT; ++i) {
        reduce_tile<PoolType::SUM, ReduceDim::REDUCE_ROW>(c_src, c_scal, i, 0, 0);
    }
    reduce_uninit();
    transpose_dest_init<true, true>();
    transpose_dest<true, true>(0);
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_dst);
    tile_regs_release();
    cb_push_back(c_dst, 1);
}

// column vector (per token) of: colsum(G) * inv_w [+ eps -> rsqrt]
template <bool RSTD>
FORCE_INLINE void combine(uint32_t c_g, uint32_t c_scal, uint32_t c_dst, uint32_t inv_w, uint32_t eps) {
    cb_wait_front(c_g, 1);
    reconfig_data_format(c_g, c_scal);
    pack_reconfig_data_format(c_dst);
    reduce_init<PoolType::SUM, ReduceDim::REDUCE_COL>(c_g, c_scal, c_dst);
    cb_reserve_back(c_dst, 1);
    tile_regs_acquire();
    reduce_tile<PoolType::SUM, ReduceDim::REDUCE_COL>(c_g, c_scal, 0, 0, 0);
    reduce_uninit();
    transpose_dest_init<true, true>();
    transpose_dest<true, true>(0);
    binop_with_scalar_tile_init();
    COL_MUL(0, inv_w);
    if constexpr (RSTD) {
        COL_ADD(0, eps);
        rsqrt_tile_init();
        COL_RSQRT(0);
    }
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_dst);
    tile_regs_release();
    cb_push_back(c_dst, 1);
    cb_pop_front(c_g, 1);
}

ADDLN_KERNEL_MAIN {
    constexpr uint32_t NT = get_compile_time_arg_val(ADDLN_CT0 + 0);
    constexpr uint32_t NX = get_compile_time_arg_val(ADDLN_CT0 + 1);
    constexpr uint32_t INV_W_BITS = get_compile_time_arg_val(ADDLN_CT0 + 2);
    constexpr uint32_t EPS_BITS = get_compile_time_arg_val(ADDLN_CT0 + 3);
    static_assert(NT <= 4, "one fp32 DST half per pass");
    constexpr uint32_t c_x = ADDLN_C_X, c_y = ADDLN_C_Y, c_scal = ADDLN_C_SCAL, c_gam = ADDLN_C_GAM, c_bet = ADDLN_C_BET;
    constexpr uint32_t c_g1 = ADDLN_C_G1, c_g2 = ADDLN_C_G2;
    constexpr uint32_t c_s = ADDLN_C_S, c_out = ADDLN_C_OUT;
    constexpr uint32_t c_p1 = ADDLN_C_P1, c_mean = 25, c_xmm = 26, c_sq = 27, c_p2 = ADDLN_C_P2, c_rstd = 29;

#ifdef HAS_Y
    constexpr uint32_t c_in = c_s;
#ifdef ADDLN_FUSED  // after the cmm phase of the same kernel: the hardware is configured, only reconfigure formats
    reconfig_data_format(c_x, c_y);
    pack_reconfig_data_format(c_s);
#else
    compute_kernel_hw_startup(c_x, c_y, c_s);
#endif
    {ZONE("A-add");
    cb_wait_front(c_x, NT);
    cb_wait_front(c_y, NT);
    add_init(c_x, c_y);
    cb_reserve_back(c_s, NT);
    tile_regs_acquire();
    for (uint32_t i = 0; i < NT; ++i) {
        add_tiles(c_x, c_y, i, i, i);
    }
    tile_regs_commit();
    tile_regs_wait();
    for (uint32_t i = 0; i < NT; ++i) {
        pack_tile(i, c_s);
    }
    tile_regs_release();
    cb_push_back(c_s, NT);
    cb_wait_front(c_s, NT);
    }
#else
    constexpr uint32_t c_in = c_x;
    compute_kernel_hw_startup(c_scal, c_x, c_p1);
    cb_wait_front(c_x, NT);
#endif
    cb_wait_front(c_scal, 1);

#ifdef ADDLN_1R
    // round 6: ONE gather round -- the per-token partials of s and s^2 travel together, var = E[s^2] - mean^2
    // (mean^2 / var <= 2.2 on the real stream, so the fp32 cancellation is harmless)
    constexpr uint32_t c_es2 = 30, c_msq = 31;
    {ZONE("B-sq");
    reconfig_data_format(c_in, c_in);
    pack_reconfig_data_format(c_sq);
    mul_init(c_in, c_in);
    cb_reserve_back(c_sq, NT);
    tile_regs_acquire();
    for (uint32_t i = 0; i < NT; ++i) {
        mul_tiles(c_in, c_in, i, i, i);
    }
    tile_regs_commit();
    tile_regs_wait();
    for (uint32_t i = 0; i < NT; ++i) {
        pack_tile(i, c_sq);
    }
    tile_regs_release();
    cb_push_back(c_sq, NT);
#ifndef ADDLN_SPLITSEND
    cb_wait_front(c_sq, NT);
#endif
    }
    {ZONE("B-p1"); row_partial<NT>(c_in, c_scal, c_p1);}
#ifdef ADDLN_SPLITSEND
    // round 9: the unpacker no longer waits for the s^2 pack before the sum partial (it reads c_in only)
    cb_wait_front(c_sq, NT);
#endif
    {ZONE("F-p2"); row_partial<NT>(c_sq, c_scal, c_p2);}
    cb_pop_front(c_sq, NT);
    {ZONE("C-g1wait"); cb_wait_front(c_g1, 1);}
    {ZONE("D-mean"); combine<false>(c_g1, c_scal, c_mean, INV_W_BITS, EPS_BITS);}
    {ZONE("G-g2wait"); cb_wait_front(c_g2, 1);}
    {ZONE("H-es2"); combine<false>(c_g2, c_scal, c_es2, INV_W_BITS, EPS_BITS);}
#ifdef ADDLN_TAIL
    // round 9: mean^2 after the E[s^2] combine (its mean round trip hidden) and d = s - mean between mean^2 and
    // rstd (hides the mean^2 round trip); only the rstd -> output round trip stays exposed
    {ZONE("H-msq");
    // mean^2
    cb_wait_front(c_mean, 1);
    reconfig_data_format(c_mean, c_mean);
    pack_reconfig_data_format(c_msq);
    mul_init(c_mean, c_mean);
    cb_reserve_back(c_msq, 1);
    tile_regs_acquire();
    mul_tiles(c_mean, c_mean, 0, 0, 0);
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_msq);
    tile_regs_release();
    cb_push_back(c_msq, 1);
    }
    {ZONE("E-d");
    // d = s - mean
    reconfig_data_format(c_in, c_mean);
    pack_reconfig_data_format(c_xmm);
    sub_bcast_cols_init(c_in, c_mean);
    cb_reserve_back(c_xmm, NT);
    tile_regs_acquire();
    for (uint32_t i = 0; i < NT; ++i) {
        sub_tiles_bcast_cols(c_in, c_mean, i, 0, i);
    }
    tile_regs_commit();
    tile_regs_wait();
    for (uint32_t i = 0; i < NT; ++i) {
        pack_tile(i, c_xmm);
    }
    tile_regs_release();
    cb_push_back(c_xmm, NT);
    cb_pop_front(c_mean, 1);
    }
    {ZONE("H-rstd");
    // rstd = rsqrt(E[s^2] - mean^2 + eps)
    cb_wait_front(c_es2, 1);
    cb_wait_front(c_msq, 1);
    reconfig_data_format(c_es2, c_msq);
    pack_reconfig_data_format(c_rstd);
    sub_init(c_es2, c_msq);
    cb_reserve_back(c_rstd, 1);
    tile_regs_acquire();
    sub_tiles(c_es2, c_msq, 0, 0, 0);
    binop_with_scalar_tile_init();
    COL_ADD(0, EPS_BITS);
    rsqrt_tile_init();
    COL_RSQRT(0);
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_rstd);
    tile_regs_release();
    cb_push_back(c_rstd, 1);
    cb_pop_front(c_es2, 1);
    cb_pop_front(c_msq, 1);
    }
    cb_wait_front(c_xmm, NT);
#else
    {ZONE("H-rstd");
    // mean^2
    cb_wait_front(c_mean, 1);
    reconfig_data_format(c_mean, c_mean);
    pack_reconfig_data_format(c_msq);
    mul_init(c_mean, c_mean);
    cb_reserve_back(c_msq, 1);
    tile_regs_acquire();
    mul_tiles(c_mean, c_mean, 0, 0, 0);
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_msq);
    tile_regs_release();
    cb_push_back(c_msq, 1);
    // rstd = rsqrt(E[s^2] - mean^2 + eps)
    cb_wait_front(c_es2, 1);
    cb_wait_front(c_msq, 1);
    reconfig_data_format(c_es2, c_msq);
    pack_reconfig_data_format(c_rstd);
    sub_init(c_es2, c_msq);
    cb_reserve_back(c_rstd, 1);
    tile_regs_acquire();
    sub_tiles(c_es2, c_msq, 0, 0, 0);
    binop_with_scalar_tile_init();
    COL_ADD(0, EPS_BITS);
    rsqrt_tile_init();
    COL_RSQRT(0);
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, c_rstd);
    tile_regs_release();
    cb_push_back(c_rstd, 1);
    cb_pop_front(c_es2, 1);
    cb_pop_front(c_msq, 1);
    }
    {ZONE("E-d");
    // d = s - mean
    reconfig_data_format(c_in, c_mean);
    pack_reconfig_data_format(c_xmm);
    sub_bcast_cols_init(c_in, c_mean);
    cb_reserve_back(c_xmm, NT);
    tile_regs_acquire();
    for (uint32_t i = 0; i < NT; ++i) {
        sub_tiles_bcast_cols(c_in, c_mean, i, 0, i);
    }
    tile_regs_commit();
    tile_regs_wait();
    for (uint32_t i = 0; i < NT; ++i) {
        pack_tile(i, c_xmm);
    }
    tile_regs_release();
    cb_push_back(c_xmm, NT);
    cb_pop_front(c_mean, 1);
    cb_wait_front(c_xmm, NT);
    }
#endif
#else
    {ZONE("B-p1"); row_partial<NT>(c_in, c_scal, c_p1);}
    {ZONE("C-g1wait"); cb_wait_front(c_g1, 1);}
    {ZONE("D-mean"); combine<false>(c_g1, c_scal, c_mean, INV_W_BITS, EPS_BITS);}
    {ZONE("E-d-sq");

    // d = s - mean
    cb_wait_front(c_mean, 1);
    reconfig_data_format(c_in, c_mean);
    pack_reconfig_data_format(c_xmm);
    sub_bcast_cols_init(c_in, c_mean);
    cb_reserve_back(c_xmm, NT);
    tile_regs_acquire();
    for (uint32_t i = 0; i < NT; ++i) {
        sub_tiles_bcast_cols(c_in, c_mean, i, 0, i);
    }
    tile_regs_commit();
    tile_regs_wait();
    for (uint32_t i = 0; i < NT; ++i) {
        pack_tile(i, c_xmm);
    }
    tile_regs_release();
    cb_push_back(c_xmm, NT);
    cb_pop_front(c_mean, 1);

    // d*d
    cb_wait_front(c_xmm, NT);
    reconfig_data_format(c_xmm, c_xmm);
    pack_reconfig_data_format(c_sq);
    mul_init(c_xmm, c_xmm);
    cb_reserve_back(c_sq, NT);
    tile_regs_acquire();
    for (uint32_t i = 0; i < NT; ++i) {
        mul_tiles(c_xmm, c_xmm, i, i, i);
    }
    tile_regs_commit();
    tile_regs_wait();
    for (uint32_t i = 0; i < NT; ++i) {
        pack_tile(i, c_sq);
    }
    tile_regs_release();
    cb_push_back(c_sq, NT);

    cb_wait_front(c_sq, NT);
    }
    {ZONE("F-p2"); row_partial<NT>(c_sq, c_scal, c_p2);}
    cb_pop_front(c_sq, NT);
    {ZONE("G-g2wait"); cb_wait_front(c_g2, 1);}
    {ZONE("H-rstd"); combine<true>(c_g2, c_scal, c_rstd, INV_W_BITS, EPS_BITS);}
#endif
    {ZONE("I-out");

    // out = d * rstd * gamma + beta  (ADDLN_NOAFF: out = d * rstd, the affine is folded into the next matmul)
    cb_wait_front(c_rstd, 1);
#ifndef ADDLN_NOAFF
    cb_wait_front(c_gam, NT);
#ifndef ADDLN_NOBETA
    cb_wait_front(c_bet, NT);
#endif
#endif
    reconfig_data_format(c_xmm, c_rstd);
    pack_reconfig_data_format(c_out);
    mul_bcast_cols_init(c_xmm, c_rstd);
    cb_reserve_back(c_out, NT);
    tile_regs_acquire();
    for (uint32_t i = 0; i < NT; ++i) {
        mul_tiles_bcast_cols(c_xmm, c_rstd, i, 0, i);
    }
#ifndef ADDLN_NOAFF
    reconfig_data_format_srcb(c_rstd, c_gam);
    mul_reuse_dest_init<EltwiseBinaryReuseDestType::DEST_TO_SRCA>(c_gam);
    for (uint32_t i = 0; i < NT; ++i) {
        mul_reuse_dest_tiles<EltwiseBinaryReuseDestType::DEST_TO_SRCA>(c_gam, i, i);
    }
#ifndef ADDLN_NOBETA
    reconfig_data_format_srcb(c_gam, c_bet);
    add_reuse_dest_init<EltwiseBinaryReuseDestType::DEST_TO_SRCA>(c_bet);
    for (uint32_t i = 0; i < NT; ++i) {
        add_reuse_dest_tiles<EltwiseBinaryReuseDestType::DEST_TO_SRCA>(c_bet, i, i);
    }
#endif
#endif
    tile_regs_commit();
    tile_regs_wait();
    for (uint32_t i = 0; i < NT; ++i) {
        pack_tile(i, c_out);
    }
    tile_regs_release();
    cb_push_back(c_out, NT);
    }
    cb_pop_front(c_rstd, 1);
    cb_pop_front(c_xmm, NT);
#ifndef ADDLN_NOAFF
    cb_pop_front(c_gam, NT);
#ifndef ADDLN_NOBETA
    cb_pop_front(c_bet, NT);
#endif
#endif
    cb_pop_front(c_in, NT);
#ifdef HAS_Y
    cb_pop_front(c_x, NT);
    cb_pop_front(c_y, NT);
#endif
    cb_pop_front(c_scal, 1);
}