// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// cmm data mover v5 ("one RISC per operand"):
//   PAR 0 (NCRISC, NOC_0): BOTH weight pipes (even K blocks -> c_w, odd -> c_w1).  Member 0 of the weight group
//     (the top row of the column) is the column's ONLY weight sender: it reads the host-packed contiguous blocks
//     from DRAM (tt_kernels.CmmWeight; one-packet reads tagged with transaction ids 1..3, the library DRAM-sharded
//     reader pattern -- never trid 0) into a staging ring of SS slots, LA blocks ahead, and multicasts each block
//     down the column (NOC_0 from the top row: no wrap-around).  The packed layout removed the per-tile issue cost
//     that bound the library's single in1 sender; the top-row sender avoids the mid-column multicast penalty.
//     Optionally the first CMM_PF pairs are already in a persistent L1 buffer (tt_kernels.CmmPrefetch).
//   PAR 1 (BRISC, NOC_1): BOTH activation pipes (contiguous owners in the row, odd blocks rotated by one member),
//     the bias read and the output write.
// compile-time / runtime args: as mover.cpp (the packed table follows the activation group; then the prefetch
// buffer address when CMM_PF).
#include "common.hpp"

#ifndef CMM_LA
#define CMM_LA 1
#endif
#ifndef CMM_SS
#define CMM_SS 3
#endif

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
}

#ifdef CMM_W2
// round 12 (xmm, CMM_W2): TWO weight senders per group -- the even K blocks are read and multicast by member 0 on NOC_0
// (NCRISC), the odd ones by member G-1 on NOC_1 (BRISC; with CMM_A_RES it has no in0 work).  Twice the DRAM readers,
// each with half the bytes, and the odd blocks' DRAM responses / multicasts use the other NoC.  The owner reads its next
// block straight into its own next ring slot (reserving two blocks ahead) while the current one is multicast in place.
template <uint32_t NBP, uint32_t S, uint32_t BT, uint32_t SEM0>
FORCE_INLINE void w2_run(Pipe<NBP, S, BT, SEM0>& p, uint32_t par, uint32_t w_addr, uint32_t wtab, uint32_t blk_bytes) {
    NoOther none;
    if (p.g.me == p.single_m) {
        auto fetch = [&](uint32_t j) {
            CZONE("WISS");
            const uint32_t d = get_arg_val<uint32_t>(wtab + 2 * j + par);
            uint64_t src = get_noc_addr_from_bank_id<true>(d & 0xff, w_addr + (d >> 8) * blk_bytes);
            uint32_t dst = p.slot_addr(j), left = blk_bytes;
            const uint32_t t = 1 + j % 3;
            noc_async_read_set_trid(t);
            while (noc_available_transactions(noc_index, t) < ((NOC_MAX_TRANSACTION_ID_COUNT + 1) / 2)) {
            }
            while (left) {
                const uint32_t n = left < NOC_MAX_BURST_SIZE ? left : NOC_MAX_BURST_SIZE;
                noc_async_read<NOC_MAX_BURST_SIZE>(src, dst, n);
                src += n;
                dst += n;
                left -= n;
            }
        };
#ifndef CMM_W2F
#define CMM_W2F 2
#endif
        // F blocks in flight: blocks 0 .. F-1 up front; once block j has landed, block j + F is issued (its slot: the
        // ring holds S blocks, F + 1 <= S) BEFORE block j is multicast, so F reads stay in flight during the multicast
        constexpr uint32_t F = CMM_W2F;
        static_assert(F + 1 <= S && F <= 3, "CMM_W2F: ring slots / trids");
        for (uint32_t j = 0; j < F && j < NBP; ++j) {
            cb_reserve_back(p.cb, (j + 1) * BT);
            fetch(j);
        }
        for (uint32_t j = 0; j < NBP; ++j) {
            {
                CZONE("WDRAM");
                noc_async_read_barrier_with_trid(1 + j % 3);
            }
            if (j + F < NBP) {
                cb_reserve_back(p.cb, (F + 1) * BT);  // slots of blocks j .. j + F (j not pushed yet) all free
                fetch(j + F);
            }
            const uint32_t slot = p.slot_addr(j);
            p.send(slot, slot, j, none);
            cb_push_back(p.cb, BT);
        }
        noc_async_read_set_trid(0);
    } else {
        for (uint32_t j = 0; j < NBP; ++j) {
            p.recv_step2(j, none);
        }
    }
    p.finish();
}
#endif

#ifdef CMM_A2N
// round 12 (CMM_A2N): keep the ready signals of two other pipes flowing while one pipe waits (NCRISC runs the two weight
// pipes AND the odd activation pipe)
template <typename A, typename B>
struct Two {
    A* a;
    B* b;
    FORCE_INLINE void signal_free() {
        a->signal_free();
        b->signal_free();
    }
};
#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()>();
#ifdef CMM_XMM
    constexpr auto o_args = TensorAccessorArgs<a_args.next_compile_time_args_offset()>();  // output (page writes)
#endif
    constexpr uint32_t BT = KB * NT;
    constexpr uint32_t NB = 2 * NBP;
#ifdef CMM_MPC
    constexpr uint32_t MPC = CMM_MPC;  // round 11 (layout k2): token tile rows per core (activation block = KB x MPC)
#else
    constexpr uint32_t MPC = 1;
#endif
    constexpr uint32_t ABT = KB * MPC;
    constexpr uint32_t c_a0 = 0, c_w0 = 1, c_bias = 2, c_wst = 3, c_ast0 = 4, c_a1 = 5, c_w1 = 6, c_ast1 = 8;
    constexpr uint32_t 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);
    LineGroup wg, ag;
    wg.parse(9);
    ag.parse(9 + 13 + 2 * wg.G);
    const uint32_t wtab = ag.args + 2 * ag.G;  // per global K block of this column: bank | slot << 8

    if constexpr (PAR == 0) {
#ifdef CMM_W2
        static_assert(CMM_LA == 1, "CMM_W2: one block ahead per sender");
#ifdef CMM_PF
        static_assert(false, "CMM_W2 without the weight prefetch");
#endif
        {
            Pipe<NBP, S, BT, 0> p0;
            p0.g = wg;
            p0.single = true;
            p0.single_m = 0;
            p0.init(c_w0, get_write_ptr(9));
            w2_run(p0, 0, w_addr, wtab, BT * p0.page);
        }
        if (false) {
#endif
        // ---------------- weights (NCRISC / NOC_0) ----------------
        Pipe<NBP, S, BT, 0> p0;
        Pipe<NBP, S, BT, S> p1;
        p0.g = wg;
        p1.g = wg;
        p0.single = true;
        p1.single = true;
        const uint32_t flags = get_write_ptr(9);
        p0.init(c_w0, flags);
        p1.init(c_w1, flags + 16 * (S + 1));
#ifdef CMM_A2N
        // round 12: the ODD activation blocks run on this RISC / NOC_0 (the even ones stay on the BRISC / NOC_1): two
        // owner hand-over chains per grid row in parallel on the two NoCs instead of one serial chain
        static_assert(MPC >= 1, "");
#if !defined(CMM_A_LOCAL) || defined(CMM_XMM) || defined(CMM_AEARLY) || defined(CMM_NO_A) || defined(CMM_NO_W) || defined(CMM_W2)
        static_assert(false, "CMM_A2N: k2 / out cmm with the local (GELU-helper) activation layout only");
#endif
        Pipe<NBP, S, ABT, 3 * S> q1;
        q1.g = ag;
        q1.g.rot = (q1.g.rot + 1) % q1.g.G;  // odd blocks: the next member (as on the BRISC before)
        q1.init(c_a1, get_write_ptr(10) + 16 * (S + 1));
        Two<decltype(p1), decltype(q1)> o_w0{&p1, &q1};
        Two<decltype(p0), decltype(q1)> o_w1{&p0, &q1};
        Two<decltype(p0), decltype(p1)> o_a1{&p0, &p1};
        auto q1_step = [&](uint32_t j) {
            if (q1.owner(j)) {
#ifdef CMM_GELU_HELP
                volatile tt_l1_ptr uint32_t* gs = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(CMM_GELU_SEM));
                while (true) {
                    invalidate_l1_cache();
                    if (*gs >= ABT) {
                        break;
                    }
                    p0.signal_free();
                    p1.signal_free();
                    q1.signal_free();
                }
#endif
                q1.step_from2(j, a_addr + q1.own_index(j) * ABT * q1.page, o_a1);
            } else {
                q1.recv_step2(j, o_a1);
            }
        };
#define CMM_OW0 o_w0
#define CMM_OW1 o_w1
#define CMM_Q1STEP(j) q1_step(j)
#else
#define CMM_OW0 p1
#define CMM_OW1 p0
#define CMM_Q1STEP(j)
#endif
        const uint32_t wpage = p0.page;
#ifdef CMM_NO_W  // timing experiment only (wrong results): no weight movement at all
        for (uint32_t j = 0; j < NB; ++j) {
            const uint32_t c = (j & 1) ? c_w1 : c_w0;
            cb_reserve_back(c, BT);
            cb_push_back(c, BT);
        }
#else
        if (wg.me == 0) {
            const uint32_t blk_bytes = BT * wpage;
            const uint32_t stg = get_write_ptr(c_wst);
#ifdef CMM_PF
            const uint32_t pf = get_arg_val<uint32_t>(wtab + NB);
            constexpr uint32_t PFB = 2 * CMM_PF;  // prefetched blocks (sequence numbers 0 .. PFB-1)
#else
            constexpr uint32_t PFB = 0;
#endif
            // block s (sequence = 2j + parity) -> its staging slot / trid
#ifndef CMM_NOSTAGE
            auto slot = [&](uint32_t s) -> uint32_t {
#ifdef CMM_PF
                if (s < PFB) {
                    return pf + s * blk_bytes;
                }
#endif
                return stg + ((s - PFB) % CMM_SS) * blk_bytes;
            };
#endif
            auto trid = [&](uint32_t s) -> uint32_t { return 1 + (s - PFB) % 3; };
#ifdef CMM_NOSTAGE
            // round 8: no staging ring -- a DRAM block lands straight in this sender's own ring slot of its pipe (the
            // next unpushed block of that pipe when LA == 1, so cb_reserve_back gives exactly its slot and waits until
            // the compute freed it), and is then multicast to the OTHER members only (Pipe::send in-place mode)
            static_assert(CMM_LA == 1, "CMM_NOSTAGE needs a lookahead of exactly one block");
            uint32_t ring[2] = {0, 0};
            auto slot = [&](uint32_t s) -> uint32_t {
#ifdef CMM_PF
                if (s < PFB) {
                    return pf + s * blk_bytes;
                }
#endif
                return ring[s & 1];
            };
#endif
            auto fetch = [&](uint32_t s) {
                if (s >= NB || s < PFB) {
                    return;
                }
#ifdef CMM_NOSTAGE
                {
                    const uint32_t cbp = (s & 1) ? c_w1 : c_w0;
                    cb_reserve_back(cbp, BT);
                    ring[s & 1] = get_write_ptr(cbp);
                }
#endif
#ifdef CMM_NO_DRAM  // timing experiment only (wrong results): weights "already in L1"
                return;
#endif
#ifdef CMM_SKIP_FIRST  // timing experiment only (wrong results): blocks < CMM_SKIP_FIRST "already in L1"
                if (s < CMM_SKIP_FIRST) {
                    return;
                }
#endif
                CZONE("WISS");
                const uint32_t d = get_arg_val<uint32_t>(wtab + s);
                uint64_t src = get_noc_addr_from_bank_id<true>(d & 0xff, w_addr + (d >> 8) * blk_bytes);
                uint32_t dst = slot(s), left = blk_bytes;
                const uint32_t t = trid(s);
                noc_async_read_set_trid(t);
                while (noc_available_transactions(noc_index, t) < ((NOC_MAX_TRANSACTION_ID_COUNT + 1) / 2)) {
                }
                while (left) {
                    const uint32_t n = left < NOC_MAX_BURST_SIZE ? left : NOC_MAX_BURST_SIZE;
                    noc_async_read<NOC_MAX_BURST_SIZE>(src, dst, n);
                    src += n;
                    dst += n;
                    left -= n;
                }
            };
            static_assert(CMM_LA + 1 <= CMM_SS && CMM_SS <= 3, "staging ring / trid ring");
            // blocks PFB .. PFB+LA-1 are fetched up front (except in-place with prefetched blocks ahead of them: their
            // pipe slot is not the next one yet), every later block LA steps ahead -- each block exactly once (round 8:
            // the old loop fetched block PFB a second time when CMM_PF > 0)
#ifdef CMM_NOSTAGE
            constexpr uint32_t PRE = PFB == 0 ? CMM_LA : 0;
#else
            constexpr uint32_t PRE = CMM_LA;
#endif
            for (uint32_t s = PFB; s < PFB + PRE; ++s) {
                fetch(s);
            }
            for (uint32_t j = 0; j < NBP; ++j) {
                for (uint32_t par = 0; par < 2; ++par) {
                    const uint32_t s = 2 * j + par;
#ifdef CMM_DUPFETCH  // A/B: the round-5..7 loop (block PFB fetched a second time when CMM_PF > 0)
                    fetch(s + CMM_LA);
#else
                    if (s + CMM_LA >= PFB + PRE) {
                        fetch(s + CMM_LA);
                    }
#endif
                    if (s >= PFB) {
                        CZONE("WDRAM");
                        noc_async_read_barrier_with_trid(trid(s));
                    }
                    if (par == 0) {
                        p0.step_from2(j, slot(s), CMM_OW0);
                    } else {
                        p1.step_from2(j, slot(s), CMM_OW1);
                        CMM_Q1STEP(j);
                    }
                }
            }
        } else {
            for (uint32_t j = 0; j < NBP; ++j) {
                p0.recv_step2(j, CMM_OW0);
                p1.recv_step2(j, CMM_OW1);
                CMM_Q1STEP(j);
            }
        }
#endif
        p0.finish();
#ifdef CMM_W2
        }
#endif
#ifdef CMM_KSPLIT
        {
            // the partner core wrote its partial of this core's token row into c_int (acked) and then incremented
            // semaphore 4 S: release the compute's PACK thread (c_gate, see compute.cpp)
            CZONE("XRECV");
            volatile tt_l1_ptr uint32_t* xs = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(4 * S));
            while (true) {
                invalidate_l1_cache();
                if (*xs >= 1) {
                    break;
                }
            }
            cb_wait_front(22, 1);
            cb_pop_front(22, 1);
        }
#endif
    } else {
        // ---------------- activation, bias, output (BRISC / NOC_1) ----------------
#ifdef CMM_A_RES
        // round 12 (xmm): this core's whole in0 (its MPC token rows, all K blocks) is already resident in L1 -- c_a0 /
        // c_a1 are backed by the two resident tensors (even / odd K blocks, written by the preceding add+LN epilogue):
        // push it all at once; only the bias is read
        {
            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;
            }
            cb_reserve_back(c_a0, NBP * ABT);
            cb_push_back(c_a0, NBP * ABT);
            cb_reserve_back(c_a1, NBP * ABT);
            cb_push_back(c_a1, NBP * ABT);
        }
#ifdef CMM_W2
        {
            Pipe<NBP, S, BT, S> p1;
            p1.g = wg;
            p1.single = true;
            p1.single_m = wg.G - 1;
            p1.init(c_w1, get_write_ptr(9) + 16 * (S + 1));
            w2_run(p1, 1, w_addr, wtab, BT * p1.page);
        }
#endif
        // the bias (trid 0) is needed only after the last K block: waited for here, off the weight stream's start
        noc_async_read_barrier();
        cb_push_back(c_bias, NT);
#else
#ifdef CMM_W2
        static_assert(false, "CMM_W2 needs CMM_A_RES (the BRISC sends the odd weight blocks)");
#endif
        Pipe<NBP, S, ABT, 2 * S> q0;
        Pipe<NBP, S, ABT, 3 * S> q1;
        q0.g = ag;
        q1.g = ag;
        q1.g.rot = (q1.g.rot + 1) % q1.g.G;  // odd blocks: the next member (one block per member when G = NB)
#ifdef CMM_XMM
        // round 11 (xmm): ONE activation sender per grid column (its bottom core, member G-1), multicasting every K block
        // up the column (NOC_1: no wrap-around); it gathers the blocks from the width-sharded input itself
        q0.single = true;
        q1.single = true;
        q0.single_m = ag.G - 1;
        q1.single_m = ag.G - 1;
#endif
        const uint32_t flags = get_write_ptr(10);
        q0.init(c_a0, flags);
#ifndef CMM_A2N  // (A2N: the odd pipe lives on the NCRISC, which owns its flags)
        q1.init(c_a1, flags + 16 * (S + 1));
#endif
        const uint32_t apage = q0.page;
        const auto a = TensorAccessor(a_args, a_addr, apage);
#ifdef CMM_A_LOCAL
        // the activation is BLOCK-sharded on this grid with ONE K block per core (tt_kernels.gelu writing the
        // fc2 input in the cmm owner layout): every member owns exactly the block in its own shard, no gather
        const uint32_t ast0 = a_addr, ast1 = a_addr;
#else
        const uint32_t ast0 = get_write_ptr(c_ast0), ast1 = get_write_ptr(c_ast1);
#endif
#if !defined(CMM_NO_A) && !defined(CMM_NO_AGATHER) && !defined(CMM_A_LOCAL)  // CMM_NO_AGATHER: timing only
        CZONE("AGATH");
        // own blocks of each parity into the staging CBs; with the k2 K split (MPC 2) a block is MPC token rows (this
        // core's row m first, then the partner row (m + 3) % 6) of KB tiles of this core's K half (KT / 2 * (m / 3) ..)
#ifdef CMM_KSPLIT
        const uint32_t kbase = (KT / 2) * (m / 3);
#else
        constexpr uint32_t kbase = 0;
#endif
        auto gather = [&](auto& q, uint32_t ast, uint32_t par) {
            for (uint32_t j = q.first_own; j < NBP && q.owner(j); ++j) {
                uint32_t dst = ast + q.own_index(j) * ABT * apage;
                const uint32_t k0 = kbase + (2 * j + par) * KB;
#ifdef CMM_XMM
                noc_async_read_set_trid(1 + 2 * j + par);  // per-block barrier before its multicast
#endif
                for (uint32_t i = 0; i < MPC; ++i) {
#ifdef CMM_XMM
                    const uint32_t mi = m + i;  // consecutive token rows
#else
                    const uint32_t mi = i == 0 ? m : (m + 3) % 6;
#endif
                    for (uint32_t k = 0; k < KB; ++k) {
                        noc_async_read(a.get_noc_addr(a_page(mi, k0 + k, KT)), dst, apage);
                        dst += apage;
                    }
                }
            }
        };
        gather(q0, ast0, 0);
        gather(q1, ast1, 1);
#ifdef CMM_XMM
        noc_async_read_set_trid(0);
#else
        {
            CZONE("AGBAR");
            noc_async_read_barrier();
        }
#endif
#endif
        const uint32_t bpage = get_local_cb_interface(c_bias).fifo_page_size;
        const auto bias = TensorAccessor(b_args, b_addr, bpage);
#if defined(CMM_BIAS_LATE) && !defined(CMM_XMM)
        // round 11: the bias is only needed after the last K block -- its DRAM reads (slow while the weight senders
        // saturate DRAM at kernel start: ~8 us for qkv's 12 tiles) no longer gate the activation pipeline
        noc_async_read_barrier();  // own activation blocks gathered
#endif
        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;
        }
#ifndef CMM_BIAS_LATE
        noc_async_read_barrier();
        cb_push_back(c_bias, NT);
#endif
#ifdef CMM_NO_A  // timing experiment only (wrong results): no activation movement at all
        for (uint32_t j = 0; j < NB; ++j) {
            const uint32_t c = (j & 1) ? c_a1 : c_a0;
            cb_reserve_back(c, ABT);  // round 11: ABT (KB x MPC) -- KB here deadlocked the k2 compute (chip 3 fault)
            cb_push_back(c, ABT);
        }
#else
#ifdef CMM_GELU_HELP
        // round 8: the GELU helper cores write this core's own activation block into its shard and count the tiles on
        // semaphore CMM_GELU_SEM (zeroed by dispatch): wait for all KB before multicasting it (keep the ready signals
        // of both pipes flowing meanwhile)
        auto gelu_wait = [&]() {
            volatile tt_l1_ptr uint32_t* gs = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(CMM_GELU_SEM));
            while (true) {
                invalidate_l1_cache();
                if (*gs >= ABT) {
                    break;
                }
                q0.signal_free();
#ifndef CMM_A2N
                q1.signal_free();
#endif
            }
        };
#else
        auto gelu_wait = [&]() {};
#endif
#ifdef CMM_AEARLY
        // round 11: an owner multicasts its block as soon as it is ready (GELU done, every peer's slot and its own slot
        // free), not only after it has received all earlier blocks (the in-order loop made the blocks a hand-over
        // chain).  The data lands in ring slot j % S of every member; the in-order loop below only pushes it.
#ifdef CMM_GELU_HELP
        auto gelu_done = [&]() -> bool {
            invalidate_l1_cache();
            return *reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(CMM_GELU_SEM)) >= ABT;
        };
#else
        auto gelu_done = [&]() -> bool { return true; };
#endif
        uint32_t nx0 = q0.first_own, nx1 = q1.first_own;  // next own block to send per pipe
        auto try_send = [&]() {
            if (nx0 < NBP && q0.owner(nx0) && gelu_done() && q0.can_send(nx0)) {
                q0.send(ast0 + q0.own_index(nx0) * ABT * apage, q0.slot_addr(nx0), nx0, q1);
                ++nx0;
            }
            if (nx1 < NBP && q1.owner(nx1) && gelu_done() && q1.can_send(nx1)) {
                q1.send(ast1 + q1.own_index(nx1) * ABT * apage, q1.slot_addr(nx1), nx1, q0);
                ++nx1;
            }
        };
        auto step = [&](auto& q, auto& o, uint32_t j, uint32_t& nx) {
            q.signal_free();
            o.signal_free();
            if (q.owner(j)) {
                while (nx <= j) {
                    try_send();
                    q.signal_free();
                    o.signal_free();
                }
                cb_reserve_back(q.cb, ABT);
                cb_push_back(q.cb, ABT);
                return;
            }
            cb_reserve_back(q.cb, ABT);
            {
                CZONE("RECV");
                volatile tt_l1_ptr uint32_t* v = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(q.valid_addr(j));
                while (true) {
                    invalidate_l1_cache();
                    if (*v >= j + 1) {
                        break;
                    }
                    try_send();
                    q.signal_free();
                    o.signal_free();
                }
            }
            cb_push_back(q.cb, ABT);
        };
        for (uint32_t j = 0; j < NBP; ++j) {
            step(q0, q1, j, nx0);
            step(q1, q0, j, nx1);
        }
#else
#if defined(CMM_XMM) && !defined(CMM_NO_AGATHER)
        auto xwait = [&](uint32_t t) { noc_async_read_barrier_with_trid(t); };
#else
        auto xwait = [&](uint32_t) {};
#endif
#ifdef CMM_A2N
        NoOther none_;
        for (uint32_t j = 0; j < NBP; ++j) {
            if (q0.owner(j)) {
                gelu_wait();
                q0.step_from2(j, ast0 + q0.own_index(j) * ABT * apage, none_);
            } else {
                q0.recv_step2(j, none_);
            }
        }
        if (false)
#endif
        for (uint32_t j = 0; j < NBP; ++j) {
            if (q0.owner(j)) {
                gelu_wait();
                xwait(1 + 2 * j);
                q0.step_from2(j, ast0 + q0.own_index(j) * ABT * apage, q1);
            } else {
                q0.recv_step2(j, q1);
            }
            if (q1.owner(j)) {
                gelu_wait();
                xwait(2 + 2 * j);
                q1.step_from2(j, ast1 + q1.own_index(j) * ABT * apage, q0);
            } else {
                q1.recv_step2(j, q0);
            }
        }
#endif
#endif
#ifdef CMM_BIAS_LATE
        noc_async_read_barrier();
        cb_push_back(c_bias, NT);
#endif
#endif  // CMM_A_RES
#ifdef CMM_KSPLIT
        {
            // this core's partial of the PARTNER's token row (compute: c_send, NT fp32 tiles) -> the partner's c_int (same
            // L1 address on every core: nothing was pushed there yet), acked, then the partner's semaphore 4 S
            CZONE("XSEND");
#ifdef CMM_PF
            constexpr uint32_t XA = 1;
#else
            constexpr uint32_t XA = 0;
#endif
            const uint32_t px = get_arg_val<uint32_t>(wtab + NB + XA), py = get_arg_val<uint32_t>(wtab + NB + XA + 1);
            constexpr uint32_t c_int = 24;
#ifdef CMM_NSB
            // wide output (compute_k2w.cpp): the partner half is tiles NT .. 2 NT - 1 of c_part (never popped)
            constexpr uint32_t c_send = 26;
            const uint32_t xbytes = NT * get_local_cb_interface(c_send).fifo_page_size;
            cb_wait_front(c_send, 2 * NT);
            const uint32_t xsrc = get_read_ptr(c_send) + xbytes;
#else
            constexpr uint32_t c_send = 19;
            const uint32_t xbytes = NT * get_local_cb_interface(c_send).fifo_page_size;
            cb_wait_front(c_send, NT);
            const uint32_t xsrc = get_read_ptr(c_send);
#endif
            noc_async_write(xsrc, get_noc_addr(px, py, get_write_ptr(c_int)), xbytes);
#ifdef CMM_XFLAG
            // round 11: library ordering (same NOC, same VC, same destination): the data flushed out of this core, then a
            // plain 4-byte write of the flag (value 1) instead of an acked barrier + atomic increment
            noc_async_writes_flushed();
            {
                volatile tt_l1_ptr uint32_t* one = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(10) + 16 * (2 * S + 2));
                *one = 1;
                noc_async_write(get_write_ptr(10) + 16 * (2 * S + 2), get_noc_addr(px, py, get_semaphore(4 * S)), 4);
                noc_async_writes_flushed();
            }
#else
            noc_async_write_barrier();
            noc_semaphore_inc(get_noc_addr(px, py, get_semaphore(4 * S)), 1);
#endif
#ifndef CMM_NSB
            cb_pop_front(c_send, NT);
#endif
        }
#endif
#ifdef CMM_XMM
        {
            // MPC x NT tiles -> pages of the interleaved output; CMM_OUT_QKV: straight into the head-split q / k / v
            // ([1, QH, QMT * 32, QDT * 32] each, addresses in args 5..7) -- the head split folded into the matmul
            const uint32_t opage = get_local_cb_interface(c_out).fifo_page_size;
            cb_wait_front(c_out, MPC * NT);
            uint32_t rp = get_read_ptr(c_out);
            for (uint32_t i = 0; i < MPC; ++i) {
                for (uint32_t c = 0; c < NT; ++c) {
                    const uint32_t mr = m + i, n = n0 + c;
#ifdef CMM_OUT_QKV
                    constexpr uint32_t QW = CMM_QH * CMM_QDT;
                    const uint32_t which = n / QW, rem = n % QW;
                    const uint32_t pg = (rem / CMM_QDT) * (CMM_QMT * CMM_QDT) + mr * CMM_QDT + rem % CMM_QDT;
                    const auto o = TensorAccessor(o_args, get_arg_val<uint32_t>(5 + which), opage);
                    noc_async_write(rp, o.get_noc_addr(pg), opage);
#else
                    const auto o = TensorAccessor(o_args, get_arg_val<uint32_t>(5), opage);
                    noc_async_write(rp, o.get_noc_addr(mr * NTOT + n), opage);
#endif
                    rp += opage;
                }
            }
            noc_async_write_barrier();
            cb_pop_front(c_out, MPC * NT);
        }
#elif !defined(CMM_ADDLN)  // fused add+LN epilogue (kernels/cmm/fused_*.cpp): the output stays in c_out for the compute
        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);
#endif
#ifdef CMM_A_RES
        noc_async_write_barrier();
#else
        q0.finish();
#endif
    }
}