Foundry/tt/kernels/resident_matmul_writer.c

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/**
 * @file resident_matmul_writer.c
 * @brief Writer third of the RESIDENT matmul triad, and the processor that PUBLISHES
 * COMPLETION -- both per core (completed_count/completed_token) and per group.
 *
 * WHY THE WRITER OWNS COMPLETION. An op is not done when the reader has fetched its
 * operands or when the FPU has finished multiplying; it is done when the result has left
 * L1. Phase 5's rule is that completion is produced BY the work item, and in a pipeline
 * that means by the LAST stage of it. So this processor owns `completed_count` and
 * `completed_token`, the reader owns `tail`, and nothing else writes either.
 *
 * AND WHY IT FETCHES ON `head` BUT PUBLISHES ON `tail`. A descriptor with no writer-side
 * work -- WRITE_WORD, a semaphore post -- could otherwise be retired before the reader had
 * performed it, which is the "completion that is not evidence of work" this protocol
 * exists to prevent, one processor over. So the reader's half is a PRECONDITION of
 * publishing completion. It is deliberately NOT a precondition of starting work, and that
 * distinction cost a run: with the writer's CURSOR following `tail`, the output CB (two
 * pages) fills after two tiles, the compute thread blocks in `cb_reserve_back`, the reader
 * then blocks because nothing is draining the input CBs, and the reader can never reach
 * the end of the descriptor to bump `tail`. Three processors each waiting on the next.
 * The only reason that was a diagnosis instead of a dead core is that every one of those
 * waits is bounded and askable.
 *
 *
 * C99 on Foundry's own `ft_tensix_dataflow.h` — no tt-metal include (plan Phase 6.6).
 * The wait mode, which used to be a `template <uint32_t kMode>` argument, is now an
 * ordinary parameter carrying the same compile-time value; `-Os` folds it identically.
 * Compile-time args:
 *   0: CB_OUT index   1: output tile size in bytes
 *   2: wait mode -- 0 upstream (the cost arm), 1 stoppable (ships), 2 bound-only (the
 *      red proof for the ask)
 * Runtime args: the ten of resident_queue_protocol.h, identical on all five processors
 * Descriptor words: a0 the parameter block, a2 the DESTINATION ROW PITCH in tiles (0 means
 *   `tile_cols`, which is what every Mt = 1 descriptor computes anyway), a5 the group
 *   sequence. See resident_matmul_block_protocol.h for why the pitch is not a parameter word.
 *
 * Build (NCRISC, NoC 1):
 *   CTARGS=16,2048,1 NOC=1 KERNEL_TEXT_BASE=0x18000 bash \
 *       foundry/tt/build_tensix_dataflow_kernel.sh \
 *       foundry/tt/kernels/resident_matmul_writer.c ncrisc /tmp/ttk_rtriad
 *
 * @status See docs/tt_bringup_log.md.
 * @issues None
 * @todo None
 */

#include <stdbool.h>
#include <stdint.h>

#include "ft_tensix_dataflow.h"
#include "resident_matmul_block_protocol.h"
#include "resident_queue_protocol.h"

#define CB_OUT ft_tensix_compile_time_arg(0)
#define OUT_TILE_BYTES ft_tensix_compile_time_arg(1)
#define WAIT_MODE ft_tensix_compile_time_arg(2)
/**
 * The writer half of one matmul descriptor: drain Mt*Nt output tiles to the destination.
 *
 * vector_add_writer.c's loop, with a stoppable wait and the drain rule. The barrier is
 * still per tile, as it is there -- this kernel is about the dispatch model, and changing
 * the writer's barrier policy in the same round would put two variables in one measurement.
 *
 * THE TILES DO NOT ARRIVE IN DESTINATION ORDER ONCE THE ROW BLOCK IS MORE THAN ONE, so the
 * counter this loop used to be is gone. The compute thread packs a whole (row block x
 * column block) section before moving on, which is `MB` short runs of `NB` tiles inside one
 * output row-band -- so the writer re-walks the reader's own three loops and computes
 * `(m, n)` from them. `row_pitch` is the tile pitch of the DESTINATION grid, which is the
 * weight's FULL Nt and not this member's slice of it, and it arrives in descriptor word a2
 * because the sixteen parameter words are full and a seventeenth would move the shared slot
 * stride (resident_matmul_block_protocol.h).
 *
 * AT MB = 1 THIS IS THE OLD LOOP, TERM FOR TERM: the row loop runs once at `m = 0`, so the
 * pitch is multiplied by zero and cannot matter, the column blocks are walked in order and
 * `n` is the running counter the old code called `tile`.
 */
static bool write_one_matmul(ftq_context *ctx, volatile FTQ_L1_PTR uint32_t *parameters,
                             uint32_t row_pitch)
{
    const uint32_t tile_rows = parameters[kMatmulTileRows];
    const uint32_t tile_cols = parameters[kMatmulTileCols];
    const uint32_t block_word = parameters[kMatmulColumnBlock];
    const uint32_t column_block = ftmm_column_block(block_word);
    const uint32_t row_block = ftmm_row_block(block_word);
    const uint32_t destination_x = parameters[kMatmulCX];
    const uint32_t destination_y = parameters[kMatmulCY];
    const uint32_t destination_base = parameters[kMatmulCBase];
    /* 0 is what every descriptor written before the word existed says, and for those it is
     * also the only value that can be right: they are all Mt = 1. */
    const uint32_t pitch = row_pitch > 0u ? row_pitch : tile_cols;
    const uint32_t probe = parameters[kMatmulBNMajor];

    for (uint32_t mt = 0; mt < tile_rows; mt += row_block) {
        const uint32_t rows = (tile_rows - mt) < row_block ? (tile_rows - mt) : row_block;

        for (uint32_t nt = 0; nt < tile_cols; nt += column_block) {
            const uint32_t columns =
                (tile_cols - nt) < column_block ? (tile_cols - nt) : column_block;

            for (uint32_t row = 0; row < rows; row++) {
                for (uint32_t column = 0; column < columns; column++) {
                    const uint32_t tile = (mt + row) * pitch + nt + column;
                    const bool have_tile =
                        ftq_cb_wait_front_stoppable(ctx, (int32_t)CB_OUT, 1, WAIT_MODE);

                    if (have_tile && (probe & kMatmulProbeNoWrite) == 0u) {
                        ft_tensix_noc_write(
                            ft_tensix_cb_read_pointer(CB_OUT),
                            ft_tensix_noc_address(destination_x, destination_y,
                                                  destination_base + tile * OUT_TILE_BYTES),
                            OUT_TILE_BYTES);
                        /* kMatmulWriterFlush: the page may be popped once its write has LEFT
                         * L1; the ack the completion post needs is taken once, below. */
                        if ((probe & kMatmulWriterFlush) != 0u)
                            ft_tensix_noc_writes_flushed_barrier();
                        else
                            ft_tensix_noc_write_barrier();
                    }
                    ft_tensix_cb_pop_front(CB_OUT, 1);
                }
            }
        }
    }
    ft_tensix_noc_write_barrier();
    return !ctx->aborting;
}

/**
 * M2: wait until every member named by the pull table has finished group sequence `sequence`.
 * Reads only (16 B NoC read of the member's control bytes 0xE0..0xEF). Stoppable and bounded:
 * a stop request or the spin bound aborts the descriptor instead of hanging the core.
 */
static bool pull_group(ftq_context *ctx, uint32_t members, uint32_t sequence)
{
    const uint32_t control = (uint32_t)(uintptr_t)ctx->control;
    volatile FTQ_L1_PTR uint32_t *table =
        (volatile FTQ_L1_PTR uint32_t *)(uintptr_t)(control + kPullTableOffset);
    volatile FTQ_L1_PTR uint32_t *scratch =
        (volatile FTQ_L1_PTR uint32_t *)(uintptr_t)(control + kPullScratchOffset);
    uint32_t member;

    if (table[0] != kPullMagic || members > table[1] || members > kPullMaxMembers) {
        ftq_abort(ctx, kErrBadParameters);
        return false;
    }
    for (member = 1u; member < members; member++) {
        const uint32_t xy = table[2u + member];
        uint32_t spins = 0u;

        for (;;) {
            ft_tensix_noc_read(ft_tensix_noc_address(xy & 0xFFFFu, xy >> 16,
                                                     control + kFieldGroupDone * 4u),
                               (uint32_t)(uintptr_t)scratch, 16u);
            ft_tensix_noc_read_barrier();
            ft_bh_l1_cache_invalidate();
            if ((int32_t)(scratch[0] - sequence) >= 0)
                break;
            if ((++spins & 1023u) == 0u &&
                (ftq_stop_requested(ctx) || spins >= ctx->cb_spin_limit)) {
                ftq_abort(ctx, kErrSemWaitTimeout);
                return false;
            }
        }
    }
    return true;
}

void kernel_main(void)
{
    uint32_t busy_started = 0u;
    ftq_context ctx;
    uint32_t completed;
    bool running = true;

    ftq_context_init(&ctx);
    ftq_status_set(&ctx, kStateBoot);

    if (!ftq_check_control(&ctx)) {
        ftq_status_set(&ctx, kStateExit);
        return;
    }

    // All five processors resume from completed_count, the ONE number that means "every
    // stage finished this descriptor". See resident_matmul_reader.c on why a relaunch is
    // only defined over a queue whose head, tail and completed_count agree.
    completed = ctx.control[kFieldCompletedCount];
    ctx.cursor = completed;
    ctx.control[12] = 0u;
    ctx.control[13] = 0u;
    ctx.control[14] = 0u;
    ctx.control[15] = 0u;
    ctx.control[61] = 0u;

    while (running) {
        // FETCH on `head`, so the writer drains the output CB while the reader is still
        // filling the input ones. Publishing waits on `tail`, below.
        if (!ftq_wait_for_work(&ctx, kFieldHead))
            break;

        ctx.aborting = false;
        ftq_status_set(&ctx, kStateExec);
        /* M1 CARD-BUSY CLOCK: this processor alone writes control words 12-15 (12 last end
         * stamp, 13 busy cycles >> 4, 14 idle cycles >> 4 between descriptors, 15 count).
         * Each delta is one descriptor or one gap, so the 32-bit wrap never bites. */
        {
            const uint32_t now = ft_bh_wall_clock_low();

            if (ctx.control[15] != 0u)
                ctx.control[14] += (now - ctx.control[12]) >> 4;
            busy_started = now;
        }

        volatile FTQ_L1_PTR uint32_t *descriptor = ftq_descriptor(&ctx);
        const uint32_t opcode = descriptor[kDescOpcode];
        const uint32_t token = descriptor[kDescToken];
        const uint32_t group_sequence = descriptor[kDescA5];
        bool done = true;

        ftq_status_opcode(&ctx, opcode);
        if (opcode == kOpMatmul) {
            ftq_status_set(&ctx, kStateCbWaitFront);
            /* a2 is the DESTINATION ROW PITCH in tiles; 0 means `tile_cols`. It is a
             * descriptor word rather than a parameter word because the parameter block is
             * full at sixteen and every resident family shares that slot stride. */
            done = write_one_matmul(&ctx, ftq_parameters(descriptor), descriptor[kDescA2]);
            ftq_status_set(&ctx, kStateExec);
        } else if (opcode == kOpStop) {
            running = false;
        }

        if (!done) {
            running = false;
            continue;
        }

        // THE PRECONDITION, and the only place it belongs: the reader must have finished
        // this descriptor before its completion is published. For a matmul that is already
        // true by construction — the last output tile cannot exist until the last input pair
        // was read — so this costs nothing there. For an opcode the writer does not
        // participate in, it is the whole of the guarantee.
        if (!ftq_wait_until_reaches(&ctx, kFieldTail, ctx.cursor + 1u)) {
            running = false;
            continue;
        }

        ctx.cursor++;
        // Token before count: a host that sees the count is entitled to trust the token it
        // reads next. Both are plain stores from ONE in-order RISC-V, which is the only
        // ordering argument this protocol makes and the only one it needs -- every other
        // control word has a different single writer (ft_tt_queue.h, THE ORDERING
        // ASSUMPTION).
        {
            const uint32_t pull = descriptor[kDescA3];

            /* a4 = the sequence every pulled member must have reached: this descriptor's
             * own for a split projection, the chain's last split one for a solo projection
             * that follows it on member 0. */
            const uint32_t pull_started = ft_bh_wall_clock_low();

            if (opcode == kOpMatmul && (pull & ~kPullRedSkip) > 1u &&
                (pull & kPullRedSkip) == 0u &&
                !pull_group(&ctx, pull & ~kPullRedSkip, descriptor[kDescA4])) {
                running = false;
                continue;
            }
            /* M8: word 61 = cycles >> 4 member 0 spent pulling (waiting on the others). */
            if (opcode == kOpMatmul && (pull & ~kPullRedSkip) > 1u)
                ctx.control[61] += (ft_bh_wall_clock_low() - pull_started) >> 4;
        }
        ctx.control[kFieldGroupDone] = group_sequence;
        {
            const uint32_t now = ft_bh_wall_clock_low();

            ctx.control[13] += (now - busy_started) >> 4;
            ctx.control[12] = now;
            ctx.control[15] += 1u;
        }
        ctx.control[kFieldCompletedToken] = token;
        ctx.control[kFieldCompletedCount] = ++completed;
        ftq_group_post(&ctx, group_sequence, token);
    }

    // The writer does NOT touch kFieldError or kFieldState: the reader owns those, and one
    // writer per word is the whole ordering argument. This processor's own reason for
    // leaving is in its status slot, which is where a five-processor diagnosis belongs.
    ftq_status_set(&ctx, kStateExit);
}