// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hattn2 writer (BRISC), round 15: the two [1, 32] tiny context tiles of head h -> row 0 of tiles 2h, 2h+1 of the
// head-split output [1, H, 32, 64] (bf16; read by the out-projection GEMV as the row-form [1, 1, 32, H*64] vector).
// It first reads head h's 12 V tiles (needed last) on its own NoC, in parallel with the reader's q / K reads.
// compile-time: TensorAccessorArgs(out), TensorAccessorArgs(v).  runtime: out_addr, h, v_addr, v_nc, v_base (see reader)
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

void kernel_main() {
    constexpr auto o_args = TensorAccessorArgs<0>();
    constexpr auto v_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
    constexpr uint32_t c_v = 2, c_out = 16, page = 2048, FB = 512, MT = 6;
    const uint32_t o_addr = get_arg_val<uint32_t>(0), h = get_arg_val<uint32_t>(1);
    const auto out = TensorAccessor(o_args, o_addr, page);
    {
        const auto v = TensorAccessor(v_args, get_arg_val<uint32_t>(2), page);
        cb_reserve_back(c_v, 2 * MT);
        uint32_t vd = get_write_ptr(c_v);
        for (uint32_t t = 0; t < 2 * MT; ++t) {
            noc_async_read(v.get_noc_addr((t / 2) * get_arg_val<uint32_t>(3) + get_arg_val<uint32_t>(4) + (t % 2)), vd, page);
            vd += page;
        }
        noc_async_read_barrier();
        cb_push_back(c_v, 2 * MT);
    }
    cb_wait_front(c_out, 2);
    const uint32_t src = get_read_ptr(c_out);
    for (uint32_t d = 0; d < 2; ++d) {
        const uint64_t dst = out.get_noc_addr(2 * h + d);
        noc_async_write(src + d * 64, dst, 32);
        noc_async_write(src + d * 64 + 32, dst + FB, 32);
    }
    noc_async_write_barrier();
    cb_pop_front(c_out, 2);
}