// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
//
// GELU (erf form) for the ViT MLP, fp32 SFPU math:
//   gelu(x) = x * (0.5 + x * P(x^2))   for |x| <= 4   (P: degree-8 fit of (Phi(x) - 0.5) / x)
//           = x                         for x > 4      (|error| <= 1.3e-4)
//           = 0                         for x < -4     (|error| <= 1.3e-4)
// (implemented branch-free: u = min(x^2, 16), CDF = clamp(0.5 + x P(u), 0, 1))
// max |error| vs the exact erf GELU: 2e-5 inside [-4, 4], 1.3e-4 at the clamps (bf16 output ULP at 4 is
// 0.03).  End-to-end on the real checkpoint (host fp32 simulation, scripts/gelu_sens.py) it is
// indistinguishable from rounding the exact GELU to bf16 (PCC 0.99999993 min over 32 crops).
// compile-time: none.  runtime: n (tiles)
#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/compute_kernel_api.h"

#ifdef TRISC_MATH
#include "sfpi.h"
namespace ckernel::sfpu {
inline void hamer_gelu_face() {
    // two rows per iteration: independent Horner chains hide the MAD latency
    for (size_t i = 0; i < 8; i += 2) {
        sfpi::vFloat x0 = sfpi::dst_reg[i];
        sfpi::vFloat x1 = sfpi::dst_reg[i + 1];
        // |x| > 4: u saturates at 16 and 0.5 + x*P(16) leaves [0, 1] -> clamped CDF of 1 (g = x) or 0 (g = 0)
        sfpi::vFloat u0 = sfpi::min(x0 * x0, 16.0f);
        sfpi::vFloat u1 = sfpi::min(x1 * x1, 16.0f);
        sfpi::vFloat a0, a1;
        { const sfpi::vFloat c8 = 8.040229593e-11f; const sfpi::vFloat c7 = -6.988785106e-09f;
        a0 = u0 * c8 + c7;
        a1 = u1 * c8 + c7; }
        { const sfpi::vFloat c = 2.712364733e-07f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = -6.289864814e-06f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = sfpi::vConstFloatPrgm0;  // 9.886907355e-05f
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = sfpi::vConstFloatPrgm1;  // -1.133757411e-03f
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = sfpi::vConstFloatPrgm2;  // 9.877118282e-03f
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = -6.641025841e-02f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = 3.989226222e-01f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        sfpi::vFloat p0 = sfpi::clamp(x0 * a0 + 0.5f, 0.0f, 1.0f);
        sfpi::vFloat p1 = sfpi::clamp(x1 * a1 + 0.5f, 0.0f, 1.0f);
        sfpi::dst_reg[i] = x0 * p0;
        sfpi::dst_reg[i + 1] = x1 * p1;
    }
}
inline void hamer_gelu_init() {
    // three of the Horner coefficients live in the programmable constant registers (saves 6 SFPLOADI per row pair)
    sfpi::vConstFloatPrgm0 = 9.886907355e-05f;
    sfpi::vConstFloatPrgm1 = -1.133757411e-03f;
    sfpi::vConstFloatPrgm2 = 9.877118282e-03f;
}
}  // namespace ckernel::sfpu
#endif

inline void hamer_gelu_tile(uint32_t idst) {
    MATH(SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, hamer_gelu_face, idst, VectorMode::RC));
}

void kernel_main() {
    const uint32_t n = get_arg_val<uint32_t>(0);
#ifdef GELU_CB_IN  // the cmm GELU helpers (kernels/cmm/help_*.cpp) use CB 20 / 21
    constexpr uint32_t cb_in = GELU_CB_IN, cb_out = GELU_CB_OUT;
#else
    constexpr uint32_t cb_in = 0, cb_out = 16;
#endif
    init_sfpu(cb_in, cb_out);
    MATH((ckernel::sfpu::hamer_gelu_init()));
    for (uint32_t i = 0; i < n; i++) {
        cb_wait_front(cb_in, 1);
        tile_regs_acquire();
        copy_tile(cb_in, 0, 0);
#ifndef HG_NOP
        hamer_gelu_tile(0);
#endif
        tile_regs_commit();
        tile_regs_wait();
        cb_reserve_back(cb_out, 1);
        pack_tile(0, cb_out);
        cb_push_back(cb_out, 1);
        cb_pop_front(cb_in, 1);
        tile_regs_release();
    }
}