// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// add+LayerNorm writer (BRISC, NOC_1): two all-gathers of the per-token partial statistics inside each
// grid row, then this core's NT normalized tiles -> the interleaved output.
// Gather (unicast only; round 2's multicast version hung chip 7 with a NOC_1-unswapped rectangle):
//   the compute packs the partial as a ROW vector (fp32 tile, row 0 = 16 floats in face 0 + 16 in face 1);
//   that 2 x 64-byte row goes to row `col` of the 1-tile gather buffer of every core of the grid row
//   (itself included).  Rows NX..31 of the buffer are zeroed locally first (disjoint from the peers'
//   writes), so the compute can column-sum the whole tile.  Then each peer's round semaphore is
//   incremented and we wait for NX-1 increments.  Semaphores are initialised (0) per launch.
// compile-time: NT, NX, TensorAccessorArgs(out)
// runtime: out_addr, out_start (first page of this core), col, sem1, sem2, then NX (x, y) peer coordinates
// OUT_WSHARD (define): the output is WIDTH-sharded on NS cores, SW tiles wide (the sharded in0 of the
//   1D-mcast qkv / fc1 matmuls); extra compile-time args SW, NS; extra runtime args: row, then NS (x, y).
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
#include "addln_ids.hpp"
#ifdef ADDLN_PROF
#include "tools/profiler/kernel_profiler.hpp"
#define WZONE(n) DeviceZoneScopedN(n)
#else
#define WZONE(n)
#endif

FORCE_INLINE void zero_tail(uint32_t base, uint32_t NX) {
    // fp32 tile: face0 rows 0..15 at [0, 1024), face1 [1024, 2048), faces 2, 3 [2048, 4096); a row = 64 B
    const uint64_t zeros = get_noc_addr(MEM_ZEROS_BASE);
    const uint32_t tail = (16 - NX) * 64;
    if (tail) {
        noc_async_read(zeros, base + NX * 64, tail);
        noc_async_read(zeros, base + 1024 + NX * 64, tail);
    }
    for (uint32_t off = 2048; off < 4096; off += MEM_ZEROS_SIZE) {
        noc_async_read(zeros, base + off, MEM_ZEROS_SIZE);
    }
}

template <uint32_t c_part, uint32_t c_gath>
FORCE_INLINE void gather_round(uint32_t col, uint32_t sem_id, uint32_t NX, uint32_t my_x, uint32_t my_y) {
    cb_reserve_back(c_gath, 1);
    const uint32_t base = get_write_ptr(c_gath);
    const uint32_t row = base + col * 64;
    cb_wait_front(c_part, 1);
    const uint32_t src = get_read_ptr(c_part);
#ifdef ADDLN_NOSYNC  // timing experiment only: own row, no peers (numerically wrong)
    noc_async_write(src, get_noc_addr(row), 64);
    noc_async_write(src + 1024, get_noc_addr(row + 1024), 64);
    noc_async_write_barrier();
    cb_pop_front(c_part, 1);
    cb_push_back(c_gath, 1);
    return;
#endif
    for (uint32_t k = 0; k < NX; ++k) {
        const uint32_t px = get_arg_val<uint32_t>(ADDLN_RT0 + 5 + 2 * k), py = get_arg_val<uint32_t>(ADDLN_RT0 + 6 + 2 * k);
        noc_async_write(src, get_noc_addr(px, py, row), 64);
        noc_async_write(src + 1024, get_noc_addr(px, py, row + 1024), 64);
    }
#ifdef ADDLN_FLUSHED  // experiment: rely on in-order delivery of the writes and the following atomic
    noc_async_writes_flushed();
#else
    noc_async_write_barrier();
#endif
    const uint32_t sem_addr = get_semaphore(sem_id);
    for (uint32_t k = 0; k < NX; ++k) {
        const uint32_t px = get_arg_val<uint32_t>(ADDLN_RT0 + 5 + 2 * k), py = get_arg_val<uint32_t>(ADDLN_RT0 + 6 + 2 * k);
        if (px == my_x && py == my_y) {
            continue;
        }
        noc_semaphore_inc(get_noc_addr(px, py, sem_addr), 1);
    }
    volatile tt_l1_ptr uint32_t* sem = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(sem_addr);
    noc_semaphore_wait(sem, NX - 1);
    noc_semaphore_set(sem, 0);
    noc_async_atomic_barrier();
    cb_pop_front(c_part, 1);
    cb_push_back(c_gath, 1);
}

#ifdef ADDLN_MCAST
// round 12: the partial row goes to the same offset of every row peer's gather buffer -> ONE multicast (loopback, the
// physical rectangle spanned by the NX peers = exactly the grid row) instead of NX unicasts (~0.6 us of issue per partial)
FORCE_INLINE uint64_t row_mcast_addr(uint32_t NX, uint32_t addr) {
    uint32_t sx = 0xffff, sy = 0xffff, ex = 0, ey = 0;
    for (uint32_t k = 0; k < NX; ++k) {
        const uint32_t px = get_arg_val<uint32_t>(ADDLN_RT0 + 5 + 2 * k), py = get_arg_val<uint32_t>(ADDLN_RT0 + 6 + 2 * k);
        sx = px < sx ? px : sx;
        sy = py < sy ? py : sy;
        ex = px > ex ? px : ex;
        ey = py > ey ? py : ey;
    }
    return noc_index == 1 ? get_noc_multicast_addr(ex, ey, sx, sy, addr) : get_noc_multicast_addr(sx, sy, ex, ey, addr);
}
#endif

#ifdef ADDLN_1R
// one round for both partials (round 6): rows `col` of G1 (sum s) and G2 (sum s^2) on every row peer, one increment
template <uint32_t c_pa, uint32_t c_ga, uint32_t c_pb, uint32_t c_gb>
FORCE_INLINE void gather_round2(uint32_t col, uint32_t sem_id, uint32_t NX, uint32_t my_x, uint32_t my_y) {
    cb_reserve_back(c_ga, 1);
    cb_reserve_back(c_gb, 1);
    const uint32_t ra = get_write_ptr(c_ga) + col * 64, rb = get_write_ptr(c_gb) + col * 64;
    const uint32_t sem_addr = get_semaphore(sem_id);
#ifdef ADDLN_SPLITSEND
    // round 9: the sum partial goes out while the compute still reduces the sum-of-squares partial
    {WZONE("W-pwait");
    cb_wait_front(c_pa, 1);
    }
    const uint32_t sa = get_read_ptr(c_pa);
    {WZONE("W-send");
#ifdef ADDLN_MCAST
    noc_async_write_multicast_loopback_src(sa, row_mcast_addr(NX, ra), 64, NX);
    noc_async_write_multicast_loopback_src(sa + 1024, row_mcast_addr(NX, ra + 1024), 64, NX);
#else
    for (uint32_t k = 0; k < NX; ++k) {
        const uint32_t px = get_arg_val<uint32_t>(ADDLN_RT0 + 5 + 2 * k), py = get_arg_val<uint32_t>(ADDLN_RT0 + 6 + 2 * k);
        noc_async_write(sa, get_noc_addr(px, py, ra), 64);
        noc_async_write(sa + 1024, get_noc_addr(px, py, ra + 1024), 64);
    }
#endif
    }
    {WZONE("W-pwait2");
    cb_wait_front(c_pb, 1);
    }
    const uint32_t sb = get_read_ptr(c_pb);
    {WZONE("W-send2");
#ifdef ADDLN_MCAST
    noc_async_write_multicast_loopback_src(sb, row_mcast_addr(NX, rb), 64, NX);
    noc_async_write_multicast_loopback_src(sb + 1024, row_mcast_addr(NX, rb + 1024), 64, NX);
#else
    for (uint32_t k = 0; k < NX; ++k) {
        const uint32_t px = get_arg_val<uint32_t>(ADDLN_RT0 + 5 + 2 * k), py = get_arg_val<uint32_t>(ADDLN_RT0 + 6 + 2 * k);
        noc_async_write(sb, get_noc_addr(px, py, rb), 64);
        noc_async_write(sb + 1024, get_noc_addr(px, py, rb + 1024), 64);
    }
#endif
    }
#else
    {WZONE("W-pwait");
    cb_wait_front(c_pa, 1);
    cb_wait_front(c_pb, 1);
    }
    const uint32_t sa = get_read_ptr(c_pa), sb = get_read_ptr(c_pb);
    {WZONE("W-send");
    for (uint32_t k = 0; k < NX; ++k) {
        const uint32_t px = get_arg_val<uint32_t>(ADDLN_RT0 + 5 + 2 * k), py = get_arg_val<uint32_t>(ADDLN_RT0 + 6 + 2 * k);
        noc_async_write(sa, get_noc_addr(px, py, ra), 64);
        noc_async_write(sa + 1024, get_noc_addr(px, py, ra + 1024), 64);
        noc_async_write(sb, get_noc_addr(px, py, rb), 64);
        noc_async_write(sb + 1024, get_noc_addr(px, py, rb + 1024), 64);
    }
    }
#endif
    {WZONE("W-bar");
    noc_async_write_barrier();
    }
    {WZONE("W-inc");
    for (uint32_t k = 0; k < NX; ++k) {
        const uint32_t px = get_arg_val<uint32_t>(ADDLN_RT0 + 5 + 2 * k), py = get_arg_val<uint32_t>(ADDLN_RT0 + 6 + 2 * k);
        if (px == my_x && py == my_y) {
            continue;
        }
        noc_semaphore_inc(get_noc_addr(px, py, sem_addr), 1);
    }
    }
    volatile tt_l1_ptr uint32_t* sem = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(sem_addr);
    {WZONE("W-swait");
    noc_semaphore_wait(sem, NX - 1);
    }
    noc_semaphore_set(sem, 0);
    noc_async_atomic_barrier();
    cb_pop_front(c_pa, 1);
    cb_pop_front(c_pb, 1);
    cb_push_back(c_ga, 1);
    cb_push_back(c_gb, 1);
}
#endif

#ifdef ADDLN_LEADER
// round 12: gather through the row LEADER (peer 0) instead of all-to-all.  Every other core writes its two partial rows
// (sum, sum of squares; face 0 / face 1 halves) into row `col` of the LEADER's gather buffers only (4 small writes instead
// of 4 NX), acks them and increments the leader's sem_a.  The leader waits for NX-1 increments, then multicasts the
// filled rows 0 .. NX-1 of both gather buffers to the rest of the row (ONE source per row: no concurrent multicasts) and
// sets everyone's sem_b to 1 with a multicast after flushing the data (same NoC / VC: the flag cannot overtake it).
template <uint32_t c_pa, uint32_t c_ga, uint32_t c_pb, uint32_t c_gb>
FORCE_INLINE void gather_leader(uint32_t col, uint32_t sem_a, uint32_t sem_b, uint32_t NX) {
    cb_reserve_back(c_ga, 1);
    cb_reserve_back(c_gb, 1);
    const uint32_t ga = get_write_ptr(c_ga), gb = get_write_ptr(c_gb);
    const uint32_t ra = ga + col * 64, rb = gb + col * 64;
    const uint32_t lx = get_arg_val<uint32_t>(ADDLN_RT0 + 5), ly = get_arg_val<uint32_t>(ADDLN_RT0 + 6);
    {WZONE("W-pwait");
    cb_wait_front(c_pa, 1);
    }
    const uint32_t sa = get_read_ptr(c_pa);
    noc_async_write(sa, get_noc_addr(lx, ly, ra), 64);
    noc_async_write(sa + 1024, get_noc_addr(lx, ly, ra + 1024), 64);
    {WZONE("W-pwait2");
    cb_wait_front(c_pb, 1);
    }
    const uint32_t sb = get_read_ptr(c_pb);
    noc_async_write(sb, get_noc_addr(lx, ly, rb), 64);
    noc_async_write(sb + 1024, get_noc_addr(lx, ly, rb + 1024), 64);
    {WZONE("W-bar");
    noc_async_write_barrier();
    }
    volatile tt_l1_ptr uint32_t* sa_p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(sem_a));
    volatile tt_l1_ptr uint32_t* sb_p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(sem_b));
    if (col == 0) {
        {WZONE("W-lwait");
        noc_semaphore_wait(sa_p, NX - 1);
        }
        WZONE("W-lmc");
        uint32_t sx = 0xffff, sy = 0xffff, ex = 0, ey = 0;
        for (uint32_t k = 0; k < NX; ++k) {
            const uint32_t px = get_arg_val<uint32_t>(ADDLN_RT0 + 5 + 2 * k), py = get_arg_val<uint32_t>(ADDLN_RT0 + 6 + 2 * k);
            sx = px < sx ? px : sx;
            sy = py < sy ? py : sy;
            ex = px > ex ? px : ex;
            ey = py > ey ? py : ey;
        }
        auto mc = [&](uint32_t addr) -> uint64_t {
            return noc_index == 1 ? get_noc_multicast_addr(ex, ey, sx, sy, addr) : get_noc_multicast_addr(sx, sy, ex, ey, addr);
        };
        // the leader sits inside the rectangle: multicast to the NX - 1 others (no loopback), as cmm's in-place sends
        noc_async_write_multicast(ga, mc(ga), NX * 64, NX - 1);
        noc_async_write_multicast(ga + 1024, mc(ga + 1024), NX * 64, NX - 1);
        noc_async_write_multicast(gb, mc(gb), NX * 64, NX - 1);
        noc_async_write_multicast(gb + 1024, mc(gb + 1024), NX * 64, NX - 1);
        noc_async_writes_flushed();
        *sb_p = 1;  // the flag value (the leader's own sem_b is the multicast source)
        noc_semaphore_set_multicast(get_semaphore(sem_b), mc(get_semaphore(sem_b)), NX - 1);
        noc_async_write_barrier();
    } else {
        {WZONE("W-inc");
        noc_semaphore_inc(get_noc_addr(lx, ly, get_semaphore(sem_a)), 1);
        }
        {WZONE("W-swait");
        noc_semaphore_wait(sb_p, 1);
        }
    }
    noc_async_atomic_barrier();
    cb_pop_front(c_pa, 1);
    cb_pop_front(c_pb, 1);
    cb_push_back(c_ga, 1);
    cb_push_back(c_gb, 1);
}
#endif

ADDLN_KERNEL_MAIN {
    constexpr uint32_t NT = get_compile_time_arg_val(ADDLN_CT0);
    constexpr uint32_t NX = get_compile_time_arg_val(ADDLN_CT0 + 1);
    constexpr auto out_args = TensorAccessorArgs<ADDLN_CT0 + 2>();
    const uint32_t out_addr = get_arg_val<uint32_t>(ADDLN_RT0 + 0);
    const uint32_t out_start = get_arg_val<uint32_t>(ADDLN_RT0 + 1);
    const uint32_t col = get_arg_val<uint32_t>(ADDLN_RT0 + 2);
    const uint32_t sem1 = get_arg_val<uint32_t>(ADDLN_RT0 + 3), sem2 = get_arg_val<uint32_t>(ADDLN_RT0 + 4);
    const uint32_t my_x = get_arg_val<uint32_t>(ADDLN_RT0 + 5 + 2 * col), my_y = get_arg_val<uint32_t>(ADDLN_RT0 + 6 + 2 * col);
    constexpr uint32_t c_g1 = ADDLN_C_G1, c_g2 = ADDLN_C_G2, c_p1 = ADDLN_C_P1, c_p2 = ADDLN_C_P2, c_out = ADDLN_C_OUT;

    zero_tail(get_write_ptr(c_g1), NX);
    zero_tail(get_write_ptr(c_g2), NX);
    noc_async_read_barrier();

#ifdef ADDLN_1R
#ifdef ADDLN_LEADER
    gather_leader<c_p1, c_g1, c_p2, c_g2>(col, sem1, sem2, NX);
#else
    gather_round2<c_p1, c_g1, c_p2, c_g2>(col, sem1, NX, my_x, my_y);
#endif
#else
    gather_round<c_p1, c_g1>(col, sem1, NX, my_x, my_y);
    gather_round<c_p2, c_g2>(col, sem2, NX, my_x, my_y);
#endif

    const uint32_t page = get_local_cb_interface(c_out).fifo_page_size;
    {WZONE("W-owait");
    cb_wait_front(c_out, NT);
    }
    uint32_t rp = get_read_ptr(c_out);
#if defined(OUT_XRES)
    // round 12: the LN output goes straight into the RESIDENT in0 of the following xmm fc1 (tt_kernels.xmm_res_alloc: two
    // tensors, even / odd K blocks, same L1 address on all 12 x 10 cores).  This core's NT tiles = K block `col` (kb = NT)
    // of token row `row`: one NT-tile write to every core whose token pair is row / 2 -- three full grid rows
    // (multicast, loopback when this core sits in one) and the row-9 cores x = p, p + 3, p + 6, p + 9 (unicasts).
    // extra runtime args after the peers: row, res0, res1, then 3 x (sx, sy, ex, ey, self_in), then 4 x (x, y)
    {
        constexpr uint32_t a0 = 5 + 2 * NX;
        const uint32_t row = get_arg_val<uint32_t>(ADDLN_RT0 + a0);
        const uint32_t res = get_arg_val<uint32_t>(ADDLN_RT0 + a0 + 1 + (col & 1));
        // shard tile (j, r, k): j = col / 2, r = row % 2, k = 0 .. NT-1; a block is 2 x NT tiles
        const uint32_t dst = res + ((col >> 1) * 2 * NT + (row & 1) * NT) * page;
        const uint32_t bytes = NT * page;
        for (uint32_t i = 0; i < 3; ++i) {
            const uint32_t b = ADDLN_RT0 + a0 + 3 + 5 * i;
            const uint32_t sx = get_arg_val<uint32_t>(b), sy = get_arg_val<uint32_t>(b + 1);
            const uint32_t ex = get_arg_val<uint32_t>(b + 2), ey = get_arg_val<uint32_t>(b + 3);
            const uint64_t ma = noc_index == 1 ? get_noc_multicast_addr(ex, ey, sx, sy, dst) : get_noc_multicast_addr(sx, sy, ex, ey, dst);
            if (get_arg_val<uint32_t>(b + 4)) {
                noc_async_write_multicast_loopback_src(rp, ma, bytes, 12);
            } else {
                noc_async_write_multicast(rp, ma, bytes, 12);
            }
        }
        for (uint32_t i = 0; i < 4; ++i) {
            const uint32_t b = ADDLN_RT0 + a0 + 18 + 2 * i;
            noc_async_write(rp, get_noc_addr(get_arg_val<uint32_t>(b), get_arg_val<uint32_t>(b + 1), dst), bytes);
        }
    }
#elif defined(OUT_WSHARD)
    constexpr uint32_t SW = get_compile_time_arg_val(out_args.next_compile_time_args_offset());
    constexpr uint32_t NS = get_compile_time_arg_val(out_args.next_compile_time_args_offset() + 1);
    constexpr uint32_t a0 = 5 + 2 * NX;
    const uint32_t row = get_arg_val<uint32_t>(ADDLN_RT0 + a0);
    for (uint32_t i = 0; i < NT; ++i) {
        const uint32_t t = col * NT + i;  // global tile column
        const uint32_t j = t / SW;        // shard core
        const uint32_t sx = get_arg_val<uint32_t>(ADDLN_RT0 + a0 + 1 + 2 * j), sy = get_arg_val<uint32_t>(ADDLN_RT0 + a0 + 2 + 2 * j);
        noc_async_write(rp, get_noc_addr(sx, sy, out_addr + (row * SW + (t - j * SW)) * page), page);
        rp += page;
    }
#else
    const auto out = TensorAccessor(out_args, out_addr, page);
    for (uint32_t i = 0; i < NT; ++i) {
        noc_async_write(rp, out.get_noc_addr(out_start + i), page);
        rp += page;
    }
#endif
    noc_async_write_barrier();
    cb_pop_front(c_out, NT);
}