Hugging Face's logo

changh95
/
hamer-p150

Keypoint Detection
blackhole
p150
tt-dit-server
tt-model-cache
tt-model-container
tenstorrent
ttnn
tt-metal
tt-nn
hand-pose
mano
3d-mesh
hamer
tt-model-catalog
Model card Files
xet
 Community
1
hamer-p150 / code /hamer /kernels /eltwise /gelu_compute.cpp
changh95's picture
changh95
Optimized build (2026-10-03): model call 3.49 ms, device 3.36 ms
58ba6df verified
Raw History Blame Contribute Delete
3.99 kB
// 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();
}
}