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 /hgemv /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
8.15 kB
// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hgemv compute, round 6 (see reader.cpp). One output tile per core.
// HG_LN (K = 1024: the whole vector is ONE dense tile): mean = sum(x) / K, d = x - mean, rstd = rsqrt(sum(d^2) / K
// + eps), h = d * rstd (scalar broadcasts; gamma / beta are folded into W / b on the host) -> c_hd (bf16),
// which the writer spreads into the row-form operand c_hrow.
// y = sum_k hrow_k @ W_k (fp32 DST) (+ bias) [GELU, the model's erf form] [+ residual] -> c_out
// compile-time: KT, inv_k bits, eps bits
#include <cstdint>
#include "api/compute/matmul.h"
#include "api/compute/reduce.h"
#include "api/compute/bcast.h"
#include "api/compute/eltwise_binary.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/rsqrt.h"
#include "api/compute/eltwise_unary/binop_with_scalar.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/reconfig_data_format.h"
#ifdef HG_PROF
#include "tools/profiler/kernel_profiler.hpp"
#define PZ(n) DeviceZoneScopedN(n)
#else
#define PZ(n)
#endif
#ifdef HG_GELU
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#ifdef TRISC_MATH
#include "sfpi.h"
namespace ckernel::sfpu {
// the erf-form GELU of kernels/eltwise/gelu_compute.cpp (degree-8 CDF polynomial, fp32 SFPU)
inline void hg_gelu_face() {
for (size_t i = 0; i < 8; i += 2) {
sfpi::vFloat x0 = sfpi::dst_reg[i];
sfpi::vFloat x1 = sfpi::dst_reg[i + 1];
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 = 9.886907355e-05f;
a0 = a0 * u0 + c;
a1 = a1 * u1 + c; }
{ const sfpi::vFloat c = -1.133757411e-03f;
a0 = a0 * u0 + c;
a1 = a1 * u1 + c; }
{ const sfpi::vFloat c = 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;
}
}
} // namespace ckernel::sfpu
#endif
inline void hg_gelu_tile(uint32_t idst) {
MATH(SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, hg_gelu_face, idst, VectorMode::RC));
}
#endif
#ifdef HG_LN
// scalar statistic of the dense tile in c_src: sum * inv_k [+ eps -> rsqrt], packed (only [0, 0] is used)
template <bool RSTD>
FORCE_INLINE void scalar_stat(uint32_t c_src, uint32_t c_scal, uint32_t c_dst, uint32_t inv_k, uint32_t eps) {
reconfig_data_format(c_src, c_scal); // REDUCE_SCALAR: data -> SrcA, scaler -> SrcB (not swapped)
pack_reconfig_data_format(c_dst);
reduce_init<PoolType::SUM, ReduceDim::REDUCE_SCALAR>(c_src, c_scal, c_dst);
cb_reserve_back(c_dst, 1);
tile_regs_acquire();
reduce_tile<PoolType::SUM, ReduceDim::REDUCE_SCALAR>(c_src, c_scal, 0, 0, 0);
reduce_uninit();
binop_with_scalar_tile_init();
mul_unary_tile(0, inv_k);
if constexpr (RSTD) {
add_unary_tile(0, eps);
rsqrt_tile_init();
rsqrt_tile(0);
}
tile_regs_commit();
tile_regs_wait();
pack_tile(0, c_dst);
tile_regs_release();
cb_push_back(c_dst, 1);
}
#endif
void kernel_main() {
constexpr uint32_t KT = get_compile_time_arg_val(0);
constexpr uint32_t INV_K_BITS = get_compile_time_arg_val(1);
constexpr uint32_t EPS_BITS = get_compile_time_arg_val(2);
constexpr uint32_t c_xd = 0, c_w = 1, c_bias = 2, c_scal = 3, c_res = 4, c_hrow = 5, c_hd = 6, c_out = 16;
constexpr uint32_t c_mean = 24, c_d = 25, c_sq = 26, c_rstd = 27, c_f32 = 28; // c_f32: fp32 format reference
#ifdef HG_LN
compute_kernel_hw_startup(c_xd, c_scal, c_mean);
{PZ("WX");
cb_wait_front(c_xd, 1);}
{PZ("LN");
cb_wait_front(c_scal, 1);
scalar_stat<false>(c_xd, c_scal, c_mean, INV_K_BITS, EPS_BITS);
// d = x - mean
cb_wait_front(c_mean, 1);
reconfig_data_format(c_xd, c_mean);
pack_reconfig_data_format(c_d);
sub_bcast_scalar_init(c_xd, c_mean);
cb_reserve_back(c_d, 1);
tile_regs_acquire();
sub_tiles_bcast_scalar(c_xd, c_mean, 0, 0, 0);
tile_regs_commit();
tile_regs_wait();
pack_tile(0, c_d);
tile_regs_release();
cb_push_back(c_d, 1);
cb_pop_front(c_mean, 1);
cb_pop_front(c_xd, 1);
// d^2
cb_wait_front(c_d, 1);
reconfig_data_format(c_d, c_d);
pack_reconfig_data_format(c_sq);
mul_init(c_d, c_d);
cb_reserve_back(c_sq, 1);
tile_regs_acquire();
mul_tiles(c_d, c_d, 0, 0, 0);
tile_regs_commit();
tile_regs_wait();
pack_tile(0, c_sq);
tile_regs_release();
cb_push_back(c_sq, 1);
cb_wait_front(c_sq, 1);
scalar_stat<true>(c_sq, c_scal, c_rstd, INV_K_BITS, EPS_BITS);
cb_pop_front(c_sq, 1);
// h = d * rstd -> bf16 dense tile for the writer
cb_wait_front(c_rstd, 1);
reconfig_data_format(c_d, c_rstd);
pack_reconfig_data_format(c_hd);
mul_bcast_scalar_init(c_d, c_rstd);
cb_reserve_back(c_hd, 1);
tile_regs_acquire();
mul_tiles_bcast_scalar(c_d, c_rstd, 0, 0, 0);
tile_regs_commit();
tile_regs_wait();
pack_tile(0, c_hd);
tile_regs_release();
cb_push_back(c_hd, 1);
#ifdef HG_TINY
// round 10: the row-major dense LN output IS the matmul operand: c_hrow = KT [1, 32] tiny tiles over c_hd's memory
// (no spread by the writer)
cb_reserve_back(c_hrow, KT);
cb_push_back(c_hrow, KT);
#endif
cb_pop_front(c_rstd, 1);
cb_pop_front(c_d, 1);
cb_pop_front(c_scal, 1);
}
#else
compute_kernel_hw_startup(c_hrow, c_w, c_out);
#endif
// y = sum_k hrow_k @ W_k
{PZ("WH");
cb_wait_front(c_hrow, KT);}
{PZ("WW");
cb_wait_front(c_w, KT);}
PZ("MM");
#ifdef HG_TINY
// round 10: [1, 32] tiny in0 / output tiles -- reprogram the tile / face geometry, not only the formats
reconfig_full_operand(c_w, c_hrow); // srcA = in1 = W, srcB = in0 = h
pack_reconfig_data_format<true>(c_out);
#else
reconfig_data_format(c_w, c_hrow); // srcA = in1 = W, srcB = in0 = h
pack_reconfig_data_format(c_out);
#endif
matmul_init(c_hrow, c_w);
tile_regs_acquire();
for (uint32_t k = 0; k < KT; ++k) {
matmul_tiles(c_hrow, c_w, k, k, 0);
}
// + bias [GELU] [+ residual]: SFPU fp32 adds in DST (operands copied in with unpack-to-dest fp32, so the matmul
// result never goes through a 19-bit source register)
cb_wait_front(c_bias, 1);
#ifdef HG_TINY
reconfig_full_operand_srca(c_bias);
copy_tile_to_dst_init_short(c_bias);
#else
copy_tile_to_dst_init_short_with_dt(c_w, c_bias);
#endif
copy_tile(c_bias, 0, 1);
add_binary_tile_init();
add_binary_tile(0, 1, 0);
#ifdef HG_GELU
hg_gelu_tile(0);
#endif
#ifdef HG_RES
cb_wait_front(c_res, 1);
#ifdef HG_TINY
reconfig_full_operand_srca(c_res);
copy_tile_to_dst_init_short(c_res);
#else
copy_tile_to_dst_init_short_with_dt(c_bias, c_res);
#endif
copy_tile(c_res, 0, 1);
add_binary_tile_init();
add_binary_tile(0, 1, 0);
#endif
tile_regs_commit();
cb_reserve_back(c_out, 1);
tile_regs_wait();
pack_tile(0, c_out);
tile_regs_release();
cb_push_back(c_out, 1);
cb_pop_front(c_hrow, KT);
cb_pop_front(c_w, KT);
cb_pop_front(c_bias, 1);
#ifdef HG_RES
cb_pop_front(c_res, 1);
#endif
}