code/tt_frnet/tt/kernels/mkh_writer.cpp

// SPDX-License-Identifier: Apache-2.0
// MK-H (OPT round 5, FRNET_POINT_MK=head; OPT_PLAN 3.1 / D19): FRHead as ONE row-streamed module megakernel, writer
// (RISCV_1): the NTO fp32 output tiles (logits) of each tile row r of this core's range to the output [.., N, 32 NTO].
// CT args: [0] per-core RT-arg count (P3), [1] NTO, then the TensorAccessorArgs of the output.
// Common RT args: [out_addr]. Per-core RT args: [r0, nr].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t NTO = get_compile_time_arg_val(1);
constexpr auto o_args = TensorAccessorArgs<2>();
constexpr uint32_t cb_out = 20;
constexpr uint32_t FTB = 4096;

void kernel_main() {
    const uint32_t r0 = get_arg_val<uint32_t>(0);
    const uint32_t nr = get_arg_val<uint32_t>(1);
    const auto o = TensorAccessor(o_args, get_common_arg_val<uint32_t>(0), FTB);
    for (uint32_t r = r0; r < r0 + nr; ++r) {
        cb_wait_front(cb_out, NTO);
        uint32_t rp = get_read_ptr(cb_out);
        for (uint32_t j = 0; j < NTO; ++j) {
            noc_async_write(rp, o.get_noc_addr(r * NTO + j), FTB);
            rp += FTB;
        }
        noc_async_writes_flushed();
        cb_pop_front(cb_out, NTO);
    }
    noc_async_write_barrier();
}