Download code/hamer/kernels/cmm/compute.cpp from changh95/hamer-p150: direct link, hf CLI and curl.
- Browser
- Download file 10.1 kB
-
https://huggingface.co/changh95/hamer-p150/resolve/main/code/hamer/kernels/cmm/compute.cpp
- Command line
-
hf download hf://changh95/hamer-p150/code/hamer/kernels/cmm/compute.cpp
-
curl -L -o compute.cpp https://huggingface.co/changh95/hamer-p150/resolve/main/code/hamer/kernels/cmm/compute.cpp
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 | |
| 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 | |
| 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)); | |
| } | |
| // 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))); | |
| MATH((throttled_mop_status = 0)); | |
| } | |
| 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))); | |
| 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); | |
| } | |
| })); | |
| MATH((llk_math_matmul<CMM_FIDV, MM_THROTTLE>(0, ct, rt))); | |
| } | |
| 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; | |
| // 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); | |
| compute_kernel_hw_startup<SrcOrder::Reverse>(c_a, c_w, c_int); | |
| // 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); | |
| constexpr uint32_t MPC = CMM_MPC; // round 11 (xmm): MPC token rows x NT columns, all of K, in full-sync DST | |
| constexpr uint32_t MPC = 1; | |
| 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); | |
| } | |
| for (uint32_t k = 0; k < KB; ++k) { | |
| CMM_MM_BLOCK(ca, cw, k, k * NT, NT, MPC, KB); | |
| } | |
| cb_pop_front(ca, KB * MPC); | |
| cb_pop_front(cw, KB * NT); | |
| } | |
| tile_regs_commit(); | |
| 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); | |
| 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); | |
| // 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; | |
| 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); | |
| } | |