hamer-p150 / code /hamer /kernels /cmm /compute.cpp
changh95's picture
Optimized build (2026-10-03): model call 3.49 ms, device 3.36 ms
58ba6df verified
Raw History Blame Contribute Delete
10.1 kB
// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// cmm compute: out[1 x NT] = sum_b a_b[1 x KB] @ w_b[KB x NT] (whole K accumulated in DST, fp32),
// packed to an fp32 intermediate, + bias (row broadcast), packed into the sharded output.
// compile-time: NB, KB, NT
#include <cstdint>
#include "api/compute/matmul.h"
#include "api/compute/bcast.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/cb_api.h"
#include "api/compute/pack.h"
#include "api/compute/reg_api.h"
#ifdef CMM_PROF
#include "tools/profiler/kernel_profiler.hpp"
#define KZONE(n) DeviceZoneScopedN(n)
#else
#define KZONE(n)
#endif
#ifdef CMM_GELU
#include "api/compute/common.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 {
// same erf-form GELU as kernels/eltwise/gelu_compute.cpp (the model's default standalone GELU since round 5)
inline void cmm_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 cmm_gelu_tile(uint32_t idst) {
MATH(SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, cmm_gelu_face, idst, VectorMode::RC));
}
#endif
#ifdef CMM_MM_FID
// round 8: the cmm matmul at an explicit fidelity (CMM_MM_FID, e.g. 2 = HiFi2) when this file is the first phase of a
// kernel compiled at another MATH_FIDELITY (kernels/cmm/fused_compute.cpp: the add+LN epilogue needs HiFi4). Same
// calls as api/compute/matmul.h matmul_block_init / matmul_block (incl. the Blackhole dynamic throttle).
constexpr ckernel::MathFidelity CMM_FIDV = static_cast<ckernel::MathFidelity>(CMM_MM_FID);
ALWI void cmm_mm_init(uint32_t a, uint32_t b, uint32_t ct, uint32_t rt, uint32_t kt) {
state_configure(b, a, __builtin_LINE());
UNPACK((llk_unpack_AB_matmul_init(a, b, 0, ct, rt, kt)));
MATH((llk_math_matmul_init<CMM_FIDV, MM_THROTTLE>(a, b, 0, ct, rt)));
#ifdef ARCH_BLACKHOLE
MATH((throttled_mop_status = 0));
#endif
}
ALWI void cmm_mm_block(uint32_t a, uint32_t b, uint32_t ia, uint32_t ib, uint32_t ct, uint32_t rt, uint32_t kt) {
state_configure(b, a, __builtin_LINE());
UNPACK((llk_unpack_AB_matmul(a, b, ia, ib, ct, rt, kt)));
#ifdef ARCH_BLACKHOLE
MATH(({
volatile uint32_t mm_throttle_en = *(throttle_ptr) % 2;
if (mm_throttle_en) {
if (throttled_mop_status != 1) {
llk_math_matmul_init<CMM_FIDV, MM_THROTTLE_MAX>(a, b, 0, ct, rt);
throttled_mop_status = 1;
}
llk_math_matmul<CMM_FIDV, MM_THROTTLE_MAX>(0, ct, rt);
} else {
if (throttled_mop_status != 0) {
llk_math_matmul_init<CMM_FIDV, MM_THROTTLE>(a, b, 0, ct, rt);
throttled_mop_status = 0;
}
llk_math_matmul<CMM_FIDV, MM_THROTTLE>(0, ct, rt);
}
}));
#else
MATH((llk_math_matmul<CMM_FIDV, MM_THROTTLE>(0, ct, rt)));
#endif
}
#define CMM_MM_INIT(a, b, ct, rt, kt) cmm_mm_init(a, b, ct, rt, kt)
#define CMM_MM_BLOCK(a, b, ia, ib, ct, rt, kt) cmm_mm_block(a, b, ia, ib, ct, rt, kt)
#else
#define CMM_MM_INIT(a, b, ct, rt, kt) matmul_block_init(a, b, 0, ct, rt, kt)
#define CMM_MM_BLOCK(a, b, ia, ib, ct, rt, kt) matmul_block(a, b, ia, ib, 0, 0, ct, rt, kt)
#endif
void kernel_main() {
constexpr uint32_t NB = get_compile_time_arg_val(0);
constexpr uint32_t KB = get_compile_time_arg_val(1);
constexpr uint32_t NT = get_compile_time_arg_val(2);
constexpr uint32_t c_a = 0, c_w = 1, c_a1 = 5, c_w1 = 6, c_bias = 2, c_out = 16, c_int = 24;
#ifdef CMM_GELU
// phase 0: GELU of this core's own activation blocks (raw c_araw -> staging c_ast, per parity), which
// its movers then multicast; runtime args: own block count of parity 0, of parity 1
{
constexpr uint32_t c_raw0 = 11, c_raw1 = 12, c_st0 = 4, c_st1 = 8;
const uint32_t n_own[2] = {get_arg_val<uint32_t>(0), get_arg_val<uint32_t>(1)};
init_sfpu(c_raw0, c_st0);
for (uint32_t par = 0; par < 2; ++par) {
const uint32_t c_raw = par ? c_raw1 : c_raw0, c_st = par ? c_st1 : c_st0;
for (uint32_t blk = 0; blk < n_own[par]; ++blk) {
{KZONE("GWAIT"); cb_wait_front(c_raw, 1);}
KZONE("GELU");
for (uint32_t t = 0; t < KB; ++t) {
cb_wait_front(c_raw, 1);
tile_regs_acquire();
copy_tile(c_raw, 0, 0);
cmm_gelu_tile(0);
tile_regs_commit();
tile_regs_wait();
cb_reserve_back(c_st, 1);
pack_tile(0, c_st);
cb_push_back(c_st, 1);
cb_pop_front(c_raw, 1);
tile_regs_release();
}
}
}
}
compute_kernel_hw_startup<SrcOrder::Reverse>(c_a, c_w, c_int);
#else
compute_kernel_hw_startup<SrcOrder::Reverse>(c_a, c_w, c_int);
#endif
#ifdef CMM_KSPLIT
// round 11 (layout k2): this core holds CMM_MPC = 2 token tile rows (own row first, then the partner's) of HALF of K;
// the 2 x NT fp32 partials stay in DST (full-sync DST, 8 fp32 tiles). The partner's half goes out through c_send
// (the BRISC writes it into the partner's c_int); this core's own half is packed with packer L1 accumulation onto
// the partner's partial, which the partner wrote into THIS core's c_int (NCRISC pops the c_gate token once it
// landed: the PACK thread blocks in cb_reserve_back(c_gate) until then).
constexpr uint32_t MPC = CMM_MPC, c_send = 19, c_gate = 22;
cb_reserve_back(c_gate, 1); // arm the gate (1-page CB, now full until the NCRISC pops it)
cb_push_back(c_gate, 1);
#elif defined(CMM_MPC)
constexpr uint32_t MPC = CMM_MPC; // round 11 (xmm): MPC token rows x NT columns, all of K, in full-sync DST
#else
constexpr uint32_t MPC = 1;
#endif
CMM_MM_INIT(c_a, c_w, NT, MPC, KB);
tile_regs_acquire();
// even K blocks arrive in (c_a, c_w) (NCRISC / NOC_0), odd ones in (c_a1, c_w1) (BRISC / NOC_1)
for (uint32_t b = 0; b < NB; ++b) {
const uint32_t ca = (b & 1) ? c_a1 : c_a, cw = (b & 1) ? c_w1 : c_w;
{
KZONE("CWA");
cb_wait_front(ca, KB * MPC);
}
{
KZONE("CWW");
cb_wait_front(cw, KB * NT);
}
#ifndef CMM_NO_MM
for (uint32_t k = 0; k < KB; ++k) {
CMM_MM_BLOCK(ca, cw, k, k * NT, NT, MPC, KB);
}
#endif
cb_pop_front(ca, KB * MPC);
cb_pop_front(cw, KB * NT);
}
tile_regs_commit();
#ifdef CMM_KSPLIT
cb_reserve_back(c_send, NT);
tile_regs_wait();
for (uint32_t i = 0; i < NT; ++i) {
pack_tile(NT + i, c_send);
}
cb_push_back(c_send, NT);
{
KZONE("CXWAIT");
cb_reserve_back(c_gate, 1); // the partner's partial is in c_int
}
cb_reserve_back(c_int, NT);
PACK((llk_pack_reconfig_l1_acc(1)));
for (uint32_t i = 0; i < NT; ++i) {
pack_tile(i, c_int);
}
PACK((llk_pack_reconfig_l1_acc(0)));
tile_regs_release();
cb_push_back(c_int, NT);
#else
cb_reserve_back(c_int, MPC * NT);
tile_regs_wait();
for (uint32_t i = 0; i < MPC * NT; ++i) {
pack_tile(i, c_int);
}
tile_regs_release();
cb_push_back(c_int, MPC * NT);
#endif
#if defined(CMM_MPC) && !defined(CMM_KSPLIT)
// xmm: + bias (the same NT bias tiles for every token row) -> c_out (MPC x NT tiles)
cb_wait_front(c_int, MPC * NT);
cb_wait_front(c_bias, NT);
reconfig_data_format(c_int, c_bias);
pack_reconfig_data_format(c_out);
add_bcast_rows_init(c_int, c_bias);
cb_reserve_back(c_out, MPC * NT);
tile_regs_acquire();
for (uint32_t i = 0; i < MPC * NT; ++i) {
add_tiles_bcast_rows(c_int, c_bias, i, i % NT, i);
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t i = 0; i < MPC * NT; ++i) {
pack_tile(i, c_out);
}
tile_regs_release();
cb_push_back(c_out, MPC * NT);
cb_pop_front(c_int, MPC * NT);
cb_pop_front(c_bias, NT);
return;
#endif
cb_wait_front(c_int, NT);
cb_wait_front(c_bias, NT);
reconfig_data_format(c_int, c_bias);
pack_reconfig_data_format(c_out);
add_bcast_rows_init(c_int, c_bias);
cb_reserve_back(c_out, NT);
tile_regs_acquire();
for (uint32_t i = 0; i < NT; ++i) {
add_tiles_bcast_rows(c_int, c_bias, i, i, i);
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t i = 0; i < NT; ++i) {
pack_tile(i, c_out);
}
tile_regs_release();
cb_push_back(c_out, NT);
cb_pop_front(c_int, NT);
cb_pop_front(c_bias, NT);
}