// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc. // SPDX-License-Identifier: Apache-2.0 // cmm data mover, one per RISC: PAR = 0 on NCRISC (NOC_0), 1 on BRISC (NOC_1). K blocks b with // b % 2 == PAR are this RISC's: block j of its pipes is global block 2j + PAR. Both RISCs run // - a weight pipe: KB x NT tiles of this core's weight column slab (interleaved DRAM, row stride NTOT), // all-gathered inside the weight group, and // - an activation pipe: 1 x KB tiles of this core's activation tile row (interleaved, row stride KT), // all-gathered inside the activation group, // so the four streams are spread over both NoCs (common.hpp). PAR 0 also reads the bias tiles, PAR 1 // writes the NT output tiles to the output shard (out_addr + ooff on core (ox, oy)). // compile-time: NBP (blocks per RISC), KB, NT, NTOT, KT, S, PAR, TensorAccessorArgs(w), (bias), (a) // runtime: w_addr, b_addr, a_addr, n0, m, out_addr, ox, oy, ooff, wgroup (13 + 2G), agroup (13 + 2G) #include "common.hpp" // activation tile (token row m, K tile kk) -> page of the activation tensor FORCE_INLINE uint32_t a_page(uint32_t m, uint32_t kk, uint32_t KT) { #ifdef CMM_A_HEADS // [1, H, M, D] head-split SDPA output (head_concat folded into this read) return (kk / CMM_A_DT) * (CMM_A_MT * CMM_A_DT) + m * CMM_A_DT + kk % CMM_A_DT; #else return m * KT + kk; #endif } void kernel_main() { constexpr uint32_t NBP = 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 NTOT = get_compile_time_arg_val(3); constexpr uint32_t KT = get_compile_time_arg_val(4); constexpr uint32_t S = get_compile_time_arg_val(5); constexpr uint32_t PAR = get_compile_time_arg_val(6); constexpr auto w_args = TensorAccessorArgs<7>(); constexpr auto b_args = TensorAccessorArgs(); constexpr auto a_args = TensorAccessorArgs(); constexpr uint32_t BT = KB * NT; constexpr uint32_t c_a = PAR ? 5 : 0, c_w = PAR ? 6 : 1, c_wst = PAR ? 7 : 3, c_ast = PAR ? 8 : 4; constexpr uint32_t c_flags = PAR ? 10 : 9, c_bias = 2, c_out = 16; const uint32_t w_addr = get_arg_val(0); const uint32_t b_addr = get_arg_val(1); const uint32_t a_addr = get_arg_val(2); const uint32_t n0 = get_arg_val(3); const uint32_t m = get_arg_val(4); Pipe wp; Pipe ap; wp.g.parse(9); ap.g.parse(9 + 13 + 2 * wp.g.G); #ifdef CMM_AROT1 // odd activation blocks owned by the odd row members (even ones by the even members): one activation // block per member instead of two on every other member (balances the CMM_GELU work) if constexpr (PAR == 1) { ap.g.rot = (ap.g.rot + 1) % ap.g.G; } #endif const uint32_t flags = get_write_ptr(c_flags); wp.init(c_w, flags); ap.init(c_a, flags + 16 * (S + 1)); const uint32_t wpage = wp.page, apage = ap.page; const auto w = TensorAccessor(w_args, w_addr, wpage); const auto a = TensorAccessor(a_args, a_addr, apage); const uint32_t wst = get_write_ptr(c_wst), ast = get_write_ptr(c_ast); #if defined(CMM_AROT1) constexpr uint32_t G_ROT_FIX_ = PAR == 1 ? 1 : 0; // ap.g.rot of PAR 1 carries the +1 already #else constexpr uint32_t G_ROT_FIX_ = 0; #endif #ifdef CMM_GELU // fused GELU on the activation (= the fc1 output): the owner reads its raw blocks into c_araw, the compute // applies the GELU into the staging CB c_ast, the owner multicasts the GELU'd block (CB-synchronised) constexpr uint32_t c_araw = PAR ? 12 : 11; const uint32_t araw = get_write_ptr(c_araw); uint32_t a_own = 0; #endif #ifdef CMM_PACKED const uint32_t wtab = ap.g.args + 2 * ap.g.G; // per global K block of this column: bank | slot << 8 #endif #ifdef CMM_PF const uint32_t pf_src = get_arg_val(wtab + 2 * NBP) + PAR * BT * wp.page; // tt_kernels.CmmPrefetch #endif // weight block j (pipe-local) of parity `par` -> staging (owner only). Plain reads + plain barriers. // (DO NOT use noc_async_read_barrier_with_trid(0): it never returned after plain noc_async_read()s // issued with trid 0 and hung chip 23, see OPT_REPORT.md round 4.) auto fetch_w = [&](uint32_t j, uint32_t par, uint32_t st) { if (j >= NBP || !wp.owner(j)) { return; } #ifdef CMM_SKIP_FIRST // timing experiment only (wrong results): pairs < CMM_SKIP_FIRST are "already in L1" if (j < CMM_SKIP_FIRST) { return; } #endif #ifdef CMM_PF // pairs < CMM_PF were prefetched into the persistent L1 buffer by the preceding op if (j < CMM_PF) { return; } #endif CZONE("WISS"); uint32_t dst = st + wp.own_index(j) * BT * wpage; #ifdef CMM_PACKED // host-packed weight (tt_kernels.CmmWeight): the block is contiguous in one DRAM bank const uint32_t d = get_arg_val(wtab + 2 * j + par); noc_async_read(get_noc_addr_from_bank_id(d & 0xff, w_addr + (d >> 8) * BT * wpage), dst, BT * wpage); #else uint32_t row = (2 * j + par) * KB * NTOT + n0; for (uint32_t k = 0; k < KB; ++k) { for (uint32_t n = 0; n < NT; ++n) { #ifndef CMM_NO_DRAM noc_async_read(w.get_noc_addr(row + n), dst, wpage); #endif dst += wpage; } row += NTOT; } #endif }; #ifdef CMM_DRAM0 // every weight DRAM read on NCRISC / NOC_0 (DRAM reads on NOC_1 measured at ~half the NOC_0 rate, // scripts/micro/dram_bench.py); NCRISC fetches BOTH blocks of a pair j (global 2j, 2j+1) into the two // staging CBs and publishes "own pairs landed" in a local semaphore that BRISC waits on before // multicasting its odd block (semaphore zeroed by dispatch every launch; only NCRISC writes it). #ifndef CMM_LA #error "CMM_DRAM0 needs CMM_LA" #endif constexpr bool DRAM_HERE = PAR == 0; const uint32_t wst1 = get_write_ptr(7); // BRISC's staging CB (same core, same L1) volatile tt_l1_ptr uint32_t* staged = reinterpret_cast(get_semaphore(4 * S)); auto fetch_pair = [&](uint32_t j) { fetch_w(j, 0, wst); fetch_w(j, 1, wst1); }; #else constexpr bool DRAM_HERE = true; auto fetch_pair = [&](uint32_t j) { fetch_w(j, PAR, wst); }; #endif #ifdef CMM_LA // lazy weight fetch: block j is read from DRAM only when the pipe is near block j, so the DRAM // serves the blocks roughly in consumption order (an eager prefetch of every block at once made the // first block land only after ~all the DRAM traffic, round 4). The first pairs go out first. constexpr uint32_t LA = CMM_LA; if constexpr (DRAM_HERE) { for (uint32_t j = 0; j < LA; ++j) { fetch_pair(j); } } #endif // own activation blocks (L1) and the bias #ifdef CMM_PRO_BARRIER { CZONE("PRO"); #endif #if defined(CMM_GELU) && defined(CMM_DRAM0) // BRISC (no DRAM reads in this mode) reads the raw own activation blocks of BOTH parities right away and // hands them to the compute, so the GELU runs while NCRISC's first weight blocks are still in DRAM flight if constexpr (PAR == 1) { for (uint32_t p = 0; p < 2; ++p) { const uint32_t rot = (ap.g.rot + ap.g.G - G_ROT_FIX_ + p) % ap.g.G; // parity p ownership (CMM_AROT1) const uint32_t q = (ap.g.me + ap.g.G - rot) % ap.g.G; const uint32_t first = (q * NBP + ap.g.G - 1) / ap.g.G; const uint32_t c_raw = p ? 12 : 11; uint32_t dst = get_write_ptr(c_raw), n = 0; for (uint32_t j = first; j < NBP && ((j * ap.g.G) / NBP + rot) % ap.g.G == ap.g.me; ++j) { const uint32_t k0 = (2 * j + p) * KB; for (uint32_t k = 0; k < KB; ++k) { noc_async_read(a.get_noc_addr(a_page(m, k0 + k, KT)), dst, apage); dst += apage; } ++n; } noc_async_read_barrier(); if (n) { cb_reserve_back(c_raw, n * KB); cb_push_back(c_raw, n * KB); } } } if constexpr (false) #endif for (uint32_t j = ap.first_own; j < NBP && ap.owner(j); ++j) { #ifdef CMM_GELU uint32_t dst = araw + ap.own_index(j) * KB * apage; ++a_own; #else uint32_t dst = ast + ap.own_index(j) * KB * apage; #endif const uint32_t k0 = (2 * j + PAR) * KB; for (uint32_t k = 0; k < KB; ++k) { noc_async_read(a.get_noc_addr(a_page(m, k0 + k, KT)), dst, apage); dst += apage; } } if constexpr (PAR == 0) { const uint32_t bpage = get_local_cb_interface(c_bias).fifo_page_size; const auto bias = TensorAccessor(b_args, b_addr, bpage); cb_reserve_back(c_bias, NT); uint32_t bp = get_write_ptr(c_bias); for (uint32_t i = 0; i < NT; ++i) { noc_async_read(bias.get_noc_addr(n0 + i), bp, bpage); bp += bpage; } } #ifdef CMM_PRO_BARRIER noc_async_read_barrier(); } #endif #ifndef CMM_LA for (uint32_t j = wp.first_own; j < NBP && wp.owner(j); ++j) { fetch_pair(j); } #endif noc_async_read_barrier(); if constexpr (PAR == 0) { cb_push_back(c_bias, NT); } #ifdef CMM_GELU if (a_own) { cb_reserve_back(c_araw, a_own * KB); cb_push_back(c_araw, a_own * KB); } #endif for (uint32_t j = 0; j < NBP; ++j) { #ifdef CMM_NO_W cb_reserve_back(c_w, BT); cb_push_back(c_w, BT); #else #ifdef CMM_LA if constexpr (DRAM_HERE) { #ifndef CMM_LA_AFTER fetch_pair(j + LA); #endif if (wp.owner(j)) { CZONE("WDRAM"); noc_async_read_barrier(); // pair j landed (also waits for the few younger own blocks in flight) #ifdef CMM_DRAM0 *staged = wp.own_index(j) + 1; #endif } } #ifdef CMM_DRAM0 if constexpr (!DRAM_HERE) { if (wp.owner(j)) { CZONE("WSTG"); const uint32_t need = wp.own_index(j) + 1; while (true) { invalidate_l1_cache(); if (*staged >= need) { break; } wp.signal_free(); ap.signal_free(); } } } #endif #endif #ifdef CMM_PF if (j < CMM_PF) { wp.step_from(j, pf_src); } else #endif wp.step(j, wst); #if defined(CMM_LA) && defined(CMM_LA_AFTER) if constexpr (DRAM_HERE) { fetch_pair(j + LA); // after block j arrived: pair j + LA lands while pairs j..j+LA-1 are consumed } #endif #endif #ifdef CMM_NO_A cb_reserve_back(c_a, KB); cb_push_back(c_a, KB); #else #ifdef CMM_GELU if (ap.owner(j)) { CZONE("AGEL"); cb_wait_front(c_ast, KB); // GELU'd by the compute } #endif ap.step(j, ast); #ifdef CMM_GELU if (ap.owner(j)) { cb_pop_front(c_ast, KB); // (sent: the data left L1 before the flag; slots are addressed absolutely) } #endif #endif } if constexpr (PAR == 1) { const uint32_t out_addr = get_arg_val(5); const uint32_t ox = get_arg_val(6), oy = get_arg_val(7), ooff = get_arg_val(8); const uint32_t opage = get_local_cb_interface(c_out).fifo_page_size; cb_wait_front(c_out, NT); noc_async_write(get_read_ptr(c_out), get_noc_addr(ox, oy, out_addr + ooff), NT * opage); noc_async_write_barrier(); cb_pop_front(c_out, NT); } wp.finish(); }