code/tt_frnet/tt/kernels/epi_compute.cpp

// SPDX-License-Identifier: Apache-2.0
// EPI compute (frnet-p150 OPT round 2, tt/epi.py): one image-branch epilogue per tile with the LLKs and the order of
// the stock programs it replaces, so both outputs are bit-identical to them:
//   y = ((p0 + p1) + p2) + ...   add_binary_tile<NearestEven>: binary_ng's fp32 SFPU add (KSplitConv's partial sums
//                                in order; the BasicBlock's ttnn.add(y, identity)), operands unpacked to DEST as fp32
//   y = hardswish(y)             (ACT 1) ttnn.hardswish's fp32 chain (hardswish_kernel.cpp, kIsFloat32):
//                                D1 = y; hardsigmoid_tile(D0); mul_binary_tile(D0, D1, D0)
//   out32 = y                    fp32 (OUT32)
//   out16 = bf16(y)              typecast_tile<Float32, Float16_b> (ttnn.typecast; OUT16): the next conv's operand
// DEST tiles: 0 = y, 1 = operand / copy. CT args: [0] T, [1] NC, [2] GX, [3] NP, [4] ACT, [5] OUT32, [6] OUT16.
#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/activations.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"

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 ACT = get_compile_time_arg_val(4);
constexpr uint32_t OUT32 = get_compile_time_arg_val(5);
constexpr uint32_t OUT16 = get_compile_time_arg_val(6);
constexpr uint32_t cb32 = 16, cb16 = 17;

FORCE_INLINE void add_part(uint32_t cb) {
    copy_init(cb);
    copy_tile(cb, 0, 1);
    add_binary_tile_init();
    add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
}

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;
    }
    compute_kernel_hw_startup(0, OUT32 ? cb32 : cb16);
    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 constexpr (NP > 3) {
            cb_wait_front(3, 1);
        }
        if constexpr (NP > 4) {
            cb_wait_front(4, 1);
        }
        tile_regs_acquire();
        copy_init(0);
        copy_tile(0, 0, 0);
        if constexpr (NP > 1) {
            add_part(1);
        }
        if constexpr (NP > 2) {
            add_part(2);
        }
        if constexpr (NP > 3) {
            add_part(3);
        }
        if constexpr (NP > 4) {
            add_part(4);
        }
        if constexpr (ACT == 1) {
            copy_dest_values_init();
            copy_dest_values<DataFormat::Float32>(0, 1);
            hardsigmoid_tile_init();
            hardsigmoid_tile(0);
            mul_binary_tile_init();
            mul_binary_tile(0, 1, 0);
        }
        if constexpr (OUT16) {
            copy_dest_values_init();
            copy_dest_values<DataFormat::Float32>(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);
        }
        tile_regs_commit();
        if constexpr (OUT32) {
            cb_reserve_back(cb32, 1);
        }
        if constexpr (OUT16) {
            cb_reserve_back(cb16, 1);
        }
        tile_regs_wait();
        if constexpr (OUT32) {
            if constexpr (OUT16) {
                pack_reconfig_data_format(cb32);
            }
            pack_tile(0, cb32);
        }
        if constexpr (OUT16) {
            if constexpr (OUT32) {
                pack_reconfig_data_format(cb16);
            }
            pack_tile(1, cb16);
        }
        tile_regs_release();
        if constexpr (OUT32) {
            cb_push_back(cb32, 1);
        }
        if constexpr (OUT16) {
            cb_push_back(cb16, 1);
        }
        cb_pop_front(0, 1);
        if constexpr (NP > 1) {
            cb_pop_front(1, 1);
        }
        if constexpr (NP > 2) {
            cb_pop_front(2, 1);
        }
        if constexpr (NP > 3) {
            cb_pop_front(3, 1);
        }
        if constexpr (NP > 4) {
            cb_pop_front(4, 1);
        }
    }
}