// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// add+LayerNorm reader (NCRISC): marks the block-sharded x (and y) shards as resident, builds the SUM
// reduce scaler tile locally (bf16 1.0 in the first row of each face; no DRAM round trip on the
// critical path), then reads this core's NT gamma / beta tiles (full 32-row tiles, interleaved DRAM),
// which the compute needs only for its last pass.
// compile-time: NT, TensorAccessorArgs(gamma), TensorAccessorArgs(beta)
// runtime: gamma_addr, beta_addr, col0 (first gamma/beta tile of this core)
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
#include "addln_ids.hpp"

ADDLN_KERNEL_MAIN {
    constexpr uint32_t NT = get_compile_time_arg_val(ADDLN_CT0);
    [[maybe_unused]] constexpr auto ga_args = TensorAccessorArgs<ADDLN_CT0 + 1>();
    [[maybe_unused]] constexpr auto be_args = TensorAccessorArgs<ga_args.next_compile_time_args_offset()>();
    const uint32_t ga_addr = get_arg_val<uint32_t>(ADDLN_RT0 + 0);
    const uint32_t be_addr = get_arg_val<uint32_t>(ADDLN_RT0 + 1);
    const uint32_t col0 = get_arg_val<uint32_t>(ADDLN_RT0 + 2);
    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;

    // sharded inputs live in the CB already (globally allocated)
    cb_reserve_back(c_x, NT);
    cb_push_back(c_x, NT);
#if defined(HAS_Y) && !defined(ADDLN_Y_INTERNAL)  // fused after cmm: the compute produces y itself
    cb_reserve_back(c_y, NT);
    cb_push_back(c_y, NT);
#endif
#ifndef ADDLN_SLOWSCAL
    // round 8: the scaler first and fast -- the compute needs it right after the add (~0.5 us in), and 512 single-word
    // stores from NCRISC took ~1.2 us of the critical path: zero the tile with NoC reads of the hardware zero region,
    // then write the 32 non-zero words
    {
        cb_reserve_back(c_scal, 1);
        const uint32_t sp = get_write_ptr(c_scal);
        const uint64_t zeros = get_noc_addr(MEM_ZEROS_BASE);
        for (uint32_t off = 0; off < 2048; off += MEM_ZEROS_SIZE) {
            noc_async_read(zeros, sp + off, MEM_ZEROS_SIZE);
        }
        noc_async_read_barrier();
        volatile tt_l1_ptr uint32_t* sc = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(sp);
        for (uint32_t f = 0; f < 4; ++f) {
            for (uint32_t j = 0; j < 8; ++j) {
                sc[f * 128 + j] = 0x3F803F80u;
            }
        }
        cb_push_back(c_scal, 1);
    }
#endif
#ifndef ADDLN_NOAFF  // round 9: with the LN affine folded into the next weights there is no gamma / beta pass
    // gamma / beta reads first (in flight while the scaler is built)
    const uint32_t gb_page = get_local_cb_interface(c_gam).fifo_page_size;
    const auto ga = TensorAccessor(ga_args, ga_addr, gb_page);
    [[maybe_unused]] const auto be = TensorAccessor(be_args, be_addr, gb_page);
    cb_reserve_back(c_gam, NT);
#ifndef ADDLN_NOBETA  // round 9: beta folded into the next matmul's bias
    cb_reserve_back(c_bet, NT);
#endif
    uint32_t gp = get_write_ptr(c_gam);
    [[maybe_unused]] uint32_t bp = get_write_ptr(c_bet);
    for (uint32_t i = 0; i < NT; ++i) {
        noc_async_read(ga.get_noc_addr(col0 + i), gp, gb_page);
#ifndef ADDLN_NOBETA
        noc_async_read(be.get_noc_addr(col0 + i), bp, gb_page);
#endif
        gp += gb_page;
        bp += gb_page;
    }
    // scaler: bf16 tile (4 faces of 16x16, 512 B each), 1.0 (0x3F80) in row 0 of every face, 0 elsewhere
#ifdef ADDLN_SLOWSCAL
    cb_reserve_back(c_scal, 1);
    volatile tt_l1_ptr uint32_t* sc = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(c_scal));
    for (uint32_t i = 0; i < 512; ++i) {
        sc[i] = 0;
    }
    for (uint32_t f = 0; f < 4; ++f) {
        for (uint32_t j = 0; j < 8; ++j) {
            sc[f * 128 + j] = 0x3F803F80u;
        }
    }
    cb_push_back(c_scal, 1);
#endif
    noc_async_read_barrier();
    cb_push_back(c_gam, NT);
#ifndef ADDLN_NOBETA
    cb_push_back(c_bet, NT);
#endif
#elif defined(ADDLN_SLOWSCAL)
#error "ADDLN_NOAFF needs the fast scaler path"
#endif
#ifdef ADDLN_PREFETCH
    // cross-op weight prefetch (tt_kernels.CmmPrefetch): this LN is latency-bound and leaves the DRAM idle, so the
    // owner cores of a later cmm_linear read its weight pairs into that op's persistent L1 buffer.
    // runtime args 3..: npf, w_addr, npf x (bank, bank offset, local dst, bytes).  Plain reads + one plain barrier.
    const uint32_t npf = get_arg_val<uint32_t>(ADDLN_RT0 + 3);
    if (npf) {
        const uint32_t w_addr = get_arg_val<uint32_t>(ADDLN_RT0 + 4);
        for (uint32_t i = 0; i < npf; ++i) {
            const uint32_t bank = get_arg_val<uint32_t>(ADDLN_RT0 + 5 + 4 * i), off = get_arg_val<uint32_t>(ADDLN_RT0 + 6 + 4 * i);
            const uint32_t ldst = get_arg_val<uint32_t>(ADDLN_RT0 + 7 + 4 * i), bytes = get_arg_val<uint32_t>(ADDLN_RT0 + 8 + 4 * i);
            noc_async_read(get_noc_addr_from_bank_id<true>(bank, w_addr + off), ldst, bytes);
        }
        noc_async_read_barrier();
    }
#endif
}