Blame for code/hamer/kernels/cmm/mover.cpp

// 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<w_args.next_compile_time_args_offset()>();
    constexpr auto a_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
    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<uint32_t>(0);
    const uint32_t b_addr = get_arg_val<uint32_t>(1);
    const uint32_t a_addr = get_arg_val<uint32_t>(2);
    const uint32_t n0 = get_arg_val<uint32_t>(3);
    const uint32_t m = get_arg_val<uint32_t>(4);
    Pipe<NBP, S, BT, PAR * 2 * S> wp;
    Pipe<NBP, S, KB, PAR * 2 * S + S> 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<uint32_t>(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<uint32_t>(wtab + 2 * j + par);
        noc_async_read(get_noc_addr_from_bank_id<true>(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<volatile tt_l1_ptr uint32_t*>(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<uint32_t>(5);
        const uint32_t ox = get_arg_val<uint32_t>(6), oy = get_arg_val<uint32_t>(7), ooff = get_arg_val<uint32_t>(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();
}