// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hgemv reader (NCRISC, NOC_0), round 6: one output tile column n per core of y = [LN](x) @ W (+ b) [GELU] [+ r]
// for the 1-token MANO head.  Vectors are "row-form" tensors [1, 1, 32, L] (TILE, interleaved L1): only row 0 of
// every tile is meaningful (face 0 row 0 = elements 0..15, face 1 row 0 = 16..31 of that tile).
//   weights: the core's packed contiguous [KT x 1] tile column (CmmWeight, GX = N/32, NT = 1, one K block), one DRAM
//            run, transaction id 1 (waited last); the small reads below use transaction id 2.  Never trid 0.
//   HG_LN:   x row 0 of the KT = 32 tiles gathered into ONE dense tile (row t <- tile t's row 0), the LayerNorm runs
//            on that tile; else x's row 0s go straight into the row-form matmul operand c_hrow.
//   HG_DIN=D: x is instead a DENSE vector with D replica tiles [1, 1, 32*D, 32] (one page read, replica n % D); without
//            LN the writer spreads it into c_hrow.
//   bias tile n (full tile), HG_RES: residual tile n row 0.
//   HG_TINY (round 10): the dense vectors are ROW-MAJOR (element 32t + j at word 32t + j of the tile buffer, so row t
//            of the "tile" is exactly a [1, 32] tiny tile), the matmul operand c_hrow is KT [1, 32] tiny tiles (an
//            alias of the dense tile's memory, or gathered contiguously), and bias / residual are [1, 32] tiny tiles.
// compile-time: KT, ES (x element bytes: 4 fp32 / 2 bf16), RES_ES, TensorAccessorArgs(x, bias[, res])
// runtime: x_addr, w_addr, w_bank, w_off, b_addr, res_addr, n
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

void kernel_main() {
    constexpr uint32_t KT = get_compile_time_arg_val(0);
    constexpr uint32_t ES = get_compile_time_arg_val(1);
    constexpr uint32_t RES_ES = get_compile_time_arg_val(2);
    constexpr auto x_args = TensorAccessorArgs<3>();
    constexpr auto b_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
#ifdef HG_RES
    constexpr auto r_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
#endif
    constexpr uint32_t c_xd = 0, c_w = 1, c_bias = 2, c_scal = 3, c_res = 4, c_hrow = 5;
    const uint32_t x_addr = get_arg_val<uint32_t>(0);
    const uint32_t w_addr = get_arg_val<uint32_t>(1), w_bank = get_arg_val<uint32_t>(2), w_off = get_arg_val<uint32_t>(3);
    const uint32_t b_addr = get_arg_val<uint32_t>(4), r_addr = get_arg_val<uint32_t>(5), n = get_arg_val<uint32_t>(6);

    noc_async_read_set_trid(2);  // the small latency-critical reads first, the weight run (trid 1) behind them
    constexpr uint32_t FB = 256 * ES;  // face bytes of x
    const uint32_t xpage = 1024 * ES;
    const auto x = TensorAccessor(x_args, x_addr, xpage);
#ifdef HG_DIN
    // dense input: replica (n % ND) of the dense vector tile (written by the producing hgemv), one 4 KB / 2 KB read
    cb_reserve_back(c_xd, 1);
#if defined(HG_TINY) && !defined(HG_LN)
    cb_reserve_back(c_hrow, KT);  // the same memory as KT [1, 32] tiny tiles (row-major dense vector)
#endif
    noc_async_read(x.get_noc_addr(n % HG_DIN), get_write_ptr(c_xd), xpage);
#elif defined(HG_LN)
    cb_reserve_back(c_xd, 1);
    const uint32_t xd = get_write_ptr(c_xd);
    for (uint32_t t = 0; t < KT; ++t) {
        const uint64_t s = x.get_noc_addr(t);
#ifdef HG_TINY
        const uint32_t d = xd + t * 32 * ES;  // row-major dense tile
        noc_async_read(s, d, 16 * ES);
        noc_async_read(s + FB, d + 16 * ES, 16 * ES);
#else
        const uint32_t d = xd + ((t / 16) * 2) * FB + (t % 16) * 16 * ES;
        noc_async_read(s, d, 16 * ES);
        noc_async_read(s + FB, d + FB, 16 * ES);
#endif
    }
#endif
#ifdef HG_LN
    // SUM scaler: bf16 1.0 in row 0 of every face
    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));
#ifdef HG_FASTX
    // round 10: zero the tile with NoC reads of the hardware zero region (512 single-word stores took ~1 us), then the
    // 32 non-zero words; x and the scaler go to the compute before the bias / residual / weight reads are issued
    {
        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, get_write_ptr(c_scal) + off, MEM_ZEROS_SIZE);
        }
        noc_async_read_barrier_with_trid(2);  // also x
    }
#else
    for (uint32_t i = 0; i < 512; ++i) {
        sc[i] = 0;
    }
#endif
    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);
#ifdef HG_FASTX
    cb_push_back(c_xd, 1);
#endif
#elif !defined(HG_DIN)
    cb_reserve_back(c_hrow, KT);
    const uint32_t hr = get_write_ptr(c_hrow);
    for (uint32_t t = 0; t < KT; ++t) {
        const uint64_t s = x.get_noc_addr(t);
#ifdef HG_TINY
        noc_async_read(s, hr + t * 32 * ES, 16 * ES);  // tiny tile t = row 0 of x tile t, contiguous
        noc_async_read(s + FB, hr + t * 32 * ES + 16 * ES, 16 * ES);
#else
        noc_async_read(s, hr + t * xpage, 16 * ES);
        noc_async_read(s + FB, hr + t * xpage + FB, 16 * ES);
#endif
    }
#endif
#if defined(HG_FASTX) && defined(HG_DIN) && !defined(HG_LN)
    noc_async_read_barrier_with_trid(2);
    cb_push_back(c_xd, 1);
#ifdef HG_TINY
    cb_push_back(c_hrow, KT);
#endif
#endif
#ifdef HG_TINY
    {
        // bias: row 0 of the row-form bias tile n -> one [1, 32] tiny tile (BES = its element bytes)
        constexpr uint32_t BES = HG_BES;
        const auto bias = TensorAccessor(b_args, b_addr, 1024 * BES);
        cb_reserve_back(c_bias, 1);
        const uint64_t s = bias.get_noc_addr(n);
        noc_async_read(s, get_write_ptr(c_bias), 16 * BES);
        noc_async_read(s + 256 * BES, get_write_ptr(c_bias) + 16 * BES, 16 * BES);
    }
#else
    const uint32_t bpage = get_local_cb_interface(c_bias).fifo_page_size;
    const auto bias = TensorAccessor(b_args, b_addr, bpage);
    cb_reserve_back(c_bias, 1);
    noc_async_read(bias.get_noc_addr(n), get_write_ptr(c_bias), bpage);
#endif
#ifdef HG_RES
    constexpr uint32_t RFB = 256 * RES_ES;
    const auto res = TensorAccessor(r_args, r_addr, 1024 * RES_ES);
    cb_reserve_back(c_res, 1);
    {
        const uint64_t s = res.get_noc_addr(n);
        noc_async_read(s, get_write_ptr(c_res), 16 * RES_ES);
#ifdef HG_TINY
        noc_async_read(s + RFB, get_write_ptr(c_res) + 16 * RES_ES, 16 * RES_ES);
#else
        noc_async_read(s + RFB, get_write_ptr(c_res) + RFB, 16 * RES_ES);
#endif
    }
#endif
#ifdef HG_XFIRST
    // round 10: x / bias / residual land (and go to the compute) BEFORE the 64 KB weight run is issued -- the weight
    // responses otherwise share this core's inbound port with the x reads (x landed at ~1.9 us median, 4.5 us max)
    noc_async_read_barrier_with_trid(2);
#if defined(HG_LN) || defined(HG_DIN)
#ifndef HG_FASTX
    cb_push_back(c_xd, 1);
#if defined(HG_TINY) && defined(HG_DIN) && !defined(HG_LN)
    cb_push_back(c_hrow, KT);
#endif
#endif
#else
    cb_push_back(c_hrow, KT);
#endif
    cb_push_back(c_bias, 1);
#ifdef HG_RES
    cb_push_back(c_res, 1);
#endif
#endif
    // weights: one contiguous run, trid 1 (issued after the small reads so x is not queued behind 64 KB)
    const uint32_t wpage = get_local_cb_interface(c_w).fifo_page_size;
    cb_reserve_back(c_w, KT);
#ifndef HG_WL1  // round 10: HG_WL1 -- the weights are already in this core's L1 (prefetched by the previous hgemv)
    {
        noc_async_read_set_trid(1);
        uint64_t src = get_noc_addr_from_bank_id<true>(w_bank, w_addr + w_off);
        uint32_t dst = get_write_ptr(c_w), left = KT * wpage;
        while (left) {
            const uint32_t m = left < NOC_MAX_BURST_SIZE ? left : NOC_MAX_BURST_SIZE;
            noc_async_read<NOC_MAX_BURST_SIZE>(src, dst, m);
            src += m;
            dst += m;
            left -= m;
        }
    }
#endif
#ifndef HG_XFIRST
    noc_async_read_barrier_with_trid(2);
#if defined(HG_LN) || defined(HG_DIN)
#ifndef HG_FASTX
    cb_push_back(c_xd, 1);
#if defined(HG_TINY) && defined(HG_DIN) && !defined(HG_LN)
    cb_push_back(c_hrow, KT);
#endif
#endif
#else
    cb_push_back(c_hrow, KT);
#endif
    cb_push_back(c_bias, 1);
#ifdef HG_RES
    cb_push_back(c_res, 1);
#endif
#endif
    // (round 10: replacing this with the plain noc_async_read_barrier() hung chip 1 in scripts/micro/hgemv_test.py
    // -- OPT_REPORT "known hangs"; keep the transaction-id barrier)
    noc_async_read_barrier_with_trid(1);
    cb_push_back(c_w, KT);
}