code/tt_frnet/tt/kernels/p2p_compute.cpp

// SPDX-License-Identifier: Apache-2.0
// P2P (OPT round 4, FRNET_P2P): pixel2point gather straight into the KCAT operand, compute: each gathered 32-row
// block of CB 0 (row-major, Kt pages) tilized into Kt bf16 tiles of CB 16 (tilize_block: an exact bf16 datacopy
// through DEST, the LLK of the embedding op's own tilized output).
// CT args: [0] per-core RT-arg count (P3), [1] Kt, [2] TZ (0 = tilize_block, 1 = fast_tilize_block, 9 = no tilize:
// timing diagnostic only, the output is garbage). Per-core RT args: [nb, 0].
#include <cstdint>

#include "api/compute/common.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/tilize.h"

constexpr uint32_t Kt = get_compile_time_arg_val(1);
constexpr uint32_t TZ = get_compile_time_arg_val(2);
constexpr uint32_t cb_rm = 0, cb_t = 16;

void kernel_main() {
    const uint32_t nb = get_arg_val<uint32_t>(0);
    if (nb == 0) {
        return;
    }
    compute_kernel_hw_startup(cb_rm, cb_t);
    if constexpr (TZ == 1) {
        fast_tilize_init(cb_rm, Kt, cb_t);
    } else if constexpr (TZ == 0) {
        tilize_init(cb_rm, Kt, cb_t);
    }
    for (uint32_t b = 0; b < nb; ++b) {
        cb_wait_front(cb_rm, Kt);
        cb_reserve_back(cb_t, Kt);
        if constexpr (TZ == 1) {
            fast_tilize_block(cb_rm, Kt, cb_t);
        } else if constexpr (TZ == 0) {
            tilize_block(cb_rm, Kt, cb_t);
        }
        cb_push_back(cb_t, Kt);
        cb_pop_front(cb_rm, Kt);
    }
    if constexpr (TZ == 1) {
        fast_tilize_uninit(cb_rm, cb_t, Kt);
    } else if constexpr (TZ == 0) {
        tilize_uninit(cb_rm, cb_t);
    }
}