code/tt_meteor/tt/kernels/terms_epi_compute.cpp

// SPDX-License-Identifier: Apache-2.0
// TERMS_EPI compute (meteor-p150 OPT_PLAN #2, tt/epi.py): the TermsConv epilogue of tt/layers.py in one pass per
// tile, with the LLKs and the order of the stock programs it replaces, so the output is bit-identical (the SceneSeg
// EPI pattern, sceneseg-p150 OPT_REPORT R1.3 / R2.3):
//   y = (p0 + p1) + p2     add_binary_tile<NearestEven>  (binary_ng's fp32 SFPU add; operands unpacked to DEST as
//                                                          fp32: UnpackToDestFp32 on CBs 0..4)
//   y = y + res            add_binary_tile<NearestEven>  (EPI_RES: ttaw residual_add(z, idt, "relu"))
//   y = relu(y)            relu_tile (EPI_RELU: binary_ng's fused post-activation of the last add, unary_op_utils)
//   hi = bf16(y), lo = bf16(y - fp32(hi))   (EPI_SPLIT: tt.layers.split_terms = typecast, typecast back (exact),
//                          fp32 SFPU subtract, typecast), packed to CBs 17 / 18 as bf16
// DEST tiles: 0 = y, 1 = operand / hi, 2 = lo. EPI_MAIN 0 skips the fp32 main output (split only).
// NP = 1 with no residual / relu is the plain split of an fp32 map (split_terms as one program).
// CT args: [0] T, [1] NC, [2] GX, [3] NP. Defines: EPI_RES, EPI_RELU, EPI_SPLIT, EPI_MAIN.
#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/relu.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"

#ifndef EPI_RES
#define EPI_RES 0
#endif
#ifndef EPI_RELU
#define EPI_RELU 0
#endif
#ifndef EPI_SPLIT
#define EPI_SPLIT 0
#endif
#ifndef EPI_MAIN
#define EPI_MAIN 1
#endif

using namespace ckernel;

constexpr uint32_t T = get_compile_time_arg_val(0);
constexpr uint32_t NC = get_compile_time_arg_val(1);
constexpr uint32_t GX = get_compile_time_arg_val(2);
constexpr uint32_t NP = get_compile_time_arg_val(3);
constexpr uint32_t cb_res = 4, cb_out = 16, cb_hi = 17, cb_lo = 18;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = T / NC, E = T % NC;
    const uint32_t t0 = core * B + (core < E ? core : E);
    const uint32_t n = B + (core < E ? 1 : 0);
    if (n == 0) {
        return;
    }
#if EPI_MAIN
    compute_kernel_hw_startup(0, cb_out);
#else
    compute_kernel_hw_startup(0, cb_hi);
#endif
    for (uint32_t t = t0; t < t0 + n; ++t) {
        cb_wait_front(0, 1);
        if constexpr (NP > 1) {
            cb_wait_front(1, 1);
        }
        if constexpr (NP > 2) {
            cb_wait_front(2, 1);
        }
#if EPI_RES
        cb_wait_front(cb_res, 1);
#endif
        tile_regs_acquire();
        copy_init(0);
        copy_tile(0, 0, 0);
        if constexpr (NP > 1) {
            copy_init(1);
            copy_tile(1, 0, 1);
            add_binary_tile_init();
            add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
        }
        if constexpr (NP > 2) {
            copy_init(2);
            copy_tile(2, 0, 1);
            add_binary_tile_init();
            add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
        }
#if EPI_RES
        copy_init(cb_res);
        copy_tile(cb_res, 0, 1);
        add_binary_tile_init();
        add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
#endif
#if EPI_RELU
        relu_tile_init();
        relu_tile(0);
#endif
#if EPI_SPLIT
        // hi (tile 1) = bf16(y); lo (tile 2) = bf16(y - hi); y stays in tile 0
        copy_dest_values_init();
        copy_dest_values<DataFormat::Float32>(0, 1);
        copy_dest_values<DataFormat::Float32>(0, 2);
        typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
        typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(1);
        sub_binary_tile_init();
        sub_binary_tile<ckernel::DstRoundingMode::NearestEven>(2, 1, 2);
        typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
        typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(2);
#endif
        tile_regs_commit();
#if EPI_MAIN
        cb_reserve_back(cb_out, 1);
#endif
#if EPI_SPLIT
        cb_reserve_back(cb_hi, 1);
        cb_reserve_back(cb_lo, 1);
#endif
        tile_regs_wait();
#if EPI_MAIN
#if EPI_SPLIT
        pack_reconfig_data_format(cb_out);
#endif
        pack_tile(0, cb_out);
#endif
#if EPI_SPLIT
        pack_reconfig_data_format(cb_hi);
        pack_tile(1, cb_hi);
        pack_tile(2, cb_lo);
#endif
        tile_regs_release();
#if EPI_MAIN
        cb_push_back(cb_out, 1);
#endif
#if EPI_SPLIT
        cb_push_back(cb_hi, 1);
        cb_push_back(cb_lo, 1);
#endif
        cb_pop_front(0, 1);
        if constexpr (NP > 1) {
            cb_pop_front(1, 1);
        }
        if constexpr (NP > 2) {
            cb_pop_front(2, 1);
        }
#if EPI_RES
        cb_pop_front(cb_res, 1);
#endif
    }
}