code/tt_sam3/tt/kernels/gelu_bb_compute.cpp

// SPDX-License-Identifier: Apache-2.0
// sam3 backbone fc1 GELU (erf form) as a custom SFPU kernel, BLK tiles per DST acquire.
//   GELU_VARIANT 0: ttnn's own accurate bf16 path (calculate_gelu<false, false>) -- bit-identical to ttnn.gelu
//   GELU_VARIANT 1: ttnn's core polynomial and identity region verbatim (same coefficients / Horner order), negative
//                   tail (-5.5426, -3.125) by ttnn's exp code (GELU_TAIL 0, bit-identical) or a degree-9 fit (GELU_TAIL 1)
//   GELU_VARIANT 3: branch-free clamped odd polynomial (fast, NOT bit-identical; opt-in SAM3_BB_GELU=poly)
// Measured on the real fc1 outputs [5184, 4736] (perf/dev/gelu_kernel_bench.py, 12x10, DRAM in / out, BLK 8):
//   ttnn.gelu exact 0.433 ms, variant 0 0.437 ms (SFPU-bound), variant 1/tail 1 0.408 ms, variant 3 0.238 ms
//   (DRAM-bound: a plain copy kernel takes 0.229 ms).
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/gelu.h"
#include "api/compute/compute_kernel_api.h"

#ifndef GELU_TAIL
#define GELU_TAIL 0
#endif
#ifdef TRISC_MATH
#include "sfpi.h"
namespace ckernel::sfpu {
#if GELU_VARIANT == 1
sfpi_inline sfpi::vFloat sam3_gelu_value(sfpi::vFloat x) {
    sfpi::vFloat x2 = x * x;
    sfpi::vFloat result = 0.0f;
    v_if(x > -5.54259443f) {
#if GELU_TAIL == 0
        {   // ttnn's exp-region code verbatim (Moroz exp_21f * corr_H, rounded to the 2^-25 grid)
            constexpr float NEG_HALF_ONE_LN2 = -0.72134752044f;
            sfpi::vFloat xlog2 = x2 * NEG_HALF_ONE_LN2 + 127.0f;
            sfpi::vInt z = _float_to_int32_for_exp_21f_(xlog2);
            sfpi::vInt exponential_part = sfpi::exexp(sfpi::as<sfpi::vFloat>(z), sfpi::ExponentMode::Biased);
            sfpi::vMag fractional_part = sfpi::exman(sfpi::as<sfpi::vFloat>(z));
            sfpi::vFloat frac = sfpi::convert<sfpi::vFloat>(fractional_part, sfpi::RoundMode::Nearest);
            frac = PolynomialEvaluator::eval(frac, 1.0017248f, 7.839635491371155e-08f, 4.791750143340323e-15f);
            sfpi::vFloat exp_val = sfpi::setexp(frac, exponential_part);
            sfpi::vFloat H = exp_val * PolynomialEvaluator::eval(x, GELU_HCORR_C0, GELU_HCORR_C1, GELU_HCORR_C2, GELU_HCORR_C3);
            constexpr float ROUND_TO_GRID = 0.375f;
            sfpi::vFloat Hs = (H + ROUND_TO_GRID) - ROUND_TO_GRID;
            result = x * Hs;
        }
#elif GELU_TAIL == 1
        {   // GELU(x) directly as a degree-9 polynomial in t = (x + c) / h over the tail (rel err ~2e-4)
            sfpi::vFloat t = x * GELU_T_SCALE + GELU_T_OFF;
            result = PolynomialEvaluator::eval(t, GELU_TP0, GELU_TP1, GELU_TP2, GELU_TP3, GELU_TP4, GELU_TP5, GELU_TP6, GELU_TP7, GELU_TP8, GELU_TP9);
        }
#endif
        v_and(x >= -3.125f);
        sfpi::vFloat odd_poly = PolynomialEvaluator::eval(
            x2, GELU_CDF_CORE_C1, GELU_CDF_CORE_C3, GELU_CDF_CORE_C5, GELU_CDF_CORE_C7, GELU_CDF_CORE_C9,
            GELU_CDF_CORE_C11, GELU_CDF_CORE_C13);
        sfpi::vFloat phi = GELU_CDF_CORE_C0 + x * odd_poly;
        result = x * phi;
        v_and(x >= 2.78125f);
        result = x;
    }
    v_endif;
    return result;
}
inline void sam3_gelu_face() {
#pragma GCC unroll 8
    for (int d = 0; d < 8; d++) {
        sfpi::vFloat in = sfpi::dst_reg[0];
        sfpi::vFloat r = sam3_gelu_value(in);
        sfpi::dst_reg[0] = sfpi::convert<sfpi::vFloat16b>(r, sfpi::RoundMode::Nearest);
        sfpi::dst_reg++;
    }
}
#endif
#if GELU_VARIANT == 3
// cheap branch-free form (moge-2 / rf-detr style): t = clamp(x, -4, 4), Phi = 0.5 + t P(t^2) (degree 7 in t^2,
// Phi(4) = 1 pinned), GELU = t Phi(t) + max(x - 4, 0); max |err| vs erf-GELU 1.35e-4 before the bf16 rounding
inline void sam3_gelu_face() {
#pragma GCC unroll 8
    for (int d = 0; d < 8; d++) {
        sfpi::vFloat x = sfpi::dst_reg[0];
        sfpi::vFloat t = sfpi::clamp(x, -4.0f, 4.0f);
        sfpi::vFloat u = t * t;
        sfpi::vFloat p = PolynomialEvaluator::eval(u, 3.9880555868e-01f, -6.6064164042e-02f, 9.5814364031e-03f, -1.0210596956e-03f,
                                                   7.6192176493e-05f, -3.7097588574e-06f, 1.0444770737e-07f, -1.2780509984e-09f);
        sfpi::vFloat phi = t * p + 0.5f;
        sfpi::vFloat r = t * phi + sfpi::max(x - 4.0f, 0.0f);
        sfpi::dst_reg[0] = sfpi::convert<sfpi::vFloat16b>(r, sfpi::RoundMode::Nearest);
        sfpi::dst_reg++;
    }
}
#endif
}  // namespace ckernel::sfpu
#endif

constexpr uint32_t cb_in = get_compile_time_arg_val(0);
constexpr uint32_t cb_out = get_compile_time_arg_val(1);
constexpr uint32_t BLK = get_compile_time_arg_val(2);

void kernel_main() {
    const uint32_t n_tiles = get_arg_val<uint32_t>(0);
    init_sfpu(cb_in, cb_out);
#if GELU_VARIANT == 0
    gelu_tile_init<false>();
#else
    MATH((::ckernel::llk_math_eltwise_unary_sfpu_init<::SfpuType::gelu>()));
#endif
    for (uint32_t t = 0; t < n_tiles; t += BLK) {
        const uint32_t b = (n_tiles - t) < BLK ? (n_tiles - t) : BLK;
        cb_wait_front(cb_in, b);
        tile_regs_acquire();
        for (uint32_t j = 0; j < b; ++j) {
            copy_tile(cb_in, j, j);
#if GELU_VARIANT == 0
            gelu_tile<false>(j);
#else  // variants 1 / 3
            MATH(SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, sam3_gelu_face, j, VectorMode::RC));
#endif
        }
        tile_regs_commit();
        tile_regs_wait();
        cb_reserve_back(cb_out, b);
        for (uint32_t j = 0; j < b; ++j) {
            pack_tile(j, cb_out);
        }
        cb_push_back(cb_out, b);
        cb_pop_front(cb_in, b);
        tile_regs_release();
    }
}