matmul — Matrix Multiply on the VPE

A GEMM C = A · B (A is M×K, B is K×N) is M·N dot products of length K: C[i][j] = dot(row i of A, column j of B). The VPE is a vector reduction engine, so dot()/dot2() express that directly, with one row and one or two columns in SRAM while the matrices stay in DRAM. B is supplied transposed (Bt, shape N×K) so column j of B is the contiguous row j of Bt.

The point of matmul/main.cpp is that the context binds no dimension, so one kernel can compute a batch of differently-sized matrices without touching its VpeContext.

void matmulRows(mu::vdma::VpeContext& ctx, const float* a, const float* bt, float* c,
                uint32_t n, uint32_t k, uint32_t rowStart, uint32_t rowEnd,
                uint64_t aBuf, uint64_t b0Buf, uint64_t b1Buf)
{
    for (uint32_t i = rowStart; i < rowEnd; ++i)
    {
        ctx.load(a + static_cast<uint64_t>(i) * k, aBuf, k);  // A row i, shared operand
        uint32_t j = 0;
        for (; j + 1 < n; j += 2)  // two columns per step via dot2
        {
            ctx.load(bt + static_cast<uint64_t>(j) * k, b0Buf, k);
            ctx.load(bt + static_cast<uint64_t>(j + 1) * k, b1Buf, k);
            auto r = ctx.dot2(aBuf, b0Buf, b1Buf, k);
            c[static_cast<uint64_t>(i) * n + j] = r.r1;
            c[static_cast<uint64_t>(i) * n + j + 1] = r.r2;
        }
        if (j < n)  // odd trailing column
        {
            ctx.load(bt + static_cast<uint64_t>(j) * k, b0Buf, k);
            c[static_cast<uint64_t>(i) * n + j] = ctx.dot(aBuf, b0Buf, k);
        }
    }
}

All four kernels are built on that loop:

Kernel	What it adds
`matmul`	One matrix, output rows split across tasks (as in knn).
`matmulBatch`	A list of square matrices of differing sizes. Buffers are allocated once, sized for the largest matrix; each matmul then runs at its own per-call dimension, so a new size costs one more loop iteration, not a new context. One task, so one VPE.
`matmulBatchParallel`	Flattens the batch into `(matrix, row-block)` units and takes every `taskCount`-th one. The enumeration is deterministic on every task, so ownership needs no shared state and no locking.
`matmulBatchParallelTiled`	Keeps `ROWTILE` A rows resident and reuses each loaded B column across all of them.

Task count

The compute-parallel width is the VPE count, not the thread count: a subsystem has 128 threads but 2 VPEs (24 SUBs → 48 VPEs), and every dot/dot2 serialises on its VPE. So raise TASKCOUNT toward 48 and dispatch time falls; past that there is nothing left to occupy. Each dispatch also pays a fixed submit + synchronize cost, so raise REPS to enlarge the batch and dilute it.

Round-robin over work units buys lock-free ownership, not an even schedule: because the batch cycles a fixed size list, a task can keep drawing the same slot. Weight units by dim² if you need real balance.

Memory-bound vs compute-bound

matmulRows streams all of B past one resident A row, so it re-reads B once per A row: an arithmetic intensity of only ~0.5 FLOP/byte, which means throughput tracks DRAM bandwidth, not the VPEs. The tiled kernel reads B once per tile of A rows instead, lifting intensity to ~0.5 · ROWTILE:

// One B column, loaded once, reduced against every resident A row of the tile.
for (uint32_t j = 0; j < n; ++j)
{
    ctx.load(bt + static_cast<uint64_t>(j) * k, bColBuf, k);   // B column j: one load
    for (uint32_t r = 0; r + 1 < rows; r += 2)                 // reuse across the tile
    {
        auto res = ctx.dot2(bColBuf, aBufs[r], aBufs[r + 1], k);
        c[(rowStart + r) * n + j]     = res.r1;
        c[(rowStart + r + 1) * n + j] = res.r2;
    }
}

The trade-off is SRAM: a bigger tile holds more buffers per task, so fewer tasks fit per VPE, and concurrency is what hides load latency. The kernel clamps ROWTILE to what the pool can hold and prints the value used.

Both parallel tests print a DRAM GB/s figure next to GFLOP/s (the minimum traffic the access pattern demands over the dispatch time, not a measured bandwidth), and PEAKBW=<GB/s> turns that into a percentage of peak. Near peak means memory-bound, so raise ROWTILE; well below peak with GFLOP/s no longer rising means compute-bound at the VPE ceiling. The example runs the plain and tiled kernels back to back so the shift shows in one run.

Allocation (bytes) vs per-call dimension (elements)

Footprint (SRAM held, hence how many tasks fit on a VPE) follows the bytes you allocate.
Work (elements processed) follows the per-call dimension you pass to each op.

A smaller dimension does not shrink the footprint or free a VPE for another task; the buffers still occupy what you allocated. To fit more concurrent tasks, allocate fewer bytes. Only vectors ever live in SRAM: even the largest matmul in the batch holds a few K-length rows, never the matrices themselves.

./bin/xarith_matmul
REPS=24 TASKCOUNT=384 SIZES=256 ./bin/xarith_matmul

Environment Variables:

Variable	Default	Description
`DEVICEID`	0	Device ID
`NUMSUB`	24	Number of subsystems
`TASKCOUNT`	48	Number of subtasks (raised to `NUMSUB` if lower)
`REPS`	4	Times `SIZES` is cycled into the batch
`SIZES`	`64,128,48,256,96,192`	Square matrix sizes making up the batch
`ROWTILE`	8	A rows resident in the tiled kernel
`USEDOT`	0	`1` reduces the tiled kernel with single `dot()` instead of fused `dot2()`
`PEAKBW`	(unset)	Device peak DRAM bandwidth in GB/s, to print achieved bandwidth as a percentage

matmul/sweep.sh drives run.sh matmul across these knobs and tabulates the throughput lines, for finding the ROWTILE/TASKCOUNT sweet spot without running each point by hand.