// SPDX-License-Identifier: Apache-2.0
// Output tail writer (moge-2 opt kernel): c_16 holds per block a 32 x (B*32) row-major block (row c =
// channel c of B*32 consecutive pixels, out_bands_compute.cpp). Pixel tile t0 of the block lies in band
// t0 / TPB; channel rows c < ROWS are written as one B*64-byte segment into page c (= channel row) of that
// band's ROW_MAJOR [ROWS, TPB*32] tensor (rows ROW0 + c with the ROW0 define). Blocks never straddle a band (TPB % B == 0, host check).
// Runtime args: band addresses [8], num_blocks, start_block. Compile args: B, ROWS, TPB, TensorAccessorArgs.
#include "api/dataflow/dataflow_api.h"
#include "api/dataflow/noc.h"
#include "api/dataflow/dataflow_buffer.h"
#include "api/tensor/noc_traits.h"

#ifndef ROW0
#define ROW0 0   // first channel row written (MOGE_OUT_GROUPS: one output tensor per head's channel rows)
#endif

void kernel_main() {
    uint32_t addr[8];
    for (uint32_t i = 0; i < 8; ++i) {
        addr[i] = get_arg_val<uint32_t>(i);
    }
    const uint32_t nblocks = get_arg_val<uint32_t>(8);
    const uint32_t start_block = get_arg_val<uint32_t>(9);
    constexpr uint32_t B = get_compile_time_arg_val(0);
    constexpr uint32_t ROWS = get_compile_time_arg_val(1);
    constexpr uint32_t TPB = get_compile_time_arg_val(2);
    constexpr auto d_args = TensorAccessorArgs<3>();
    constexpr uint32_t cb_id = 16;
    constexpr uint32_t seg = B * 32 * 2;
    Noc noc;
    DataflowBuffer dfb(cb_id);
    for (uint32_t b = 0; b < nblocks; ++b) {
        const uint32_t t0 = (start_block + b) * B;
        const uint32_t band = t0 / TPB;
        const uint32_t off = (t0 - band * TPB) * 32 * 2;
        const auto s = TensorAccessor(d_args, addr[band]);
        dfb.wait_front(B);
        for (uint32_t c = 0; c < ROWS; ++c) {
            noc.async_write(dfb, s, seg, {.offset_bytes = (ROW0 + c) * seg}, {.page_id = c, .offset_bytes = off});
        }
        noc.async_writes_flushed();
        dfb.pop_front(B);
    }
    noc.async_write_barrier();
}