// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hgemv writer (BRISC, NOC_1), round 6.  HG_LN: the normalized dense tile (row t = elements 32t..32t+31, bf16) is
// spread into the row-form matmul operand (tile t row 0), by local NoC copies.  Then row 0 of the output tile goes
// to tile n of the row-form output tensor.
// HG_DOUT=D: row 0 is also written as row n of D replica tiles of a dense output vector [1, 1, 32*D, 32].
// compile-time: KT, OES (output element bytes), TensorAccessorArgs(out)[, TensorAccessorArgs(dense out)];
// runtime: out_addr, n[, dense_addr]
#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 OES = get_compile_time_arg_val(1);
    constexpr auto o_args = TensorAccessorArgs<2>();
    constexpr uint32_t c_hd = 6, c_hrow = 5, c_out = 16;
    const uint32_t o_addr = get_arg_val<uint32_t>(0), n = get_arg_val<uint32_t>(1);
#if (defined(HG_LN) || defined(HG_DIN)) && !defined(HG_TINY)
    {
        // dense tile (row t = elements 32t..32t+31) -> row-form operand (tile t row 0): the LN output (bf16), or the
        // dense input itself without LN (element size HES)
#ifdef HG_LN
        constexpr uint32_t HES = 2, c_src = c_hd;
#else
        constexpr uint32_t HES = HG_IES, c_src = 0;
#endif
        constexpr uint32_t HFB = 256 * HES, HPAGE = 1024 * HES;
        cb_wait_front(c_src, 1);
        cb_reserve_back(c_hrow, KT);
        const uint32_t hd = get_read_ptr(c_src), hr = get_write_ptr(c_hrow);
        for (uint32_t t = 0; t < KT; ++t) {
            const uint32_t s = hd + ((t / 16) * 2) * HFB + (t % 16) * 16 * HES;
            noc_async_read(get_noc_addr(s), hr + t * HPAGE, 16 * HES);
            noc_async_read(get_noc_addr(s + HFB), hr + t * HPAGE + HFB, 16 * HES);
        }
        noc_async_read_barrier();
        cb_push_back(c_hrow, KT);
        cb_pop_front(c_src, 1);
    }
#endif
    constexpr uint32_t OFB = 256 * OES;
    const auto out = TensorAccessor(o_args, o_addr, 1024 * OES);
    cb_wait_front(c_out, 1);
    const uint32_t src = get_read_ptr(c_out);
    const uint64_t dst = out.get_noc_addr(n);
#ifdef HG_TINY
    // round 10: c_out is a [1, 32] tiny tile: its two face rows are contiguous
    constexpr uint32_t SRC1 = 16 * OES;
#else
    constexpr uint32_t SRC1 = OFB;
#endif
#ifdef HG_REP
    // replicate row 0 into all 32 rows (consumers that process whole tiles, e.g. SDPA on the query)
    for (uint32_t r = 0; r < 32; ++r) {
        const uint32_t off = ((r / 16) * 2) * OFB + (r % 16) * 16 * OES;
        noc_async_write(src, dst + off, 16 * OES);
        noc_async_write(src + SRC1, dst + off + OFB, 16 * OES);
    }
#else
    noc_async_write(src, dst, 16 * OES);
    noc_async_write(src + SRC1, dst + OFB, 16 * OES);
#endif
#ifdef HG_DOUT
    {
        // dense replicas: row n of each of the HG_DOUT replica tiles of the dense output vector
        constexpr auto d_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
        const auto dout = TensorAccessor(d_args, get_arg_val<uint32_t>(2), 1024 * OES);
#ifdef HG_TINY
        // round 10: row-major dense replicas: row n = 32 contiguous elements = this core's tiny output tile
        for (uint32_t c = 0; c < HG_DOUT; ++c) {
            noc_async_write(src, dout.get_noc_addr(c) + n * 32 * OES, 32 * OES);
        }
#else
        const uint32_t off = ((n / 16) * 2) * OFB + (n % 16) * 16 * OES;
        for (uint32_t c = 0; c < HG_DOUT; ++c) {
            const uint64_t d = dout.get_noc_addr(c) + off;
            noc_async_write(src, d, 16 * OES);
            noc_async_write(src + OFB, d + OFB, 16 * OES);
        }
#endif
    }
#endif
    noc_async_write_barrier();
#ifdef HG_SIGNAL
    // round 15: this program also runs the head attention (kernels/hattn2, HA2_WAIT) on other cores: the output column
    // (a query column) is in L1 (acked above) -> increment the attention core's semaphore 0
    noc_semaphore_inc(get_noc_addr(get_arg_val<uint32_t>(HG_SIGNAL), get_arg_val<uint32_t>(HG_SIGNAL + 1), get_semaphore(0)), 1);
    noc_async_atomic_barrier();
#endif
    cb_pop_front(c_out, 1);
}