code/tt_frnet/tt/kernels/kcat_compute.cpp

// SPDX-License-Identifier: Apache-2.0
// KCAT operand build (tt/kcat.py), compute: per fp32 x tile, x_hi = bf16(x) (typecast_tile Float32 -> Float16_b,
// the LLK of the stock ttnn.typecast), x_lo = bf16(x - x_hi) (sub_binary_tile in fp32: exact; then the same
// typecast), both packed as bf16 (exact: the DEST values are bf16-representable). The same values as the stock
// typecast -> typecast -> subtract -> typecast chain of PointLinear._split (bit-identity checked on the device).
// Mode 2 (OPT round 3, HEAD-SKIP): x = x + skip first (add_binary_tile<NearestEven> on the fp32 tiles of CBs 0 and
// 1 in DEST: binary_ng's fp32 SFPU add, the LLK of the stock ttnn.add(y, skip) it replaces; as tt/epi.py), then the
// same split. CT args: [0] per-core RT-arg count (P3), [1] mode (0 = fp32 split; 1 = bf16 input: nothing to do here,
// the writer reads CB 0 itself; 2 = fp32 x + fp32 skip, then the split). Per-core RT args: [n, 0].
#include <cstdint>

#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/typecast.h"
#include "api/compute/tile_move_copy.h"

using namespace ckernel;

constexpr uint32_t mode = get_compile_time_arg_val(1);
constexpr uint32_t cb_x = 0, cb_s = 1, cb_hi = 16, cb_lo = 17;

void kernel_main() {
    if constexpr (mode == 1) {
        return;
    } else {
        const uint32_t n = get_arg_val<uint32_t>(0);
        if (n == 0) {
            return;
        }
        compute_kernel_hw_startup(cb_x, cb_hi);
        for (uint32_t t = 0; t < n; ++t) {
            cb_wait_front(cb_x, 1);
            tile_regs_acquire();
            copy_init(cb_x);
            copy_tile(cb_x, 0, 0);
            if constexpr (mode == 2) {
                cb_wait_front(cb_s, 1);
                copy_init(cb_s);
                copy_tile(cb_s, 0, 1);
                add_binary_tile_init();
                add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
                copy_dest_values_init();
                copy_dest_values<DataFormat::Float32>(0, 1);
            } else {
                copy_tile(cb_x, 0, 1);
            }
            typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
            typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(0);
            sub_binary_tile_init();
            sub_binary_tile<ckernel::DstRoundingMode::NearestEven>(1, 0, 1);
            typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
            typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(1);
            cb_reserve_back(cb_hi, 1);
            cb_reserve_back(cb_lo, 1);
            tile_regs_commit();
            tile_regs_wait();
            pack_tile(0, cb_hi);
            pack_tile(1, cb_lo);
            tile_regs_release();
            cb_push_back(cb_hi, 1);
            cb_push_back(cb_lo, 1);
            cb_pop_front(cb_x, 1);
            if constexpr (mode == 2) {
                cb_pop_front(cb_s, 1);
            }
        }
    }
}