code/tt_frnet/tt/kernels/p2p_reader.cpp

// SPDX-License-Identifier: Apache-2.0
// P2P (OPT round 4, FRNET_P2P): pixel2point gather straight into the KCAT operand, reader (RISCV_0).
// For this core's contiguous range of 32-row blocks: the block's 32 uint32 pixel indices (from the ROW_MAJOR [1, N]
// index tensor, one page of N * 4 bytes), then the 32 bf16 table rows (ROW_MAJOR [V, C], one page of C * 2 bytes per
// pixel) gathered into one row-major block of CB 0 (Kt pages of 2048 bytes = 32 rows x C bf16), which the compute
// kernel tilizes. BB blocks are read per barrier (CB 0 holds 2 BB blocks); the indices are read IB blocks at a time.
// CT args: [0] per-core RT-arg count (P3), [1] Kt (C / 32), [2] row bytes (C * 2), [3] BB, [4] IB, [5] index page
// bytes (N * 4), [6] DIAG (0; 1 = timing diagnostic without table reads), then the TensorAccessorArgs of the table and of the index.
// Common RT args: [table_addr, index_addr]. Per-core RT args: [b0, nb] (first block, block count).
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t Kt = get_compile_time_arg_val(1);
constexpr uint32_t RB = get_compile_time_arg_val(2);
constexpr uint32_t BB = get_compile_time_arg_val(3);
constexpr uint32_t IB = get_compile_time_arg_val(4);
constexpr uint32_t IPB = get_compile_time_arg_val(5);
constexpr uint32_t DIAG = get_compile_time_arg_val(6);   // timing diagnostic: 1 = no table reads (garbage output)
constexpr auto t_args = TensorAccessorArgs<7>();
constexpr auto i_args = TensorAccessorArgs<t_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_rm = 0, cb_idx = 1;
constexpr uint32_t BLOCK_BYTES = 32 * RB;

void kernel_main() {
    const uint32_t b0 = get_arg_val<uint32_t>(0);
    const uint32_t nb = get_arg_val<uint32_t>(1);
    if (nb == 0) {
        return;
    }
    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);
    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 uint32_t end = b0 + nb;
    uint32_t ib_end = b0;   // indices of blocks [ib_start, ib_end) are in ibuf
    uint32_t ib_start = b0;
    for (uint32_t b = b0; b < end; b += BB) {
        const uint32_t k = (end - b) < BB ? (end - b) : BB;
        if (b + k > ib_end) {   // refill: IB blocks of indices from block b (BB divides IB: never straddles)
            ib_start = b;
            const uint32_t m = (end - b) < IB ? (end - b) : IB;
            noc_async_read(index.get_noc_addr(0, b * 128), ibuf, m * 128);
            noc_async_read_barrier();
            ib_end = b + m;
        }
        cb_reserve_back(cb_rm, k * Kt);
        uint32_t wp = get_write_ptr(cb_rm);
        const uint32_t base = (b - ib_start) * 32;
        if constexpr (DIAG != 1) {
            for (uint32_t j = 0; j < 32 * k; ++j) {
                noc_async_read(table.get_noc_addr(idx[base + j]), wp, RB);
                wp += RB;
            }
        }
        noc_async_read_barrier();
        cb_push_back(cb_rm, k * Kt);
    }
}