code/tt_frnet/tt/kernels/mkh_reader.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, reader
// (RISCV_0). Once: every layer's K-concatenated two-term weight W' (bf16 TILE, K' x Nt tiles, row-major tile order)
// into CB 2 (layer after layer, never popped) and every bias row (fp32 TILE, Nt tiles) into CB 3. Then per tile row
// r of this core's contiguous range:
//   - the 32 pixel indices of the rows (ROW_MAJOR uint32 [1, N], one page) -- read IB blocks at a time -- and the
//     32 bf16 rows of the gather table (ROW_MAJOR [V, C], one C * 2 byte page per pixel) into CB 0 (KT0 pages =
//     one row-major 32-row block; the compute tilizes it): the pixel2point(fuse2) gather;
//   - for each layer with a skip, its Nt fp32 tiles of row r (q, pf2, pf1, pf0) into CB 4, in layer order.
// CT args: [0] per-core RT-arg count (P3), [1] NL, [2] KT0 (C / 32), [3] index page bytes (N * 4), [4] IB,
// [5 + 3 l ..] per layer l < MAXL: K' tiles, Nt, SKIP; then the TensorAccessorArgs of the table, of the index, and one
// shared by all W', one by all biases, one by all skips (DRAM interleaved; tt/mk_head.py asserts it).
// Common RT args: [table, index, W_0 .. W_4, B_0 .. B_4, S_0 .. S_4] (unused slots 0).
// Per-core RT args: [r0, nr] (first tile row, tile-row count).
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t MAXL = 5;
constexpr uint32_t NL = get_compile_time_arg_val(1);
constexpr uint32_t KT0 = get_compile_time_arg_val(2);
constexpr uint32_t IPB = get_compile_time_arg_val(3);
constexpr uint32_t IB = get_compile_time_arg_val(4);
constexpr uint32_t KP[MAXL] = {get_compile_time_arg_val(5), get_compile_time_arg_val(8), get_compile_time_arg_val(11),
                               get_compile_time_arg_val(14), get_compile_time_arg_val(17)};
constexpr uint32_t NT[MAXL] = {get_compile_time_arg_val(6), get_compile_time_arg_val(9), get_compile_time_arg_val(12),
                               get_compile_time_arg_val(15), get_compile_time_arg_val(18)};
constexpr uint32_t SK[MAXL] = {get_compile_time_arg_val(7), get_compile_time_arg_val(10), get_compile_time_arg_val(13),
                               get_compile_time_arg_val(16), get_compile_time_arg_val(19)};
constexpr uint32_t ST = SK[0] * NT[0] + SK[1] * NT[1] + SK[2] * NT[2] + SK[3] * NT[3] + SK[4] * NT[4];
constexpr auto t_args = TensorAccessorArgs<20>();
constexpr auto i_args = TensorAccessorArgs<t_args.next_compile_time_args_offset()>();
constexpr auto w_args = TensorAccessorArgs<i_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();
constexpr auto s_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_rm = 0, cb_idx = 1, cb_w = 2, cb_b = 3, cb_s = 4;
constexpr uint32_t RB = KT0 * 64;   // table row bytes (C bf16)
constexpr uint32_t WTB = 2048, 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);
    if (nr == 0) {
        return;
    }
    // weights and biases, once
    for (uint32_t l = 0; l < NL; ++l) {
        const uint32_t nw = KP[l] * NT[l];
        const auto w = TensorAccessor(w_args, get_common_arg_val<uint32_t>(2 + l), WTB);
        cb_reserve_back(cb_w, nw);
        uint32_t wp = get_write_ptr(cb_w);
        for (uint32_t t = 0; t < nw; ++t) {
            noc_async_read(w.get_noc_addr(t), wp, WTB);
            wp += WTB;
        }
        const auto b = TensorAccessor(b_args, get_common_arg_val<uint32_t>(2 + MAXL + l), FTB);
        cb_reserve_back(cb_b, NT[l]);
        uint32_t bp = get_write_ptr(cb_b);
        for (uint32_t t = 0; t < NT[l]; ++t) {
            noc_async_read(b.get_noc_addr(t), bp, FTB);
            bp += FTB;
        }
        noc_async_read_barrier();
        cb_push_back(cb_w, nw);
        cb_push_back(cb_b, NT[l]);
    }
    const auto table = TensorAccessor(t_args, get_common_arg_val<uint32_t>(0), RB);
    const auto index = TensorAccessor(i_args, get_common_arg_val<uint32_t>(1), IPB);
    cb_reserve_back(cb_idx, 1);
    volatile tt_l1_ptr uint32_t* idx = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_idx));
    const uint32_t ibuf = get_write_ptr(cb_idx);
    const uint32_t end = r0 + nr;
    uint32_t ib_start = r0, ib_end = r0;
    for (uint32_t r = r0; r < end; ++r) {
        if (r >= ib_end) {
            ib_start = r;
            const uint32_t m = (end - r) < IB ? (end - r) : IB;
            noc_async_read(index.get_noc_addr(0, r * 128), ibuf, m * 128);
            noc_async_read_barrier();
            ib_end = r + m;
        }
        cb_reserve_back(cb_rm, KT0);
        uint32_t wp = get_write_ptr(cb_rm);
        const uint32_t base = (r - ib_start) * 32;
        for (uint32_t j = 0; j < 32; ++j) {
            noc_async_read(table.get_noc_addr(idx[base + j]), wp, RB);
            wp += RB;
        }
        if constexpr (ST > 0) {
            cb_reserve_back(cb_s, ST);   // the row's skip tiles of every layer (CB 4 holds 2 ST: no wrap inside)
            uint32_t sp = get_write_ptr(cb_s);
            for (uint32_t l = 0; l < NL; ++l) {
                if (SK[l] == 0) {
                    continue;
                }
                const auto s = TensorAccessor(s_args, get_common_arg_val<uint32_t>(2 + 2 * MAXL + l), FTB);
                for (uint32_t j = 0; j < NT[l]; ++j) {
                    noc_async_read(s.get_noc_addr(r * NT[l] + j), sp, FTB);
                    sp += FTB;
                }
            }
        }
        noc_async_read_barrier();
        cb_push_back(cb_rm, KT0);
        if constexpr (ST > 0) {
            cb_push_back(cb_s, ST);
        }
    }
}