code/tt_moge/kernels/gelu_poly_compute.cpp

// SPDX-License-Identifier: Apache-2.0
// GELU(x) = x * Phi(x) as a branch-free polynomial (moge-2 opt kernel, used through ttnn.generic_op).
//   t = clamp(x, -4, 4);  u = t^2 / 16;  Phi(t) = 0.5 + t * P(u)  (P: degree 7, weighted minimax fit of the
//   normal CDF on [0, 4] with Phi(4) = 1 pinned);  GELU = x * Phi(t)  (x > 4: x; x < -4: |y| < 1.3e-7 |x|)
// The four highest-order coefficients are fp16 (one SFPLOADI each, refit after quantization), the three
// lowest fp32 in the programmable constant registers. 27 SFPU instructions per row (v1: 32).
// Max |error| vs the exact (erf) GELU 2.1e-4 (fp32 evaluation, before the bf16 round-to-nearest).
// Reads TILE bf16 pages from c_0, writes c_2. Runtime arg 0: number of tiles for this core.
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/dataflow/dataflow_buffer.h"
#if defined(TRISC_MATH)
#include "sfpi.h"
#include "llk_math_eltwise_unary_sfpu_macros.h"
namespace ckernel::sfpu {
inline void moge_gelu_init() {
    sfpi::vConstFloatPrgm0 = 0.3983510434627533f;
    sfpi::vConstFloatPrgm1 = -1.0425554513931274f;
    sfpi::vConstFloatPrgm2 = 2.3193752765655518f;
}
inline void moge_gelu_poly() {
#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;
        u = u * 0.0625f;                                                    // exact power-of-2 scaling
        sfpi::vFloat p = u * sfpi::vFloat(sfpi::sFloat16a(0xafb6u)) + sfpi::vFloat(sfpi::sFloat16a(0x3ae0u));   // -0.1204833984375, 0.859375
        p = p * u + sfpi::vFloat(sfpi::sFloat16a(0xc0e2u));                                          // -2.44140625
        p = p * u + sfpi::vFloat(sfpi::sFloat16a(0x438fu));                                          // 3.779296875
        p = p * u + sfpi::vFloat(sfpi::sFloat16a(0xc341u));                                          // -3.626953125
        p = p * u + sfpi::vConstFloatPrgm2;
        p = p * u + sfpi::vConstFloatPrgm1;
        p = p * u + sfpi::vConstFloatPrgm0;
        sfpi::vFloat phi = t * p + 0.5f;
        sfpi::vFloat r = x * phi;
        sfpi::dst_reg[0] = sfpi::convert<sfpi::vFloat16b>(r, sfpi::RoundMode::Nearest);
        sfpi::dst_reg++;
    }
}
}  // namespace ckernel::sfpu
#endif

void kernel_main() {
    uint32_t num_tiles = get_arg_val<uint32_t>(0);
    constexpr uint32_t BLOCK = get_compile_time_arg_val(0);   // tiles per DST acquire (<= 8, bf16 dest)
    constexpr auto cb_input = tt::CBIndex::c_0;
    constexpr auto cb_output = tt::CBIndex::c_2;
    DataflowBuffer dfb_in(cb_input);
    DataflowBuffer dfb_out(cb_output);
    init_sfpu(cb_input, cb_output);
    MATH((::ckernel::llk_math_eltwise_unary_sfpu_init<::SfpuType::gelu>(::ckernel::sfpu::moge_gelu_init)));
    uint32_t done = 0;
    while (done < num_tiles) {
        const uint32_t n = (num_tiles - done) < BLOCK ? (num_tiles - done) : BLOCK;
        dfb_in.wait_front(n);
        dfb_out.reserve_back(n);
        tile_regs_acquire();
        for (uint32_t j = 0; j < n; ++j) {
            copy_tile(cb_input, j, j);
            MATH((_llk_math_eltwise_unary_sfpu_params_(::ckernel::sfpu::moge_gelu_poly, j, VectorMode::RC)));
        }
        tile_regs_commit();
        tile_regs_wait();
        for (uint32_t j = 0; j < n; ++j) {
            pack_tile(j, cb_output);
        }
        tile_regs_release();
        dfb_in.pop_front(n);
        dfb_out.push_back(n);
        done += n;
    }
}