// SPDX-License-Identifier: Apache-2.0
// rowgather_k (tt/rowgather_k.py, OPT_PLAN R2 "SCA glue"): out[n] = table[index[n]] (a zero row when index[n] >= R,
// e.g. the sentinel) for TILE tables, without the untilize / zero-row pad / embedding / tilize of the RM gather.
//
//   table bf16 TILE [1, 1, R, D], L1 interleaved (DRAM would need 64-byte aligned reads: watcher finding): row q = row q % 32 of tile row q / 32 (two 32-byte face rows
//         per 32-column tile: faces 2 (r / 16) and + 1, row r % 16);
//   index uint32 ROW_MAJOR [1, N] (one page);
//   out   bf16 TILE [1, 1, N, D] (rows >= N zero).
// Work item = one 32-row block (DT tiles); items split over the cores (contiguous ranges) and, per core, alternately
// over the two data-movement RISC-Vs (SLOT 0 / 1, each its own CB): read 32 indices, 2 DT reads per row straight into
// the output tiles' face rows (zeros page for a zero row), one barrier, then DT tile writes. Data movement only; no
// semaphores, no multicast; common runtime args only [table, index, out] (P3).
#include <stdint.h>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t SLOT = get_compile_time_arg_val(0);
constexpr uint32_t cb_out = get_compile_time_arg_val(1);
constexpr uint32_t cb_idx = get_compile_time_arg_val(2);
constexpr uint32_t R = get_compile_time_arg_val(3);       // table rows (logical)
constexpr uint32_t N = get_compile_time_arg_val(4);       // output rows (logical)
constexpr uint32_t DT = get_compile_time_arg_val(5);      // 32-column tiles per row
constexpr uint32_t NWORK = get_compile_time_arg_val(6);   // 32-row output blocks
constexpr uint32_t NC = get_compile_time_arg_val(7);      // cores
constexpr uint32_t GX = get_compile_time_arg_val(8);      // grid width
constexpr auto t_args = TensorAccessorArgs<9>();
constexpr auto i_args = TensorAccessorArgs<t_args.next_compile_time_args_offset()>();
constexpr auto o_args = TensorAccessorArgs<i_args.next_compile_time_args_offset()>();

constexpr uint32_t TB16 = 2048;

void kernel_main() {
    const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (ci >= NC) {
        return;
    }
    const uint32_t w0 = ci * NWORK / NC;
    const uint32_t w1 = (ci + 1) * NWORK / NC;
    const auto t_acc = TensorAccessor(t_args, get_common_arg_val<uint32_t>(0), TB16);
    const auto i_acc = TensorAccessor(i_args, get_common_arg_val<uint32_t>(1), N * 4);
    const auto o_acc = TensorAccessor(o_args, get_common_arg_val<uint32_t>(2), TB16);
    const uint64_t zeros = get_noc_addr(MEM_ZEROS_BASE);
    cb_reserve_back(cb_idx, 1);
    const uint32_t ibuf = get_write_ptr(cb_idx);
    volatile tt_l1_ptr uint32_t* idx = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(ibuf);
    const uint64_t ipage = i_acc.get_noc_addr(0);
    cb_reserve_back(cb_out, DT);                                // one scratch block, reused after each write barrier
    const uint32_t ob = get_write_ptr(cb_out);
    for (uint32_t w = w0 + SLOT; w < w1; w += 2) {
        // the block's 32 indices (128 B at offset 128 w; the last block may run past N: those rows are zero and the
        // bytes read past the page are ignored)
        const uint32_t nvalid = (w + 1) * 32 <= N ? 32 : N - w * 32;
        noc_async_read(ipage + w * 128, ibuf, 128);
        noc_async_read_barrier();
        for (uint32_t r = 0; r < 32; ++r) {
            const uint32_t roff = (((r >> 4) << 1) << 9) + ((r & 15) << 5);
            const uint32_t q = r < nvalid ? idx[r] : R;
            if (q >= R) {
                for (uint32_t t = 0; t < DT; ++t) {
                    noc_async_read(zeros, ob + t * TB16 + roff, 32);
                    noc_async_read(zeros, ob + t * TB16 + roff + 512, 32);
                }
                continue;
            }
            const uint32_t sr = q & 31;
            const uint32_t soff = (((sr >> 4) << 1) << 9) + ((sr & 15) << 5);
            const uint32_t st = (q >> 5) * DT;
            for (uint32_t t = 0; t < DT; ++t) {
                const uint64_t s = t_acc.get_noc_addr(st + t, soff);
                noc_async_read(s, ob + t * TB16 + roff, 32);
                noc_async_read(s + 512, ob + t * TB16 + roff + 512, 32);
            }
        }
        noc_async_read_barrier();
        for (uint32_t t = 0; t < DT; ++t) {
            noc_async_write(ob + t * TB16, o_acc.get_noc_addr(w * DT + t), TB16);
        }
        noc_async_write_barrier();
    }
}