code/tt_frnet/tt/kernels/mkh_compute.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, compute.
// Per tile row (32 points) of this core's range, every intermediate stays in L1 / DEST:
//   g = tilize(gathered fuse2 rows)                                    (tilize_block: exact bf16 datacopy, as P2P)
//   layer 0:  y = act([g | g] @ [W_hi; W_lo] + b)                      (bf16 input: x_lo = 0)
//   layer l:  y = act([x_hi | x_hi | x_lo] @ [W_hi; W_lo; W_hi] + b)   x = y_{l-1} + skip_{l-1}
// with exactly the arithmetic of the stock path it replaces (KCAT operand kernel + ttnn.linear 1D program, OPT rounds
// 1-4): the K-concatenated matmul accumulated over K' in fp32 DEST in K' order (matmul_block, HiFi4; the stock
// matmul's cross-block reload is exact, UnpackToDestFp32), packed fp32, the bias added by the FPU row broadcast of
// the packed partials (add_tiles_bcast_rows: the stock fused-bias epilogue), ReLU, then in DEST: + skip (copy_tile
// with UnpackToDestFp32 + add_binary_tile<NearestEven>, the HEAD-SKIP add) and the two-term split of the sum
// (typecast_tile RNE, sub_binary_tile, typecast_tile: kcat_compute.cpp), packed bf16 as the next layer's x_hi / x_lo.
// The last layer (logits) is packed fp32 to CB 20 for the writer.
// CT args: [0] per-core RT-arg count (P3), [1] NL, [2] KT0, [3] (unused), [4] (unused), [5 + 3 l ..] per layer:
// K' tiles, Nt, SKIP; [20 + l] RELU per layer; [25] SB (output tiles per DEST batch of the matmul, <= 4).
// Per-core RT args: [nr, 0].
#include <cstdint>

#include "api/compute/bcast.h"
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/copy_dest_values.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/relu.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/matmul.h"
#include "api/compute/pack.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/tilize.h"

using namespace ckernel;

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 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 RL[MAXL] = {get_compile_time_arg_val(20), get_compile_time_arg_val(21), get_compile_time_arg_val(22),
                               get_compile_time_arg_val(23), get_compile_time_arg_val(24)};
constexpr uint32_t SB = get_compile_time_arg_val(25);
constexpr uint32_t EB = get_compile_time_arg_val(26);   // epilogue tiles per DEST acquire (2 slots each; fp32 DEST: <= 2)

constexpr uint32_t cb_rm = 0, cb_w = 2, cb_b = 3, cb_s = 4, cb_g = 5;
constexpr uint32_t cb_hi[2] = {16, 18}, cb_lo[2] = {17, 19};
constexpr uint32_t cb_out = 20, cb_i = 24;

void kernel_main() {
    const uint32_t nr = get_arg_val<uint32_t>(0);
    if (nr == 0) {
        return;
    }
    uint32_t woff[MAXL], boff[MAXL];
    {
        uint32_t w = 0, b = 0;
        for (uint32_t l = 0; l < NL; ++l) {
            woff[l] = w;
            boff[l] = b;
            w += KP[l] * NT[l];
            b += NT[l];
        }
        compute_kernel_hw_startup<SrcOrder::Reverse>(cb_g, cb_w, cb_i);
        cb_wait_front(cb_w, w);
        cb_wait_front(cb_b, b);
    }
    for (uint32_t r = 0; r < nr; ++r) {
        // pixel2point(fuse2): the gathered row-major block -> KT0 bf16 tiles
        reconfig_data_format_srca(cb_rm);
        pack_reconfig_data_format(cb_g);
        tilize_init(cb_rm, KT0, cb_g);
        cb_wait_front(cb_rm, KT0);
        cb_reserve_back(cb_g, KT0);
        tilize_block(cb_rm, KT0, cb_g);
        cb_push_back(cb_g, KT0);
        cb_pop_front(cb_rm, KT0);
        tilize_uninit(cb_rm, cb_g);

        for (uint32_t l = 0; l < NL; ++l) {
            const uint32_t nt = NT[l], kp = KP[l];
            const uint32_t a = l == 0 ? cb_g : cb_hi[(l - 1) & 1];
            const uint32_t lo = l == 0 ? cb_g : cb_lo[(l - 1) & 1];
            const uint32_t kt = l == 0 ? kp / 2 : kp / 3;
            const bool last = l + 1 == NL;
            const uint32_t sb = nt < SB ? nt : SB;
            // y = X' @ W' (fp32 DEST over all K' tiles) -> fp32 partials in CB 24
            cb_wait_front(a, kt);
            if (l > 0) {
                cb_wait_front(lo, kt);
            }
            reconfig_data_format<SrcOrder::Reverse>(a, cb_w);
            pack_reconfig_data_format(cb_i);
            matmul_block_init(a, cb_w, 0, sb, 1, kp);
            for (uint32_t n0 = 0; n0 < nt; n0 += sb) {
                tile_regs_acquire();
                for (uint32_t k = 0; k < kp; ++k) {
                    uint32_t cb = a, it = k;
                    if (k >= 2 * kt) {
                        cb = lo;
                        it = k - 2 * kt;
                    } else if (k >= kt) {
                        it = k - kt;
                    }
                    matmul_block(cb, cb_w, it, woff[l] + k * nt + n0, 0, 0, sb, 1, kp);
                }
                tile_regs_commit();
                cb_reserve_back(cb_i, sb);
                tile_regs_wait();
                for (uint32_t i = 0; i < sb; ++i) {
                    pack_tile(i, cb_i);
                }
                tile_regs_release();
                cb_push_back(cb_i, sb);
            }
            cb_pop_front(a, kt);
            if (l > 0) {
                cb_pop_front(lo, kt);
            }
            // epilogue per output tile: + bias (FPU row broadcast), ReLU, then + skip and the split, or fp32 out
            const uint32_t ohi = cb_hi[l & 1], olo = cb_lo[l & 1];
            if (SK[l]) {
                cb_wait_front(cb_s, nt);
            }
            for (uint32_t n = 0; n < nt; n += EB) {
                const uint32_t ge = (nt - n) < EB ? (nt - n) : EB;   // tiles of this DEST group (2 slots each)
                cb_wait_front(cb_i, ge);
                reconfig_data_format(cb_i, cb_b);
                add_bcast_rows_init(cb_i, cb_b);
                tile_regs_acquire();
                for (uint32_t e = 0; e < ge; ++e) {
                    add_tiles_bcast_rows(cb_i, cb_b, e, boff[l] + n + e, 2 * e);
                }
                if (RL[l]) {
                    relu_tile_init();
                    for (uint32_t e = 0; e < ge; ++e) {
                        relu_tile(2 * e);
                    }
                }
                if (last) {
                    tile_regs_commit();
                    cb_reserve_back(cb_out, ge);
                    pack_reconfig_data_format(cb_out);
                    tile_regs_wait();
                    for (uint32_t e = 0; e < ge; ++e) {
                        pack_tile(2 * e, cb_out);
                    }
                    tile_regs_release();
                    cb_push_back(cb_out, ge);
                } else {
                    if (SK[l]) {
                        reconfig_data_format_srca(cb_s);
                        copy_init(cb_s);
                        for (uint32_t e = 0; e < ge; ++e) {
                            copy_tile(cb_s, n + e, 2 * e + 1);
                        }
                        add_binary_tile_init();
                        for (uint32_t e = 0; e < ge; ++e) {
                            add_binary_tile<ckernel::DstRoundingMode::NearestEven>(2 * e, 2 * e + 1, 2 * e);
                        }
                    }
                    copy_dest_values_init();
                    for (uint32_t e = 0; e < ge; ++e) {
                        copy_dest_values<DataFormat::Float32>(2 * e, 2 * e + 1);
                    }
                    typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
                    for (uint32_t e = 0; e < ge; ++e) {
                        typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(2 * e);
                    }
                    sub_binary_tile_init();
                    for (uint32_t e = 0; e < ge; ++e) {
                        sub_binary_tile<ckernel::DstRoundingMode::NearestEven>(2 * e + 1, 2 * e, 2 * e + 1);
                    }
                    typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
                    for (uint32_t e = 0; e < ge; ++e) {
                        typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(2 * e + 1);
                    }
                    tile_regs_commit();
                    cb_reserve_back(ohi, ge);
                    cb_reserve_back(olo, ge);
                    pack_reconfig_data_format(ohi);
                    tile_regs_wait();
                    for (uint32_t e = 0; e < ge; ++e) {
                        pack_tile(2 * e, ohi);
                        pack_tile(2 * e + 1, olo);
                    }
                    tile_regs_release();
                    cb_push_back(ohi, ge);
                    cb_push_back(olo, ge);
                }
                cb_pop_front(cb_i, ge);
            }
            if (SK[l]) {
                cb_pop_front(cb_s, nt);
            }
        }
    }
}