code/tt_moge/kernels/writer_batched.cpp

// SPDX-License-Identifier: Apache-2.0
// Interleaved page writer from c_2, BLOCK pages per flush (moge-2 opt kernel).
// Runtime args: dst_addr, num_pages, start_id. Compile args: BLOCK, TensorAccessorArgs.
#include "api/dataflow/dataflow_api.h"
#include "api/dataflow/noc.h"
#include "api/dataflow/dataflow_buffer.h"
#include "api/tensor/noc_traits.h"

void kernel_main() {
    const uint32_t dst_addr = get_arg_val<uint32_t>(0);
    const uint32_t num_pages = get_arg_val<uint32_t>(1);
    const uint32_t start_id = get_arg_val<uint32_t>(2);
    constexpr uint32_t BLOCK = get_compile_time_arg_val(0);
    constexpr auto dst_args = TensorAccessorArgs<1>();
    constexpr uint32_t cb_id = 2;
    const uint32_t page_bytes = get_local_cb_interface(cb_id).fifo_page_size;
    const auto s = TensorAccessor(dst_args, dst_addr);
    Noc noc;
    DataflowBuffer dfb(cb_id);
    uint32_t i = start_id;
    const uint32_t end_id = start_id + num_pages;
    while (i < end_id) {
        const uint32_t n = (end_id - i) < BLOCK ? (end_id - i) : BLOCK;
        dfb.wait_front(n);
        for (uint32_t j = 0; j < n; ++j) {
            noc.async_write(dfb, s, page_bytes, {.offset_bytes = j * page_bytes}, {.page_id = i + j});
        }
        noc.async_writes_flushed();
        dfb.pop_front(n);
        i += n;
    }
    noc.async_write_barrier();
}