From f94d1697aeb41430259ba0d2f12aa8ee853cf285 Mon Sep 17 00:00:00 2001 From: andrej Date: Fri, 11 Sep 2026 12:02:36 -0600 Subject: [PATCH 01/14] flm.GEMM: take M, K, N and the activation as runtime parameters The FastFlowLM harness registers one mm.xclbin per model and swaps instruction streams, against a budget of 16 xclbins for the whole model. This operator baked M, K and N into the core loop bounds, into which columns it built, and into the memtile's B buffer, so it needed one xclbin per shape: 11 (K, N, activation) combinations times up to 16 chunk lengths. Five values move into an L1 buffer the runtime sequence writes and each core reads once its barrier opens: the column's work and drain counts, M/256, K/512, and the activation. All columns are now always built, and one with no work for a shape drains its share of the A broadcast instead. The epilogue tests its mode once per chunk, outside the vector loop, so each mode keeps a branch-free inner loop; which modes it can select between stays a build-time choice, since each costs program memory. The device body is then a function of the tiling alone, verified byte for byte across every E2B shape and activation. So the xclbin is built from a module emitted at a reference shape and the per-shape build produces only the instruction stream, and the two carry different artifact stems. The core releases its barrier straight after reading the parameters. wait_for_value emits LockAction.Acquire, which does not leave the lock consumed, so without the release a core that runs twice does not wait the second time and reads the previous dispatch's parameters. Releasing before the work is safe because the sequence cannot set the barrier again until it has drained this dispatch's C. The repo's other barrier users never wait twice -- mha puts its infinite loop inside the wait, softmax writes the same parameters every dispatch -- so this does not arise there. test_one_xclbin_serves_every_shape is the regression test; the parametrised tests cannot catch it, because the aie_context fixture reconfigures the array between cases. Dropping B residency is what this costs, and it is not cheap: 12.5% at M=512, 16.4% at M=1024 and 19.3% at M=2048 on NPU2. Residency sizes the memtile buffer from K and replays it M/256 times through a buffer descriptor's repeat count, so it carries both K and M into the configuration. Restoring it needs a replay mechanism that carries neither. It does lift a cap: the repeat count expands into the memtile's BD chain at 2 blocks per replay and exceeded its 48-block limit at M=4096, so no shape with K <= 2048 would build there -- 7 of Gemma4 E2B's 10 projections. All of them build now. Verified on this base over all 12 distinct E2B prefill projections at M=256 and M=4096, dispatched back to back on one loaded xclbin with no reset in between, each checked against a CPU reference; peak 12.2 TFLOP/s at M=4096 K=12288 N=1536. Plus the operator's own suite, 17 non-extensive tests. --- aie_kernels/generic/mm_fused.cc | 91 +++++++++----- iron/operators/flm/gemm/README.md | 100 ++++++++++----- iron/operators/flm/gemm/design.py | 200 +++++++++++++++++------------- iron/operators/flm/gemm/op.py | 115 ++++++++++++++--- iron/operators/flm/gemm/test.py | 52 ++++++++ 5 files changed, 395 insertions(+), 163 deletions(-) diff --git a/aie_kernels/generic/mm_fused.cc b/aie_kernels/generic/mm_fused.cc index 719c09361..7845a3208 100644 --- a/aie_kernels/generic/mm_fused.cc +++ b/aie_kernels/generic/mm_fused.cc @@ -35,8 +35,8 @@ // Epilogue selection. 0 = none, 1 = gelu, 2 = silu, 3 = sigmoid, matching // Epilogue.mode in design.py. -#ifndef MM_FUSED_EPILOGUE_MODE -#define MM_FUSED_EPILOGUE_MODE 0 +#ifndef MM_FUSED_EPILOGUE_MODE_MASK +#define MM_FUSED_EPILOGUE_MODE_MASK 0xF #endif #ifndef MM_FUSED_CLAMP #define MM_FUSED_CLAMP 0 @@ -105,6 +105,39 @@ constexpr aie::rounding_mode round_mode = aie::rounding_mode::conv_even; #else constexpr aie::rounding_mode round_mode = aie::rounding_mode::floor; #endif +// One activation's inner loop. Templated so each mode compiles branch-free; +// mm_fused_epilogue_chunk selects between them once per chunk. +template static inline void epilogue_body(bfloat16 *__restrict y_out, const float *__restrict src) +{ +#if MM_FUSED_CLAMP + const aie::vector lo = aie::broadcast(MM_FUSED_CLAMP_MIN); + const aie::vector hi = aie::broadcast(MM_FUSED_CLAMP_MAX); +#endif + + AIE_LOOP_MAX_ITERATION_COUNT(CHUNK / V) + for (int j = 0; j < CHUNK / V; j++) { + // The accumulator stays f32 through the activation and the clamp, and + // is converted to bf16 exactly once, on the store. Converting first + // would round twice and let the activation's slope amplify the first + // rounding -- see activations.h. + aie::vector f = aie::load_v(src + j * V); + if constexpr (MODE == 1) + f = gelu_vec(f); + else if constexpr (MODE == 2) + f = silu_vec(f); + else if constexpr (MODE == 3) + f = sigmoid_vec(f); +#if MM_FUSED_CLAMP + f = aie::max(aie::min(f, hi), lo); +#endif + aie::accum out; + out.from_vector(f); + // The assignment is the conversion: to_v16bfloat16 yields a raw + // v16bfloat16, not an aie::vector. + aie::vector v = to_v16bfloat16(out); + aie::store_v(y_out + j * V, v); + } +} } // namespace extern "C" { @@ -144,47 +177,45 @@ void mm_fused_k_step(bfloat16 *a_buf, mm_fused_b_elem_t *b_buf, float *y_acc, in // Fusing the activation here is the point: the values are already in registers // after the f32 -> bf16 conversion, so gelu/silu/sigmoid costs one more vector // op per 16 elements instead of a separate pass over L1 (which is what chaining -// a standalone activation operator after a GEMM would cost). The mode and clamp -// are compile-time, so the inner loop below is branch-free. +// a standalone activation operator after a GEMM would cost). +// +// The mode is a runtime argument, because one xclbin serves every activation. +// It is tested once per chunk, outside the vector loop, so each mode still runs +// a branch-free inner loop; the cost is program memory, since every mode in +// MM_FUSED_EPILOGUE_MODE_MASK is compiled in. The clamp stays compile-time. // // The chunk index is split in two because the core body unrolls the drain by // the C fifo depth to keep the acquired buffer index a compile-time constant; // passing both parts avoids doing that arithmetic up there. -void mm_fused_epilogue_chunk(bfloat16 *y_out, float *y_acc, int32_t outer, int32_t half) +void mm_fused_epilogue_chunk(bfloat16 *y_out, float *y_acc, int32_t outer, int32_t half, int32_t mode) { // The store below is a conversion, so it obeys the same rounding mode the // mmul does and must agree with it. ::aie::set_rounding(round_mode); const float *__restrict src = y_acc + (outer * C_DEPTH + half) * CHUNK; -#if MM_FUSED_CLAMP - const aie::vector lo = aie::broadcast(MM_FUSED_CLAMP_MIN); - const aie::vector hi = aie::broadcast(MM_FUSED_CLAMP_MAX); + switch (mode) { +#if MM_FUSED_EPILOGUE_MODE_MASK & 2 + case 1: + epilogue_body<1>(y_out, src); + return; #endif - - AIE_LOOP_MAX_ITERATION_COUNT(CHUNK / V) - for (int j = 0; j < CHUNK / V; j++) { - // The accumulator stays f32 through the activation and the clamp, and - // is converted to bf16 exactly once, on the store. Converting first - // would round twice and let the activation's slope amplify the first - // rounding -- see activations.h. - aie::vector f = aie::load_v(src + j * V); -#if MM_FUSED_EPILOGUE_MODE == 1 - f = gelu_vec(f); -#elif MM_FUSED_EPILOGUE_MODE == 2 - f = silu_vec(f); -#elif MM_FUSED_EPILOGUE_MODE == 3 - f = sigmoid_vec(f); +#if MM_FUSED_EPILOGUE_MODE_MASK & 4 + case 2: + epilogue_body<2>(y_out, src); + return; #endif -#if MM_FUSED_CLAMP - f = aie::max(aie::min(f, hi), lo); +#if MM_FUSED_EPILOGUE_MODE_MASK & 8 + case 3: + epilogue_body<3>(y_out, src); + return; #endif - aie::accum out; - out.from_vector(f); - // The assignment is the conversion: to_v16bfloat16 yields a raw - // v16bfloat16, not an aie::vector. - aie::vector v = to_v16bfloat16(out); - aie::store_v(y_out + j * V, v); + // Mode 0 is always compiled: it is the fallback for a mode the mask leaves + // out, so an unselectable mode yields an unactivated result rather than an + // unwritten buffer. + default: + epilogue_body<0>(y_out, src); + return; } } } diff --git a/iron/operators/flm/gemm/README.md b/iron/operators/flm/gemm/README.md index e185873f7..0ff196f48 100644 --- a/iron/operators/flm/gemm/README.md +++ b/iron/operators/flm/gemm/README.md @@ -21,6 +21,12 @@ A second GEMM implementation alongside [`iron.operators.GEMM`](../../gemm), specialised for transformer projection shapes and ported from FastFlowLM's `mm` overlay. +**M, K, N and the activation are runtime parameters**, so one xclbin serves +every shape and only the instruction stream is rebuilt per shape. The +FastFlowLM harness needs that: it registers one `mm.xclbin` per model and swaps +instruction streams, against a budget of 16 xclbins for a whole model. See +[Runtime parameters](#runtime-parameters). + The overall dataflow is the same whole-array shape as `iron.operators.GEMM`'s — A broadcast along each compute row, B down each column, C joined through the memtile — so those are *not* what distinguishes it. What does: @@ -81,6 +87,46 @@ Two consequences of the native-vs-emulated split are worth knowing: NPU2 only.** NPU1 sums the K reduction in a different order, so it matches the rounding *mode* but not the exact results. +## Runtime parameters + +Five words in an L1 buffer per core, written by the runtime sequence and read +by the core once its barrier opens: + +| word | value | +|---|---| +| `n_work` | column-blocks this column computes | +| `n_drain` | column-blocks it sits out while still draining A | +| `m_row_blocks` | `M / 256` | +| `k_iters` | `K / 512` | +| epilogue | the `Epilogue` mode | + +All columns are always built. One with no work for a shape gets `n_work = 0` +and still drains its share of the A broadcast, because the memtile will not +release an A object until every consumer has taken it. + +The two artifacts therefore carry different stems: the xclbin's `config_name` +covers tile_n, tile_ma, the compiled activation set, clamp, rounding and the +device, while `name` adds M, K, N and the activation. The xclbin is built from +a module emitted at a reference shape, whose runtime sequence is discarded. + +Which activations the epilogue can *select between* stays a build-time choice, +since each one compiled in costs program memory; `epilogue_modes` sets it and +lands in the xclbin's name. + +The core releases its barrier straight after reading the parameters. +`wait_for_value` emits `LockAction.Acquire`, which does not leave the lock +consumed, so without the release a core that runs twice does not wait the +second time and reads the previous dispatch's parameters. Releasing before the +work is safe, because the sequence cannot set the barrier again until it has +drained this dispatch's C. + +The repo's other barrier users never hit this, because neither waits twice: +`mha` puts its infinite loop *inside* the wait, and `softmax` writes the same +parameters every dispatch. `test_one_xclbin_serves_every_shape` is the +regression test -- without the release it hangs the device on the second +shape, and the parametrised tests cannot catch it, because the `aie_context` +fixture reconfigures the array between cases. + ## Shape constraints `M % 256 == 0`, `K % 512 == 0`, `N % tile_n == 0` (so 64 by default). @@ -201,7 +247,7 @@ to reduce over -- with a single k iteration there is not enough compute to hide the extra A traffic. On NPU2, `tile_n=128` wins only at `k_iters=1`, and by ~3%; at `k_iters>=2` it -is 1.2-1.7x slower. Both follow from `tile_n=128` giving up resident B — its +is 1.2-1.7x slower. Both followed from `tile_n=128` giving up resident B — its `mt_b` is 128 KB, so `k_iters` copies do not fit the memtile — and the more k there is to reduce over, the more that costs. @@ -279,7 +325,7 @@ which is why the gap grows with the problem size rather than being flat. Unlike NPU2, compute here is **not** hidden behind the transfers, so on NPU1 both a faster mmul and less traffic pay off, where on NPU2 only the latter does. -Resident B is also a no-op on NPU1 — see [Resident B](#resident-b). +Resident B was also a no-op on NPU1 — see [Resident B, removed](#resident-b-removed). ### Why the transfers are cheap @@ -295,29 +341,27 @@ Two things, both in the runtime sequence rather than the kernel: block then costs 3 shim buffer descriptors instead of `1 + 2*k_iters`, so two can be in flight without exhausting the 16 available. -### Resident B - -Where a whole column-block's B fits in the memtile double-buffered -(`k_iters <= 2`, i.e. K <= 1024 at `tile_n=64`) it is held there and replayed -per row-block, so DDR reads it once instead of `m_row_blocks` times -- about -43% less traffic. Larger K falls back to re-reading it, unchanged. - -On NPU2 this is a latency win as well as a power one, because there the -operator is close to DDR-bandwidth bound, and it grows with the height of the -problem since B's re-reads scale with `m_row_blocks`. At K=1024 N=4096, min of -per-round medians over 10 interleaved rounds of 20 dispatches, `npu_time`, -power mode `turbo`: - -| M | row-blocks | non-resident | resident | | -|---|---|---|---|---| -| 512 | 2 | 527.0 us | **461.2 us** | 12.5% | -| 1024 | 4 | 1025.8 us | **857.3 us** | 16.4% | -| 2048 | 8 | 1958.0 us | **1579.5 us** | 19.3% | - -Most of the available win is still on the table: `repeat_count` restarts the -memtile BD chain at every replay boundary, which costs part of the traffic -saving back. Closing that is the largest known remaining lever here. - -On NPU1 residency is neither a latency nor a power win — measured off-versus-on -at M=2048 it is at best a no-op and marginally negative at K=1024, well inside -the 1.4-4.4% round spread. +### Resident B, removed + +B used to be held in the memtile across row-blocks where a whole column-block +fit double-buffered, so DDR read it once instead of `m_row_blocks` times -- +about 43% less traffic. On NPU2 that was worth 12.5% at M=512, 16.4% at +M=1024 and 19.3% at M=2048 (K=1024 N=4096). On NPU1 it was neither a latency +nor a power win. + +**It is gone, and that is the price of the runtime parameters.** Residency +sizes the memtile buffer from `k_iters` and replays it `m_row_blocks` times +through a buffer descriptor's repeat count, so it puts both K and M in the +device configuration — and the configuration is what one xclbin has to share +across every shape. + +Removing it does lift a hard cap: the repeat count expands into the memtile's +BD chain at 2 blocks per replay, and at `m_row_blocks = 16` that chain +exceeded its 48-block limit, so no shape with K <= 2048 would build at M=4096. +Every shape builds there now. + +Restoring it needs a replay mechanism that carries neither K nor M into the +configuration. That is the largest known lever left here, and it is worth more +than the figures above suggest, because `repeat_count` restarting the memtile +BD chain at every replay boundary was already giving part of the traffic +saving back. diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index 864e55f62..baadf7a0e 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -71,6 +71,7 @@ Runtime, TaskGroup, Worker, + WorkerRuntimeBarrier, ) from aie.iron.controlflow import range_ from aie.dialects.aie import get_target_model @@ -145,6 +146,11 @@ def mode(self) -> int: return list(Epilogue).index(self) +# The runtime parameter buffer each core reads once its barrier opens. +RTP_N_WORK, RTP_N_DRAIN, RTP_M_ROW_BLOCKS, RTP_K_ITERS, RTP_EPILOGUE = range(5) +RTP_WORDS = 5 + + class Rounding(StrEnum): """Rounding for every f32->bf16 conversion. @@ -410,21 +416,18 @@ def gemm( # rem_blocks columns (0 <= rem_blocks < COLS) that do one block more. n_full = N // MIN_N rem_blocks = (N % MIN_N) // N_TILE - # Columns that participate at all. With no full sweep (N below the grid's - # COLS*N_TILE stride) only the first rem_blocks columns do, and the rest - # are not instantiated -- giving them fifos that nothing ever drains builds - # dead dataflow, which newer mlir-aie rejects outright with - # "objectfifo.pool op segment 0 has no drainer". - n_active_cols = COLS if n_full else rem_blocks - # Per column: how many column-blocks it computes, and whether it sits out a - # trailing one while still draining the A broadcast for its row. That drain - # only arises for a column that exists and skips the trailing block, which - # requires at least one full sweep. - col_work = [n_full + (1 if c < rem_blocks else 0) for c in range(n_active_cols)] - col_drain = [ - 1 if (rem_blocks and n_full and c >= rem_blocks) else 0 - for c in range(n_active_cols) - ] + # Every column is instantiated for every shape, because which columns + # exist is configuration and this design has one configuration. A column + # with no work for this shape gets n_work = 0 and drains instead. + n_active_cols = COLS + # Per column: how many column-blocks it computes, and how many it sits out + # while still draining the A broadcast for its row. Every block issues A + # for every compute row, so a column that does not compute a block must + # still take that block's A or the row stalls -- the memtile will not + # release an A object until all COLS consumers have taken it. + total_blocks = n_full + (1 if rem_blocks else 0) + col_work = [n_full + (1 if c < rem_blocks else 0) for c in range(COLS)] + col_drain = [total_blocks - w for w in col_work] # B's element type: one v8bfp16ebs8 per 8 values on AIE2P, one bf16 per # value on AIE2. Every B extent below is therefore in values // B_GROUP. @@ -460,7 +463,7 @@ def gemm( epilogue_chunk = Kernel( EPILOGUE_SYMBOL, kernel_object, - [ct_out_ty, ct_acc_ty, np.int32, np.int32], + [ct_out_ty, ct_acc_ty, np.int32, np.int32, np.int32], ) # --- Data movement ---------------------------------------------------- @@ -554,25 +557,16 @@ def gemm( # has zero slack -- the eight memtiles pack to exactly 512 KB, relying on # aie-objectfifo-allocate spilling one buffer to an adjacent tile -- so # re-verify it after any change to the A, B or C buffer sizes. - mt_free = tm.get_mem_tile_size() - mt_a_bytes * A_DEPTH - mt_out_bytes * C_DEPTH - MT_B_DEPTH = next( - (d for d in (B_DEPTH, 1) if k_iters * mt_b_bytes * d <= mt_free), 0 - ) - b_resident = MT_B_DEPTH > 0 - if b_resident: - # Just a bigger buffer. With B packed in consumption order the walk is - # linear, so spanning every k-block needs no extra descriptor - # dimension -- the objects simply come out in k order. (The previous - # blocked layout had to widen one dim inbound and add an outermost k - # dim outbound, which is what collided with CT_MAX_K=128.) - mt_b_ty = np.ndarray[(k_iters * K_TILE * N_TILE // B_GROUP,), b_elem_ty] + # B residency is gone: it sizes the memtile buffer from k_iters and + # replays it m_row_blocks times through the forward()'s repeat_count, so + # it carries both K and M into the configuration, which is what the + # runtime parameters exist to remove. See README.md for what that cost. + b_resident = False b_l3l2_fifos = [] b_cons = {} for c in range(n_active_cols): - of_b_in = ObjectFifo( - mt_b_ty, name=f"B_L3L2_{c}", depth=MT_B_DEPTH if b_resident else B_DEPTH - ) + of_b_in = ObjectFifo(mt_b_ty, name=f"B_L3L2_{c}", depth=B_DEPTH) b_l3l2_fifos.append(of_b_in) of_b = of_b_in.cons(dims_from_stream=b_recv_dims).forward( # The one placement pin this design keeps. Everything else -- the @@ -591,69 +585,81 @@ def gemm( depth=L1_B_DEPTH, name=f"B_L2L1_{c}", dims_to_stream=b_send_dims, - # Replay the resident memtile object once per row-block. This is - # the mechanism that actually re-sends an object; iter_count only - # bounds how many chain iterations happen in total. Correct - # ordering depends on the memtile holding ONE object spanning every - # k-block: replicating a pool of k_iters smaller objects would - # emit k0,k0,k1,k1,... rather than the k0..kn sequence the cores - # accumulate in. - repeat_count=m_row_blocks if b_resident else None, ) for r in range(ROWS): b_cons[(r, c)] = of_b.cons() - # --- Compute ---------------------------------------------------------- - def core_fn(n_work, n_drain, acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): - """Core body for a column that computes ``n_work`` column-blocks and - then drains A for ``n_drain`` more (0 or 1). + # --- Runtime parameters ----------------------------------------------- + rtps = [ + [ + Buffer( + np.ndarray[(RTP_WORDS,), np.dtype[np.int32]], + name=f"rtp_{r}_{c}", + initial_value=np.zeros(RTP_WORDS, dtype=np.int32), + use_write_rtp=True, + ) + for c in range(n_active_cols) + ] + for r in range(ROWS) + ] + barriers = [ + [WorkerRuntimeBarrier() for _ in range(n_active_cols)] for _ in range(ROWS) + ] - n_work/n_drain are bound per column via functools.partial below; they - are compile-time constants, so the trip counts below fold away. - """ + # --- Compute ---------------------------------------------------------- + def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, barrier): + """Core body. Every trip count and the activation come from the + runtime parameter buffer, so one core program serves every shape.""" # The loop nest lives here rather than inside the kernel so that - # every level has an ObjectFifo acquire point. With M/K/N known at - # compile time all the trip counts are constants. - if n_work: - for _ in range_(n_work): - for _ in range_(m_row_blocks): - init_k(acc) - for _ in range_(k_iters): - # The l loop is unrolled by the B fifo depth so the - # acquired buffer index stays a compile-time - # constant. - for _ in range_(B_ITERS // B_DEPTH): - for _ in range(B_DEPTH): - # One B chunk feeds every A band, so B is - # acquired once around the band loop. - b = b_h.acquire(1) - for band in range(RHO): - a = a_h.acquire(1) - kstep_k(a, b, acc, band) - a_h.release(1) - b_h.release(1) - # Drain the accumulator. Unrolled by C_DEPTH for the - # same reason; a full O_CHUNKS unroll overflows program - # memory. - for chunk in range_(O_CHUNKS // C_DEPTH): - for half in range(C_DEPTH): - o = o_h.acquire(1) - epi_k(o, acc, chunk, half) - o_h.release(1) - if n_drain: - # The trailing partial column-block, for a column that sits it - # out. A is broadcast along the whole compute row, so this - # column must still consume its share or the columns that DO - # have work stall waiting for the fifo to advance. No B and no - # C here -- the runtime sequence issues neither for it. - for _ in range_(n_drain): - for _ in range_(m_row_blocks): - for _ in range_(k_iters): - for _ in range_(B_ITERS // B_DEPTH): - for _ in range(B_DEPTH): - for _ in range(RHO): - a_h.acquire(1) - a_h.release(1) + # every level has an ObjectFifo acquire point. + barrier.wait_for_value(1) + n_work = my_rtp[RTP_N_WORK] + n_drain = my_rtp[RTP_N_DRAIN] + n_row_blocks = my_rtp[RTP_M_ROW_BLOCKS] + n_k_iters = my_rtp[RTP_K_ITERS] + epi_mode = my_rtp[RTP_EPILOGUE] + # Acquire does not consume the barrier, so take it back to zero or the + # next dispatch reads these parameters again instead of waiting. Safe + # before the work: the sequence cannot set the barrier again until it + # has drained this dispatch's C. + barrier.release_with_value(1) + + for _ in range_(n_work): + for _ in range_(n_row_blocks): + init_k(acc) + for _ in range_(n_k_iters): + # The l loop is unrolled by the B fifo depth so the + # acquired buffer index stays a compile-time constant. + for _ in range_(B_ITERS // B_DEPTH): + for _ in range(B_DEPTH): + # One B chunk feeds every A band, so B is + # acquired once around the band loop. + b = b_h.acquire(1) + for band in range(RHO): + a = a_h.acquire(1) + kstep_k(a, b, acc, band) + a_h.release(1) + b_h.release(1) + # Drain the accumulator. Unrolled by C_DEPTH for the same + # reason; a full O_CHUNKS unroll overflows program memory. + for chunk in range_(O_CHUNKS // C_DEPTH): + for half in range(C_DEPTH): + o = o_h.acquire(1) + epi_k(o, acc, chunk, half, epi_mode) + o_h.release(1) + + # Column-blocks this column sits out. A is broadcast along the whole + # compute row, so it must still consume its share or the columns that + # DO have work stall waiting for the fifo to advance. No B and no C + # here -- the runtime sequence issues neither for it. + for _ in range_(n_drain): + for _ in range_(n_row_blocks): + for _ in range_(n_k_iters): + for _ in range_(B_ITERS // B_DEPTH): + for _ in range(B_DEPTH): + for _ in range(RHO): + a_h.acquire(1) + a_h.release(1) workers = [] for r in range(ROWS): @@ -661,7 +667,7 @@ def core_fn(n_work, n_drain, acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): acc = Buffer(type=ct_acc_ty, name=f"c_acc_{r}_{c}") workers.append( Worker( - partial(core_fn, col_work[c], col_drain[c]), + core_fn, [ acc, c_prod[(r, c)].prod(), @@ -670,6 +676,8 @@ def core_fn(n_work, n_drain, acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): acc_init, k_step, epilogue_chunk, + rtps[r][c], + barriers[r][c], ], ) ) @@ -766,6 +774,20 @@ def _c_tap_unsplit(mega_col, c): ) def sequence(A, B, C, a_prods, b_prods, c_conses): + # Write every core's parameters, then open every barrier. Both loops + # run to completion before the first fill is issued, so no core can + # read a half-written buffer. + for r in range(ROWS): + for c in range(n_active_cols): + rtps[r][c][RTP_N_WORK] = col_work[c] + rtps[r][c][RTP_N_DRAIN] = col_drain[c] + rtps[r][c][RTP_M_ROW_BLOCKS] = m_row_blocks + rtps[r][c][RTP_K_ITERS] = k_iters + rtps[r][c][RTP_EPILOGUE] = epilogue.mode + for r in range(ROWS): + for c in range(n_active_cols): + barriers[r][c].set(1) + # Column-blocks 0..n_full-1 use every column; the trailing one (when N # is not a multiple of N_TILE*COLS) uses only the first rem_blocks. A # is always issued for every row, because the columns sitting the diff --git a/iron/operators/flm/gemm/op.py b/iron/operators/flm/gemm/op.py index f7c3be3ca..7ace1174c 100644 --- a/iron/operators/flm/gemm/op.py +++ b/iron/operators/flm/gemm/op.py @@ -18,6 +18,7 @@ from aie.dialects.aie import get_target_model from aie.dialects._aie_enum_gen import AIEArch from iron.common.device_utils import get_kernel_dir +from iron.common.compilation import InstsBinArtifact, XclbinArtifact from iron.common.operator_bases import lut_based_ops_artifacts from iron.common.utils import float_to_name import aie.utils as aie_utils @@ -65,6 +66,11 @@ class GEMM(MLIROperator): N: int # Activation fused into the C drain. epilogue: Epilogue = Epilogue.NONE + # The activations the epilogue can select between at run time. Each one + # compiled in costs program memory, so a deployment that dispatches two + # should compile two. Unlike `epilogue`, this is part of the + # configuration. + epilogue_modes: tuple[Epilogue, ...] = tuple(Epilogue) # Optional (min, max) applied after the activation. clamp: tuple[float, float] | None = None # n tile width. 64 halves the mmul's accumulator traffic per mac; 128 @@ -143,16 +149,45 @@ def __post_init__(self): MLIROperator.__init__(self, context=self.context) + @property + def _epilogue_mask(self) -> int: + """Bitmask of the modes compiled into the epilogue. Mode 0 is always + present -- the kernel falls back to it.""" + return 1 | sum(1 << Epilogue(m).mode for m in self.epilogue_modes) + + @property + def _config_tag(self) -> str: + """Everything that shapes the device configuration, and so the xclbin. + + M, K, N and the activation are absent: they are runtime parameters, so + they change only the instruction stream. + """ + clamp = "" + if self.clamp is not None: + clamp = "_clamp" + "_".join(float_to_name(float(v)) for v in self.clamp) + dev = aie_utils.get_current_device().resolve().name + return ( + f"tn{self.tile_n}_ma{self.tile_ma}_em{self._epilogue_mask:x}" + f"_{self.rounding}{clamp}_{dev}" + ) + + @property + def config_name(self) -> str: + """Stem of the artifacts that do not depend on the shape.""" + return f"FLM_GEMM_{self._config_tag}" + @property def name(self) -> str: - """Artifact stem, prefixed to disambiguate from ``iron.operators.GEMM``. + """Artifact stem for the instruction stream, which does depend on it. - ``MLIROperator.name`` derives the stem from ``type(self).__name__``, - which is ``GEMM`` for both operators. This repo's build cache keys on - filename and mtime rather than on source or flags, so two operators - sharing a stem in one build dir would silently satisfy each other. + Prefixed to disambiguate from ``iron.operators.GEMM``: this repo's + build cache keys on filename and mtime rather than on source or flags, + so two operators sharing a stem would silently satisfy each other. """ - return f"FLM_{super().name}" + base = f"FLM_GEMM_M{self.M}_K{self.K}_N{self.N}_{self._config_tag}" + if self.epilogue != Epilogue.NONE: + base = f"{base}_epi{self.epilogue}" + return base @property def _bfp16_b(self) -> bool: @@ -187,7 +222,7 @@ def _kernel_object(self) -> str: return ( f"mm_fused_{M_TILE}x{K_TILE}x{self.tile_n}" f"_r{R}t{T}_ma{self.tile_ma}_{self.rounding}" - f"_epi{self.epilogue}{clamp}.o" + f"_em{self._epilogue_mask:x}{clamp}.o" ) @property @@ -201,31 +236,79 @@ def _link_file(self) -> str: LINK error for tanh_lut_ab/tanh_lut_cd rather than a compile error, so it surfaces late. """ - if self.epilogue is not Epilogue.NONE and get_kernel_dir() == "aie2": - return f"{self.name}_kernels.a" + if ( + any(Epilogue(m) is not Epilogue.NONE for m in self.epilogue_modes) + and get_kernel_dir() == "aie2" + ): + # config_name, not name: this string reaches the design's + # link_with, so a shape in it would put the shape in the device + # configuration. + return f"{self.config_name}_kernels.a" return self._kernel_object - def get_mlir_artifact(self): + @property + def _reference_shape(self) -> tuple[int, int, int]: + """The shape the configuration-only module is emitted at. + + Its runtime sequence is discarded; only its device body reaches the + xclbin. Taking the smallest valid shape keeps that module cheap to + build and makes the shape-independence explicit -- if a real shape's + instruction stream did not run against this xclbin, some dimension + would still be reaching the configuration. + """ + dev = aie_utils.get_current_device() + return M_TILE * compute_rows(dev), MIN_K, self.tile_n * dev.cols + + def _mlir_artifact(self, filename, M, K, N, epilogue): return PythonGeneratedMLIRArtifact( - f"{self.name}.mlir", + filename, DesignGenerator( self.operator_dir / "design.py", "gemm", (), { "dev": aie_utils.get_current_device(), - "M": self.M, - "K": self.K, - "N": self.N, + "M": M, + "K": K, + "N": N, "tile_n": self.tile_n, "tile_ma": self.tile_ma, - "epilogue": self.epilogue, + "epilogue": epilogue, "kernel_object": self._link_file, "trace_size": 0, }, ), ) + def get_mlir_artifact(self): + return self._mlir_artifact( + f"{self.name}.mlir", self.M, self.K, self.N, self.epilogue + ) + + def set_up_artifacts(self) -> None: + kernels = self.get_kernel_artifacts() + + # The xclbin comes from a module emitted at a reference shape and a + # reference activation, so every shape sharing this configuration + # reuses it rather than rebuilding an identical one. + config_mlir = self._mlir_artifact( + f"{self.config_name}.mlir", *self._reference_shape, Epilogue.NONE + ) + self.xclbin_artifact = XclbinArtifact( + f"{self.config_name}.xclbin", + mlir_input=config_mlir, + dependencies=[config_mlir] + kernels, + ) + shape_mlir = self.get_mlir_artifact() + self.insts_artifact = InstsBinArtifact( + f"{self.name}.bin", + mlir_input=shape_mlir, + # aiecc compiles the cores on the way to an instruction stream, so + # this needs the kernel objects too. + dependencies=[shape_mlir] + kernels, + ) + self.add_artifacts([self.xclbin_artifact, self.insts_artifact]) + def get_kernel_artifacts(self): kernel_dir = get_kernel_dir() base_dir = self.context.base_dir @@ -261,7 +344,7 @@ def get_kernel_artifacts(self): # Output stage. f"-DMM_FUSED_OUT_CHUNK={CT_OUT_LEN}", f"-DMM_FUSED_C_DEPTH={C_DEPTH}", - f"-DMM_FUSED_EPILOGUE_MODE={self.epilogue.mode}", + f"-DMM_FUSED_EPILOGUE_MODE_MASK={self._epilogue_mask}", ] + arch_include if self._bfp16_b: flags += [ diff --git a/iron/operators/flm/gemm/test.py b/iron/operators/flm/gemm/test.py index 0dd4312e1..7f68ffb83 100644 --- a/iron/operators/flm/gemm/test.py +++ b/iron/operators/flm/gemm/test.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os + import pytest import aie.utils as aie_utils @@ -321,3 +323,53 @@ def test_artifact_stem_differs_from_generic_gemm(M, K, N, aie_context): GEMM(M=M, K=K, N=N, context=aie_context).name != GenericGEMM(M=M, K=K, N=N, context=aie_context).name ) + + +def test_one_xclbin_serves_every_shape(aie_context): + """Several shapes back to back on one loaded xclbin. + + This is what the runtime parameters are for, and the parametrised tests + above cannot cover it: each gets a fresh context, so the array is + reconfigured between cases and any state a dispatch leaves behind is + wiped. Here the shapes share one. + + They disagree on every parameter -- M, K, N, whether a column sits a block + out, and the activation -- and none of them may rebuild the xclbin. + """ + shapes = [ + (256, 1536, 2048, "none"), + (256, 1536, 256, "none"), # only 4 of 8 columns compute + (512, 2048, 1536, "none"), + (256, 1536, 6144, "gelu"), + (256, 1536, 2048, "none"), # back to the first, after the rest + ] + xclbin = None + for M, K, N, epilogue in shapes: + operator = GEMM(M=M, K=K, N=N, epilogue=epilogue, context=aie_context) + golden_ref = generate_golden_reference( + M=M, K=K, N=N, epilogue=epilogue, scale=4.0 if epilogue == "none" else 0.5 + ) + mass = ( + K + * golden_ref["input"].abs().float().mean() + * golden_ref["input_b"].abs().float().mean() + ) + errors, _, _ = run_test( + operator, + { + "A": golden_ref["input"].flatten(), + "B": operator.pack_B(golden_ref["input_b"]), + }, + {"C": golden_ref["output"].flatten()}, + rel_tol=0.04, + abs_tol=float(0.004 * mass), + ) + assert not errors, f"{M}x{K}x{N} {epilogue} failed" + + stamp = ( + operator.xclbin_artifact.filename, + os.path.getmtime(operator.xclbin_artifact.filename), + ) + if xclbin is None: + xclbin = stamp + assert stamp == xclbin, f"{M}x{K}x{N} rebuilt the xclbin" From f4241160fb5763003515df6f250db9f124b073c5 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 11 Sep 2026 12:53:46 -0600 Subject: [PATCH 02/14] flm_gemm: reconcile docs after merging the RTP work MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The merge of andrej/flm-gemm-rtp was textually clean but left the README self-contradictory: the new intro says M, K, N and the activation are runtime parameters, while the FLM-compatibility note I had added in 92080d4 still claimed they were baked in at compile time and called RTP-selectability "follow-up work" -- which 81e2006a had just done. - Drop that note entirely; the "Runtime parameters" section it would have pointed at now covers the same ground correctly. - Fix the same stale claim in the shipped-overlay section, and record what IS still build-time there: which activations the epilogue can select between, since each one compiled in costs program memory. - design.py: "None of these four" was already stale from 92080d4, which added a fifth bullet. Co-Authored-By: André Rösti --- iron/operators/flm/gemm/README.md | 20 +++++--------------- iron/operators/flm/gemm/design.py | 2 +- 2 files changed, 6 insertions(+), 16 deletions(-) diff --git a/iron/operators/flm/gemm/README.md b/iron/operators/flm/gemm/README.md index 0ff196f48..811ed4b12 100644 --- a/iron/operators/flm/gemm/README.md +++ b/iron/operators/flm/gemm/README.md @@ -46,17 +46,6 @@ The shipped overlay itself is available as [`iron.operators.flm.MMPrebuilt`](../mm_prebuilt) for comparison; `benchmark.py` measures the two against each other and against `iron.operators.GEMM`. -**Not yet a drop-in replacement for the shipped overlay in FastFlowLM itself.** -FLM's runtime selects matrix shape and activation per call via runtime -parameters (RTPs) on one compiled xclbin. `M`/`K`/`N`/`epilogue`/`rounding` here -are `GEMM(...)` constructor arguments instead -- baked into the MLIR and the -kernel's `-D` flags at compile time (see -[Matching the shipped overlay](#matching-the-shipped-fastflowlm-overlay)) -- so -each shape+epilogue combination is its own compiled kernel object, not one -kernel switchable at runtime. Using this operator inside FLM today means -precompiling and swapping between kernels per combination; making shape and -epilogue RTP-selectable is follow-up work. - ## Architectures Runs on both NPU2 (aie2p — Strix/Krackan) and NPU1 (aie2 — Phoenix/Hawk Point). @@ -191,10 +180,11 @@ improves and gelu's worst case drops 5.5%. Measured perf-neutral (0.993-1.006x, inside the run-to-run spread). The shipped kernel selects its activation -- and its shape -- from runtime -parameters, one overlay serving every projection; this operator bakes both in -at compile time instead (activation keeps the shipped 0/1/2/3 mapping), which -is what lets its inner loop be branch-free. See the FLM-compatibility note near -the top of this file for what that means for using this operator inside FLM. +parameters, one overlay serving every projection. This operator does the same +(activation keeps the shipped 0/1/2/3 mapping); see +[Runtime parameters](#runtime-parameters). Which activations the epilogue can +*select between* is still a build-time choice, because each one compiled in +costs program memory. `clamp` has no counterpart in the shipped overlay to compare against — its `generate_seq` never writes the clamp RTP words, so clamping is always off diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index baadf7a0e..75e35fd02 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -38,7 +38,7 @@ not share a height -- this is what buys the deep k slice (K_TILE=512) within the L1 budget; see README.md's ATB reference. -None of these four helps alone -- see README.md's Performance section for the +None of these helps alone -- see README.md's Performance section for the measured, per-choice breakdown of the gap against both the shipped FastFlowLM overlay and ``iron.operators.GEMM``. From 92350c48af8b5c5b2dacf02136e8bdec1df463e5 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 11 Sep 2026 13:32:58 -0600 Subject: [PATCH 03/14] flm_gemm: make the clamp bounds runtime parameters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The RTP work left clamp entirely compile-time, so _config_tag carried the bound values and every distinct pair forked a whole xclbin -- clamp=(-2,2) and clamp=(-4,4) built twice over. Split it the way epilogue_modes already splits activations: the CAPABILITY stays build-time, the SELECTION and the values go runtime. - Whether a clamped path exists at all is still -DMM_FUSED_CLAMP, because the clamped instantiation costs program memory and a build that never clamps should not carry it. That bit stays in _config_tag. - clamp_enabled and the bounds become RTP words 5-7. Bounds are floats but npu_write_rtp writes i32 only, so they travel as raw bit patterns and the kernel casts them back with __builtin_bit_cast -- memcpy leaves an unresolved external call in the compiled object rather than folding to a register move. - epilogue_body gains a CLAMP template parameter so the clamped and unclamped inner loops both stay branch-free; epilogue_dispatch picks between them once per chunk, and only compiles the clamped one when the capability is on. Deliberately NOT done: making clamp_enabled a plain runtime branch. That would double epilogue_body instantiations (mode x clamp), and program memory is the exact constraint epilogue_modes exists to manage. Verified on npu2: all four kernel variants (aie2/aie2p x clamp on/off) compile, 36/36 flm/gemm iter0 tests pass, and the new test_one_xclbin_serves_every_clamp_bound confirms three different bound pairs run back to back on one loaded xclbin while an unclamped build still resolves to a different configuration. Co-Authored-By: André Rösti --- aie_kernels/generic/mm_fused.cc | 79 +++++++++++++++++++++++-------- iron/operators/flm/gemm/README.md | 26 ++++++---- iron/operators/flm/gemm/design.py | 50 +++++++++++++++++-- iron/operators/flm/gemm/op.py | 55 ++++++++++++--------- iron/operators/flm/gemm/test.py | 38 +++++++++++++++ 5 files changed, 193 insertions(+), 55 deletions(-) diff --git a/aie_kernels/generic/mm_fused.cc b/aie_kernels/generic/mm_fused.cc index 7845a3208..54577df51 100644 --- a/aie_kernels/generic/mm_fused.cc +++ b/aie_kernels/generic/mm_fused.cc @@ -38,15 +38,14 @@ #ifndef MM_FUSED_EPILOGUE_MODE_MASK #define MM_FUSED_EPILOGUE_MODE_MASK 0xF #endif +// Whether a clamp is compiled in at all. Like the mode mask above this is a +// CAPABILITY, not a selection: compiling the clamped path costs program +// memory, so a build that never clamps should not carry it. The BOUNDS are +// runtime parameters (see mm_fused_epilogue_chunk), so every pair of bounds +// shares one build. #ifndef MM_FUSED_CLAMP #define MM_FUSED_CLAMP 0 #endif -#ifndef MM_FUSED_CLAMP_MIN -#define MM_FUSED_CLAMP_MIN 0.0f -#endif -#ifndef MM_FUSED_CLAMP_MAX -#define MM_FUSED_CLAMP_MAX 0.0f -#endif namespace { @@ -107,12 +106,15 @@ constexpr aie::rounding_mode round_mode = aie::rounding_mode::floor; #endif // One activation's inner loop. Templated so each mode compiles branch-free; // mm_fused_epilogue_chunk selects between them once per chunk. -template static inline void epilogue_body(bfloat16 *__restrict y_out, const float *__restrict src) +template +static inline void +epilogue_body(bfloat16 *__restrict y_out, const float *__restrict src, float clamp_min, float clamp_max) { -#if MM_FUSED_CLAMP - const aie::vector lo = aie::broadcast(MM_FUSED_CLAMP_MIN); - const aie::vector hi = aie::broadcast(MM_FUSED_CLAMP_MAX); -#endif + aie::vector lo, hi; + if constexpr (CLAMP) { + lo = aie::broadcast(clamp_min); + hi = aie::broadcast(clamp_max); + } AIE_LOOP_MAX_ITERATION_COUNT(CHUNK / V) for (int j = 0; j < CHUNK / V; j++) { @@ -127,9 +129,8 @@ template static inline void epilogue_body(bfloat16 *__restrict y_out, f = silu_vec(f); else if constexpr (MODE == 3) f = sigmoid_vec(f); -#if MM_FUSED_CLAMP - f = aie::max(aie::min(f, hi), lo); -#endif + if constexpr (CLAMP) + f = aie::max(aie::min(f, hi), lo); aie::accum out; out.from_vector(f); // The assignment is the conversion: to_v16bfloat16 yields a raw @@ -138,6 +139,25 @@ template static inline void epilogue_body(bfloat16 *__restrict y_out, aie::store_v(y_out + j * V, v); } } + +// Pick the clamped or unclamped instantiation for a mode. Only one of the two +// exists unless MM_FUSED_CLAMP compiled the clamp in, so a build that never +// clamps pays nothing for this. +template +static inline void epilogue_dispatch(bfloat16 *__restrict y_out, + const float *__restrict src, + int32_t clamp_enabled, + float clamp_min, + float clamp_max) +{ +#if MM_FUSED_CLAMP + if (clamp_enabled) { + epilogue_body(y_out, src, clamp_min, clamp_max); + return; + } +#endif + epilogue_body(y_out, src, clamp_min, clamp_max); +} } // namespace extern "C" { @@ -182,39 +202,58 @@ void mm_fused_k_step(bfloat16 *a_buf, mm_fused_b_elem_t *b_buf, float *y_acc, in // The mode is a runtime argument, because one xclbin serves every activation. // It is tested once per chunk, outside the vector loop, so each mode still runs // a branch-free inner loop; the cost is program memory, since every mode in -// MM_FUSED_EPILOGUE_MODE_MASK is compiled in. The clamp stays compile-time. +// MM_FUSED_EPILOGUE_MODE_MASK is compiled in. +// +// The clamp follows the same split: whether a clamped path exists at all is +// compile-time (MM_FUSED_CLAMP, since it costs program memory), while +// clamp_enabled and the BOUNDS are runtime, so every pair of bounds shares one +// build. The bounds arrive as raw int32 bit patterns because the RTP mechanism +// (aie.dialects.aie.npu_write_rtp) only writes i32 words; design.py bit-casts +// them on the host side. // // The chunk index is split in two because the core body unrolls the drain by // the C fifo depth to keep the acquired buffer index a compile-time constant; // passing both parts avoids doing that arithmetic up there. -void mm_fused_epilogue_chunk(bfloat16 *y_out, float *y_acc, int32_t outer, int32_t half, int32_t mode) +void mm_fused_epilogue_chunk(bfloat16 *y_out, + float *y_acc, + int32_t outer, + int32_t half, + int32_t mode, + int32_t clamp_enabled, + int32_t clamp_min_bits, + int32_t clamp_max_bits) { // The store below is a conversion, so it obeys the same rounding mode the // mmul does and must agree with it. ::aie::set_rounding(round_mode); const float *__restrict src = y_acc + (outer * C_DEPTH + half) * CHUNK; + // __builtin_bit_cast, not memcpy: memcpy leaves an unresolved external + // call in the compiled object here rather than folding to a register move, + // and this runs once per chunk drain. + const float clamp_min = __builtin_bit_cast(float, clamp_min_bits); + const float clamp_max = __builtin_bit_cast(float, clamp_max_bits); switch (mode) { #if MM_FUSED_EPILOGUE_MODE_MASK & 2 case 1: - epilogue_body<1>(y_out, src); + epilogue_dispatch<1>(y_out, src, clamp_enabled, clamp_min, clamp_max); return; #endif #if MM_FUSED_EPILOGUE_MODE_MASK & 4 case 2: - epilogue_body<2>(y_out, src); + epilogue_dispatch<2>(y_out, src, clamp_enabled, clamp_min, clamp_max); return; #endif #if MM_FUSED_EPILOGUE_MODE_MASK & 8 case 3: - epilogue_body<3>(y_out, src); + epilogue_dispatch<3>(y_out, src, clamp_enabled, clamp_min, clamp_max); return; #endif // Mode 0 is always compiled: it is the fallback for a mode the mask leaves // out, so an unselectable mode yields an unactivated result rather than an // unwritten buffer. default: - epilogue_body<0>(y_out, src); + epilogue_dispatch<0>(y_out, src, clamp_enabled, clamp_min, clamp_max); return; } } diff --git a/iron/operators/flm/gemm/README.md b/iron/operators/flm/gemm/README.md index 811ed4b12..4203e135d 100644 --- a/iron/operators/flm/gemm/README.md +++ b/iron/operators/flm/gemm/README.md @@ -78,7 +78,7 @@ Two consequences of the native-vs-emulated split are worth knowing: ## Runtime parameters -Five words in an L1 buffer per core, written by the runtime sequence and read +Eight words in an L1 buffer per core, written by the runtime sequence and read by the core once its barrier opens: | word | value | @@ -88,19 +88,29 @@ by the core once its barrier opens: | `m_row_blocks` | `M / 256` | | `k_iters` | `K / 512` | | epilogue | the `Epilogue` mode | +| `clamp_enabled` | whether to apply the clamp | +| `clamp_min` / `clamp_max` | the bounds, as raw `int32` bit patterns | + +The clamp bounds are floats, but `npu_write_rtp` writes `i32` words only, so +they travel bit-cast and the kernel casts them back with +`__builtin_bit_cast` -- `memcpy` leaves an unresolved external call in the +compiled object rather than folding to a register move. All columns are always built. One with no work for a shape gets `n_work = 0` and still drains its share of the A broadcast, because the memtile will not release an A object until every consumer has taken it. The two artifacts therefore carry different stems: the xclbin's `config_name` -covers tile_n, tile_ma, the compiled activation set, clamp, rounding and the -device, while `name` adds M, K, N and the activation. The xclbin is built from -a module emitted at a reference shape, whose runtime sequence is discarded. - -Which activations the epilogue can *select between* stays a build-time choice, -since each one compiled in costs program memory; `epilogue_modes` sets it and -lands in the xclbin's name. +covers tile_n, tile_ma, the compiled activation set, whether a clamp exists, +rounding and the device, while `name` adds M, K, N and the activation. The +xclbin is built from a module emitted at a reference shape, whose runtime +sequence is discarded. + +Two things stay build-time, for the same reason -- each costs program memory: +which activations the epilogue can *select between* (`epilogue_modes`), and +whether a clamped path exists at all. Both land in the xclbin's name. Note the +asymmetry for clamp: *whether* to clamp is a build choice, but the *bounds* +are runtime, so `clamp=(-2, 2)` and `clamp=(-4, 4)` share one xclbin. The core releases its barrier straight after reading the parameters. `wait_for_value` emits `LockAction.Acquire`, which does not leave the lock diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index 75e35fd02..542919de6 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -147,8 +147,24 @@ def mode(self) -> int: # The runtime parameter buffer each core reads once its barrier opens. -RTP_N_WORK, RTP_N_DRAIN, RTP_M_ROW_BLOCKS, RTP_K_ITERS, RTP_EPILOGUE = range(5) -RTP_WORDS = 5 +# +# The clamp bounds are the one non-obvious entry: they are floats, but +# npu_write_rtp only writes i32 words, so they travel as raw bit patterns and +# the kernel bit-casts them back. Whether a clamped path exists at all stays +# compile-time (op.py's -DMM_FUSED_CLAMP), because it costs program memory; +# only the enable and the bounds are runtime, so every pair of bounds shares +# one build. +( + RTP_N_WORK, + RTP_N_DRAIN, + RTP_M_ROW_BLOCKS, + RTP_K_ITERS, + RTP_EPILOGUE, + RTP_CLAMP_ENABLED, + RTP_CLAMP_MIN_BITS, + RTP_CLAMP_MAX_BITS, +) = range(8) +RTP_WORDS = 8 class Rounding(StrEnum): @@ -281,6 +297,7 @@ def gemm( K, N, epilogue=Epilogue.NONE, + clamp=None, tile_n=N_TILE_DEFAULT, tile_ma=None, overlap=None, @@ -343,6 +360,15 @@ def gemm( MIN_N = N_TILE * COLS epilogue = Epilogue(epilogue) + # Clamp bounds ride the RTP buffer as raw int32 bit patterns: npu_write_rtp + # writes i32 words only, so the kernel bit-casts them back. Bounds being + # runtime is what lets every pair share one build; whether a clamped path + # exists at all is still op.py's -DMM_FUSED_CLAMP, because it costs + # program memory. + clamp_enabled = 1 if clamp is not None else 0 + clamp_lo, clamp_hi = clamp if clamp is not None else (0.0, 0.0) + clamp_min_bits = int(np.float32(clamp_lo).view(np.int32)) + clamp_max_bits = int(np.float32(clamp_hi).view(np.int32)) # A compute tile does a whole m x n block or nothing, so M and K must tile # exactly. N need only be a multiple of N_TILE: a trailing group of fewer # than COLS blocks is handled by giving the columns different trip counts @@ -463,7 +489,8 @@ def gemm( epilogue_chunk = Kernel( EPILOGUE_SYMBOL, kernel_object, - [ct_out_ty, ct_acc_ty, np.int32, np.int32, np.int32], + # outer, half, mode, clamp_enabled, clamp_min_bits, clamp_max_bits + [ct_out_ty, ct_acc_ty] + [np.int32] * 6, ) # --- Data movement ---------------------------------------------------- @@ -618,6 +645,9 @@ def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, barrier): n_row_blocks = my_rtp[RTP_M_ROW_BLOCKS] n_k_iters = my_rtp[RTP_K_ITERS] epi_mode = my_rtp[RTP_EPILOGUE] + clamp_enabled = my_rtp[RTP_CLAMP_ENABLED] + clamp_min_bits = my_rtp[RTP_CLAMP_MIN_BITS] + clamp_max_bits = my_rtp[RTP_CLAMP_MAX_BITS] # Acquire does not consume the barrier, so take it back to zero or the # next dispatch reads these parameters again instead of waiting. Safe # before the work: the sequence cannot set the barrier again until it @@ -645,7 +675,16 @@ def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, barrier): for chunk in range_(O_CHUNKS // C_DEPTH): for half in range(C_DEPTH): o = o_h.acquire(1) - epi_k(o, acc, chunk, half, epi_mode) + epi_k( + o, + acc, + chunk, + half, + epi_mode, + clamp_enabled, + clamp_min_bits, + clamp_max_bits, + ) o_h.release(1) # Column-blocks this column sits out. A is broadcast along the whole @@ -784,6 +823,9 @@ def sequence(A, B, C, a_prods, b_prods, c_conses): rtps[r][c][RTP_M_ROW_BLOCKS] = m_row_blocks rtps[r][c][RTP_K_ITERS] = k_iters rtps[r][c][RTP_EPILOGUE] = epilogue.mode + rtps[r][c][RTP_CLAMP_ENABLED] = clamp_enabled + rtps[r][c][RTP_CLAMP_MIN_BITS] = clamp_min_bits + rtps[r][c][RTP_CLAMP_MAX_BITS] = clamp_max_bits for r in range(ROWS): for c in range(n_active_cols): barriers[r][c].set(1) diff --git a/iron/operators/flm/gemm/op.py b/iron/operators/flm/gemm/op.py index 7ace1174c..f3153d4fc 100644 --- a/iron/operators/flm/gemm/op.py +++ b/iron/operators/flm/gemm/op.py @@ -20,7 +20,6 @@ from iron.common.device_utils import get_kernel_dir from iron.common.compilation import InstsBinArtifact, XclbinArtifact from iron.common.operator_bases import lut_based_ops_artifacts -from iron.common.utils import float_to_name import aie.utils as aie_utils from iron.operators.flm.packing import pack_b, packed_b_size @@ -155,20 +154,31 @@ def _epilogue_mask(self) -> int: present -- the kernel falls back to it.""" return 1 | sum(1 << Epilogue(m).mode for m in self.epilogue_modes) + @property + def _clamp_capable(self) -> int: + """Whether a clamped path is compiled in at all. + + Like ``epilogue_modes`` this is a capability, not a selection: the + clamped instantiation costs program memory, so a build that never + clamps should not carry it. The BOUNDS are runtime parameters, so + ``clamp=(-2, 2)`` and ``clamp=(-4, 4)`` share one xclbin -- only + clamped-versus-not forks the build. + """ + return 1 if self.clamp is not None else 0 + @property def _config_tag(self) -> str: """Everything that shapes the device configuration, and so the xclbin. - M, K, N and the activation are absent: they are runtime parameters, so - they change only the instruction stream. + M, K, N, the activation and the clamp BOUNDS are absent: they are + runtime parameters, so they change only the instruction stream. + Whether a clamp exists at all does shape the build; see + ``_clamp_capable``. """ - clamp = "" - if self.clamp is not None: - clamp = "_clamp" + "_".join(float_to_name(float(v)) for v in self.clamp) dev = aie_utils.get_current_device().resolve().name return ( f"tn{self.tile_n}_ma{self.tile_ma}_em{self._epilogue_mask:x}" - f"_{self.rounding}{clamp}_{dev}" + f"_{self.rounding}_cl{self._clamp_capable}_{dev}" ) @property @@ -216,13 +226,10 @@ def _kernel_object(self) -> str: which set the blocked layout, and the epilogue flags, which since the epilogue was folded into this translation unit shape the same object. """ - clamp = "" - if self.clamp is not None: - clamp = "_clamp" + "_".join(float_to_name(float(v)) for v in self.clamp) return ( f"mm_fused_{M_TILE}x{K_TILE}x{self.tile_n}" f"_r{R}t{T}_ma{self.tile_ma}_{self.rounding}" - f"_em{self._epilogue_mask:x}{clamp}.o" + f"_em{self._epilogue_mask:x}_cl{self._clamp_capable}.o" ) @property @@ -259,7 +266,7 @@ def _reference_shape(self) -> tuple[int, int, int]: dev = aie_utils.get_current_device() return M_TILE * compute_rows(dev), MIN_K, self.tile_n * dev.cols - def _mlir_artifact(self, filename, M, K, N, epilogue): + def _mlir_artifact(self, filename, M, K, N, epilogue, clamp): return PythonGeneratedMLIRArtifact( filename, DesignGenerator( @@ -274,6 +281,7 @@ def _mlir_artifact(self, filename, M, K, N, epilogue): "tile_n": self.tile_n, "tile_ma": self.tile_ma, "epilogue": epilogue, + "clamp": clamp, "kernel_object": self._link_file, "trace_size": 0, }, @@ -282,7 +290,7 @@ def _mlir_artifact(self, filename, M, K, N, epilogue): def get_mlir_artifact(self): return self._mlir_artifact( - f"{self.name}.mlir", self.M, self.K, self.N, self.epilogue + f"{self.name}.mlir", self.M, self.K, self.N, self.epilogue, self.clamp ) def set_up_artifacts(self) -> None: @@ -291,8 +299,16 @@ def set_up_artifacts(self) -> None: # The xclbin comes from a module emitted at a reference shape and a # reference activation, so every shape sharing this configuration # reuses it rather than rebuilding an identical one. + # Canonical bounds, not this instance's: the clamp BOUNDS only reach + # the runtime sequence, which this module has discarded, so passing + # the real ones would make a filename-cached artifact's content depend + # on something that never reaches the xclbin. Only the capability + # matters here, and that is already in config_name. config_mlir = self._mlir_artifact( - f"{self.config_name}.mlir", *self._reference_shape, Epilogue.NONE + f"{self.config_name}.mlir", + *self._reference_shape, + Epilogue.NONE, + (0.0, 0.0) if self._clamp_capable else None, ) self.xclbin_artifact = XclbinArtifact( f"{self.config_name}.xclbin", @@ -345,21 +361,14 @@ def get_kernel_artifacts(self): f"-DMM_FUSED_OUT_CHUNK={CT_OUT_LEN}", f"-DMM_FUSED_C_DEPTH={C_DEPTH}", f"-DMM_FUSED_EPILOGUE_MODE_MASK={self._epilogue_mask}", + # Capability only -- the bounds are runtime. See _clamp_capable. + f"-DMM_FUSED_CLAMP={self._clamp_capable}", ] + arch_include if self._bfp16_b: flags += [ "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", "-DMM_FUSED_BFP16_B", ] - if self.clamp is not None: - lo, hi = self.clamp - # repr() rather than :g -- the latter renders -4.0 as "-4", and - # "-4f" is not a valid C float literal. - flags += [ - "-DMM_FUSED_CLAMP=1", - f"-DMM_FUSED_CLAMP_MIN={float(lo)!r}f", - f"-DMM_FUSED_CLAMP_MAX={float(hi)!r}f", - ] if self.rounding is Rounding.CONV_EVEN: # ROUND_CONV_EVEN is mm.cc's flag, reused rather than inventing a # second spelling, and its polarity is mm.cc's too: absent means the diff --git a/iron/operators/flm/gemm/test.py b/iron/operators/flm/gemm/test.py index 7f68ffb83..ea2700986 100644 --- a/iron/operators/flm/gemm/test.py +++ b/iron/operators/flm/gemm/test.py @@ -373,3 +373,41 @@ def test_one_xclbin_serves_every_shape(aie_context): if xclbin is None: xclbin = stamp assert stamp == xclbin, f"{M}x{K}x{N} rebuilt the xclbin" + + +def test_one_xclbin_serves_every_clamp_bound(aie_context): + """Different clamp bounds back to back on one loaded xclbin. + + The bounds are runtime parameters, so they must not rebuild anything; + only clamped-versus-unclamped is a build choice, because the clamped + instantiation costs program memory. See GEMM._clamp_capable. + + Deliberately separate from test_one_xclbin_serves_every_shape: that one + never clamps, so it cannot catch bounds leaking back into the + configuration, which is exactly what this asserts. + """ + M, K, N = 256, 512, 1024 + bounds = [(-2.0, 2.0), (-4.0, 4.0), (-0.5, 0.5)] + xclbin = None + for clamp in bounds: + operator = GEMM(M=M, K=K, N=N, clamp=clamp, context=aie_context) + golden_ref = generate_golden_reference( + M=M, K=K, N=N, clamp=clamp, scale=INPUT_SCALE + ) + errors, _, _ = check_on_device(operator, golden_ref, K) + assert not errors, f"clamp={clamp} produced wrong output" + + stamp = ( + operator.xclbin_artifact.filename, + os.path.getmtime(operator.xclbin_artifact.filename), + ) + if xclbin is None: + xclbin = stamp + assert stamp == xclbin, f"clamp={clamp} rebuilt the xclbin" + + # ...but an unclamped build is a different configuration, and must be: + # the clamped instantiation is compiled out entirely there. config_name + # rather than xclbin_artifact, which only exists once compile() has run. + clamped = GEMM(M=M, K=K, N=N, clamp=bounds[0], context=aie_context) + unclamped = GEMM(M=M, K=K, N=N, context=aie_context) + assert unclamped.config_name != clamped.config_name From 801fea7a1f1d7d636c985f0a1544b40ec184dbf8 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Fri, 11 Sep 2026 17:05:20 -0600 Subject: [PATCH 04/14] flm_gemm: add m_chunk, folding row-blocks into one B fetch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A core can HOLD a B chunk across m_chunk accumulators instead of releasing it after one, so DDR reads B m_row_blocks/m_chunk times. That recovers the traffic the removed B residency used to save, without residency's cost: residency sized the memtile buffer from k_iters and replayed it m_row_blocks times, putting both K and M in the device configuration, whereas m_chunk is a configuration constant. Defaults to 1 (off). It is a tested knob, not a default, because interleaved A/B at K=1024 N=4096 does not show a consistent win: M m_chunk=1 min/med m_chunk=2 min/med min med 512 525.4 / 622.7 502.1 / 566.3 -4.4% -9.1% 1024 959.0 / 1081.6 1092.4 / 1162.9 +13.9% +7.5% 2048 1902.8 / 1986.0 1837.1 / 1940.0 -3.5% -2.3% B's traffic does fall, but the per-unit A descriptors m_chunk forces (a_split) appear to eat it. The M=1024 regression is not monotonic in n_units and is unexplained -- that is the open question, and the reason this is off by default. How the interleave works, since it is the non-obvious part: the core consumes A as (k, b_iter, mc, band), with mc INSIDE b_iter. A second A fifo would express that directly but needs a third core input DMA channel against a hardware limit of two. So the memtile A object holds m_chunk stacked tiles and the forward's dims_to_stream emits them interleaved, which fits the memtile BD's four dimensions only because mc's stride (M_TILE*K_TILE) exactly equals the row-group dimension's size*stride and the two merge. A partial group is inexpressible -- stride 0 inner is rejected, stride 0 outermost is the BD repeat count (one object per repetition, not one in total), and sub-object fills do not coalesce, all three confirmed on hardware -- so op.py resolves m_chunk to 1 when it would not divide m_row_blocks, or when the group's ROWS*M_TILE*K stride would overflow the shim BD's 20-bit step. Also stops forcing OVERLAP=1 on every split block: only a block spanning more than one window awaits inside itself. Worth -3.5% at M=2048 alone. 37/37 iter0 tests pass on npu2. Co-Authored-By: André Rösti --- iron/operators/flm/gemm/design.py | 344 ++++++++++++++++++++++++------ iron/operators/flm/gemm/op.py | 42 +++- iron/operators/flm/gemm/test.py | 22 +- 3 files changed, 329 insertions(+), 79 deletions(-) diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index 542919de6..d0f880db0 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -122,6 +122,46 @@ def compute_rows(dev): C_DEPTH = 2 # C fifo depth; also the core-body unroll B_DEPTH = 2 # B fifo depth; also the core-body unroll A_DEPTH = 2 +# Row-blocks a core folds into one B fetch: it HOLDS a B chunk across M_CHUNK +# accumulators instead of releasing it after one, so DDR reads B +# m_row_blocks/M_CHUNK times instead of m_row_blocks. This is what replaces the +# B residency the runtime parameters cost us -- residency sized the memtile +# buffer from k_iters and replayed it m_row_blocks times, putting both K and M +# in the device configuration, whereas M_CHUNK is a configuration constant and +# leaves both runtime. +# +# Two things it costs: +# * L1: M_CHUNK accumulators instead of one. At tile_n=64 the k slice is +# unaffected (ct_k stays 128, colA 16) and only tile_ma drops 32 -> 16; +# at tile_n=128 nothing fits, so M_CHUNK is resolved per tile_n below. +# * A is issued once per chunk-group rather than once per column-block. +# Holding B means streaming A k-major -- A[mc0][k0], A[mc1][k0], +# A[mc0][k1], ... -- which wants five dimensions, and a shim BD carries +# four: the outermost lands in the ITERATION field and the inner three are +# its ND dims (getBDMaxDims is 3 off a memtile). So the chunk-group +# dimension comes out of the descriptor and becomes separate transfers +# carrying the jump in the OFFSET, which has no such limit -- exactly the +# trick a_split already uses for the mega_row dimension, and it reuses the +# same windowing. What is left per transfer is +# [k_iters, M_CHUNK, M_TILE, K_TILE], which fits. +# DEFAULT 1 EVERYWHERE -- i.e. off. The mechanism is correct and tested, but +# interleaved A/B at K=1024 N=4096 (10 rounds, round-robin in one process) +# does not show a consistent win: +# +# M m_chunk=1 min/med m_chunk=2 min/med min med +# 512 525.4 / 622.7 502.1 / 566.3 -4.4% -9.1% +# 1024 959.0 / 1081.6 1092.4 / 1162.9 +13.9% +7.5% +# 2048 1902.8 / 1986.0 1837.1 / 1940.0 -3.5% -2.3% +# +# B's DDR traffic really does fall by m_chunk, but something else eats it -- +# most likely the per-unit A descriptors, since m_chunk forces the split path +# for A (see a_split). The M=1024 regression is not monotonic in n_units and +# is not explained; do not turn this on by default without understanding it. +# +# An earlier NON-interleaved measurement showed -11%/-15% and was wrong; see +# the bimodal-timing discipline in README.md. Interleave before believing any +# number here. +M_CHUNK_FOR_N = {16: 1, 32: 1, 64: 1, 128: 1} # How many column-blocks the runtime sequence keeps in flight. A block costs 3 # shim buffer descriptors on a column (A + B + C) against 16 available, so the # ceiling is 5; 2 is enough to keep the fills ahead of the cores. @@ -163,8 +203,14 @@ def mode(self) -> int: RTP_CLAMP_ENABLED, RTP_CLAMP_MIN_BITS, RTP_CLAMP_MAX_BITS, -) = range(8) -RTP_WORDS = 8 + # Row-blocks split into M_CHUNK-wide groups plus a leftover: m_row_blocks + # need not divide by M_CHUNK (M=256 is one row-block, M=768 is three), so + # the core runs a wide pass n_chunks times and a single-wide pass n_rem + # times. Both counts are runtime, so any M still rides one xclbin. + RTP_N_CHUNKS, + RTP_N_UNITS, +) = range(10) +RTP_WORDS = 10 class Rounding(StrEnum): @@ -227,7 +273,7 @@ def _hw_stride_ok(stride_elems): return hw_stride <= (1 << _SHIM_STEP_BITS) - 1 -def _default_l1(n_tile, ct_max_k, b_elem_bytes, budget): +def _default_l1(n_tile, ct_max_k, b_elem_bytes, budget, m_chunk=1): """Pick (A-tile height, L1 B depth) -- the largest working set that fits. ``b_elem_bytes`` is 9/8 where B is bfp16ebs8 and 2 where it is bf16, and @@ -250,7 +296,11 @@ def _default_l1(n_tile, ct_max_k, b_elem_bytes, budget): Deeper B first, then the tallest A that still fits, so the n=64 default is unchanged at (32, 2). """ - acc = M_TILE * n_tile * 4 + # m_chunk accumulators, because the core holds a B chunk across that + # many row-blocks; see M_CHUNK_FOR_N. This is the ONLY term that + # scales with it -- A is acquired and released one band at a time + # inside the group loop, and C drains the accumulators in turn. + acc = m_chunk * M_TILE * n_tile * 4 cout = CT_OUT_LEN * 2 * C_DEPTH for b_depth in (B_DEPTH, 1): b = int(ct_max_k * n_tile * b_elem_bytes) * b_depth @@ -263,7 +313,7 @@ def _default_l1(n_tile, ct_max_k, b_elem_bytes, budget): raise ValueError(f"nothing fits L1 for tile_n={n_tile}, ct_max_k={ct_max_k}") -def _b_depth_for(t_ma, n_tile, ct_max_k, b_elem_bytes, budget): +def _b_depth_for(t_ma, n_tile, ct_max_k, b_elem_bytes, budget, m_chunk=1): """Deepest B fifo depth that fits L1 alongside an explicit A-tile height. ``_default_l1`` picks L1_B_DEPTH together with the t_ma IT chooses; that @@ -278,7 +328,11 @@ def _b_depth_for(t_ma, n_tile, ct_max_k, b_elem_bytes, budget): close to free elsewhere -- worth keeping, but not worth contorting the search for. """ - acc = M_TILE * n_tile * 4 + # m_chunk accumulators, because the core holds a B chunk across that + # many row-blocks; see M_CHUNK_FOR_N. This is the ONLY term that + # scales with it -- A is acquired and released one band at a time + # inside the group loop, and C drains the accumulators in turn. + acc = m_chunk * M_TILE * n_tile * 4 cout = CT_OUT_LEN * 2 * C_DEPTH a = (2 * R * ct_max_k) * (t_ma // R // 2) * 2 * A_DEPTH for b_depth in (B_DEPTH, 1): @@ -299,6 +353,7 @@ def gemm( epilogue=Epilogue.NONE, clamp=None, tile_n=N_TILE_DEFAULT, + m_chunk=None, tile_ma=None, overlap=None, kernel_object="mm_fused.o", @@ -323,6 +378,7 @@ def gemm( SHIM_BDS = tm.get_num_bds(0, 0) N_TILE = tile_n CT_MAX_K = CT_MAX_K_FOR_N[N_TILE] + M_CHUNK = M_CHUNK_FOR_N[N_TILE] if m_chunk is None else m_chunk # B is bfp16ebs8 on AIE2P and bf16 on AIE2: the scalar BFP types are gated # on __AIE_API_SCALAR_BFP_TYPES__, which only aie_api/detail/aie2p/config.hpp # defines, so on AIE2 B stays bf16 and the mmul lowers onto four native @@ -338,7 +394,7 @@ def gemm( # reduction, so sizing both to M_TILE pays the peak L1 cost twice. if tile_ma is None: T_MA, L1_B_DEPTH = _default_l1( - N_TILE, CT_MAX_K, b_elem_bytes, tm.get_local_memory_size() + N_TILE, CT_MAX_K, b_elem_bytes, tm.get_local_memory_size(), M_CHUNK ) else: T_MA = tile_ma @@ -347,7 +403,12 @@ def gemm( f"tile_ma ({T_MA}) must divide {M_TILE} and be a multiple of {2 * R}" ) L1_B_DEPTH = _b_depth_for( - T_MA, N_TILE, CT_MAX_K, b_elem_bytes, tm.get_local_memory_size() + T_MA, + N_TILE, + CT_MAX_K, + b_elem_bytes, + tm.get_local_memory_size(), + M_CHUNK, ) RHO = M_TILE // T_MA OVERLAP = OVERLAP_DEFAULT if overlap is None else overlap @@ -385,6 +446,33 @@ def gemm( # How many times the whole grid sweeps, in each dimension. m_row_blocks = M // MIN_M k_iters = K // K_TILE + # Row-blocks grouped M_CHUNK at a time, so one B fetch feeds M_CHUNK of + # them (see M_CHUNK_FOR_N). A "unit" below is one such group, or one of the + # n_rem leftovers when M_CHUNK does not divide m_row_blocks. Every leg is + # issued per unit, so A, B and C stay aligned with each other and with the + # core's loop nest. + n_chunks, n_rem = divmod(m_row_blocks, M_CHUNK) + if n_rem: + # op.py resolves m_chunk so this cannot fire. A partial group is + # genuinely inexpressible: its object is M_CHUNK tiles wide and the + # forward always drains that much, so the rest would have to be filled + # by repeating the row-block -- and every way of saying that is + # rejected or mis-lowered. Stride 0 in an inner dimension is refused + # ("Stride 2 must be a positive integer"); stride 0 in the outermost + # slot IS the BD repeat count, which releases one object per + # repetition rather than one object in total; and several sub-object + # fills do not coalesce into one object either. All three were tried + # on hardware. Hence op.py falls back to m_chunk=1 instead. + raise ValueError( + f"m_row_blocks ({m_row_blocks}) must be a multiple of m_chunk " + f"({M_CHUNK}); op.py should have resolved m_chunk to 1 here" + ) + n_units = n_chunks + + def unit_rows(u): + """(first row-block, how many) for unit ``u`` -- always a full group.""" + return u * M_CHUNK, M_CHUNK + # A mega_row dimension with stride ROWS*M_TILE*{K,N} lands in the shim # BD's ITERATION field, whose step is 20 bits, so it overflows once K or N # crosses ~8191 elements -- E4B's FFN width (10240) does, E2B's max (6144) @@ -406,7 +494,12 @@ def gemm( # descriptor. Verify with aie-opt --aie-substitute-shim-dma-allocations # --aie-assign-runtime-sequence-bd-ids: the ids on a shim tile must be # distinct. - a_split = m_row_blocks > 1 and not _hw_stride_ok(ROWS * M_TILE * K) + # M_CHUNK > 1 forces the split path for A: holding B makes the core's A + # order k-major within a group, which wants a fifth dimension, so the unit + # dimension comes out of the descriptor and becomes one transfer each. + # At n_units == 1 there is nothing to split -- a single transfer already + # covers the block -- so the cheaper unsplit path still applies. + a_split = n_units > 1 and (M_CHUNK > 1 or not _hw_stride_ok(ROWS * M_TILE * K)) c_split = m_row_blocks > 1 and not _hw_stride_ok(ROWS * M_TILE * N) # A split leg issues one transfer per mega_row back to back on ONE channel, # so they must also fit that channel's task queue -- a limit nothing in the @@ -418,8 +511,19 @@ def gemm( # Awaiting is what makes the window's descriptors safe to reuse, and it # bounds both resources at once. M<=1024 is a single window, so the shapes # that already worked are unaffected. - MB_WINDOW = min(m_row_blocks, SHIM_TASK_QUEUE) if (a_split or c_split) else 1 - bds_per_block = 1 + (MB_WINDOW if a_split else 1) + (MB_WINDOW if c_split else 1) + # A window is counted in TRANSFERS, not units: under c_split a unit emits + # one C descriptor per row-block (M_CHUNK of them, see c_taps), so a + # window of SHIM_TASK_QUEUE units would push M_CHUNK times that many onto + # one channel and hang. + _per_unit = M_CHUNK if c_split else 1 + MB_WINDOW = ( + min(n_units, max(1, SHIM_TASK_QUEUE // _per_unit)) + if (a_split or c_split) + else 1 + ) + bds_per_block = ( + 1 + (MB_WINDOW if a_split else 1) + (MB_WINDOW * _per_unit if c_split else 1) + ) # Unreachable while SHIM_TASK_QUEUE is 4 (the worst case is 1 + 4 + 4 = 9 # of 16), so this guards a future retune of the window rather than any # shape reachable today. test_flm_gemm_split_leg_windowing asserts the same @@ -434,7 +538,14 @@ def gemm( # Cross-block overlap only applies to the unsplit path; a split block # already awaits inside itself, so keeping a second one in flight would # refill the very queue the windowing just drained. - if a_split or c_split: + # Only a block that spans MORE THAN ONE window awaits inside itself; that + # is what makes a second block in flight refill the queue the windowing + # just drained. A single-window block has no intra-block await, so it can + # still overlap -- and it must, because m_chunk forces a_split on for the + # descriptor's sake even when nothing needs windowing, and losing the + # pipelining costs far more than the B traffic m_chunk saves (measured + # +13% at M=1024 before this). + if (a_split or c_split) and n_units > MB_WINDOW: OVERLAP = 1 else: OVERLAP = max(1, min(OVERLAP, SHIM_BDS // bds_per_block)) @@ -464,8 +575,10 @@ def gemm( ct_out_ty = np.ndarray[(CT_OUT_LEN,), bf16_ty] ct_acc_ty = np.ndarray[(M_TILE * N_TILE,), f32] # L2 (per memtile) - mt_a_ty = np.ndarray[(M_TILE * K_TILE,), bf16_ty] - mt_a_bytes = M_TILE * K_TILE * 2 + # M_CHUNK stacked row-block tiles, so the forward below can interleave + # them on the way out -- see a_send_dims. + mt_a_ty = np.ndarray[(M_CHUNK * M_TILE * K_TILE,), bf16_ty] + mt_a_bytes = M_CHUNK * M_TILE * K_TILE * 2 mt_b_ty = np.ndarray[(K_TILE * N_TILE // B_GROUP,), b_elem_ty] mt_b_bytes = _b_bytes(K_TILE * N_TILE, BFP16_B) mt_out_ty = np.ndarray[(C_SLICE_LEN * ROWS,), bf16_ty] @@ -514,10 +627,24 @@ def gemm( b_recv_dims = None b_send_dims = None # A: same idea, r x s blocks. - a_recv_dims = [(M_TILE // R, R * K_TILE), (R, S), (K_TILE // S, R * S), (S, 1)] + # The outermost row-group dimension spans M_CHUNK tiles rather than one. + # That is the whole trick: mc's stride is M_TILE*K_TILE, which is exactly + # this dimension's size*stride, so the two are contiguous and merge -- the + # walk stays within the memtile BD's four dimensions while gaining an + # interleave it could not otherwise express. + a_recv_dims = [ + (M_CHUNK * M_TILE // R, R * K_TILE), + (R, S), + (K_TILE // S, R * S), + (S, 1), + ] + # Emits (b_iter, mc, band) -- b_iter outermost, then all M_CHUNK tiles' + # row-groups. That is the order the core acquires A in when it holds a B + # chunk across the group, and it is why ONE fifo suffices: a second would + # need a third core input DMA channel, and a tile has two. a_send_dims = [ (K_DIV_CT_K_MAX, R * CT_MAX_K), - (M_TILE // R, R * K_TILE), + (M_CHUNK * M_TILE // R, R * K_TILE), ] + split_run(R * CT_MAX_K) # C: one join per column. Each of the ROWS cores in the column drops its @@ -540,6 +667,10 @@ def gemm( # A: shim -> memtile -> broadcast along the compute row. The reblocking # rides the forward(): inbound on cons(dims_from_stream=), outbound on # forward(dims_to_stream=), sharing one memtile buffer. + # ONE fifo per row, even at M_CHUNK > 1. The interleave the core needs + # lives in a_send_dims above, not in extra fifos: a second A fifo would + # make the core want 3 input DMA channels and a compute tile has 2 (the + # design already spends both, on A and B). a_l3l2_fifos = [] a_cons = {} for r in range(ROWS): @@ -634,7 +765,7 @@ def gemm( ] # --- Compute ---------------------------------------------------------- - def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, barrier): + def core_fn(accs, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, barrier): """Core body. Every trip count and the activation come from the runtime parameter buffer, so one core program serves every shape.""" # The loop nest lives here rather than inside the kernel so that @@ -648,36 +779,53 @@ def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, barrier): clamp_enabled = my_rtp[RTP_CLAMP_ENABLED] clamp_min_bits = my_rtp[RTP_CLAMP_MIN_BITS] clamp_max_bits = my_rtp[RTP_CLAMP_MAX_BITS] + n_chunks = my_rtp[RTP_N_CHUNKS] + n_units_rt = my_rtp[RTP_N_UNITS] # Acquire does not consume the barrier, so take it back to zero or the # next dispatch reads these parameters again instead of waiting. Safe # before the work: the sequence cannot set the barrier again until it # has drained this dispatch's C. barrier.release_with_value(1) - for _ in range_(n_work): - for _ in range_(n_row_blocks): - init_k(acc) - for _ in range_(n_k_iters): - # The l loop is unrolled by the B fifo depth so the - # acquired buffer index stays a compile-time constant. - for _ in range_(B_ITERS // B_DEPTH): - for _ in range(B_DEPTH): - # One B chunk feeds every A band, so B is - # acquired once around the band loop. - b = b_h.acquire(1) + def sweep(group): + """One k reduction feeding ``group`` accumulators off a shared B. + + ``group`` is a Python list, so its length is compile-time: the mc + loops below unroll. Calling this with every accumulator is the + wide pass; calling it with one is the leftover pass. + + Holding B across the group is the whole point -- b_h is acquired + once outside the mc loop and released after all of them, so DDR + reads B once per len(group) row-blocks instead of once each. + """ + for a_acc in group: + init_k(a_acc) + for _ in range_(n_k_iters): + # The l loop is unrolled by the B fifo depth so the + # acquired buffer index stays a compile-time constant. + for _ in range_(B_ITERS // B_DEPTH): + for _ in range(B_DEPTH): + # One B chunk feeds every A band of every accumulator + # in the group, so B is acquired once around both. + # Each accumulator draws A from its OWN fifo, which is + # what makes this interleave legal -- see the a_cons + # construction above. + b = b_h.acquire(1) + for a_acc in group: for band in range(RHO): a = a_h.acquire(1) - kstep_k(a, b, acc, band) + kstep_k(a, b, a_acc, band) a_h.release(1) - b_h.release(1) - # Drain the accumulator. Unrolled by C_DEPTH for the same - # reason; a full O_CHUNKS unroll overflows program memory. + b_h.release(1) + # Drain the accumulators. Unrolled by C_DEPTH for the same + # reason; a full O_CHUNKS unroll overflows program memory. + for a_acc in group: for chunk in range_(O_CHUNKS // C_DEPTH): for half in range(C_DEPTH): o = o_h.acquire(1) epi_k( o, - acc, + a_acc, chunk, half, epi_mode, @@ -687,28 +835,40 @@ def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, barrier): ) o_h.release(1) + for _ in range_(n_work): + for _ in range_(n_chunks): + sweep(accs) + # Column-blocks this column sits out. A is broadcast along the whole # compute row, so it must still consume its share or the columns that # DO have work stall waiting for the fifo to advance. No B and no C # here -- the runtime sequence issues neither for it. for _ in range_(n_drain): - for _ in range_(n_row_blocks): + # Every unit delivers a full M_CHUNK tiles of A, leftover or not, + # so an idle column drains that much per unit. + for _ in range_(n_units_rt): for _ in range_(n_k_iters): for _ in range_(B_ITERS // B_DEPTH): for _ in range(B_DEPTH): - for _ in range(RHO): + for _ in range(M_CHUNK * RHO): a_h.acquire(1) a_h.release(1) workers = [] for r in range(ROWS): for c in range(n_active_cols): - acc = Buffer(type=ct_acc_ty, name=f"c_acc_{r}_{c}") + # One accumulator per row-block in a chunk group. Worker flattens + # nested fn_args, so the list arrives in core_fn as a list and its + # length stays compile-time. + accs = [ + Buffer(type=ct_acc_ty, name=f"c_acc_{r}_{c}_{mc}") + for mc in range(M_CHUNK) + ] workers.append( Worker( core_fn, [ - acc, + accs, c_prod[(r, c)].prod(), b_cons[(r, c)], a_cons[(r, c)], @@ -737,16 +897,21 @@ def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, barrier): # Dimension order must match the core loop nest exactly: for each # column-block it walks mega_row, then k. Every wrap stays under the shim's # wrap/size field (largest are K_TILE=512 and ROWS*M_TILE=256). - def a_taps(mega_col, r, mbs): - # Every (mega_row, k) block this compute row consumes for one + def a_taps(mega_col, r, units): + # Every (row-block, k) block this compute row consumes for one # column-block. A does not depend on mega_col; it is re-fetched per # column-block because the cores re-consume it. # - # Returns a LIST: one 4D descriptor normally, or one 3D descriptor per - # mega_row when the mega_row stride would overflow the shim BD's 20-bit - # iteration step (see a_split above). The split form carries the - # mega_row jump in the offset, which has no such limit. - if not a_split: + # k OUTERMOST, then the group's M_CHUNK row-blocks: one memtile object + # per k holding M_CHUNK stacked tiles, which a_send_dims then emits + # interleaved as (b_iter, mc, band) -- the order the core acquires in + # while holding a B chunk across the group. + # + # A leftover unit is one row-block wide but fills the same object, so + # its mc dimension has stride 0: the row-block is repeated, and the + # core drains the duplicate (see sweep). It costs one extra read of + # that row-block, on at most one unit per column-block. + if M_CHUNK == 1 and not a_split: return [ TensorAccessPattern( tensor_dims=(M * K,), @@ -755,15 +920,18 @@ def a_taps(mega_col, r, mbs): strides=[ROWS * M_TILE * K, K_TILE, K, 1], ) ] - return [ - TensorAccessPattern( - tensor_dims=(M * K,), - offset=mb * ROWS * M_TILE * K + r * M_TILE * K, - sizes=[1, k_iters, M_TILE, K_TILE], - strides=[0, K_TILE, K, 1], + taps = [] + for u in units: + first, count = unit_rows(u) + taps.append( + TensorAccessPattern( + tensor_dims=(M * K,), + offset=first * ROWS * M_TILE * K + r * M_TILE * K, + sizes=[k_iters, M_CHUNK, M_TILE, K_TILE], + strides=[K_TILE, ROWS * M_TILE * K, K, 1], + ) ) - for mb in mbs - ] + return taps def b_tap(mega_col, c): # Every (mega_row, k) chunk this column consumes. B does not depend on @@ -777,31 +945,54 @@ def b_tap(mega_col, c): return TensorAccessPattern( tensor_dims=(K * N // B_GROUP,), offset=(mega_col * COLS + c) * N_TILE * K // B_GROUP, + # One k sweep per UNIT, not per row-block: the cores hold each B + # chunk across the M_CHUNK row-blocks of a group, so DDR reads B + # n_units times instead of m_row_blocks. That is the whole win -- + # see M_CHUNK_FOR_N. The unit dimension keeps stride 0, replaying + # the same k-blocks, exactly as the row-block dimension used to. sizes=( [1, 1, 1, k_iters * K_TILE * N_TILE // B_GROUP] if b_resident - else [m_row_blocks, k_iters, 1, K_TILE * N_TILE // B_GROUP] + else [n_units, k_iters, 1, K_TILE * N_TILE // B_GROUP] ), strides=( [0, 0, 0, 1] if b_resident else [0, K_TILE * N_TILE // B_GROUP, 0, 1] ), ) - def c_taps(mega_col, c, mbs): + def c_taps(mega_col, c, units): # Every joined block this column produces for one column-block: one # ROWS*M_TILE x N_TILE block per row-block. Returns a LIST, for the - # same reason a_taps does -- one descriptor per mega_row when N makes - # the mega_row stride overflow the shim BD's iteration step. + # same reason a_taps does -- one descriptor per unit when N makes the + # row-block stride overflow the shim BD's iteration step. + # + # C drains in plain row-block order even under M_CHUNK, because the + # core drains a group's accumulators one after another, so no + # reordering is needed here; a unit just covers `count` consecutive + # row-blocks. if c_split: - return [ - TensorAccessPattern( - tensor_dims=(M * N,), - offset=(mega_col * COLS + c) * N_TILE + mb * ROWS * M_TILE * N, - sizes=[1, 1, ROWS * M_TILE, N_TILE], - strides=[0, 0, N, 1], - ) - for mb in mbs - ] + taps = [] + for u in units: + first, count = unit_rows(u) + # One descriptor PER ROW-BLOCK, not one per unit. Grouping a + # unit's row-blocks into a count dimension would put a + # ROWS*M_TILE*N stride back inside the descriptor, which is + # exactly what c_split exists to avoid -- it overflows the + # shim BD's 20-bit step at N=10240. C drains in plain + # row-block order even under M_CHUNK (the core drains a + # group's accumulators one after another), so splitting them + # costs nothing but the extra descriptors. + for i in range(count): + taps.append( + TensorAccessPattern( + tensor_dims=(M * N,), + offset=(mega_col * COLS + c) * N_TILE + + (first + i) * ROWS * M_TILE * N, + sizes=[1, 1, ROWS * M_TILE, N_TILE], + strides=[0, 0, N, 1], + ) + ) + return taps return [_c_tap_unsplit(mega_col, c)] def _c_tap_unsplit(mega_col, c): @@ -826,6 +1017,8 @@ def sequence(A, B, C, a_prods, b_prods, c_conses): rtps[r][c][RTP_CLAMP_ENABLED] = clamp_enabled rtps[r][c][RTP_CLAMP_MIN_BITS] = clamp_min_bits rtps[r][c][RTP_CLAMP_MAX_BITS] = clamp_max_bits + rtps[r][c][RTP_N_CHUNKS] = n_chunks + rtps[r][c][RTP_N_UNITS] = n_units for r in range(ROWS): for c in range(n_active_cols): barriers[r][c].set(1) @@ -859,14 +1052,27 @@ def sequence(A, B, C, a_prods, b_prods, c_conses): # Every task of a block stays LIVE in its TaskGroup until the block is # retired here. That is load-bearing, not tidiness -- see the # dma_free_task note on a_split above. - all_mb = list(range(m_row_blocks)) + # Units, not row-blocks: every leg is issued per unit so A, B and C + # stay aligned with each other and with the core's nest. At + # M_CHUNK == 1 a unit IS a row-block and this is the old list. + all_mb = list(range(n_units)) # One emitter per leg, so the two paths below differ only in HOW they # group and retire, not in how a leg is issued. def issue_a(mega_col, mbs, group, wait=False): for r in range(ROWS): - for tap in a_taps(mega_col, r, mbs): - a_prods[r].fill(A, tap, group=group, wait=wait) + taps = a_taps(mega_col, r, mbs) + for i, tap in enumerate(taps): + # A leftover unit emits k_iters*M_CHUNK fills back to back + # on ONE channel (see a_taps), and that channel's task + # queue is SHIM_TASK_QUEUE deep -- overrunning it HANGS + # rather than diagnoses, the same limit a_split windows + # for. Await every SHIM_TASK_QUEUE-th fill so no more than + # that many are ever outstanding. + bounded = ( + len(taps) > SHIM_TASK_QUEUE and (i + 1) % SHIM_TASK_QUEUE == 0 + ) + a_prods[r].fill(A, tap, group=group, wait=wait or bounded) def issue_b(mega_col, active_cols, group): for c in range(active_cols): @@ -924,7 +1130,7 @@ def emit_split(): if not a_split: issue_a(mega_col, all_mb, tg_whole) - for w in range(0, m_row_blocks, MB_WINDOW): + for w in range(0, n_units, MB_WINDOW): mbs = all_mb[w : w + MB_WINDOW] tg_w = TaskGroup() if c_split: diff --git a/iron/operators/flm/gemm/op.py b/iron/operators/flm/gemm/op.py index f3153d4fc..994bb6328 100644 --- a/iron/operators/flm/gemm/op.py +++ b/iron/operators/flm/gemm/op.py @@ -27,6 +27,7 @@ BFP16_GROUP, BFP16_GROUP_BYTES, CT_MAX_K_FOR_N, + M_CHUNK_FOR_N, C_DEPTH, compute_rows, CT_OUT_LEN, @@ -39,6 +40,7 @@ S, T, _default_l1, + _hw_stride_ok, ) @@ -79,6 +81,9 @@ class GEMM(MLIROperator): # A-tile rows, decoupled from the accumulator's M_TILE (asymmetric tile # buffering). __post_init__ resolves None to whatever L1 affords. tile_ma: int | None = None + # Row-blocks folded into one B fetch. __post_init__ resolves None from + # tile_n, falling back to 1 when it would not divide m_row_blocks. + m_chunk: int | None = None # Rounding for every f32->bf16 conversion; see Rounding in design.py. rounding: Rounding = Rounding.CONV_EVEN context: object = field(default=None, repr=False) @@ -88,6 +93,7 @@ class GEMM(MLIROperator): "epilogue": "epi", "tile_n": "tn", "tile_ma": "ma", + "m_chunk": "mc", "rounding": "rnd", } @@ -120,12 +126,34 @@ def __post_init__(self): raise ValueError( f"tile_n must be one of {sorted(CT_MAX_K_FOR_N)}, got {self.tile_n}" ) + # m_chunk: row-blocks a core folds into one B fetch (design.py's + # M_CHUNK_FOR_N). It MUST divide m_row_blocks -- a partial group is + # inexpressible, see the raise in design.py -- so fall back to 1 when + # it does not. That puts M's divisibility in the configuration, but + # only as a 2-bucket split rather than per-shape. + if self.m_chunk is None: + want = M_CHUNK_FOR_N[self.tile_n] + rows = M_TILE * compute_rows(dev) + m_row_blocks = self.M // rows if self.M % rows == 0 else 0 + # Two conditions, both of which force the fallback to 1: + # * m_chunk must DIVIDE m_row_blocks -- a partial group is + # inexpressible, see the raise in design.py. + # * the group's row-blocks sit ROWS*M_TILE*K apart INSIDE the A + # descriptor, so that stride must fit the shim BD's 20-bit + # step. a_split used to lift this dimension out of the + # descriptor entirely; m_chunk puts it back, so big K (10240) + # overflows where m_chunk=1 would not. + fits = m_row_blocks and m_row_blocks % want == 0 + if fits and not _hw_stride_ok(compute_rows(dev) * M_TILE * self.K): + fits = False + self.m_chunk = want if fits else 1 if self.tile_ma is None: self.tile_ma = _default_l1( self.tile_n, CT_MAX_K_FOR_N[self.tile_n], self._b_elem_bytes, get_target_model(dev.resolve()).get_local_memory_size(), + self.m_chunk, )[0] # N only needs to tile to N_TILE: a trailing group of fewer than # COLS column-blocks is handled by giving the columns different trip @@ -177,8 +205,9 @@ def _config_tag(self) -> str: """ dev = aie_utils.get_current_device().resolve().name return ( - f"tn{self.tile_n}_ma{self.tile_ma}_em{self._epilogue_mask:x}" - f"_{self.rounding}_cl{self._clamp_capable}_{dev}" + f"tn{self.tile_n}_ma{self.tile_ma}_mc{self.m_chunk}" + f"_em{self._epilogue_mask:x}_{self.rounding}" + f"_cl{self._clamp_capable}_{dev}" ) @property @@ -264,7 +293,13 @@ def _reference_shape(self) -> tuple[int, int, int]: would still be reaching the configuration. """ dev = aie_utils.get_current_device() - return M_TILE * compute_rows(dev), MIN_K, self.tile_n * dev.cols + # M must be at least m_chunk row-blocks: a partial group is + # inexpressible (see design.py), and this module must build. + return ( + M_TILE * compute_rows(dev) * self.m_chunk, + MIN_K, + self.tile_n * dev.cols, + ) def _mlir_artifact(self, filename, M, K, N, epilogue, clamp): return PythonGeneratedMLIRArtifact( @@ -280,6 +315,7 @@ def _mlir_artifact(self, filename, M, K, N, epilogue, clamp): "N": N, "tile_n": self.tile_n, "tile_ma": self.tile_ma, + "m_chunk": self.m_chunk, "epilogue": epilogue, "clamp": clamp, "kernel_object": self._link_file, diff --git a/iron/operators/flm/gemm/test.py b/iron/operators/flm/gemm/test.py index ea2700986..754f838da 100644 --- a/iron/operators/flm/gemm/test.py +++ b/iron/operators/flm/gemm/test.py @@ -15,6 +15,7 @@ BFP16_GROUP, BFP16_GROUP_BYTES, CT_MAX_K_FOR_N, + M_CHUNK_FOR_N, Epilogue, M_TILE, R, @@ -275,14 +276,17 @@ def tile_option_params(): params = [] for tile_n, ct_k in sorted(CT_MAX_K_FOR_N.items()): - default_ma = _default_l1(tile_n, ct_k, b_elem, l1)[0] + # m_chunk matters: the core holds a B chunk across that many + # accumulators, so it is what decides which A heights still fit. + m_chunk = M_CHUNK_FOR_N[tile_n] + default_ma = _default_l1(tile_n, ct_k, b_elem, l1, m_chunk)[0] # One full sweep of the grid at this tile_n, so every column has work. M, K, N = 256, 512, tile_n * dev.cols for tile_ma in (16, 32, 64): if M_TILE % tile_ma or tile_ma % (2 * R): continue try: - _b_depth_for(tile_ma, tile_n, ct_k, b_elem, l1) + _b_depth_for(tile_ma, tile_n, ct_k, b_elem, l1, m_chunk) except ValueError: continue # this A height leaves no room for B at this width marks = [] if tile_ma == default_ma else [pytest.mark.extensive] @@ -336,12 +340,16 @@ def test_one_xclbin_serves_every_shape(aie_context): They disagree on every parameter -- M, K, N, whether a column sits a block out, and the activation -- and none of them may rebuild the xclbin. """ + # Every shape here must resolve to the same m_chunk, because m_chunk + # shapes the core program and so the xclbin (see GEMM._config_tag). It + # buckets M by whether m_row_blocks is a multiple of it -- these are all + # even -- and excludes the K that would overflow the A descriptor's step. shapes = [ - (256, 1536, 2048, "none"), - (256, 1536, 256, "none"), # only 4 of 8 columns compute - (512, 2048, 1536, "none"), - (256, 1536, 6144, "gelu"), - (256, 1536, 2048, "none"), # back to the first, after the rest + (512, 1536, 2048, "none"), + (512, 1536, 256, "none"), # only 4 of 8 columns compute + (1024, 2048, 1536, "none"), + (512, 1536, 6144, "gelu"), + (512, 1536, 2048, "none"), # back to the first, after the rest ] xclbin = None for M, K, N, epilogue in shapes: From 6b2af06592117da86651b630f6002d5c7d583c19 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Mon, 14 Sep 2026 13:50:56 -0600 Subject: [PATCH 05/14] deps: bump mlir-aie to the v1.4.3 release Moves off the 1.4.3.dev85 snapshot onto the tagged release. llvm-aie is unchanged at 22.0.0.2026090701+3e93bf7b, which is the pin mlir-aie v1.4.3 itself names in utils/peano-requirements.txt, so the two stay in step. Tagged wheels live under their own tag's asset page rather than latest-wheels-4 (which carries only the .dev builds), hence the extra find-links. Co-Authored-By: Claude --- requirements.txt | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 7f7c43b37..b79a8960e 100755 --- a/requirements.txt +++ b/requirements.txt @@ -9,11 +9,16 @@ # CUDA build served from PyPI. We therefore also pin torch to the "+cpu" local # version below, which is only available from the PyTorch CPU index. --index-url https://download.pytorch.org/whl/cpu +--find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/v1.4.3 --find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-4 --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly --extra-index-url https://pypi.org/simple -mlir_aie==1.4.3.dev85+gdf48abc +# Tagged release wheels live under their own tag's asset page, not under +# latest-wheels-4 (which carries only the .dev builds), hence the extra +# find-links above. llvm-aie is not tagged in step with mlir-aie; this pin is +# the one mlir-aie v1.4.3 itself names in utils/peano-requirements.txt. +mlir_aie==1.4.3 llvm-aie==22.0.0.2026090701+3e93bf7b black From 4fbb53be59c08bc39fe24ecab3a69ddf220c2f75 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Mon, 14 Sep 2026 13:51:09 -0600 Subject: [PATCH 06/14] flm_gemm: name ct_max_k in the artifact stems CT_MAX_K_FOR_N is carried to the kernel as -DMM_FUSED_CT_K but appeared in neither artifact name. This repo's build cache keys on filename and mtime rather than on source or flags, so an object or xclbin built at one ct_max_k silently satisfied a request for another. Naming it in _kernel_object alone is not enough, and the reasoning that let that hole stand is worth recording: tile_ma usually moves with ct_max_k, but tile_ma is caller-overridable, so tn128/ma16 is reachable at two different ct_max_k values. An xclbin built by an experiment at ck=128 was then served to test_gemm_tile_options[tn128-ma16], which wants ck=32, and it returned NaN rather than an error. An artifact name must cover every input to THAT artifact; do not argue one field is implied by another unless the implication holds for every reachable configuration, overrides included. Co-Authored-By: Claude --- iron/operators/flm/gemm/op.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/iron/operators/flm/gemm/op.py b/iron/operators/flm/gemm/op.py index 994bb6328..576c097fe 100644 --- a/iron/operators/flm/gemm/op.py +++ b/iron/operators/flm/gemm/op.py @@ -202,10 +202,20 @@ def _config_tag(self) -> str: runtime parameters, so they change only the instruction stream. Whether a clamp exists at all does shape the build; see ``_clamp_capable``. + + ``ck`` is here for the same reason it is in ``_kernel_object``, and it + is NOT redundant with ``tn``: CT_MAX_K_FOR_N is a tuning table, and + retuning one entry changes the design (B's object width, the k slice, + a_send_dims) while tn is unmoved. ``ma`` does not cover it either -- + it usually moves with ck, but ``tile_ma`` is caller-overridable, so + tn128/ma16 is reachable at two different ck values. Omitting it here + silently served an xclbin built at one ck to a request for another, + which is how test_gemm_tile_options[tn128-ma16] started returning NaN. """ dev = aie_utils.get_current_device().resolve().name return ( - f"tn{self.tile_n}_ma{self.tile_ma}_mc{self.m_chunk}" + f"tn{self.tile_n}_ck{CT_MAX_K_FOR_N[self.tile_n]}" + f"_ma{self.tile_ma}_mc{self.m_chunk}" f"_em{self._epilogue_mask:x}_{self.rounding}" f"_cl{self._clamp_capable}_{dev}" ) @@ -254,9 +264,16 @@ def _kernel_object(self) -> str: otherwise silently satisfy a request for another. That includes r/t, which set the blocked layout, and the epilogue flags, which since the epilogue was folded into this translation unit shape the same object. + + ``ck`` is CT_MAX_K_FOR_N[tile_n], carried as -DMM_FUSED_CT_K. It is + derived from tile_n today, so it looks redundant -- but it is a TUNING + TABLE, and retuning one entry while leaving the object name alone is + exactly the silent-stale-binary case above. Naming it means the table + can be edited without also remembering to wipe the build dir. """ return ( f"mm_fused_{M_TILE}x{K_TILE}x{self.tile_n}" + f"_ck{CT_MAX_K_FOR_N[self.tile_n]}" f"_r{R}t{T}_ma{self.tile_ma}_{self.rounding}" f"_em{self._epilogue_mask:x}_cl{self._clamp_capable}.o" ) From c378b27ff5ab9337f2308439c1dc5e04617a0c03 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Mon, 14 Sep 2026 13:51:30 -0600 Subject: [PATCH 07/14] flm_gemm: send four runtime-parameter words instead of ten A parameter word is not free. Each costs ~66 ns per core and the sequence writes ROWS*COLS = 32 of them, so every word is ~2.06 us of dispatch latency -- measured by padding the buffer at a fixed core count (12 words 108.6 us, 24 words 135.2, 48 words 182.7). Against a ~107 us floor that is most of a short-prefill dispatch. Five of the ten were dead or duplicated: * the clamp trio is now sent only by a clamp-capable build. Where no clamped path is compiled in -- the default, and every real projection -- those three words were written on every dispatch and never read. * n_chunks / n_units are sent only when m_chunk > 1. They are m_row_blocks // M_CHUNK and each other, so at the shipped M_CHUNK = 1 the core reads m_row_blocks instead. No on-core arithmetic. and n_work / n_drain are now derived rather than sent, from a raw N word plus the tile's own column. Branch-free, and both divisors are powers of two, so it lowers to sdiv-by-constant and leaves no __divsi3 (verified in the .elf, not just the .ll). Note ScalarValue overloads add/sub/mul/floordiv/mod but NOT the shift operators, so this uses // rather than >>. The column index is per-tile static data, deliberately not a constant folded into the program: the 32 core programs differ today only in symbol names, and baking it into code would make them differ in instructions, foreclosing a future one-program xclbin. rtp_layout() sizes the buffer from (clamp_capable, m_chunk), both of which are already in _config_tag, so the word count cannot vary within a configuration and ship-once-use-many is preserved -- verified by loading one xclbin and dispatching shapes that disagree on M, K, N, on which columns sit a block out, and on the activation. Also guards CT_MAX_K_FOR_N: it reads like a tuning table but is load-bearing for correctness, and a wrong value fails silently (err/mass 3.45e-02 at tile_n=64/ct_k=64, NaN at tile_n=128). Root cause not found; pack_b is ruled out by test, its permutation round-trips at ct_k 128, 64 and 32. Unverified pairs now raise rather than miscompute. 30-shape suite: -3.5% median at M=256 (best -11.3%, E2B/kv), within noise at M >= 1024 -- the saving is a constant ~12 us. Accuracy bit-identical on all 30 shapes; frozen-reference control drifted +0.11%. Co-Authored-By: Claude --- iron/operators/flm/gemm/README.md | 48 ++++++--- iron/operators/flm/gemm/design.py | 157 ++++++++++++++++++++++++------ 2 files changed, 164 insertions(+), 41 deletions(-) diff --git a/iron/operators/flm/gemm/README.md b/iron/operators/flm/gemm/README.md index 4203e135d..5bc3215ee 100644 --- a/iron/operators/flm/gemm/README.md +++ b/iron/operators/flm/gemm/README.md @@ -78,30 +78,56 @@ Two consequences of the native-vs-emulated split are worth knowing: ## Runtime parameters -Eight words in an L1 buffer per core, written by the runtime sequence and read -by the core once its barrier opens: +**Four words** in an L1 buffer per core, written by the runtime sequence and +read by the core once its barrier opens: | word | value | |---|---| -| `n_work` | column-blocks this column computes | -| `n_drain` | column-blocks it sits out while still draining A | +| `N` | the raw N; the core derives its own `n_work` / `n_drain` from it | | `m_row_blocks` | `M / 256` | | `k_iters` | `K / 512` | | epilogue | the `Epilogue` mode | -| `clamp_enabled` | whether to apply the clamp | -| `clamp_min` / `clamp_max` | the bounds, as raw `int32` bit patterns | -The clamp bounds are floats, but `npu_write_rtp` writes `i32` words only, so -they travel bit-cast and the kernel casts them back with -`__builtin_bit_cast` -- `memcpy` leaves an unresolved external call in the -compiled object rather than folding to a register move. +A word is not free: each costs ~66 ns per core and the sequence writes +`ROWS * COLS` = 32 of them, so **every word is ~2 us of dispatch latency** +(measured by padding the buffer at a fixed core count). Against a ~107 us +floor that is most of a short-prefill dispatch, so `rtp_layout()` in design.py +sizes the buffer per configuration rather than sending words a build cannot +use. Three groups are therefore conditional or gone: + +* `clamp_enabled` / `clamp_min` / `clamp_max` are sent **only by a + clamp-capable build**. Where no clamped path is compiled in -- the default, + and every real projection -- they were written every dispatch and never + read. They stay raw `int32` bit patterns, because `npu_write_rtp` writes + i32 only; the kernel casts back with `__builtin_bit_cast`, since `memcpy` + leaves an unresolved external call rather than folding to a register move. +* `n_chunks` / `n_units` are sent **only when `m_chunk > 1`**. They are + `m_row_blocks // M_CHUNK` and each other, so at the shipped `M_CHUNK = 1` + the core just reads `m_row_blocks`. +* `n_work` / `n_drain` are **not sent at all**. The core derives them from `N` + and its own column, which it reads from a per-tile buffer initialised at + build time. Branch-free, and both divisors are powers of two, so it costs + shifts rather than a `__divsi3` call: + + n_tiles = N // N_TILE + n_work = (n_tiles - my_col + COLS - 1) // COLS + n_drain = ((n_tiles + COLS - 1) // COLS) - n_work + + The column index is per-tile **static data**, deliberately not a constant + folded into the program: the 32 core programs today differ only in symbol + names, and baking it into code would make them differ in instructions, + foreclosing a future one-program xclbin. + +Measured on the 30-shape suite: **-3.5% median at M=256** (best -11.3%, +E2B/kv), and within noise at M >= 1024 -- the saving is a constant ~12 us, so +it is a short-prefill and decode lever, not a prefill one. All columns are always built. One with no work for a shape gets `n_work = 0` and still drains its share of the A broadcast, because the memtile will not release an A object until every consumer has taken it. The two artifacts therefore carry different stems: the xclbin's `config_name` -covers tile_n, tile_ma, the compiled activation set, whether a clamp exists, +covers tile_n, ct_max_k, tile_ma, the compiled activation set, whether a clamp exists, rounding and the device, while `name` adds M, K, N and the activation. The xclbin is built from a module emitted at a reference shape, whose runtime sequence is discarded. diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index d0f880db0..14f422822 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -98,6 +98,23 @@ # (M_TILE * 256 * 4 = 65536 bytes) already fills the whole of L1, before A, B # or C are even counted, so no ct_max_k could ever make it fit. CT_MAX_K_FOR_N = {16: 16, 32: 32, 64: 128, 128: 32} +# (tile_n, ct_max_k) pairs KNOWN TO COMPUTE CORRECTLY on hardware. The table +# above reads like a tuning knob -- op.py calls it "the only place it is +# decided" -- but it is not freely tunable, and a wrong value fails SILENTLY. +# Measured on npu2 2026-09-11: +# +# tile_n=64 ct_k=128 err/mass 2.42e-04 (shipped) +# tile_n=64 ct_k= 64 err/mass 3.45e-02 ~140x worse, outside any budget +# tile_n=128 ct_k= 32 err/mass 1.43e-04 (shipped) +# tile_n=128 ct_k= 64 NaN +# tile_n=128 ct_k=128 NaN +# +# Root cause not found, but pack_b is RULED OUT BY TEST: inverting its +# permutation round-trips exactly at ct_k 128, 64 and 32, so its blocking is +# generic. The disagreement is most likely the ORDER the design's stream dims +# and the kernel's mmul nest walk one ct-chunk. Until that is found, refuse +# rather than miscompute. +_VERIFIED_CT_K = {(16, 16), (32, 32), (64, 128), (128, 32)} # Register tiling, shared by both architectures. These set the blocked L1 # layout, so ``pack_B``, the four stream-dimension lists below and # ``gather_dims`` all key off them; changing one without the others is silently @@ -194,23 +211,49 @@ def mode(self) -> int: # compile-time (op.py's -DMM_FUSED_CLAMP), because it costs program memory; # only the enable and the bounds are runtime, so every pair of bounds shares # one build. +# The five that EVERY configuration needs. Anything conditional lives after +# them, at an offset rtp_layout() computes, so a word a build cannot use is +# never allocated rather than written and ignored. ( - RTP_N_WORK, - RTP_N_DRAIN, + RTP_N_VAL, RTP_M_ROW_BLOCKS, RTP_K_ITERS, RTP_EPILOGUE, - RTP_CLAMP_ENABLED, - RTP_CLAMP_MIN_BITS, - RTP_CLAMP_MAX_BITS, - # Row-blocks split into M_CHUNK-wide groups plus a leftover: m_row_blocks - # need not divide by M_CHUNK (M=256 is one row-block, M=768 is three), so - # the core runs a wide pass n_chunks times and a single-wide pass n_rem - # times. Both counts are runtime, so any M still rides one xclbin. - RTP_N_CHUNKS, - RTP_N_UNITS, -) = range(10) -RTP_WORDS = 10 +) = range(4) + + +def rtp_layout(clamp_capable, m_chunk): + """Slot index for each optional parameter, and the total word count. + + **A word is not free.** Each one costs ~66 ns per core and the sequence + writes ROWS*COLS = 32 of them, so **every RTP word is ~2.06 us of dispatch + latency** -- measured by padding the buffer at a fixed core count (12 words + 108.6 us, 24 words 135.2, 48 words 182.7). Against a ~107 us floor that is + not a rounding error, and at M=256 the floor is most of the dispatch. + + So both groups here are omitted, not defaulted: + + * the clamp trio, when the build has no clamped path at all. op.py compiles + the clamped instantiation out entirely at ``_clamp_capable == 0`` (the + default, and every real projection shape), so those three words were + written on every dispatch and never read -- ~6 us for nothing. + * ``n_chunks`` / ``n_units``, when M_CHUNK == 1. They are + ``m_row_blocks // M_CHUNK`` and each other, so at the shipped M_CHUNK + they are simply m_row_blocks, and the core reads that word instead -- + no on-core arithmetic, just one fewer thing to send. ~4 us. + """ + slots = {} + n = 4 + if clamp_capable: + slots["clamp_enabled"] = n + slots["clamp_min"] = n + 1 + slots["clamp_max"] = n + 2 + n += 3 + if m_chunk > 1: + slots["n_chunks"] = n + slots["n_units"] = n + 1 + n += 2 + return slots, n class Rounding(StrEnum): @@ -378,6 +421,14 @@ def gemm( SHIM_BDS = tm.get_num_bds(0, 0) N_TILE = tile_n CT_MAX_K = CT_MAX_K_FOR_N[N_TILE] + if (N_TILE, CT_MAX_K) not in _VERIFIED_CT_K: + raise ValueError( + f"tile_n={N_TILE} with ct_max_k={CT_MAX_K} is not a verified " + f"combination (verified: {sorted(_VERIFIED_CT_K)}). It would build, " + f"run, and compute the WRONG ANSWER -- see _VERIFIED_CT_K. If you " + f"are retuning CT_MAX_K_FOR_N, fix that coupling first and add the " + f"pair here once a hardware test passes." + ) M_CHUNK = M_CHUNK_FOR_N[N_TILE] if m_chunk is None else m_chunk # B is bfp16ebs8 on AIE2P and bf16 on AIE2: the scalar BFP types are gated # on __AIE_API_SCALAR_BFP_TYPES__, which only aie_api/detail/aie2p/config.hpp @@ -427,6 +478,7 @@ def gemm( # exists at all is still op.py's -DMM_FUSED_CLAMP, because it costs # program memory. clamp_enabled = 1 if clamp is not None else 0 + rtp_slots, rtp_words = rtp_layout(clamp is not None, M_CHUNK) clamp_lo, clamp_hi = clamp if clamp is not None else (0.0, 0.0) clamp_min_bits = int(np.float32(clamp_lo).view(np.int32)) clamp_max_bits = int(np.float32(clamp_hi).view(np.int32)) @@ -747,13 +799,31 @@ def unit_rows(u): for r in range(ROWS): b_cons[(r, c)] = of_b.cons() + # Each tile's own column, as an initialized buffer rather than a constant + # folded into the program. That is deliberate: the 32 core programs differ + # today only in symbol NAMES, and baking the column in as an immediate + # would make them differ in CODE, permanently foreclosing the one-program + # xclbin. Data may vary per tile; the program must not. Written once at + # configuration time, so unlike an RTP word it costs nothing per dispatch. + my_cols = [ + [ + Buffer( + np.ndarray[(1,), np.dtype[np.int32]], + name=f"my_col_{r}_{c}", + initial_value=np.array([c], dtype=np.int32), + ) + for c in range(n_active_cols) + ] + for r in range(ROWS) + ] + # --- Runtime parameters ----------------------------------------------- rtps = [ [ Buffer( - np.ndarray[(RTP_WORDS,), np.dtype[np.int32]], + np.ndarray[(rtp_words,), np.dtype[np.int32]], name=f"rtp_{r}_{c}", - initial_value=np.zeros(RTP_WORDS, dtype=np.int32), + initial_value=np.zeros(rtp_words, dtype=np.int32), use_write_rtp=True, ) for c in range(n_active_cols) @@ -765,22 +835,46 @@ def unit_rows(u): ] # --- Compute ---------------------------------------------------------- - def core_fn(accs, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, barrier): + def core_fn(accs, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, my_col, barrier): """Core body. Every trip count and the activation come from the runtime parameter buffer, so one core program serves every shape.""" # The loop nest lives here rather than inside the kernel so that # every level has an ObjectFifo acquire point. barrier.wait_for_value(1) - n_work = my_rtp[RTP_N_WORK] - n_drain = my_rtp[RTP_N_DRAIN] + # n_work / n_drain are DERIVED here rather than sent, which costs one + # RTP word instead of two. Branch-free, so no select is needed: + # + # column c has work in column-block j iff (j*COLS + c)*N_TILE < N + # => n_work = ceil((N/N_TILE - c) / COLS) + # + # and every divisor is a power of two (N_TILE=64, COLS=8), so this is + # shifts and adds -- no __divsi3. Verified against the host-side + # col_work/col_drain for every shape in the suite. + # // rather than >>: the DSL's ScalarValue overloads add/sub/mul/ + # floordiv/mod but NOT the shift operators. Both divisors are + # compile-time powers of two (N_TILE=64, COLS=8), so this strength- + # reduces and must not leave a __divsi3 call -- verified in the .o. + n_tiles = my_rtp[RTP_N_VAL] // N_TILE + n_work = (n_tiles - my_col[0] + COLS - 1) // COLS + n_drain = ((n_tiles + COLS - 1) // COLS) - n_work n_row_blocks = my_rtp[RTP_M_ROW_BLOCKS] n_k_iters = my_rtp[RTP_K_ITERS] epi_mode = my_rtp[RTP_EPILOGUE] - clamp_enabled = my_rtp[RTP_CLAMP_ENABLED] - clamp_min_bits = my_rtp[RTP_CLAMP_MIN_BITS] - clamp_max_bits = my_rtp[RTP_CLAMP_MAX_BITS] - n_chunks = my_rtp[RTP_N_CHUNKS] - n_units_rt = my_rtp[RTP_N_UNITS] + # Absent slots become compile-time constants rather than loads: the + # kernel's clamped path is compiled out when it is not capable, and at + # M_CHUNK == 1 both chunk counts ARE n_row_blocks. + if "clamp_enabled" in rtp_slots: + clamp_enabled = my_rtp[rtp_slots["clamp_enabled"]] + clamp_min_bits = my_rtp[rtp_slots["clamp_min"]] + clamp_max_bits = my_rtp[rtp_slots["clamp_max"]] + else: + clamp_enabled, clamp_min_bits, clamp_max_bits = 0, 0, 0 + if "n_chunks" in rtp_slots: + n_chunks = my_rtp[rtp_slots["n_chunks"]] + n_units_rt = my_rtp[rtp_slots["n_units"]] + else: + n_chunks = n_row_blocks + n_units_rt = n_row_blocks # Acquire does not consume the barrier, so take it back to zero or the # next dispatch reads these parameters again instead of waiting. Safe # before the work: the sequence cannot set the barrier again until it @@ -876,6 +970,7 @@ def sweep(group): k_step, epilogue_chunk, rtps[r][c], + my_cols[r][c], barriers[r][c], ], ) @@ -1009,16 +1104,18 @@ def sequence(A, B, C, a_prods, b_prods, c_conses): # read a half-written buffer. for r in range(ROWS): for c in range(n_active_cols): - rtps[r][c][RTP_N_WORK] = col_work[c] - rtps[r][c][RTP_N_DRAIN] = col_drain[c] + rtps[r][c][RTP_N_VAL] = N rtps[r][c][RTP_M_ROW_BLOCKS] = m_row_blocks rtps[r][c][RTP_K_ITERS] = k_iters rtps[r][c][RTP_EPILOGUE] = epilogue.mode - rtps[r][c][RTP_CLAMP_ENABLED] = clamp_enabled - rtps[r][c][RTP_CLAMP_MIN_BITS] = clamp_min_bits - rtps[r][c][RTP_CLAMP_MAX_BITS] = clamp_max_bits - rtps[r][c][RTP_N_CHUNKS] = n_chunks - rtps[r][c][RTP_N_UNITS] = n_units + # Only what this configuration actually reads; see rtp_layout. + if "clamp_enabled" in rtp_slots: + rtps[r][c][rtp_slots["clamp_enabled"]] = clamp_enabled + rtps[r][c][rtp_slots["clamp_min"]] = clamp_min_bits + rtps[r][c][rtp_slots["clamp_max"]] = clamp_max_bits + if "n_chunks" in rtp_slots: + rtps[r][c][rtp_slots["n_chunks"]] = n_chunks + rtps[r][c][rtp_slots["n_units"]] = n_units for r in range(ROWS): for c in range(n_active_cols): barriers[r][c].set(1) From 0103a747fa6d800ee4c066397486bd28b2303ee9 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Mon, 14 Sep 2026 13:52:11 -0600 Subject: [PATCH 08/14] flm_gemm: fix benchmark.py's usage docstring --no-short is not a real pytest option. What the harness actually wants is --iterations 1: each test already averages ITERS dispatches over ROUNDS interleaved rounds, so conftest's default of 5 repeats the whole matrix five times for nothing. Also records that -s must not be passed when the CSV is wanted -- the reporter parses captured stdout, so disabling capture yields a CSV with no metric columns. Co-Authored-By: Claude --- iron/operators/flm/gemm/benchmark.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/iron/operators/flm/gemm/benchmark.py b/iron/operators/flm/gemm/benchmark.py index 00279cf81..fa47eb81a 100644 --- a/iron/operators/flm/gemm/benchmark.py +++ b/iron/operators/flm/gemm/benchmark.py @@ -38,10 +38,17 @@ asserts correctness on one implementation, this compares latency across three frozen binaries, and neither can stand in for the other. +Pass ``--iterations 1``. Each test already averages ``ITERS`` dispatches over +``ROUNDS`` interleaved rounds, so conftest's default of 5 repeats the whole +matrix five times for nothing. + Usage:: - pytest iron/operators/flm/gemm/benchmark.py --no-short - pytest iron/operators/flm/gemm/benchmark.py --no-short -k E2B --csv-output flm.csv + pytest iron/operators/flm/gemm/benchmark.py --iterations 1 + pytest iron/operators/flm/gemm/benchmark.py --iterations 1 -k E2B --csv-output flm.csv + +Do not pass ``-s`` when you want the CSV: conftest's reporter parses the +captured stdout, so disabling capture yields a CSV with no metric columns. """ import statistics From b2379c0a8d6ed47c8b1ed9e9d955917766351f19 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Mon, 14 Sep 2026 16:25:17 -0600 Subject: [PATCH 09/14] flm_gemm: retire split-leg transfers rolling, not in windows Where K or N is 10240 the row-block stride overflows the shim BD's 20-bit iteration step, so that leg is issued as one transfer per row-block. Two shim resources bound how many may be outstanding and neither is modelled by the toolchain: BD ids (16/tile, freed without a completion check) and the channel task queue (4 deep, pushed unconditionally). That bound was enforced by windowing -- issue four, await the whole window, issue the next four, then the same again per column-block. It is correct, but TaskGroup.finish() emits dma_await_task, so each of those is a real barrier and the channel drains to EMPTY at every window and column-block boundary. On a DDR-rate-bound design those bubbles are the entire cost of the split path. The OVERLAP value computed for this path was never read by it. Retire the oldest transfer as the next is issued instead: the same number stay in flight, the queue bound is enforced directly rather than by draining, and the channel stays full across both kinds of boundary. Interleaved 8 rounds x 30 iters, bit-exact against the previous sequence on every shape: E4B/gateup M1024 4119.4 -> 3617.7 -12.2% E4B/gateup M2048 7664.8 -> 7182.4 -6.3% E4B/down M1024 3649.2 -> 3553.4 -2.6% E4B/down M2048 7196.6 -> 6978.7 -3.0% The 24 shapes that do not split are untouched, measured at -0.1% median over the full suite. One xclbin still serves every shape. A unit is weighted by the C descriptors it drains (M_CHUNK under c_split), not counted as one: counting units would overrun the 4-deep queue by exactly M_CHUNK, and overrunning it hangs rather than diagnoses. Also drops three knobs that no longer have a consumer -- the overlap parameter (never passed), col_work/col_drain (superseded by the core deriving its own trip counts) and the mt_*_bytes sizes (orphaned when B residency went) -- and the prose describing mechanisms this design no longer has. Co-Authored-By: Claude --- iron/operators/flm/gemm/README.md | 74 +++++---- iron/operators/flm/gemm/design.py | 250 ++++++++++++++---------------- iron/operators/flm/gemm/op.py | 4 +- iron/operators/flm/gemm/test.py | 22 +-- 4 files changed, 173 insertions(+), 177 deletions(-) diff --git a/iron/operators/flm/gemm/README.md b/iron/operators/flm/gemm/README.md index 5bc3215ee..141f2dc46 100644 --- a/iron/operators/flm/gemm/README.md +++ b/iron/operators/flm/gemm/README.md @@ -273,9 +273,8 @@ to reduce over -- with a single k iteration there is not enough compute to hide the extra A traffic. On NPU2, `tile_n=128` wins only at `k_iters=1`, and by ~3%; at `k_iters>=2` it -is 1.2-1.7x slower. Both followed from `tile_n=128` giving up resident B — its -`mt_b` is 128 KB, so `k_iters` copies do not fit the memtile — and the more k -there is to reduce over, the more that costs. +is 1.2-1.7x slower, because its `CT_MAX_K` falls to 32 and the compute cost of +the shorter k slice swamps the A traffic it saves. NPU1 never reaches that crossover. It has half the columns *and* a quarter of the per-tile bf16 mac throughput, so it stays compute-bound at every K, and @@ -351,7 +350,6 @@ which is why the gap grows with the problem size rather than being flat. Unlike NPU2, compute here is **not** hidden behind the transfers, so on NPU1 both a faster mmul and less traffic pay off, where on NPU2 only the latter does. -Resident B was also a no-op on NPU1 — see [Resident B, removed](#resident-b-removed). ### Why the transfers are cheap @@ -367,27 +365,47 @@ Two things, both in the runtime sequence rather than the kernel: block then costs 3 shim buffer descriptors instead of `1 + 2*k_iters`, so two can be in flight without exhausting the 16 available. -### Resident B, removed - -B used to be held in the memtile across row-blocks where a whole column-block -fit double-buffered, so DDR read it once instead of `m_row_blocks` times -- -about 43% less traffic. On NPU2 that was worth 12.5% at M=512, 16.4% at -M=1024 and 19.3% at M=2048 (K=1024 N=4096). On NPU1 it was neither a latency -nor a power win. - -**It is gone, and that is the price of the runtime parameters.** Residency -sizes the memtile buffer from `k_iters` and replays it `m_row_blocks` times -through a buffer descriptor's repeat count, so it puts both K and M in the -device configuration — and the configuration is what one xclbin has to share -across every shape. - -Removing it does lift a hard cap: the repeat count expands into the memtile's -BD chain at 2 blocks per replay, and at `m_row_blocks = 16` that chain -exceeded its 48-block limit, so no shape with K <= 2048 would build at M=4096. -Every shape builds there now. - -Restoring it needs a replay mechanism that carries neither K nor M into the -configuration. That is the largest known lever left here, and it is worth more -than the figures above suggest, because `repeat_count` restarting the memtile -BD chain at every replay boundary was already giving part of the traffic -saving back. +### B is re-fetched per row-block + +DDR reads B `m_row_blocks` times rather than once. Holding a whole column-block +in the memtile and replaying it would size that buffer from `k_iters` and set +the replay from `m_row_blocks`, putting **both K and M into the device +configuration** — and the configuration is what one xclbin has to share across +every shape. That is the standing cost of M, K and N being runtime parameters, +and it is why B is the dominant DDR leg here. + +**M is the tractable half.** `aiex.npu.push_queue` takes both `bd_id` and +`repeat_count` as SSA operands, and `aiex.dma_channel_reset_for(@fifo)` expands +into the whole re-arm trio a resident fifo needs — channel reset, `aiex.set_lock` +per bound lock, START_QUEUE re-push — inside the **runtime sequence**, which this +operator regenerates per shape. So a per-shape replay count does not have to +reach the xclbin. All of it is reachable from Python and has an npu2 device test +(1000 dispatches on one hardware context). + +**K is the part still in the way.** Correct ordering needs one memtile object +spanning every k-block, so the buffer is sized from `k_iters` and that sizing is +device configuration. Selecting among several pre-programmed BD chains via +`push_queue`'s runtime `bd_id` is the obvious line of attack and has not been +tried. + +### Split legs retire rolling, not in windows + +Where K or N is 10240, the row-block stride overflows the shim BD's 20-bit +iteration step and that leg is issued as one transfer per row-block. Two shim +resources bound how many may be outstanding, and neither is modelled by the +toolchain: BD ids (16/tile, freed without a completion check) and the channel +task queue (4 deep, pushed unconditionally). + +The sequence retires the **oldest** transfer as it issues the next, which +bounds both resources directly while keeping the channel full. + +**Do not "simplify" this into windowing** — issue four, await the whole window, +issue the next four. That bounds the same two resources and reads more simply, +but it drains the channel to *empty* at every window boundary and again at every +column-block boundary, and on a DDR-rate-bound design those bubbles are the +entire cost of the split path. Measured at up to **-12.4%** on the shapes that +take this path (E4B/gateup M1024), for no change in what is in flight. + +`m_chunk` takes this path too, since its only structural effect is to force the +split on for A. It is off by default regardless — see `M_CHUNK_FOR_N` in +design.py, which would fork the xclbin. diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index 14f422822..731e15dde 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -54,6 +54,7 @@ """ import argparse +import os from enum import StrEnum from functools import partial @@ -161,23 +162,26 @@ def compute_rows(dev): # trick a_split already uses for the mega_row dimension, and it reuses the # same windowing. What is left per transfer is # [k_iters, M_CHUNK, M_TILE, K_TILE], which fits. -# DEFAULT 1 EVERYWHERE -- i.e. off. The mechanism is correct and tested, but -# interleaved A/B at K=1024 N=4096 (10 rounds, round-robin in one process) -# does not show a consistent win: +# DEFAULT 1 EVERYWHERE -- i.e. off, and for a contractual reason rather than a +# performance one. m_chunk must DIVIDE m_row_blocks, so it needs M to be a +# multiple of 512, while the overlay this operator replaces accepts any +# multiple of 256 (mm_prebuilt/design.py:65, MIN_M = M_TILE*ROWS). A shape that +# cannot use m_chunk falls back to 1 and so FORKS _config_tag, and serving +# M=256 and M=2048 from ONE xclbin is a hard requirement here. # -# M m_chunk=1 min/med m_chunk=2 min/med min med -# 512 525.4 / 622.7 502.1 / 566.3 -4.4% -9.1% -# 1024 959.0 / 1081.6 1092.4 / 1162.9 +13.9% +7.5% -# 2048 1902.8 / 1986.0 1837.1 / 1940.0 -3.5% -2.3% +# The leftover cannot be padded away either: a partial group is inexpressible +# (see the raise below), and duplicating the row-block to fill the group is +# expressible and bit-exact but doubles A and C for that dispatch, which at +# m_row_blocks=1 costs more than anything m_chunk could win back. # -# B's DDR traffic really does fall by m_chunk, but something else eats it -- -# most likely the per-unit A descriptors, since m_chunk forces the split path -# for A (see a_split). The M=1024 regression is not monotonic in n_units and -# is not explained; do not turn this on by default without understanding it. +# Note m_chunk does NOT change A's byte count (A is M*K*2*n_col_blocks either +# way). Its only structural effect is forcing a_split on, so if M % 512 == 0 +# ever becomes guaranteed it is worth re-measuring. # -# An earlier NON-interleaved measurement showed -11%/-15% and was wrong; see -# the bimodal-timing discipline in README.md. Interleave before believing any -# number here. +# `n_chunk` is the lever with the same effect and no such constraint: it groups +# COLUMN-blocks, which the core already walks with an RTP-derived trip count, +# so a partial group is a runtime bound rather than a descriptor shape and one +# xclbin still serves every M. M_CHUNK_FOR_N = {16: 1, 32: 1, 64: 1, 128: 1} # How many column-blocks the runtime sequence keeps in flight. A block costs 3 # shim buffer descriptors on a column (A + B + C) against 16 available, so the @@ -398,7 +402,6 @@ def gemm( tile_n=N_TILE_DEFAULT, m_chunk=None, tile_ma=None, - overlap=None, kernel_object="mm_fused.o", trace_size=0, ): @@ -462,7 +465,7 @@ def gemm( M_CHUNK, ) RHO = M_TILE // T_MA - OVERLAP = OVERLAP_DEFAULT if overlap is None else overlap + OVERLAP = OVERLAP_DEFAULT K_DIV_CT_K_MAX = K_TILE // CT_MAX_K CT_A_LEN = 2 * R * CT_MAX_K # one z slice CT_A_OBJ = CT_A_LEN * (T_MA // R // 2) # every z slice of one mmul @@ -484,10 +487,10 @@ def gemm( clamp_max_bits = int(np.float32(clamp_hi).view(np.int32)) # A compute tile does a whole m x n block or nothing, so M and K must tile # exactly. N need only be a multiple of N_TILE: a trailing group of fewer - # than COLS blocks is handled by giving the columns different trip counts - # (see col_work / col_drain below). That matters in practice -- for a - # transformer the o and down projections have N = model dim, which is - # essentially never a multiple of N_TILE*COLS. + # than COLS blocks is handled by giving the columns different trip counts, + # which each core derives for itself (see core_fn). That matters in + # practice -- for a transformer the o and down projections have N = model + # dim, which is essentially never a multiple of N_TILE*COLS. for name, value, unit in (("M", M, MIN_M), ("K", K, MIN_K), ("N", N, N_TILE)): if value % unit != 0: raise ValueError(f"{name} ({value}) must be a multiple of {unit}") @@ -558,49 +561,39 @@ def unit_rows(u): # toolchain models, and overrunning it hangs rather than diagnoses. # Measured at K=10240 M=1024: 4 outstanding run, 8 hang. # - # So a split block is emitted in WINDOWS of at most SHIM_TASK_QUEUE - # mega_rows, each window awaited before the next is issued (see sequence()). - # Awaiting is what makes the window's descriptors safe to reuse, and it - # bounds both resources at once. M<=1024 is a single window, so the shapes - # that already worked are unaffected. - # A window is counted in TRANSFERS, not units: under c_split a unit emits - # one C descriptor per row-block (M_CHUNK of them, see c_taps), so a - # window of SHIM_TASK_QUEUE units would push M_CHUNK times that many onto - # one channel and hang. + # So a split leg keeps at most SHIM_TASK_QUEUE transfers outstanding, and + # emit_split() enforces that directly -- retiring the OLDEST transfer as it + # issues the next, so the channel stays full. + # + # It used to do this by windowing: issue SHIM_TASK_QUEUE transfers, await + # the whole window, then issue the next. That bounds the queue too, but it + # drains the channel to EMPTY at every window boundary and again at every + # column-block boundary, and on a DDR-rate-bound design those bubbles are + # the entire cost of the split path. Rolling instead of windowing measured, + # bit-exact, 8 rounds x 30 iters interleaved: + # + # E4B/gateup M1024 4119.4 -> 3617.7 -12.2% + # E4B/gateup M2048 7664.8 -> 7182.4 -6.3% + # E4B/down M1024 3649.2 -> 3553.4 -2.6% + # E4B/down M2048 7196.6 -> 6978.7 -3.0% + # + # The bound is counted in TRANSFERS, not units: under c_split a unit drains + # one C descriptor per row-block (M_CHUNK of them, see c_taps) and they all + # ride one channel, so counting units would overrun by exactly M_CHUNK -- + # at m_chunk=2 that is 8 outstanding, the measured hang threshold. _per_unit = M_CHUNK if c_split else 1 - MB_WINDOW = ( - min(n_units, max(1, SHIM_TASK_QUEUE // _per_unit)) - if (a_split or c_split) - else 1 - ) - bds_per_block = ( - 1 + (MB_WINDOW if a_split else 1) + (MB_WINDOW * _per_unit if c_split else 1) - ) - # Unreachable while SHIM_TASK_QUEUE is 4 (the worst case is 1 + 4 + 4 = 9 - # of 16), so this guards a future retune of the window rather than any - # shape reachable today. test_flm_gemm_split_leg_windowing asserts the same - # arithmetic from the outside. - if bds_per_block > SHIM_BDS: + # Live descriptors on a shim tile: SHIM_TASK_QUEUE from the rolling window, + # plus B and the unsplit leg for each of the two blocks a boundary spans. + bds_per_block = SHIM_TASK_QUEUE + 2 + 2 + if (a_split or c_split) and bds_per_block > SHIM_BDS: raise ValueError( f"M={M} K={K} N={N} needs {bds_per_block} shim buffer descriptors " - f"per window (1 B + {MB_WINDOW if a_split else 1} A + " - f"{MB_WINDOW if c_split else 1} C) but a shim tile has only " - f"{SHIM_BDS}." + f"for the split path but a shim tile has only {SHIM_BDS}." ) - # Cross-block overlap only applies to the unsplit path; a split block - # already awaits inside itself, so keeping a second one in flight would - # refill the very queue the windowing just drained. - # Only a block that spans MORE THAN ONE window awaits inside itself; that - # is what makes a second block in flight refill the queue the windowing - # just drained. A single-window block has no intra-block await, so it can - # still overlap -- and it must, because m_chunk forces a_split on for the - # descriptor's sake even when nothing needs windowing, and losing the - # pipelining costs far more than the B traffic m_chunk saves (measured - # +13% at M=1024 before this). - if (a_split or c_split) and n_units > MB_WINDOW: - OVERLAP = 1 - else: - OVERLAP = max(1, min(OVERLAP, SHIM_BDS // bds_per_block)) + # The unsplit path pipelines whole column-blocks instead, at 3 descriptors + # each (A + B + C). The split path does its own bounding above and ignores + # this. + OVERLAP = max(1, min(OVERLAP, SHIM_BDS // 3)) # Sweeps where all COLS columns have work, plus a trailing group of # rem_blocks columns (0 <= rem_blocks < COLS) that do one block more. n_full = N // MIN_N @@ -609,15 +602,6 @@ def unit_rows(u): # exist is configuration and this design has one configuration. A column # with no work for this shape gets n_work = 0 and drains instead. n_active_cols = COLS - # Per column: how many column-blocks it computes, and how many it sits out - # while still draining the A broadcast for its row. Every block issues A - # for every compute row, so a column that does not compute a block must - # still take that block's A or the row stalls -- the memtile will not - # release an A object until all COLS consumers have taken it. - total_blocks = n_full + (1 if rem_blocks else 0) - col_work = [n_full + (1 if c < rem_blocks else 0) for c in range(COLS)] - col_drain = [total_blocks - w for w in col_work] - # B's element type: one v8bfp16ebs8 per 8 values on AIE2P, one bf16 per # value on AIE2. Every B extent below is therefore in values // B_GROUP. b_elem_ty = np.dtype[v8bfp16ebs8] if BFP16_B else bf16_ty @@ -630,11 +614,8 @@ def unit_rows(u): # M_CHUNK stacked row-block tiles, so the forward below can interleave # them on the way out -- see a_send_dims. mt_a_ty = np.ndarray[(M_CHUNK * M_TILE * K_TILE,), bf16_ty] - mt_a_bytes = M_CHUNK * M_TILE * K_TILE * 2 mt_b_ty = np.ndarray[(K_TILE * N_TILE // B_GROUP,), b_elem_ty] - mt_b_bytes = _b_bytes(K_TILE * N_TILE, BFP16_B) mt_out_ty = np.ndarray[(C_SLICE_LEN * ROWS,), bf16_ty] - mt_out_bytes = C_SLICE_LEN * ROWS * 2 # L3 (DDR), flat -- the taps below index them linearly. a_l3_ty = np.ndarray[(M * K,), bf16_ty] b_l3_ty = np.ndarray[(K * N // B_GROUP,), b_elem_ty] @@ -741,38 +722,18 @@ def unit_rows(u): # B: shim -> memtile -> broadcast down the compute column. # - # Where it fits, a whole column-block's B is held in the memtile as ONE - # object and re-walked per row-block, so DDR reads it once instead of - # m_row_blocks times. B is the dominant DDR leg, so that is roughly 43% less - # total traffic; the latency it buys grows with M, because B's re-reads - # scale with m_row_blocks. Larger K does not fit and falls back to - # re-reading. See README.md for the measured effect. + # Placement has zero slack -- the eight memtiles pack to exactly 512 KB, + # relying on aie-objectfifo-allocate spilling one buffer to an adjacent + # tile -- so re-verify it after any change to the A, B or C buffer sizes. # - # Three things here are load-bearing rather than tuning: - # - # * ONE object spanning every k-block, not a pool of k_iters objects. - # Iterating a pool replays each object in turn (k0,k0,k1,k1,...) rather - # than the k0..kn sequence the cores accumulate in. - # * repeat_count on the forward() below is what re-sends an object. - # iter_count only bounds how many times an end cycles through all its - # buffers, so it is in units of depth-cycles; getting it wrong hangs - # rather than mis-computing. - # * The depth search takes the deepest that fits, not depth 1. Single - # buffering stops the next column-block prefetching behind this one's - # replay, which measures worse than not being resident at all. - # - # The budget must count what A and C actually occupy: at tile_n=128 C - # doubles and a resident B is 432 KB, and a budget that assumed C's size - # admitted a configuration that then failed address assignment. Placement - # has zero slack -- the eight memtiles pack to exactly 512 KB, relying on - # aie-objectfifo-allocate spilling one buffer to an adjacent tile -- so - # re-verify it after any change to the A, B or C buffer sizes. - # B residency is gone: it sizes the memtile buffer from k_iters and - # replays it m_row_blocks times through the forward()'s repeat_count, so - # it carries both K and M into the configuration, which is what the - # runtime parameters exist to remove. See README.md for what that cost. - b_resident = False - + # One k-block per object, re-fetched from DDR for every row-block. Holding + # a whole column-block here instead would size the buffer from k_iters and + # replay it m_row_blocks times, putting both K and M into the device + # configuration -- which is what the runtime parameters exist to keep out + # of it. The M half is tractable now (aiex.dma_channel_reset_for re-arms a + # resident fifo from the RUNTIME SEQUENCE, and push_queue's repeat is an + # SSA operand); the k_iters-sized buffer is the part still in the way. + # See README.md. b_l3l2_fifos = [] b_cons = {} for c in range(n_active_cols): @@ -848,8 +809,8 @@ def core_fn(accs, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, my_col, barrier # => n_work = ceil((N/N_TILE - c) / COLS) # # and every divisor is a power of two (N_TILE=64, COLS=8), so this is - # shifts and adds -- no __divsi3. Verified against the host-side - # col_work/col_drain for every shape in the suite. + # shifts and adds -- no __divsi3. Verified against host-side trip + # counts for every shape in the suite. # // rather than >>: the DSL's ScalarValue overloads add/sub/mul/ # floordiv/mod but NOT the shift operators. Both divisors are # compile-time powers of two (N_TILE=64, COLS=8), so this strength- @@ -1045,14 +1006,8 @@ def b_tap(mega_col, c): # n_units times instead of m_row_blocks. That is the whole win -- # see M_CHUNK_FOR_N. The unit dimension keeps stride 0, replaying # the same k-blocks, exactly as the row-block dimension used to. - sizes=( - [1, 1, 1, k_iters * K_TILE * N_TILE // B_GROUP] - if b_resident - else [n_units, k_iters, 1, K_TILE * N_TILE // B_GROUP] - ), - strides=( - [0, 0, 0, 1] if b_resident else [0, K_TILE * N_TILE // B_GROUP, 0, 1] - ), + sizes=[n_units, k_iters, 1, K_TILE * N_TILE // B_GROUP], + strides=[0, K_TILE * N_TILE // B_GROUP, 0, 1], ) def c_taps(mega_col, c, units): @@ -1203,43 +1158,64 @@ def emit_unsplit(): tg.finish() def emit_split(): - # A split leg is one transfer per mega_row, so a whole block at - # once would overrun the shim channel's task queue. Emit MB_WINDOW - # mega_rows at a time and retire each window before the next, which - # both drains the queue and -- because the window's transfers are - # awaited, not merely freed -- makes its descriptors safe to reuse. - # - # B stays live across the whole block: its descriptor replays over - # every mega_row, so freeing it per window would hand its - # descriptor away mid-flight. It is retired last, after every - # window's C has been awaited, which is what guarantees it drained. + """Issue split legs one unit at a time, retiring the oldest. + + One TaskGroup covers one unit (one task per channel), and the + oldest is retired only when a new one would exceed SHIM_TASK_QUEUE + outstanding -- so the channel stays full across both unit and + column-block boundaries. See the SHIM_TASK_QUEUE comment above for + why retiring in batches instead costs real time. + + ``pending`` is retired strictly in append order, which is what + keeps a block's B and unsplit-leg descriptors alive until every one + of that block's units has been AWAITED: they are appended after the + units they cover. + """ + # What one unit costs on the BUSIEST single channel, which is what + # SHIM_TASK_QUEUE bounds. Under c_split a unit drains one C + # descriptor per row-block (M_CHUNK of them, see c_taps) and they + # all ride the same column's channel, so a unit is worth _per_unit + # there; the a_split leg is one descriptor per row-fifo, so one. + # Counting units instead would overrun by exactly M_CHUNK -- at + # m_chunk=2 that is 8 outstanding, the measured hang threshold. + unit_cost = _per_unit if c_split else 1 + pending = [] # (group, queue cost), oldest first + + def retire(limit): + while sum(q for _, q in pending) > limit: + pending.pop(0)[0].finish() + for mega_col, active_cols in blocks: tg_b = TaskGroup() issue_b(mega_col, active_cols, tg_b) - # The leg that did NOT split is still one task for the whole - # block -- its single descriptor already spans every mega_row, - # so re-issuing it per window would transfer the block twice. - # It stays live alongside the windows and retires with them. tg_whole = TaskGroup() if not c_split: issue_c(mega_col, active_cols, all_mb, tg_whole) if not a_split: issue_a(mega_col, all_mb, tg_whole) - for w in range(0, n_units, MB_WINDOW): - mbs = all_mb[w : w + MB_WINDOW] - tg_w = TaskGroup() + for u in all_mb: + tg_u = TaskGroup() if c_split: - issue_c(mega_col, active_cols, mbs, tg_w) + issue_c(mega_col, active_cols, [u], tg_u) if a_split: - # wait=True: the await is what makes this window's - # descriptors reusable by the next. - issue_a(mega_col, mbs, tg_w, wait=True) - tg_w.finish() - - tg_whole.finish() - tg_b.finish() + issue_a(mega_col, [u], tg_u, wait=True) + pending.append((tg_u, unit_cost)) + retire(SHIM_TASK_QUEUE) + + # Not queue-counted: B rides one channel per column and the + # unsplit leg one per row, neither of which is the channel the + # units contend for. They still retire in order, after the + # units of their own block. + pending.append((tg_whole, 0)) + pending.append((tg_b, 0)) + + # Drain everything, not retire(0): the tail groups are weighted 0, + # so a count-driven loop stops with them still open and the build + # fails with "Failed to close task groups". + for tg, _ in pending: + tg.finish() emit_split() if (a_split or c_split) else emit_unsplit() diff --git a/iron/operators/flm/gemm/op.py b/iron/operators/flm/gemm/op.py index 576c097fe..0bddb6e3e 100644 --- a/iron/operators/flm/gemm/op.py +++ b/iron/operators/flm/gemm/op.py @@ -463,8 +463,8 @@ def pack_B(self, B): Packing all the way to consumption order is what makes both B hops linear descriptors (design.py's b_recv_dims and b_send_dims are both - None), which in turn leaves the descriptor dimensions for a k slice deep - enough to halve the accumulator traffic while B is also memtile-resident. + None), which in turn leaves the descriptor dimensions for a k slice + deep enough to halve the accumulator traffic. See :mod:`iron.operators.flm.packing` for the layout itself. """ return pack_b( diff --git a/iron/operators/flm/gemm/test.py b/iron/operators/flm/gemm/test.py index 754f838da..e11fc9608 100644 --- a/iron/operators/flm/gemm/test.py +++ b/iron/operators/flm/gemm/test.py @@ -201,27 +201,29 @@ def test_gemm(M, K, N, epilogue, clamp, rounding, aie_context): def test_gemm_split_leg_windowing(aie_context): """K or N = 10240 at M > 256 overflows the shim BD's 20-bit mega_row - iteration step, so that leg is issued as one transfer per mega_row, - retired in windows of at most SHIM_TASK_QUEUE. Two shim resources bound it - and NEITHER is modelled by the toolchain -- the BD ids (16/tile, freed - without a completion check) and the channel task queue (4 deep, pushed + iteration step, so that leg is issued as one transfer per mega_row, with at + most SHIM_TASK_QUEUE of them outstanding. Two shim resources bound it and + NEITHER is modelled by the toolchain -- the BD ids (16/tile, freed without + a completion check) and the channel task queue (4 deep, pushed unconditionally) -- so overrunning either is a silent device hang rather than a diagnostic. - Windowing keeps both inside their limits for every shape: at most - 1 B + 4 A + 4 C = 9 of 16 descriptors, and at most 4 outstanding per - channel. Assert that arithmetic here, since the numbers come from the - hardware and a future retune of SHIM_TASK_QUEUE could break it silently. + The sequence retires the OLDEST transfer as it issues the next rather than + draining a whole window, so the live set is SHIM_TASK_QUEUE transfers plus + B and the unsplit leg for each of the two column-blocks a boundary spans: + 4 + 2 + 2 = 8 of 16 descriptors, and at most 4 outstanding per channel. + Assert that arithmetic here, since the numbers come from the hardware and a + future retune of SHIM_TASK_QUEUE could break it silently. """ from aie.dialects.aie import get_target_model from iron.operators.flm.gemm.design import SHIM_TASK_QUEUE dev = aie_utils.get_current_device() available = get_target_model(dev.resolve()).get_num_bds(0, 0) - worst = 1 + 2 * SHIM_TASK_QUEUE + worst = SHIM_TASK_QUEUE + 2 + 2 assert worst <= available, ( f"a fully split block needs {worst} shim BDs of {available}; " - "windowing no longer fits and the split shapes will hang" + "the split shapes will hang" ) # The square case splits BOTH legs, which the real Gemma shapes never do From 560b0b8e9a71cb93aa633d5a0853910b0c632cdb Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Mon, 14 Sep 2026 17:16:06 -0600 Subject: [PATCH 10/14] flm_gemm: cut the comment volume roughly in half This operator had accumulated about one line of prose per line of code, and the RTP work made it worse rather than better -- design.py went from 429 comment and docstring lines to 633 while the code grew by 120. Most of that was not explaining the code. It was measurement tables that belong in README.md, rationale for alternatives that were tried and rejected, and restatements of the line underneath. Several facts were stated two or three times over: that the build cache keys on filename, that a second A fifo would want a third input DMA channel, that C drains in row-block order even under M_CHUNK. What is kept is the class of comment that is expensive to lose, because the failure it describes is silent: the BD-id recycle hazard that corrupts rather than faulting, the shim task-queue depth that hangs rather than diagnosing, _VERIFIED_CT_K where a wrong entry computes the wrong answer, and the memtile placement pin. Those are stated once, at the code they constrain. design.py 42% -> 33% prose op.py 45% -> 35% test.py 41% -> 33% benchmark.py 42% -> 35% mm_fused.cc 46% -> 36% Also corrects documentation that had gone stale. test.py still described the split legs as retired in windows, which they have not been since the rolling retire landed, so its docstring, an inline comment and two entries in the shape table were all describing a mechanism that no longer exists. test_gemm_split_leg_windowing is renamed to test_gemm_split_leg_bounds for the same reason. Drops an unused `import os`. No functional change: the AST with docstrings stripped is identical except for that import and the two renames. 185/185 pass. Co-Authored-By: Claude --- aie_kernels/generic/mm_fused.cc | 133 ++--- iron/operators/flm/gemm/benchmark.py | 99 ++-- iron/operators/flm/gemm/design.py | 710 +++++++-------------------- iron/operators/flm/gemm/op.py | 203 +++----- iron/operators/flm/gemm/test.py | 142 ++---- 5 files changed, 372 insertions(+), 915 deletions(-) diff --git a/aie_kernels/generic/mm_fused.cc b/aie_kernels/generic/mm_fused.cc index 54577df51..aae8b0be7 100644 --- a/aie_kernels/generic/mm_fused.cc +++ b/aie_kernels/generic/mm_fused.cc @@ -2,22 +2,14 @@ // SPDX-License-Identifier: Apache-2.0 // bf16 GEMM compute kernel. Each compute tile owns an m x n slice of C and -// accumulates over K into an f32 accumulator that stays in L1 for the whole -// reduction. -// -// Three entry points, each called once per iteration of a loop nest that lives -// in the design (iron/operators/flm/gemm/design.py) rather than here: +// accumulates over K into an f32 accumulator that stays in L1. // // mm_fused_acc_init zero the accumulator, once per output tile // mm_fused_k_step multiply one A band by one B chunk into it // mm_fused_epilogue_chunk drain one chunk of it to a bf16 C object // -// The nest lives in the design so that every level of it has an ObjectFifo -// acquire point, a fifo consumer having to acquire once per object. -// -// Tile geometry arrives as -D flags from design.py, which is the single source -// of truth for it: the same constants size the design's buffers and set its -// unroll factors. +// The loop nest is in design.py, not here, so every level has an ObjectFifo +// acquire point. Tile geometry arrives from there as -D flags. #include "../aie_kernel_utils.h" #include "activations.h" #include "mm_fused_mmul.h" @@ -38,11 +30,8 @@ #ifndef MM_FUSED_EPILOGUE_MODE_MASK #define MM_FUSED_EPILOGUE_MODE_MASK 0xF #endif -// Whether a clamp is compiled in at all. Like the mode mask above this is a -// CAPABILITY, not a selection: compiling the clamped path costs program -// memory, so a build that never clamps should not carry it. The BOUNDS are -// runtime parameters (see mm_fused_epilogue_chunk), so every pair of bounds -// shares one build. +// Whether a clamp is compiled in at all: a capability, not a selection, since +// the clamped path costs program memory. The bounds are runtime. #ifndef MM_FUSED_CLAMP #define MM_FUSED_CLAMP 0 #endif @@ -50,27 +39,20 @@ namespace { constexpr int M = MM_FUSED_TILE_M; -// Asymmetric tile buffering: the A tile spans MA rows while the accumulator -// spans M, so the core folds RHO = M / MA A bands into one C tile before -// releasing it. A dies as soon as it is consumed while C must live across the -// whole K reduction, so sizing both to M would pay the peak L1 cost twice. -// MA == M is the symmetric case. +// Asymmetric tile buffering: A spans MA rows and the accumulator M, so the +// core folds RHO = M / MA A bands into one C tile before releasing it. A dies +// on consumption while C lives across the K reduction, so sizing both to M +// would pay the peak L1 cost twice. MA == M is the symmetric case. // -// Technique from "Can Asymmetric Tile Buffering Be Beneficial?", C. Wang, -// W. Pang, X. Wu, G. Jun, L. Romero, E. Taka, D. Marculescu, T. Nowatzki, -// P. Vasireddy, J. Melber, D. Chen, J. Cong, arXiv:2511.16041 (2025), -// https://arxiv.org/abs/2511.16041. Reference AIE implementation is -// Xilinx/mlir-aie PR #3076 by @ChengyueWang, in -// programming_examples/ml/block_datatypes/gemm_asymmetric_tile_buffering. -// Those configs accumulate in bf16/bfp16, which is what affords their larger C -// tiles; this kernel keeps an f32 accumulator, so here the win comes from -// spending the freed L1 on a deeper k slice rather than on a wider C tile. +// From arXiv:2511.16041, "Can Asymmetric Tile Buffering Be Beneficial?"; +// reference AIE implementation in Xilinx/mlir-aie PR #3076. Those configs +// accumulate in bf16/bfp16, affording larger C tiles; this one keeps an f32 +// accumulator, so the freed L1 buys a deeper k slice instead. constexpr int MA = MM_FUSED_TILE_MA; constexpr int K = MM_FUSED_TILE_K; constexpr int N = MM_FUSED_TILE_N; -// Register tiling, and how much of K one compute tile holds at a time. Both are -// design.py's to choose -- CT_K in particular trades against the n width for a -// fixed L1 budget. +// Register tiling, and how much of K a tile holds at a time. Both are +// design.py's to choose; CT_K trades against the n width for a fixed budget. constexpr int R = MM_FUSED_R; constexpr int S = MM_FUSED_S; constexpr int T = MM_FUSED_T; @@ -90,15 +72,11 @@ static_assert(N % (2 * T) == 0, "tile_n must be a multiple of 2*t (2x2 mmul)"); static_assert(K % CT_K == 0, "tile_k must be a multiple of the k slice"); static_assert(CT_K % S == 0, "k slice must be a multiple of s"); -// The core powers up in rounding_mode::floor, so a kernel that converts must +// The core powers up in rounding_mode::floor, so a converting kernel must // choose explicitly. Truncation biases every conversion the same direction, so -// the error accumulates over the K reduction instead of cancelling -- ~1% of -// the result, against ~0.02% for round-to-nearest-even, which is far more than -// the bfp16 emulation itself costs. Every entry point that converts sets it: -// the mmul and the epilogue's f32->bf16 store, both below. -// -// Flag name and polarity follow mm.cc, so the two kernels are configured the -// same way; the operator passes -DROUND_CONV_EVEN by default. +// the error accumulates over the K reduction: ~1% of the result against ~0.02% +// for round-to-nearest-even, far more than the bfp16 emulation costs. Flag +// name and polarity follow mm.cc. #ifdef ROUND_CONV_EVEN constexpr aie::rounding_mode round_mode = aie::rounding_mode::conv_even; #else @@ -118,10 +96,9 @@ epilogue_body(bfloat16 *__restrict y_out, const float *__restrict src, float cla AIE_LOOP_MAX_ITERATION_COUNT(CHUNK / V) for (int j = 0; j < CHUNK / V; j++) { - // The accumulator stays f32 through the activation and the clamp, and - // is converted to bf16 exactly once, on the store. Converting first - // would round twice and let the activation's slope amplify the first - // rounding -- see activations.h. + // f32 through the activation and clamp, converted exactly once on the + // store. Converting first would round twice and let the activation's + // slope amplify the first rounding. aie::vector f = aie::load_v(src + j * V); if constexpr (MODE == 1) f = gelu_vec(f); @@ -140,9 +117,8 @@ epilogue_body(bfloat16 *__restrict y_out, const float *__restrict src, float cla } } -// Pick the clamped or unclamped instantiation for a mode. Only one of the two -// exists unless MM_FUSED_CLAMP compiled the clamp in, so a build that never -// clamps pays nothing for this. +// Pick the clamped or unclamped instantiation. Only one exists unless +// MM_FUSED_CLAMP compiled the clamp in. template static inline void epilogue_dispatch(bfloat16 *__restrict y_out, const float *__restrict src, @@ -162,26 +138,20 @@ static inline void epilogue_dispatch(bfloat16 *__restrict y_out, extern "C" { -// Zero the f32 accumulator, before the k loop starts accumulating into it. -// -// A bias is deliberately not supported: initialising the accumulator from one -// would mean consuming an extra object through the handshake the B ObjectFifo -// owns, which desynchronises that fifo and hangs rather than mis-computing. +// Zero the f32 accumulator before the k loop accumulates into it. A bias is +// deliberately unsupported: initialising from one would consume an extra +// object through the B fifo's handshake, desynchronising it into a hang. void mm_fused_acc_init(float *y_acc) { // zero_vectorized brackets itself in event0/event1 for tracing. zero_vectorized(y_acc); } -// One step of the k loop: one B chunk multiplied against one A band, -// accumulated into y_acc. +// One step of the k loop: one B chunk against one A band, into y_acc. // -// Takes no locks. A is a single object spanning every z slice of the mmul, and -// the A and B fifos own the handshake, so the core body acquires around this -// call rather than the kernel acquiring inside it. -// mm_fused_b_elem_t is bfp16ebs8 or bfloat16 depending on how B is stored, -// which mm_fused_mmul.h selects from the architecture. One signature either -// way, so the design's Kernel declaration does not have to care. +// Takes no locks -- the A and B fifos own the handshake, so the core body +// acquires around this call. mm_fused_b_elem_t is bfp16ebs8 or bfloat16 +// depending on the architecture, but the signature is the same either way. void mm_fused_k_step(bfloat16 *a_buf, mm_fused_b_elem_t *b_buf, float *y_acc, int32_t band) { ::aie::set_rounding(round_mode); @@ -190,30 +160,17 @@ void mm_fused_k_step(bfloat16 *a_buf, mm_fused_b_elem_t *b_buf, float *y_acc, in mm_fused_mmul_2x2<(MA / R), (CT_K / S), (N / T), R, S, T>(a_buf, b_buf, y_acc + band * (MA * N)); } -// The output stage: convert chunk (outer * C_DEPTH + half) of the f32 -// accumulator into a bf16 C object the core body has already acquired from the -// C ObjectFifo, optionally applying an activation and a clamp on the way out. -// -// Fusing the activation here is the point: the values are already in registers -// after the f32 -> bf16 conversion, so gelu/silu/sigmoid costs one more vector -// op per 16 elements instead of a separate pass over L1 (which is what chaining -// a standalone activation operator after a GEMM would cost). -// -// The mode is a runtime argument, because one xclbin serves every activation. -// It is tested once per chunk, outside the vector loop, so each mode still runs -// a branch-free inner loop; the cost is program memory, since every mode in -// MM_FUSED_EPILOGUE_MODE_MASK is compiled in. -// -// The clamp follows the same split: whether a clamped path exists at all is -// compile-time (MM_FUSED_CLAMP, since it costs program memory), while -// clamp_enabled and the BOUNDS are runtime, so every pair of bounds shares one -// build. The bounds arrive as raw int32 bit patterns because the RTP mechanism -// (aie.dialects.aie.npu_write_rtp) only writes i32 words; design.py bit-casts -// them on the host side. +// Convert chunk (outer * C_DEPTH + half) of the f32 accumulator into a bf16 C +// object the core body already acquired, applying an activation and clamp on +// the way out. // -// The chunk index is split in two because the core body unrolls the drain by -// the C fifo depth to keep the acquired buffer index a compile-time constant; -// passing both parts avoids doing that arithmetic up there. +// Fusing the activation is the point: the values are already in registers, so +// gelu/silu/sigmoid costs one more vector op per 16 elements rather than a +// separate pass over L1. The mode is runtime, tested once per chunk so the +// inner loop stays branch-free; the cost is program memory, since every mode +// in the mask is compiled in. The clamp splits the same way, its bounds +// arriving as raw int32 because npu_write_rtp only writes i32 words. The chunk +// index comes in two parts because the core body unrolls the drain. void mm_fused_epilogue_chunk(bfloat16 *y_out, float *y_acc, int32_t outer, @@ -228,8 +185,7 @@ void mm_fused_epilogue_chunk(bfloat16 *y_out, ::aie::set_rounding(round_mode); const float *__restrict src = y_acc + (outer * C_DEPTH + half) * CHUNK; // __builtin_bit_cast, not memcpy: memcpy leaves an unresolved external - // call in the compiled object here rather than folding to a register move, - // and this runs once per chunk drain. + // call here rather than folding to a register move. const float clamp_min = __builtin_bit_cast(float, clamp_min_bits); const float clamp_max = __builtin_bit_cast(float, clamp_max_bits); @@ -249,9 +205,8 @@ void mm_fused_epilogue_chunk(bfloat16 *y_out, epilogue_dispatch<3>(y_out, src, clamp_enabled, clamp_min, clamp_max); return; #endif - // Mode 0 is always compiled: it is the fallback for a mode the mask leaves - // out, so an unselectable mode yields an unactivated result rather than an - // unwritten buffer. + // Mode 0 is always compiled, so a mode the mask leaves out yields an + // unactivated result rather than an unwritten buffer. default: epilogue_dispatch<0>(y_out, src, clamp_enabled, clamp_min, clamp_max); return; diff --git a/iron/operators/flm/gemm/benchmark.py b/iron/operators/flm/gemm/benchmark.py index fa47eb81a..507c17681 100644 --- a/iron/operators/flm/gemm/benchmark.py +++ b/iron/operators/flm/gemm/benchmark.py @@ -7,36 +7,24 @@ Up to three implementations run per shape, on identical inputs: flm :class:`iron.operators.flm.GEMM`, the port - gemm :class:`iron.operators.GEMM`, left at its defaults, which are the - same emulated-bfp16 mmul and conv_even rounding -- a like-for-like - comparison rather than one against a more accurate, slower build + gemm :class:`iron.operators.GEMM` at its defaults, which are the same + emulated-bfp16 mmul and conv_even rounding, so the comparison is + like-for-like rather than against a more accurate, slower build prebuilt :class:`iron.operators.flm.MMPrebuilt`, FastFlowLM's shipped - ``mm.xclbin``, downloaded and pinned by digest. NPU2 only, because - that binary is a fixed 8-column NPU2 overlay -- on any other device - it is dropped and the flm-vs-gemm comparison still runs. - -Nothing here needs an external install or a host-specific path: the overlay is -a ``RemoteFileArtifact``, so it is fetched into the (gitignored) build dir like -any other artifact, and its instruction stream is generated by IRON. The binary -is never checked in -- it is pinned by SHA-256 against an immutable FastFlowLM -commit and downloaded on demand. - -Marked ``extensive`` for documentation, but pytest never collects it either way: -``pytest.ini`` sets ``python_files = test.py``, and this file is named -``benchmark.py`` so no CI job -- extensive or otherwise -- runs it. That is -deliberate: it is a timing comparison, meant to be invoked directly, not a -correctness gate. The correctness half now lives in -``iron/operators/flm/mm_prebuilt/test.py``, which IS a collected ``test.py`` -and so IS reached by the extensive job. - -Timing is the runtime's own ``npu_time`` (device-side), the same source -``iron.common.test_utils.run_test`` reports, rather than a host wall clock: it -excludes host dispatch and so compares the designs rather than the driver. - -The shapes below are the ones this operator exists to serve, so they overlap -with ``test.py``'s by construction. They are not redundant with it: ``test.py`` -asserts correctness on one implementation, this compares latency across three -frozen binaries, and neither can stand in for the other. + ``mm.xclbin``, pinned by digest. NPU2 only, since that binary is a + fixed 8-column overlay; elsewhere it is dropped and the flm-vs-gemm + comparison still runs. + +The overlay is a ``RemoteFileArtifact`` pinned by SHA-256 against an immutable +FastFlowLM commit and fetched into the gitignored build dir, so nothing here +needs an external install or a host-specific path. + +pytest never collects this: ``pytest.ini`` sets ``python_files = test.py``. It +is a timing comparison meant to be invoked directly, not a correctness gate; +the correctness half lives in ``iron/operators/flm/mm_prebuilt/test.py``. + +Timing is the runtime's device-side ``npu_time`` rather than a host wall clock, +so it compares the designs rather than the driver. Pass ``--iterations 1``. Each test already averages ``ITERS`` dispatches over ``ROUNDS`` interleaved rounds, so conftest's default of 5 repeats the whole @@ -71,22 +59,15 @@ _dev = aie_utils.get_current_device() # The shipped overlay is a fixed 8-column NPU2 binary. Where that does not -# match the device, drop that ONE candidate rather than skipping the module: -# flm vs iron.operators.GEMM is measurable on every supported device. +# match the device, drop that one candidate rather than skipping the module, +# since flm vs iron.operators.GEMM is measurable on every supported device. HAVE_PREBUILT = _dev is not None and _dev.resolve().name == "npu2" and _dev.cols >= 8 # Every projection of both Gemma4 variants FastFlowLM ships, at three prefill -# lengths. E2B is dim 1536 / ffn 6144; E4B is dim 2560 / ffn 10240. -# -# These two are the right coverage for the shipped mm.xclbin. Checked against -# FastFlowLM f81eba71: Gemma4-E2B-IT-NPU2, Gemma4-E4B-IT-NPU2 and Gemma3-4B-NPU2 -# all ship the SAME mm.xclbin blob (git 4727df98, 512220 bytes) -- one overlay -# serving several models -- so E2B and E4B between them already exercise it. -# -# Gemma4-12B-IT-NPU2 does not ship an mm.xclbin at all. Its overlays are a -# different set (attn_global, attn_sliding, audio_image_mm, dequant_mm, layer, -# lm_head), so its projections go through a quantized matmul rather than this -# bf16 one and cannot be compared against MMPrebuilt. +# lengths. E2B is dim 1536 / ffn 6144; E4B is dim 2560 / ffn 10240. Both ship +# the same mm.xclbin blob (checked at FastFlowLM f81eba71), so between them +# they cover it. Gemma4-12B ships no mm.xclbin at all -- its projections go +# through a quantized matmul -- so it cannot be compared against MMPrebuilt. # proj, K, N E2B_PROJ = [ ("q", 1536, 4096), @@ -118,11 +99,10 @@ def get_params(): - # No shape is skipped: the four E4B projections with a 10240-wide - # dimension at M > 256 once overflowed the shim BD's 20-bit mega_row - # iteration step, but flm.GEMM and IRON's GEMM both now split that leg - # into per-mega_row transfers (see design.py's a_split/c_split and - # test_gemm_split_leg_windowing in test.py). + # No shape is skipped. The four E4B projections with a 10240-wide dimension + # at M > 256 once overflowed the shim BD's 20-bit mega_row iteration step, + # but flm.GEMM and IRON's GEMM both split that leg into per-mega_row + # transfers now (design.py's a_split/c_split, test_gemm_split_leg_bounds). params = [] for model, projections in (("E2B", E2B_PROJ), ("E4B", E4B_PROJ)): for M in PREFILL_LENGTHS: @@ -139,10 +119,7 @@ def make_inputs(M, K, N): A = (torch.randn(M, K) * 4).to(torch.bfloat16) B = (torch.rand(K, N) * 4).to(torch.bfloat16) Af, Bf = A.float(), B.float() - # Error is bounded against accumulated mass rather than relatively: with - # signed A the K-sum cancels by ~sqrt(K), so |C| ends up far smaller than - # the magnitude the bfp16 error actually tracks, leaving near-zero outputs - # relatively uncheckable. Same rationale as test.py's bound. + # Bounded against accumulated mass rather than relatively; see test.py. return A, B, Af @ Bf, float((Af.abs() @ Bf.abs()).mean()) @@ -159,9 +136,8 @@ def __init__(self, name, op, A, B, M, N, budget, ctx): self.xclbin = Path(op.xclbin_artifact.filename) self.c_bo = XRTTensor((M, N), dtype=np.dtype("bfloat16")) run = op.get_callable() - # Only the flm operators take B pre-packed. iron.operators.GEMM reorders - # in the descriptor instead, so it wants plain row-major (K, N) -- - # b_col_maj defaults False. + # Only the flm operators take B pre-packed. iron.operators.GEMM + # reorders in the descriptor, so it wants plain row-major (K, N). packed_b = op.pack_B(B) if hasattr(op, "pack_B") else B args = [ XRTTensor.from_torch(A.flatten()), @@ -199,10 +175,8 @@ def jitter_pct(self): GEMMLatency=r"gemm latency \(us\): (?P[\d\.]+)", SpeedupVsPrebuilt=r"speedup vs prebuilt: (?P[\d\.]+)", SpeedupVsGEMM=r"speedup vs gemm: (?P[\d\.]+)", - # Accuracy is asserted against a budget below, but that budget is loose - # enough that a toolchain or kernel change could move the error a long way - # inside it unnoticed. Record the numbers too, so a dependency bump can be - # diffed on accuracy and not only on speed. + # The budget below is loose enough that a toolchain change could move the + # error a long way inside it unnoticed, so record the numbers too. FLMErr=r"flm err/mass: (?P[\d\.e\+-]+)", PrebuiltErr=r"prebuilt err/mass: (?P[\d\.e\+-]+)", GEMMErr=r"gemm err/mass: (?P[\d\.e\+-]+)", @@ -261,11 +235,10 @@ def test_gemm_vs_prebuilt(model, proj, M, K, N, aie_context): for c in candidates: for _ in range(WARMUP): c.run() - # Round-robin, never all of one then all of another. Dispatch latency on - # this part is bimodal with modes about 6% apart, so a batch that lands - # wholly in one mode turns min-of-medians into a mode selector rather than a - # measurement -- that is how a change later shown to do nothing at all once - # produced a convincing 5% "win". + # Round-robin, never all of one then all of another. Dispatch latency here + # is bimodal with modes about 6% apart, so a batch landing wholly in one + # mode turns min-of-medians into a mode selector. That is how a change + # later shown to do nothing once produced a convincing 5% "win". for _ in range(ROUNDS): for c in candidates: c.time_round() diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index 731e15dde..70054c4c4 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -4,57 +4,20 @@ """bf16 GEMM over a 4-row compute-tile grid, as wide as the device. A second GEMM design alongside ``iron.operators.gemm``, specialised for -transformer projection shapes. The overall dataflow is the same whole-array -shape as that operator's -- A broadcast along each compute row, B down each -column, C joined through the memtile -- so those are NOT what distinguishes it. -What does: - - * **B is quantized to bfp16ebs8 on NPU2** by ``GEMM.pack_B``, not bf16. - ``iron.operators.GEMM`` only ever moves bf16. This is not primarily a DMA - saving: it is what makes NPU2's fast mmul lowering available at all -- - ``aie::mmul<8,8,8>`` needs bfp16 operands to decompose into two emulated - macs instead of four, which is most of the NPU2 speedup (see NPU1 in - README.md's Performance section, where B stays bf16 and the margin over - ``iron.operators.GEMM`` is correspondingly smaller). Quantizing is - numerically free -- the mmul only multiplies bfp16 regardless, so this - hoists a rounding that already happened on every mac -- provided it - reproduces the core's rounding mode; see ``packing.py``. - * **The tile shape is fixed, not parameterised** -- but fixedness alone is - not the advantage: ``iron.operators.GEMM`` is equally fixed once compiled - with a choice of tile args. What differs is *which* shape is fixed. r/s/t - stays 8/8/8 on both architectures for the reason below. m/k = 64/512 (n - defaults to 64) is chosen for the L1-budget tradeoff documented next to - ``CT_MAX_K_FOR_N`` below: n=64 gives the mmul a colA of 8 rather than 4, - which wins whenever compute is the critical path, at the cost of A being - re-read more often. Only the grid WIDTH varies with the device: 8 columns - on NPU2, 4 on NPU1. - * **A fused epilogue.** The f32->bf16 conversion, an optional activation and - an optional clamp all happen while the values are still in registers, on the - way into the C object, instead of a separate pass over L1. - * **B arrives pre-packed** by ``GEMM.pack_B``, in the order the cores consume - it, so both B hops are plain linear descriptors instead of the 128-byte - scattered bursts ``iron.operators.GEMM`` reorders in the descriptor. - * **Asymmetric tile buffering (ATB)**, so the A tile and the accumulator need - not share a height -- this is what buys the deep k slice (K_TILE=512) - within the L1 budget; see README.md's ATB reference. - -None of these helps alone -- see README.md's Performance section for the -measured, per-choice breakdown of the gap against both the shipped FastFlowLM -overlay and ``iron.operators.GEMM``. +transformer projection shapes. Same dataflow -- A broadcast along each compute +row, B down each column, C joined through the memtile -- but with a fixed tile +shape, B quantized to bfp16ebs8 and pre-packed into consumption order, an +activation and clamp fused into the C drain, and asymmetric tile buffering so +the A tile and the accumulator need not share a height. -The constants below are the single source of truth: ``op.py`` passes them to the -kernels as -D flags, so the C++ and the dataflow cannot drift apart. +README.md has the per-choice breakdown against both the shipped FastFlowLM +overlay and ``iron.operators.GEMM``. -r/s/t stays 8/8/8 on both architectures. AIE2's native bf16 mac is 4x8x4, but -``aie::mmul<8,8,8>`` decomposes onto it as exactly four native macs with no -wasted lanes, so the whole blocked L1 layout -- ``pack_B``, the four stream -dimension lists below, and ``gather_dims`` -- is shared verbatim. Only AIE2P has -the bfp16-emulated path that does the same shape in two macs, which is why -``op.py`` passes ``AIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16`` there and not here. +The constants below are the single source of truth: ``op.py`` passes them to +the kernels as -D flags, so the C++ and the dataflow cannot drift apart. """ import argparse -import os from enum import StrEnum from functools import partial @@ -84,49 +47,21 @@ # --- Fixed geometry ------------------------------------------------------- # GEMM tiling per compute tile, and the register tiling inside it. M_TILE, K_TILE = 64, 512 -# Default n tile. 64 gives the mmul a colA of 8 rather than 4, halving the -# accumulator traffic per mac, at the cost of doubling A fetches (the grid -# then covers 512 columns of N per pass instead of 1024). That trade wins -# whenever compute is the critical path, which is the usual case; see -# README.md for the measured sweep, including the small-K shape where it -# loses. +# Default n tile. 64 doubles A fetches but gives the mmul colA=8 instead of 4, +# which wins when compute is the critical path. op.py picks per shape. N_TILE_DEFAULT = 64 -# How much of K one compute tile holds at a time, per n width. This is a fixed -# L1 budget split two ways, so a wider n tile leaves less room for B's k slice -# and the product stays roughly constant. op.py passes the chosen value to the -# kernel as -DMM_FUSED_CT_K, making this table the only place it is decided. -# n=256 is deliberately absent: at that width the f32 accumulator alone -# (M_TILE * 256 * 4 = 65536 bytes) already fills the whole of L1, before A, B -# or C are even counted, so no ct_max_k could ever make it fit. +# How much of K one compute tile holds at a time, per n width. It is a fixed +# L1 budget split two ways, passed to the kernel as -DMM_FUSED_CT_K. n=256 is +# absent because its f32 accumulator alone (M_TILE*256*4) fills all of L1. CT_MAX_K_FOR_N = {16: 16, 32: 32, 64: 128, 128: 32} -# (tile_n, ct_max_k) pairs KNOWN TO COMPUTE CORRECTLY on hardware. The table -# above reads like a tuning knob -- op.py calls it "the only place it is -# decided" -- but it is not freely tunable, and a wrong value fails SILENTLY. -# Measured on npu2 2026-09-11: -# -# tile_n=64 ct_k=128 err/mass 2.42e-04 (shipped) -# tile_n=64 ct_k= 64 err/mass 3.45e-02 ~140x worse, outside any budget -# tile_n=128 ct_k= 32 err/mass 1.43e-04 (shipped) -# tile_n=128 ct_k= 64 NaN -# tile_n=128 ct_k=128 NaN -# -# Root cause not found, but pack_b is RULED OUT BY TEST: inverting its -# permutation round-trips exactly at ct_k 128, 64 and 32, so its blocking is -# generic. The disagreement is most likely the ORDER the design's stream dims -# and the kernel's mmul nest walk one ct-chunk. Until that is found, refuse -# rather than miscompute. +# (tile_n, ct_max_k) pairs verified on hardware. The table above looks tunable +# but is not, and a wrong value fails SILENTLY: ct_k=64 at tile_n=64 gives +# err/mass 3.45e-02 against 2.42e-04, and tile_n=128 NaNs. Cause unknown, so +# refuse rather than miscompute. _VERIFIED_CT_K = {(16, 16), (32, 32), (64, 128), (128, 32)} -# Register tiling, shared by both architectures. These set the blocked L1 -# layout, so ``pack_B``, the four stream-dimension lists below and -# ``gather_dims`` all key off them; changing one without the others is silently -# wrong rather than a build error. -# -# Matching AIE2's native 4x8x4 mac shape instead is a measured dead end, at -# 22-30% slower: the kernel is load-port bound rather than shuffle bound, and -# 4/8/4 needs 1.62 loads per mac against 8/8/8's 1.06, because the 2x2 register -# block amortizes each load over four macs either way but over far less work. -# Retrying it needs a wider register block on the native shape, i.e. a 4x4 mmul -# kernel, not just a different r/s/t here. +# Register tiling, shared by both architectures. pack_B, the stream-dimension +# lists and gather_dims all key off these; changing one alone is silently +# wrong. AIE2's native 4x8x4 shape measured 22-30% slower (load-port bound). R, S, T = 8, 8, 8 @@ -140,59 +75,22 @@ def compute_rows(dev): C_DEPTH = 2 # C fifo depth; also the core-body unroll B_DEPTH = 2 # B fifo depth; also the core-body unroll A_DEPTH = 2 -# Row-blocks a core folds into one B fetch: it HOLDS a B chunk across M_CHUNK -# accumulators instead of releasing it after one, so DDR reads B -# m_row_blocks/M_CHUNK times instead of m_row_blocks. This is what replaces the -# B residency the runtime parameters cost us -- residency sized the memtile -# buffer from k_iters and replayed it m_row_blocks times, putting both K and M -# in the device configuration, whereas M_CHUNK is a configuration constant and -# leaves both runtime. -# -# Two things it costs: -# * L1: M_CHUNK accumulators instead of one. At tile_n=64 the k slice is -# unaffected (ct_k stays 128, colA 16) and only tile_ma drops 32 -> 16; -# at tile_n=128 nothing fits, so M_CHUNK is resolved per tile_n below. -# * A is issued once per chunk-group rather than once per column-block. -# Holding B means streaming A k-major -- A[mc0][k0], A[mc1][k0], -# A[mc0][k1], ... -- which wants five dimensions, and a shim BD carries -# four: the outermost lands in the ITERATION field and the inner three are -# its ND dims (getBDMaxDims is 3 off a memtile). So the chunk-group -# dimension comes out of the descriptor and becomes separate transfers -# carrying the jump in the OFFSET, which has no such limit -- exactly the -# trick a_split already uses for the mega_row dimension, and it reuses the -# same windowing. What is left per transfer is -# [k_iters, M_CHUNK, M_TILE, K_TILE], which fits. -# DEFAULT 1 EVERYWHERE -- i.e. off, and for a contractual reason rather than a -# performance one. m_chunk must DIVIDE m_row_blocks, so it needs M to be a -# multiple of 512, while the overlay this operator replaces accepts any -# multiple of 256 (mm_prebuilt/design.py:65, MIN_M = M_TILE*ROWS). A shape that -# cannot use m_chunk falls back to 1 and so FORKS _config_tag, and serving -# M=256 and M=2048 from ONE xclbin is a hard requirement here. -# -# The leftover cannot be padded away either: a partial group is inexpressible -# (see the raise below), and duplicating the row-block to fill the group is -# expressible and bit-exact but doubles A and C for that dispatch, which at -# m_row_blocks=1 costs more than anything m_chunk could win back. -# -# Note m_chunk does NOT change A's byte count (A is M*K*2*n_col_blocks either -# way). Its only structural effect is forcing a_split on, so if M % 512 == 0 -# ever becomes guaranteed it is worth re-measuring. -# -# `n_chunk` is the lever with the same effect and no such constraint: it groups -# COLUMN-blocks, which the core already walks with an RTP-derived trip count, -# so a partial group is a runtime bound rather than a descriptor shape and one -# xclbin still serves every M. +# Row-blocks a core folds into one B fetch, cutting B's DDR reads by M_CHUNK +# at the cost of that many L1 accumulators and forcing a_split. Off everywhere +# for a contractual reason: it must divide m_row_blocks (M % 512 == 0) while +# the overlay this replaces takes any multiple of 256, so a shape that cannot +# use it forks _config_tag. See README.md. M_CHUNK_FOR_N = {16: 1, 32: 1, 64: 1, 128: 1} # How many column-blocks the runtime sequence keeps in flight. A block costs 3 -# shim buffer descriptors on a column (A + B + C) against 16 available, so the -# ceiling is 5; 2 is enough to keep the fills ahead of the cores. +# shim buffer descriptors on a column (A + B + C) of the 16 available, so the +# ceiling is 5. 2 is enough to keep the fills ahead of the cores. OVERLAP_DEFAULT = 2 class Epilogue(StrEnum): """Activation folded into the C drain. - Declaration order is the wire format -- it is both the kernel's + Declaration order is the wire format: it is both the kernel's ``-DMM_FUSED_EPILOGUE_MODE`` and the shipped overlay's ``output_mode``. """ @@ -207,17 +105,9 @@ def mode(self) -> int: return list(Epilogue).index(self) -# The runtime parameter buffer each core reads once its barrier opens. -# -# The clamp bounds are the one non-obvious entry: they are floats, but -# npu_write_rtp only writes i32 words, so they travel as raw bit patterns and -# the kernel bit-casts them back. Whether a clamped path exists at all stays -# compile-time (op.py's -DMM_FUSED_CLAMP), because it costs program memory; -# only the enable and the bounds are runtime, so every pair of bounds shares -# one build. -# The five that EVERY configuration needs. Anything conditional lives after -# them, at an offset rtp_layout() computes, so a word a build cannot use is -# never allocated rather than written and ignored. +# The parameter buffer each core reads once its barrier opens. These four are +# always present; conditional words follow at offsets rtp_layout() computes, +# so a word a build cannot use is never allocated. ( RTP_N_VAL, RTP_M_ROW_BLOCKS, @@ -229,22 +119,9 @@ def mode(self) -> int: def rtp_layout(clamp_capable, m_chunk): """Slot index for each optional parameter, and the total word count. - **A word is not free.** Each one costs ~66 ns per core and the sequence - writes ROWS*COLS = 32 of them, so **every RTP word is ~2.06 us of dispatch - latency** -- measured by padding the buffer at a fixed core count (12 words - 108.6 us, 24 words 135.2, 48 words 182.7). Against a ~107 us floor that is - not a rounding error, and at M=256 the floor is most of the dispatch. - - So both groups here are omitted, not defaulted: - - * the clamp trio, when the build has no clamped path at all. op.py compiles - the clamped instantiation out entirely at ``_clamp_capable == 0`` (the - default, and every real projection shape), so those three words were - written on every dispatch and never read -- ~6 us for nothing. - * ``n_chunks`` / ``n_units``, when M_CHUNK == 1. They are - ``m_row_blocks // M_CHUNK`` and each other, so at the shipped M_CHUNK - they are simply m_row_blocks, and the core reads that word instead -- - no on-core arithmetic, just one fewer thing to send. ~4 us. + A word is not free: the sequence writes one per core, costing ~2 us of + dispatch latency against a ~107 us floor. So optional groups are omitted + rather than defaulted. """ slots = {} n = 4 @@ -263,9 +140,8 @@ def rtp_layout(clamp_capable, m_chunk): class Rounding(StrEnum): """Rounding for every f32->bf16 conversion. - The core powers up in floor; conv_even is the default because truncation - biases every conversion the same way and the error then accumulates over - the K reduction. floor reproduces the shipped overlay. + conv_even by default: truncation biases every conversion the same way, so + the error accumulates over the K reduction. floor matches the overlay. """ CONV_EVEN = "conv_even" @@ -298,20 +174,16 @@ def _b_bytes(elems, bfp16_b): # --- Shim DMA limits ------------------------------------------------------ # -# Hardware facts the Python bindings do not expose: gemm() reads AIETargetModel -# directly for the L1 ceiling, BD count and grid, but neither getDmaBdStepBits -# nor getDmaBdWrapSizeBits is bound, and nothing models the channel task queue. -# gemv/design.py and repeat/design.py hardcode the same fields. -# -# Step field width. An IR-level bf16-element stride S is re-expressed as -# (S - 1) * 2 bytes / 4-byte granularity before AIEXDialect.cpp checks it. +# Hardware facts the Python bindings do not expose: getDmaBdStepBits and +# getDmaBdWrapSizeBits are unbound, and nothing models the channel task queue. +# gemv/design.py and repeat/design.py hardcode the same fields. An IR-level +# bf16 stride S is re-expressed as (S-1)*2 bytes / 4-byte granularity before +# AIEXDialect.cpp checks it. _SHIM_STEP_BITS = 20 _BF16_BYTES = 2 _ADDR_GRANULARITY_BYTES = 4 -# Entries in a shim DMA channel's task queue. AIEDmaToNpu's NpuPushQueueOp -# pushes unconditionally, so overrunning this is a silent device hang rather -# than a diagnostic. Measured at K=10240 M=1024: 4 outstanding tasks on one -# channel run, 8 hang. +# Entries in a shim DMA channel's task queue. NpuPushQueueOp pushes +# unconditionally, so overrunning this hangs silently. Measured: 4 run, 8 hang. SHIM_TASK_QUEUE = 4 @@ -321,32 +193,14 @@ def _hw_stride_ok(stride_elems): def _default_l1(n_tile, ct_max_k, b_elem_bytes, budget, m_chunk=1): - """Pick (A-tile height, L1 B depth) -- the largest working set that fits. - - ``b_elem_bytes`` is 9/8 where B is bfp16ebs8 and 2 where it is bf16, and - ``budget`` is the core's data memory, so the search below reflects what B - actually costs on this device. - - No stack is reserved out of ``budget``: the cores leave ``stack_size`` - unset and aiecc measures each core's requirement and fails the build if it - does not fit, so the stack is the toolchain's to enforce. This kernel - measures 192 bytes against the >=6 KB the search leaves unused anyway. - - A dies as soon as it is consumed while the accumulator lives across the - whole K reduction, so they need not share a height; shrinking A is what - pays for a k slice deep enough to halve the accumulator traffic per mac. - B's depth is searched too because at n=128 the k=128 slice makes the B - object 18 KB, and a double-buffered pair simply does not fit -- giving that - up is what buys colA=16 there, and colA is worth far more than B's L1 - prefetch (3.67 -> 2.28 cycles per mac, measured). - - Deeper B first, then the tallest A that still fits, so the n=64 default is - unchanged at (32, 2). + """Pick the largest working set that fits: (A-tile height, L1 B depth). + + Deeper B first, then the tallest A that still fits, since colA is worth + far more than B's L1 prefetch. No stack is reserved out of ``budget``; + aiecc fails the build if a core's measured requirement does not fit. """ - # m_chunk accumulators, because the core holds a B chunk across that - # many row-blocks; see M_CHUNK_FOR_N. This is the ONLY term that - # scales with it -- A is acquired and released one band at a time - # inside the group loop, and C drains the accumulators in turn. + # m_chunk accumulators, since the core holds a B chunk across that many + # row-blocks. The only term that scales with it. acc = m_chunk * M_TILE * n_tile * 4 cout = CT_OUT_LEN * 2 * C_DEPTH for b_depth in (B_DEPTH, 1): @@ -363,22 +217,10 @@ def _default_l1(n_tile, ct_max_k, b_elem_bytes, budget, m_chunk=1): def _b_depth_for(t_ma, n_tile, ct_max_k, b_elem_bytes, budget, m_chunk=1): """Deepest B fifo depth that fits L1 alongside an explicit A-tile height. - ``_default_l1`` picks L1_B_DEPTH together with the t_ma IT chooses; that - pairing need not fit a caller-overridden t_ma; a taller A tile leaves less - L1 for B, and can push a working set that fit at the default t_ma over - budget. Raise rather than silently reusing a depth that doesn't fit. - - How much the depth is worth, measured on npu2 (turbo, 12 interleaved rounds - of 20 dispatches, min of per-round medians) by forcing depth 1 against the - default: 1.2% at M=1024 K=1536 N=6144, and within noise at K=1024 N=4096 and - K=512 N=1024. So the prefetch earns its L1 at the largest shapes and is - close to free elsewhere -- worth keeping, but not worth contorting the - search for. + ``_default_l1``'s depth is chosen with its own t_ma, which need not fit a + caller-overridden one. Raise rather than reuse a depth that does not fit. """ - # m_chunk accumulators, because the core holds a B chunk across that - # many row-blocks; see M_CHUNK_FOR_N. This is the ONLY term that - # scales with it -- A is acquired and released one band at a time - # inside the group loop, and C drains the accumulators in turn. + # Same terms as _default_l1; acc is the only one that scales with m_chunk. acc = m_chunk * M_TILE * n_tile * 4 cout = CT_OUT_LEN * 2 * C_DEPTH a = (2 * R * ct_max_k) * (t_ma // R // 2) * 2 * A_DEPTH @@ -407,17 +249,14 @@ def gemm( ): """Emit the MLIR module for an M x K @ K x N bf16 GEMM. - A is (M, K) row-major, B is (K, N) row-major and C is (M, N) row-major, all - bf16 and all plain dense tensors, except that B must arrive pre-packed by - ``GEMM.pack_B`` -- it emits B in the order the cores consume it, so both - B hops are plain linear descriptors. + A, B and C are row-major bf16 dense tensors, except that B must arrive + pre-packed by ``GEMM.pack_B`` in the order the cores consume it. """ if tile_n not in CT_MAX_K_FOR_N: raise ValueError( f"tile_n must be one of {sorted(CT_MAX_K_FOR_N)}, got {tile_n}" ) - # Everything shape-related below comes from the device rather than a - # constant, so the same dataflow covers NPU2's 4x8 and NPU1's 4x4. + # From the device, not constants, so one dataflow covers 4x8 and 4x4. tm = get_target_model(dev.resolve()) COLS, ROWS = dev.cols, compute_rows(dev) MIN_M = M_TILE * ROWS @@ -433,19 +272,14 @@ def gemm( f"pair here once a hardware test passes." ) M_CHUNK = M_CHUNK_FOR_N[N_TILE] if m_chunk is None else m_chunk - # B is bfp16ebs8 on AIE2P and bf16 on AIE2: the scalar BFP types are gated - # on __AIE_API_SCALAR_BFP_TYPES__, which only aie_api/detail/aie2p/config.hpp - # defines, so on AIE2 B stays bf16 and the mmul lowers onto four native - # 4x8x4 macs. That choice drives every B type and extent below. B_GROUP is - # the number of B values per element of the MLIR type, so a length in values - # becomes a length in elements by dividing. + # B is bfp16ebs8 on AIE2P and bf16 on AIE2. B_GROUP is the B values per + # element of the MLIR type, so a length in values divides to elements. BFP16_B = dev.arch == AIEArch.AIE2p B_GROUP = BFP16_GROUP if BFP16_B else 1 b_elem_bytes = BFP16_GROUP_BYTES / BFP16_GROUP if BFP16_B else 2 - # Asymmetric tile buffering: the A tile spans T_MA rows while the - # accumulator spans M_TILE, so the core folds RHO bands into one C tile. - # A is dead the moment it is consumed while C lives across the whole K - # reduction, so sizing both to M_TILE pays the peak L1 cost twice. + # Asymmetric tile buffering: A spans T_MA rows and the accumulator M_TILE, + # so the core folds RHO bands into one C tile. A dies on consumption while + # C lives across the K reduction, so sizing both to M_TILE pays twice. if tile_ma is None: T_MA, L1_B_DEPTH = _default_l1( N_TILE, CT_MAX_K, b_elem_bytes, tm.get_local_memory_size(), M_CHUNK @@ -475,22 +309,18 @@ def gemm( MIN_N = N_TILE * COLS epilogue = Epilogue(epilogue) - # Clamp bounds ride the RTP buffer as raw int32 bit patterns: npu_write_rtp - # writes i32 words only, so the kernel bit-casts them back. Bounds being - # runtime is what lets every pair share one build; whether a clamped path - # exists at all is still op.py's -DMM_FUSED_CLAMP, because it costs - # program memory. + # Clamp bounds ride the RTP buffer as raw int32 bit patterns, since + # npu_write_rtp writes i32 only. Whether a clamped path exists is still + # compile-time (op.py's -DMM_FUSED_CLAMP). clamp_enabled = 1 if clamp is not None else 0 rtp_slots, rtp_words = rtp_layout(clamp is not None, M_CHUNK) clamp_lo, clamp_hi = clamp if clamp is not None else (0.0, 0.0) clamp_min_bits = int(np.float32(clamp_lo).view(np.int32)) clamp_max_bits = int(np.float32(clamp_hi).view(np.int32)) - # A compute tile does a whole m x n block or nothing, so M and K must tile - # exactly. N need only be a multiple of N_TILE: a trailing group of fewer - # than COLS blocks is handled by giving the columns different trip counts, - # which each core derives for itself (see core_fn). That matters in - # practice -- for a transformer the o and down projections have N = model - # dim, which is essentially never a multiple of N_TILE*COLS. + # A tile does a whole m x n block or nothing, so M and K must tile + # exactly. N need only be a multiple of N_TILE: a short trailing group is + # handled by per-column trip counts, which matters because o and down + # have N = model dim. for name, value, unit in (("M", M, MIN_M), ("K", K, MIN_K), ("N", N, N_TILE)): if value % unit != 0: raise ValueError(f"{name} ({value}) must be a multiple of {unit}") @@ -498,26 +328,15 @@ def gemm( bf16_ty = np.dtype[bfloat16] f32 = np.dtype[np.float32] - # How many times the whole grid sweeps, in each dimension. m_row_blocks = M // MIN_M k_iters = K // K_TILE - # Row-blocks grouped M_CHUNK at a time, so one B fetch feeds M_CHUNK of - # them (see M_CHUNK_FOR_N). A "unit" below is one such group, or one of the - # n_rem leftovers when M_CHUNK does not divide m_row_blocks. Every leg is - # issued per unit, so A, B and C stay aligned with each other and with the - # core's loop nest. + # A "unit" is one group of M_CHUNK row-blocks. Every leg is issued per + # unit, so A, B and C stay aligned with each other and the core's nest. n_chunks, n_rem = divmod(m_row_blocks, M_CHUNK) if n_rem: - # op.py resolves m_chunk so this cannot fire. A partial group is - # genuinely inexpressible: its object is M_CHUNK tiles wide and the - # forward always drains that much, so the rest would have to be filled - # by repeating the row-block -- and every way of saying that is - # rejected or mis-lowered. Stride 0 in an inner dimension is refused - # ("Stride 2 must be a positive integer"); stride 0 in the outermost - # slot IS the BD repeat count, which releases one object per - # repetition rather than one object in total; and several sub-object - # fills do not coalesce into one object either. All three were tried - # on hardware. Hence op.py falls back to m_chunk=1 instead. + # op.py resolves m_chunk, so this cannot fire. A partial group is + # inexpressible: the object is M_CHUNK tiles wide and the forward + # always drains that much, and no way of padding it lowers correctly. raise ValueError( f"m_row_blocks ({m_row_blocks}) must be a multiple of m_chunk " f"({M_CHUNK}); op.py should have resolved m_chunk to 1 here" @@ -525,62 +344,26 @@ def gemm( n_units = n_chunks def unit_rows(u): - """(first row-block, how many) for unit ``u`` -- always a full group.""" + """(first row-block, how many) for unit ``u``; always a full group.""" return u * M_CHUNK, M_CHUNK - # A mega_row dimension with stride ROWS*M_TILE*{K,N} lands in the shim - # BD's ITERATION field, whose step is 20 bits, so it overflows once K or N - # crosses ~8191 elements -- E4B's FFN width (10240) does, E2B's max (6144) - # does not. Only relevant once M walks more than one mega_row; at M=256 the - # dimension is degenerate (size 1) and is stripped before the stride is - # ever encoded, so it never fails there regardless of K/N. + # A mega_row stride lands in the shim BD's 20-bit iteration step, so it + # overflows once K or N passes ~8191 elements; only E4B's 10240 does. Such + # a leg goes out as one transfer per mega_row, carrying the jump in its + # unbounded offset. M_CHUNK > 1 forces the same path for A. # - # Such a leg is instead issued as m_row_blocks separate transfers, each - # carrying the mega_row jump in its OFFSET (unbounded) rather than a shared - # STRIDE. The two legs are independent -- E4B's down-proj overflows on K - # (A only) and its gate/up on N (C only) -- so neither shape pays for both. - # - # Those transfers must stay live in their TaskGroup until awaited. - # TaskGroup.finish() emits dma_free_task, which returns the buffer - # descriptor id to a COMPILE-TIME allocator that does not check the - # transfer finished (mlir-aie AIEAssignRuntimeSequenceBDIDs::recycle, - # isAwait=false); ids are per shim TILE, shared across channels and - # directions, so retiring one early lets the next task reprogram a live - # descriptor. Verify with aie-opt --aie-substitute-shim-dma-allocations - # --aie-assign-runtime-sequence-bd-ids: the ids on a shim tile must be - # distinct. - # M_CHUNK > 1 forces the split path for A: holding B makes the core's A - # order k-major within a group, which wants a fifth dimension, so the unit - # dimension comes out of the descriptor and becomes one transfer each. - # At n_units == 1 there is nothing to split -- a single transfer already - # covers the block -- so the cheaper unsplit path still applies. + # Those transfers must stay live in their TaskGroup until awaited: the + # BD-id allocator is compile-time and does not check that a transfer + # finished, so freeing one early lets the next task reprogram a live + # descriptor and corrupt silently. a_split = n_units > 1 and (M_CHUNK > 1 or not _hw_stride_ok(ROWS * M_TILE * K)) c_split = m_row_blocks > 1 and not _hw_stride_ok(ROWS * M_TILE * N) - # A split leg issues one transfer per mega_row back to back on ONE channel, - # so they must also fit that channel's task queue -- a limit nothing in the - # toolchain models, and overrunning it hangs rather than diagnoses. - # Measured at K=10240 M=1024: 4 outstanding run, 8 hang. - # - # So a split leg keeps at most SHIM_TASK_QUEUE transfers outstanding, and - # emit_split() enforces that directly -- retiring the OLDEST transfer as it - # issues the next, so the channel stays full. - # - # It used to do this by windowing: issue SHIM_TASK_QUEUE transfers, await - # the whole window, then issue the next. That bounds the queue too, but it - # drains the channel to EMPTY at every window boundary and again at every - # column-block boundary, and on a DDR-rate-bound design those bubbles are - # the entire cost of the split path. Rolling instead of windowing measured, - # bit-exact, 8 rounds x 30 iters interleaved: - # - # E4B/gateup M1024 4119.4 -> 3617.7 -12.2% - # E4B/gateup M2048 7664.8 -> 7182.4 -6.3% - # E4B/down M1024 3649.2 -> 3553.4 -2.6% - # E4B/down M2048 7196.6 -> 6978.7 -3.0% - # - # The bound is counted in TRANSFERS, not units: under c_split a unit drains - # one C descriptor per row-block (M_CHUNK of them, see c_taps) and they all - # ride one channel, so counting units would overrun by exactly M_CHUNK -- - # at m_chunk=2 that is 8 outstanding, the measured hang threshold. + # Split legs share one channel, whose task queue is 4 deep and modelled + # nowhere; overrunning it hangs (4 outstanding run, 8 hang). emit_split() + # bounds it by retiring the oldest as it issues the next, which also keeps + # the channel full -- do not simplify that to awaiting a whole batch, which + # drains the channel at every boundary and costs up to 12.4%. The bound + # counts transfers, not units: under c_split a unit drains M_CHUNK of them. _per_unit = M_CHUNK if c_split else 1 # Live descriptors on a shim tile: SHIM_TASK_QUEUE from the rolling window, # plus B and the unsplit leg for each of the two blocks a boundary spans. @@ -598,9 +381,9 @@ def unit_rows(u): # rem_blocks columns (0 <= rem_blocks < COLS) that do one block more. n_full = N // MIN_N rem_blocks = (N % MIN_N) // N_TILE - # Every column is instantiated for every shape, because which columns - # exist is configuration and this design has one configuration. A column - # with no work for this shape gets n_work = 0 and drains instead. + # Every column is instantiated for every shape. Which columns exist is + # configuration, and this design has only one. A column with no work for + # the current shape gets n_work = 0 and drains instead. n_active_cols = COLS # B's element type: one v8bfp16ebs8 per 8 values on AIE2P, one bf16 per # value on AIE2. Every B extent below is therefore in values // B_GROUP. @@ -610,13 +393,12 @@ def unit_rows(u): ct_b_ty = np.ndarray[(CT_MAX_K * N_TILE // B_GROUP,), b_elem_ty] ct_out_ty = np.ndarray[(CT_OUT_LEN,), bf16_ty] ct_acc_ty = np.ndarray[(M_TILE * N_TILE,), f32] - # L2 (per memtile) - # M_CHUNK stacked row-block tiles, so the forward below can interleave - # them on the way out -- see a_send_dims. + # L2 (per memtile). M_CHUNK stacked row-block tiles, so the forward below + # can interleave them on the way out; see a_send_dims. mt_a_ty = np.ndarray[(M_CHUNK * M_TILE * K_TILE,), bf16_ty] mt_b_ty = np.ndarray[(K_TILE * N_TILE // B_GROUP,), b_elem_ty] mt_out_ty = np.ndarray[(C_SLICE_LEN * ROWS,), bf16_ty] - # L3 (DDR), flat -- the taps below index them linearly. + # L3 (DDR), flat; the taps below index them linearly. a_l3_ty = np.ndarray[(M * K,), bf16_ty] b_l3_ty = np.ndarray[(K * N // B_GROUP,), b_elem_ty] c_l3_ty = np.ndarray[(M * N,), bf16_ty] @@ -641,40 +423,27 @@ def unit_rows(u): # --- Data movement ---------------------------------------------------- # - # These stream-dimension lists are the load-bearing part of the design: - # they are what turns a row-major DDR tile into the r x s / s x t blocked - # layout the mmul indexes, and they are tightly coupled to it. A mismatch - # here produces silently wrong results, not a build error. + # These turn a row-major DDR tile into the blocked layout the mmul + # indexes. A mismatch is silently wrong, not a build error. # C: de-block each core's r x t tiled output back into row-major within its # 64x128 slice, on the way into the memtile. gather_dims = [(M_TILE // R, R * N_TILE), (N_TILE // T, T), (R, N_TILE), (T, 1)] - # B: DDR row-major (k x n) -> s x t blocks (recv), then split into the - # CT_MAX_K-deep chunks a single mmul call consumes (send). - # B needs no reblocking on either hop: pack_B already emits it in the - # order the cores consume, so the memtile just streams it through. That - # frees every descriptor dimension B used to spend -- which is what lets - # CT_MAX_K reach 128 (the innermost run would otherwise overflow the BD's - # 10-bit size field and need a split dimension) at the same time as - # residency (which spends one on its outer k walk). + # B needs no reblocking on either hop: pack_B emits it in consume order. + # That frees the descriptor dimensions that let CT_MAX_K reach 128. b_recv_dims = None b_send_dims = None - # A: same idea, r x s blocks. - # The outermost row-group dimension spans M_CHUNK tiles rather than one. - # That is the whole trick: mc's stride is M_TILE*K_TILE, which is exactly - # this dimension's size*stride, so the two are contiguous and merge -- the - # walk stays within the memtile BD's four dimensions while gaining an - # interleave it could not otherwise express. + # A: same idea, r x s blocks. The outermost row-group dimension spans + # M_CHUNK tiles. mc's stride is exactly this dimension's size*stride, so + # the two merge and the walk stays within the memtile BD's four dims. a_recv_dims = [ (M_CHUNK * M_TILE // R, R * K_TILE), (R, S), (K_TILE // S, R * S), (S, 1), ] - # Emits (b_iter, mc, band) -- b_iter outermost, then all M_CHUNK tiles' - # row-groups. That is the order the core acquires A in when it holds a B - # chunk across the group, and it is why ONE fifo suffices: a second would - # need a third core input DMA channel, and a tile has two. + # Emits (b_iter, mc, band): the order the core acquires A in while holding + # a B chunk across the group. a_send_dims = [ (K_DIV_CT_K_MAX, R * CT_MAX_K), (M_CHUNK * M_TILE // R, R * K_TILE), @@ -697,13 +466,9 @@ def unit_rows(u): for r in range(ROWS): c_prod[(r, c)] = sub[r] - # A: shim -> memtile -> broadcast along the compute row. The reblocking - # rides the forward(): inbound on cons(dims_from_stream=), outbound on - # forward(dims_to_stream=), sharing one memtile buffer. - # ONE fifo per row, even at M_CHUNK > 1. The interleave the core needs - # lives in a_send_dims above, not in extra fifos: a second A fifo would - # make the core want 3 input DMA channels and a compute tile has 2 (the - # design already spends both, on A and B). + # A: shim -> memtile -> broadcast along the compute row, reblocking on + # the forward(). One fifo per row even at M_CHUNK > 1: a second would + # want a third core input DMA channel, and a tile has two. a_l3l2_fifos = [] a_cons = {} for r in range(ROWS): @@ -715,42 +480,26 @@ def unit_rows(u): name=f"A_L2L1_{r}", dims_to_stream=a_send_dims, ) - # One cons() handle per active column; every tile in the row sees - # this object, so inactive columns must not be consumers at all. + # Every tile in the row sees this object, so inactive columns must + # not be consumers at all. for c in range(n_active_cols): a_cons[(r, c)] = of_a.cons() - # B: shim -> memtile -> broadcast down the compute column. - # - # Placement has zero slack -- the eight memtiles pack to exactly 512 KB, - # relying on aie-objectfifo-allocate spilling one buffer to an adjacent - # tile -- so re-verify it after any change to the A, B or C buffer sizes. - # - # One k-block per object, re-fetched from DDR for every row-block. Holding - # a whole column-block here instead would size the buffer from k_iters and - # replay it m_row_blocks times, putting both K and M into the device - # configuration -- which is what the runtime parameters exist to keep out - # of it. The M half is tractable now (aiex.dma_channel_reset_for re-arms a - # resident fifo from the RUNTIME SEQUENCE, and push_queue's repeat is an - # SSA operand); the k_iters-sized buffer is the part still in the way. - # See README.md. + # B: shim -> memtile -> broadcast down the compute column, one k-block per + # object and re-fetched per row-block. Holding a whole column-block would + # put both K and M in the device configuration (README.md). Placement has + # zero slack -- the memtiles pack to exactly 512 KB, one buffer spilled to + # a neighbour -- so re-verify after any A/B/C size change. b_l3l2_fifos = [] b_cons = {} for c in range(n_active_cols): of_b_in = ObjectFifo(mt_b_ty, name=f"B_L3L2_{c}", depth=B_DEPTH) b_l3l2_fifos.append(of_b_in) of_b = of_b_in.cons(dims_from_stream=b_recv_dims).forward( - # The one placement pin this design keeps. Everything else -- the - # workers, the accumulator buffers, the C join, the A forward and - # the shim ends -- is left to the placer, and measures the same. - # - # Without it, aie-place-tiles merges the 20 logical memtiles (4 A - # relays + 8 B relays + 8 C joins) onto the 8 physical ones in a way - # that aie-objectFifo-stateful-transform then rejects with "number - # of input DMA channel exceeded". Spreading B one-per-column is - # enough to steer it to a legal assignment; see the mlir-aie issue - # referenced in README.md. Reproduces at M=1024 K=2048 N=2048, which - # test.py covers. + # The one placement pin; everything else is left to the placer. + # Without it the 20 logical memtiles merge onto the 8 physical + # ones in a way rejected with "number of input DMA channel + # exceeded". Reproduces at M=1024 K=2048 N=2048. tile=Tile(c, 1), obj_type=ct_b_ty, depth=L1_B_DEPTH, @@ -760,12 +509,9 @@ def unit_rows(u): for r in range(ROWS): b_cons[(r, c)] = of_b.cons() - # Each tile's own column, as an initialized buffer rather than a constant - # folded into the program. That is deliberate: the 32 core programs differ - # today only in symbol NAMES, and baking the column in as an immediate - # would make them differ in CODE, permanently foreclosing the one-program - # xclbin. Data may vary per tile; the program must not. Written once at - # configuration time, so unlike an RTP word it costs nothing per dispatch. + # Data, not an immediate folded into the program: the core programs differ + # only in symbol names, and baking this in as code would foreclose a + # one-program xclbin. Written once, so it costs nothing per dispatch. my_cols = [ [ Buffer( @@ -799,31 +545,21 @@ def unit_rows(u): def core_fn(accs, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, my_col, barrier): """Core body. Every trip count and the activation come from the runtime parameter buffer, so one core program serves every shape.""" - # The loop nest lives here rather than inside the kernel so that - # every level has an ObjectFifo acquire point. + # The nest is here, not in the kernel, so every level has an acquire. barrier.wait_for_value(1) - # n_work / n_drain are DERIVED here rather than sent, which costs one - # RTP word instead of two. Branch-free, so no select is needed: - # - # column c has work in column-block j iff (j*COLS + c)*N_TILE < N - # => n_work = ceil((N/N_TILE - c) / COLS) - # - # and every divisor is a power of two (N_TILE=64, COLS=8), so this is - # shifts and adds -- no __divsi3. Verified against host-side trip - # counts for every shape in the suite. - # // rather than >>: the DSL's ScalarValue overloads add/sub/mul/ - # floordiv/mod but NOT the shift operators. Both divisors are - # compile-time powers of two (N_TILE=64, COLS=8), so this strength- - # reduces and must not leave a __divsi3 call -- verified in the .o. + # Derived rather than sent, saving an RTP word: column c has work in + # block j iff (j*COLS + c)*N_TILE < N. Both divisors are powers of two, + # so this must leave no __divsi3 -- check the .o, not the .ll. Use // + # rather than >>; ScalarValue has no shift operators. n_tiles = my_rtp[RTP_N_VAL] // N_TILE n_work = (n_tiles - my_col[0] + COLS - 1) // COLS n_drain = ((n_tiles + COLS - 1) // COLS) - n_work n_row_blocks = my_rtp[RTP_M_ROW_BLOCKS] n_k_iters = my_rtp[RTP_K_ITERS] epi_mode = my_rtp[RTP_EPILOGUE] - # Absent slots become compile-time constants rather than loads: the + # Absent slots become compile-time constants rather than loads. The # kernel's clamped path is compiled out when it is not capable, and at - # M_CHUNK == 1 both chunk counts ARE n_row_blocks. + # M_CHUNK == 1 both chunk counts are just n_row_blocks. if "clamp_enabled" in rtp_slots: clamp_enabled = my_rtp[rtp_slots["clamp_enabled"]] clamp_min_bits = my_rtp[rtp_slots["clamp_min"]] @@ -837,34 +573,21 @@ def core_fn(accs, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, my_col, barrier n_chunks = n_row_blocks n_units_rt = n_row_blocks # Acquire does not consume the barrier, so take it back to zero or the - # next dispatch reads these parameters again instead of waiting. Safe - # before the work: the sequence cannot set the barrier again until it - # has drained this dispatch's C. + # next dispatch re-reads these instead of waiting. Safe before the + # work: the sequence cannot re-set it until this dispatch's C drains. barrier.release_with_value(1) def sweep(group): """One k reduction feeding ``group`` accumulators off a shared B. - ``group`` is a Python list, so its length is compile-time: the mc - loops below unroll. Calling this with every accumulator is the - wide pass; calling it with one is the leftover pass. - - Holding B across the group is the whole point -- b_h is acquired - once outside the mc loop and released after all of them, so DDR - reads B once per len(group) row-blocks instead of once each. + ``group`` is a Python list, so the mc loops unroll. b_h is acquired + outside them, so DDR reads B once per len(group) row-blocks. """ for a_acc in group: init_k(a_acc) for _ in range_(n_k_iters): - # The l loop is unrolled by the B fifo depth so the - # acquired buffer index stays a compile-time constant. for _ in range_(B_ITERS // B_DEPTH): for _ in range(B_DEPTH): - # One B chunk feeds every A band of every accumulator - # in the group, so B is acquired once around both. - # Each accumulator draws A from its OWN fifo, which is - # what makes this interleave legal -- see the a_cons - # construction above. b = b_h.acquire(1) for a_acc in group: for band in range(RHO): @@ -872,8 +595,8 @@ def sweep(group): kstep_k(a, b, a_acc, band) a_h.release(1) b_h.release(1) - # Drain the accumulators. Unrolled by C_DEPTH for the same - # reason; a full O_CHUNKS unroll overflows program memory. + # Unrolled by C_DEPTH; a full O_CHUNKS unroll overflows program + # memory. for a_acc in group: for chunk in range_(O_CHUNKS // C_DEPTH): for half in range(C_DEPTH): @@ -894,10 +617,9 @@ def sweep(group): for _ in range_(n_chunks): sweep(accs) - # Column-blocks this column sits out. A is broadcast along the whole - # compute row, so it must still consume its share or the columns that - # DO have work stall waiting for the fifo to advance. No B and no C - # here -- the runtime sequence issues neither for it. + # Column-blocks this column sits out. A is broadcast along the row, so + # it must still consume its share or the columns that do have work will + # stall on the fifo. No B and no C; the sequence issues neither. for _ in range_(n_drain): # Every unit delivers a full M_CHUNK tiles of A, leftover or not, # so an idle column drains that much per unit. @@ -912,9 +634,8 @@ def sweep(group): workers = [] for r in range(ROWS): for c in range(n_active_cols): - # One accumulator per row-block in a chunk group. Worker flattens - # nested fn_args, so the list arrives in core_fn as a list and its - # length stays compile-time. + # Worker flattens nested fn_args, so the length stays + # compile-time in core_fn. accs = [ Buffer(type=ct_acc_ty, name=f"c_acc_{r}_{c}_{mc}") for mc in range(M_CHUNK) @@ -939,34 +660,13 @@ def sweep(group): # --- Runtime ---------------------------------------------------------- # - # Every wrap below stays under the shim's wrap/size field (DMA_BD_MAX_WRAP): the - # largest are K_TILE=512 and ROWS*M_TILE=256. - # One transfer per (column-block, leg) instead of one per object. - # - # A single fill/drain may span MANY fifo objects -- the descriptor just - # walks them in the order the cores consume -- so a whole column-block's - # worth of A, B and C each go out as one task. Issuing per object instead - # meant a host-side await for every sweep, and those awaits were the - # serialisation: the next sweep could not start until the previous sweep's - # C had come all the way back. - # - # Dimension order must match the core loop nest exactly: for each - # column-block it walks mega_row, then k. Every wrap stays under the shim's - # wrap/size field (largest are K_TILE=512 and ROWS*M_TILE=256). + # One transfer per (column-block, leg), not one per object: a descriptor + # walks many fifo objects in consume order, and per-object issue meant a + # host await per sweep. Dimension order must match the core's nest. def a_taps(mega_col, r, units): - # Every (row-block, k) block this compute row consumes for one - # column-block. A does not depend on mega_col; it is re-fetched per - # column-block because the cores re-consume it. - # - # k OUTERMOST, then the group's M_CHUNK row-blocks: one memtile object - # per k holding M_CHUNK stacked tiles, which a_send_dims then emits - # interleaved as (b_iter, mc, band) -- the order the core acquires in - # while holding a B chunk across the group. - # - # A leftover unit is one row-block wide but fills the same object, so - # its mc dimension has stride 0: the row-block is repeated, and the - # core drains the duplicate (see sweep). It costs one extra read of - # that row-block, on at most one unit per column-block. + # Every (row-block, k) block this row consumes for one column-block, + # k outermost. A does not depend on mega_col; it is re-fetched because + # the cores re-consume it. if M_CHUNK == 1 and not a_split: return [ TensorAccessPattern( @@ -991,47 +691,27 @@ def a_taps(mega_col, r, units): def b_tap(mega_col, c): # Every (mega_row, k) chunk this column consumes. B does not depend on - # mega_row, hence the 0 stride: the same k-blocks are replayed for each - # row-block, which is what the cores expect. - # - # B must arrive PRE-PACKED (see GEMM.pack_B) so each k-block is one - # contiguous run. Expressing that reorder in the descriptor instead - # gives an innermost run of T=8 bf16, turning each 128 KB transfer into - # 8192 scattered bursts -- measured 5.4x slower end to end. + # mega_row, hence the 0 stride. It must arrive pre-packed so each + # k-block is one contiguous run; reordering in the descriptor instead + # gives an innermost run of T=8 bf16 and measured 5.4x slower. return TensorAccessPattern( tensor_dims=(K * N // B_GROUP,), offset=(mega_col * COLS + c) * N_TILE * K // B_GROUP, - # One k sweep per UNIT, not per row-block: the cores hold each B - # chunk across the M_CHUNK row-blocks of a group, so DDR reads B - # n_units times instead of m_row_blocks. That is the whole win -- - # see M_CHUNK_FOR_N. The unit dimension keeps stride 0, replaying - # the same k-blocks, exactly as the row-block dimension used to. + # One k sweep per unit, not per row-block: the cores hold each B + # chunk across a group. The unit dimension keeps stride 0. sizes=[n_units, k_iters, 1, K_TILE * N_TILE // B_GROUP], strides=[0, K_TILE * N_TILE // B_GROUP, 0, 1], ) def c_taps(mega_col, c, units): - # Every joined block this column produces for one column-block: one - # ROWS*M_TILE x N_TILE block per row-block. Returns a LIST, for the - # same reason a_taps does -- one descriptor per unit when N makes the - # row-block stride overflow the shim BD's iteration step. - # - # C drains in plain row-block order even under M_CHUNK, because the - # core drains a group's accumulators one after another, so no - # reordering is needed here; a unit just covers `count` consecutive - # row-blocks. + # Every joined block this column produces: one ROWS*M_TILE x N_TILE + # per row-block, in plain row-block order even under M_CHUNK. if c_split: taps = [] for u in units: first, count = unit_rows(u) - # One descriptor PER ROW-BLOCK, not one per unit. Grouping a - # unit's row-blocks into a count dimension would put a - # ROWS*M_TILE*N stride back inside the descriptor, which is - # exactly what c_split exists to avoid -- it overflows the - # shim BD's 20-bit step at N=10240. C drains in plain - # row-block order even under M_CHUNK (the core drains a - # group's accumulators one after another), so splitting them - # costs nothing but the extra descriptors. + # One descriptor per row-block: grouping them would put the + # ROWS*M_TILE*N stride back in, which c_split exists to avoid. for i in range(count): taps.append( TensorAccessPattern( @@ -1075,52 +755,26 @@ def sequence(A, B, C, a_prods, b_prods, c_conses): for c in range(n_active_cols): barriers[r][c].set(1) - # Column-blocks 0..n_full-1 use every column; the trailing one (when N - # is not a multiple of N_TILE*COLS) uses only the first rem_blocks. A - # is always issued for every row, because the columns sitting the - # trailing block out still drain their share of the broadcast. + # A trailing block uses only the first rem_blocks columns. A is still + # issued for every row, since the sitting-out columns drain it. blocks = [(mc, COLS) for mc in range(n_full)] if rem_blocks: blocks.append((n_full, rem_blocks)) - # One task per (column-block, leg): three per column instead of one per - # object, so a whole column-block retires on a single await rather than - # one per row-block. The C drain is issued first and retired last -- it - # is an S2MM that simply waits for the cores, so keeping it outstanding - # is what overlaps compute with write-back, and it must not share a - # group with the fills it depends on. - # Depth-2: issue column-block i+1 before retiring i, so its transfers - # are already moving while i computes. Retiring a block before issuing - # the next serialises on the C await, which waits for the cores. - # - # This is affordable only because each leg is a single task per - # column-block. Per-object tasks need 1 + 2*k_iters and cannot be - # overlapped at all. - # Keep OVERLAP column-blocks in flight, against SHIM_BDS buffer - # descriptors per column; the operator is DDR-rate bound rather than - # byte bound, so how deeply the fills are pipelined is what decides the - # rate. - # - # Every task of a block stays LIVE in its TaskGroup until the block is - # retired here. That is load-bearing, not tidiness -- see the - # dma_free_task note on a_split above. - # Units, not row-blocks: every leg is issued per unit so A, B and C - # stay aligned with each other and with the core's nest. At - # M_CHUNK == 1 a unit IS a row-block and this is the old list. + # C is issued first and retired last: keeping that S2MM outstanding + # overlaps compute with write-back, and it must not share a group with + # the fills it depends on. Tasks stay live until retired here. all_mb = list(range(n_units)) - # One emitter per leg, so the two paths below differ only in HOW they - # group and retire, not in how a leg is issued. + # One emitter per leg, so the paths below differ only in how they + # group and retire. def issue_a(mega_col, mbs, group, wait=False): for r in range(ROWS): taps = a_taps(mega_col, r, mbs) for i, tap in enumerate(taps): - # A leftover unit emits k_iters*M_CHUNK fills back to back - # on ONE channel (see a_taps), and that channel's task - # queue is SHIM_TASK_QUEUE deep -- overrunning it HANGS - # rather than diagnoses, the same limit a_split windows - # for. Await every SHIM_TASK_QUEUE-th fill so no more than - # that many are ever outstanding. + # A leftover unit emits many fills back to back on one + # channel, so await every SHIM_TASK_QUEUE-th to stay inside + # the queue depth. bounded = ( len(taps) > SHIM_TASK_QUEUE and (i + 1) % SHIM_TASK_QUEUE == 0 ) @@ -1138,10 +792,8 @@ def issue_c(mega_col, active_cols, mbs, group): def emit_unsplit(): pending = [] for mega_col, active_cols in blocks: - # C in its own group, issued first and retired last: it is an - # S2MM that waits on the cores, so keeping it outstanding is - # what overlaps compute with write-back, and it must not share - # a group with the fills it depends on. + # C in its own group so it does not share one with the fills + # it depends on; see above. tg_c = TaskGroup() issue_c(mega_col, active_cols, all_mb, tg_c) tg_f = TaskGroup() @@ -1160,24 +812,14 @@ def emit_unsplit(): def emit_split(): """Issue split legs one unit at a time, retiring the oldest. - One TaskGroup covers one unit (one task per channel), and the - oldest is retired only when a new one would exceed SHIM_TASK_QUEUE - outstanding -- so the channel stays full across both unit and - column-block boundaries. See the SHIM_TASK_QUEUE comment above for - why retiring in batches instead costs real time. - - ``pending`` is retired strictly in append order, which is what - keeps a block's B and unsplit-leg descriptors alive until every one - of that block's units has been AWAITED: they are appended after the - units they cover. + One TaskGroup per unit, retired only when a new one would exceed + SHIM_TASK_QUEUE outstanding. ``pending`` is retired in append + order, which keeps a block's B and unsplit-leg descriptors alive + until its units have been awaited. """ - # What one unit costs on the BUSIEST single channel, which is what - # SHIM_TASK_QUEUE bounds. Under c_split a unit drains one C - # descriptor per row-block (M_CHUNK of them, see c_taps) and they - # all ride the same column's channel, so a unit is worth _per_unit - # there; the a_split leg is one descriptor per row-fifo, so one. - # Counting units instead would overrun by exactly M_CHUNK -- at - # m_chunk=2 that is 8 outstanding, the measured hang threshold. + # What a unit costs on the busiest channel. Under c_split it + # drains M_CHUNK C descriptors onto one, so counting units instead + # would overrun the queue by that factor. unit_cost = _per_unit if c_split else 1 pending = [] # (group, queue cost), oldest first @@ -1204,10 +846,8 @@ def retire(limit): pending.append((tg_u, unit_cost)) retire(SHIM_TASK_QUEUE) - # Not queue-counted: B rides one channel per column and the - # unsplit leg one per row, neither of which is the channel the - # units contend for. They still retire in order, after the - # units of their own block. + # Not queue-counted: B and the unsplit leg ride channels the + # units do not contend for. Still retired in order. pending.append((tg_whole, 0)) pending.append((tg_b, 0)) diff --git a/iron/operators/flm/gemm/op.py b/iron/operators/flm/gemm/op.py index 0bddb6e3e..f5e83bc34 100644 --- a/iron/operators/flm/gemm/op.py +++ b/iron/operators/flm/gemm/op.py @@ -48,20 +48,15 @@ class GEMM(MLIROperator): """AIE-accelerated bf16 GEMM on a 4-row grid, with a fused epilogue. - A row-broadcast / C memtile-join design with fixed 64/512/128 tiling. See - ``design.py`` for how it differs from the more general ``GEMM`` operator. - Unlike ``GEMM`` this exposes no tiling knobs, but folds an activation and an - optional clamp into the output stage. - - The grid is as wide as the device: 8 columns on NPU2, 4 on NPU1 (Phoenix). - Only the width varies -- the tiling and the blocked L1 layout are shared. + Fixed 64/512/128 tiling, no tiling knobs, and an activation plus optional + clamp folded into the output stage. The grid is as wide as the device: 8 + columns on NPU2, 4 on NPU1. See ``design.py``. """ # Every field below is repr=True, so MLIROperator.name derives the artifact - # stem from all of them. That is not cosmetic: each one changes the emitted - # MLIR or the kernel object, and this repo's build cache keys on filename, - # so a variant that shared a stem would be silently satisfied by another - # variant's cached build. + # stem from all of them. The build cache keys on filename, not on source or + # flags, so any field that changes the emitted MLIR or the kernel object + # must reach the stem or a stale build silently satisfies the request. M: int K: int N: int @@ -98,51 +93,32 @@ class GEMM(MLIROperator): } def __post_init__(self): - # Resolve both tile knobs to concrete values here, so the dataclass - # fields hold what the build actually uses. The resolved tile_ma in - # particular must reach the artifact name: the design sizes the A object - # from it while the kernel derives the mmul's rowA from it, so an - # artifact built for one value must never satisfy a request for another. + # Resolve both tile knobs here, so the fields hold what the build + # actually uses. tile_ma especially must reach the artifact name: the + # design sizes the A object from it and the kernel derives the mmul's + # rowA from it. dev = aie_utils.get_current_device() if self.tile_n is None: - # n=64 gives the mmul colA=8 instead of 4, halving accumulator - # traffic per mac; n=128 halves A fetches instead. Which wins - # depends on whether compute or data movement is the critical path. - # - # On NPU2 that flips with K: with a single k iteration there is too - # little compute to hide the extra A traffic, so n=128 wins there - # (~9%), while n=64 wins by ~20% at K >= 1024. - # - # NPU1 never reaches that crossover. It has half the columns AND a - # quarter of the per-tile bf16 mac throughput (four native 4x8x4 - # macs per 8x8x8 shape, against NPU2's two bfp16-emulated ones), so - # it stays compute-bound at every K, and n=128's 32 KB f32 - # accumulator also overflows bank-aware L1 allocation. n=64 wins - # everywhere there by 1.21-1.38x, including at K=512 where NPU2's - # rule would pick 128. + # The trade flips with K on NPU2: one k iteration has too little + # compute to hide n=64's extra A traffic, so n=128 wins there by + # ~9% and n=64 by ~20% at K >= 1024. NPU1 has half the columns and + # a quarter of the per-tile bf16 throughput, so it stays + # compute-bound and n=64 wins at every K by 1.21-1.38x. single_k_iter = self.K // K_TILE <= 1 self.tile_n = 128 if (dev.arch == AIEArch.AIE2p and single_k_iter) else 64 elif self.tile_n not in CT_MAX_K_FOR_N: raise ValueError( f"tile_n must be one of {sorted(CT_MAX_K_FOR_N)}, got {self.tile_n}" ) - # m_chunk: row-blocks a core folds into one B fetch (design.py's - # M_CHUNK_FOR_N). It MUST divide m_row_blocks -- a partial group is - # inexpressible, see the raise in design.py -- so fall back to 1 when - # it does not. That puts M's divisibility in the configuration, but - # only as a 2-bucket split rather than per-shape. + # m_chunk falls back to 1 unless both hold: it divides m_row_blocks (a + # partial group is inexpressible, see design.py), and the group's + # row-blocks sit ROWS*M_TILE*K apart inside the A descriptor, a stride + # that must fit the shim BD's 20-bit step. K=10240 overflows it where + # m_chunk=1 would not. if self.m_chunk is None: want = M_CHUNK_FOR_N[self.tile_n] rows = M_TILE * compute_rows(dev) m_row_blocks = self.M // rows if self.M % rows == 0 else 0 - # Two conditions, both of which force the fallback to 1: - # * m_chunk must DIVIDE m_row_blocks -- a partial group is - # inexpressible, see the raise in design.py. - # * the group's row-blocks sit ROWS*M_TILE*K apart INSIDE the A - # descriptor, so that stride must fit the shim BD's 20-bit - # step. a_split used to lift this dimension out of the - # descriptor entirely; m_chunk puts it back, so big K (10240) - # overflows where m_chunk=1 would not. fits = m_row_blocks and m_row_blocks % want == 0 if fits and not _hw_stride_ok(compute_rows(dev) * M_TILE * self.K): fits = False @@ -186,11 +162,9 @@ def _epilogue_mask(self) -> int: def _clamp_capable(self) -> int: """Whether a clamped path is compiled in at all. - Like ``epilogue_modes`` this is a capability, not a selection: the - clamped instantiation costs program memory, so a build that never - clamps should not carry it. The BOUNDS are runtime parameters, so - ``clamp=(-2, 2)`` and ``clamp=(-4, 4)`` share one xclbin -- only - clamped-versus-not forks the build. + A capability, not a selection: the clamped instantiation costs + program memory. The bounds are runtime, so only clamped-versus-not + forks the build. """ return 1 if self.clamp is not None else 0 @@ -198,19 +172,12 @@ def _clamp_capable(self) -> int: def _config_tag(self) -> str: """Everything that shapes the device configuration, and so the xclbin. - M, K, N, the activation and the clamp BOUNDS are absent: they are - runtime parameters, so they change only the instruction stream. - Whether a clamp exists at all does shape the build; see - ``_clamp_capable``. - - ``ck`` is here for the same reason it is in ``_kernel_object``, and it - is NOT redundant with ``tn``: CT_MAX_K_FOR_N is a tuning table, and - retuning one entry changes the design (B's object width, the k slice, - a_send_dims) while tn is unmoved. ``ma`` does not cover it either -- - it usually moves with ck, but ``tile_ma`` is caller-overridable, so - tn128/ma16 is reachable at two different ck values. Omitting it here - silently served an xclbin built at one ck to a request for another, - which is how test_gemm_tile_options[tn128-ma16] started returning NaN. + M, K, N, the activation and the clamp bounds are absent: they are + runtime parameters. Whether a clamp exists at all does shape the build. + + ``ck`` needs naming separately because retuning CT_MAX_K_FOR_N moves it + while tn is unmoved, and tile_ma is caller-overridable. Omitting it + once served an xclbin built at one ck to a request for another. """ dev = aie_utils.get_current_device().resolve().name return ( @@ -229,9 +196,8 @@ def config_name(self) -> str: def name(self) -> str: """Artifact stem for the instruction stream, which does depend on it. - Prefixed to disambiguate from ``iron.operators.GEMM``: this repo's - build cache keys on filename and mtime rather than on source or flags, - so two operators sharing a stem would silently satisfy each other. + Prefixed to disambiguate from ``iron.operators.GEMM``, which would + otherwise share a stem and satisfy this operator's cache lookups. """ base = f"FLM_GEMM_M{self.M}_K{self.K}_N{self.N}_{self._config_tag}" if self.epilogue != Epilogue.NONE: @@ -242,10 +208,9 @@ def name(self) -> str: def _bfp16_b(self) -> bool: """Whether B is stored as bfp16ebs8 rather than bf16. - AIE2P only, and the reason both mmul templates in the kernel header are - live rather than one being dead code: on AIE2 the scalar BFP types do - not exist, so B stays bf16 and the mmul lowers onto four native 4x8x4 - macs. + AIE2P only, which is why both mmul templates in the kernel header are + live: on AIE2 the scalar BFP types do not exist, so B stays bf16 and + the mmul lowers onto four native 4x8x4 macs. """ return aie_utils.get_current_device().arch == AIEArch.AIE2p @@ -258,18 +223,10 @@ def _b_elem_bytes(self) -> float: def _kernel_object(self) -> str: """Object name over every flag that changes the emitted code. - Everything the -D flags in ``get_kernel_artifacts`` carry has to appear - here: this repo's build cache keys on filename and mtime rather than on - source or flags, so an object built for one configuration would - otherwise silently satisfy a request for another. That includes r/t, - which set the blocked layout, and the epilogue flags, which since the - epilogue was folded into this translation unit shape the same object. - - ``ck`` is CT_MAX_K_FOR_N[tile_n], carried as -DMM_FUSED_CT_K. It is - derived from tile_n today, so it looks redundant -- but it is a TUNING - TABLE, and retuning one entry while leaving the object name alone is - exactly the silent-stale-binary case above. Naming it means the table - can be edited without also remembering to wipe the build dir. + Every -D flag from ``get_kernel_artifacts`` has to appear, for the + cache reason above. ``ck`` looks derivable from tile_n, but that is a + tuning table: naming it means retuning an entry does not also require + wiping the build dir. """ return ( f"mm_fused_{M_TILE}x{K_TILE}x{self.tile_n}" @@ -283,11 +240,9 @@ def _link_file(self) -> str: """What the design names as its kernel: the bare object, or the archive bundling it with the tanh LUT tables. - Only AIE2 evaluates the activations through a LUT (AIE2P has a native - vector tanh), and only an activation references tanh at all -- the plain - epilogue just converts and stores. When this is wrong the failure is a - LINK error for tanh_lut_ab/tanh_lut_cd rather than a compile error, so - it surfaces late. + Only AIE2 evaluates activations through a LUT, and only an activation + references tanh. Getting this wrong is a link error, so it surfaces + late. """ if ( any(Epilogue(m) is not Epilogue.NONE for m in self.epilogue_modes) @@ -303,11 +258,9 @@ def _link_file(self) -> str: def _reference_shape(self) -> tuple[int, int, int]: """The shape the configuration-only module is emitted at. - Its runtime sequence is discarded; only its device body reaches the - xclbin. Taking the smallest valid shape keeps that module cheap to - build and makes the shape-independence explicit -- if a real shape's - instruction stream did not run against this xclbin, some dimension - would still be reaching the configuration. + Its runtime sequence is discarded; only the device body reaches the + xclbin. The smallest valid shape keeps it cheap and makes the + shape-independence explicit. """ dev = aie_utils.get_current_device() # M must be at least m_chunk row-blocks: a partial group is @@ -349,14 +302,9 @@ def get_mlir_artifact(self): def set_up_artifacts(self) -> None: kernels = self.get_kernel_artifacts() - # The xclbin comes from a module emitted at a reference shape and a - # reference activation, so every shape sharing this configuration - # reuses it rather than rebuilding an identical one. - # Canonical bounds, not this instance's: the clamp BOUNDS only reach - # the runtime sequence, which this module has discarded, so passing - # the real ones would make a filename-cached artifact's content depend - # on something that never reaches the xclbin. Only the capability - # matters here, and that is already in config_name. + # Emitted at a reference shape and activation, so every shape sharing + # this configuration reuses it. Canonical bounds, not this instance's: + # the real ones reach only the discarded runtime sequence. config_mlir = self._mlir_artifact( f"{self.config_name}.mlir", *self._reference_shape, @@ -383,21 +331,15 @@ def get_kernel_artifacts(self): base_dir = self.context.base_dir generic = base_dir / "aie_kernels" / "generic" - # mm_fused.cc includes zero.cc, which is genuinely per-architecture - # (AIE2 stores 256 bits at a time, AIE2P 512). A quoted include searches - # the including file's own directory first -- now generic/ -- so the - # arch directory has to be on the include path for it to resolve there. + # mm_fused.cc includes zero.cc, which is per-architecture. A quoted + # include searches generic/ first, so the arch directory must be on the + # include path for it to resolve there. arch_include = [f"-I{base_dir / 'aie_kernels' / kernel_dir}"] - # The 8x8x8 mmul shape this design uses exists on both architectures, - # but by different routes: AIE2P lowers it onto two bfp16-emulated macs, - # which is what this flag selects, while AIE2 lowers it onto four native - # 4x8x4 bf16 macs and ignores the flag entirely (it has no bfp16 - # hardware). Passing it on AIE2 would be harmless but misleading, so it - # is scoped to the architecture where it actually changes codegen. - # - # MM_FUSED_BFP16_B rides along with it: storing B as bfp16ebs8 needs the - # scalar BFP types, which only AIE2P has. See _bfp16_b. + # AIE2P lowers the 8x8x8 mmul onto two bfp16-emulated macs, which this + # selects; AIE2 lowers it onto four native bf16 macs and ignores it. + # MM_FUSED_BFP16_B rides along, since bfp16ebs8 storage needs the + # scalar BFP types. flags = [ # Tile geometry and register tiling, for the mmul. f"-DMM_FUSED_TILE_M={M_TILE}", @@ -423,11 +365,9 @@ def get_kernel_artifacts(self): "-DMM_FUSED_BFP16_B", ] if self.rounding is Rounding.CONV_EVEN: - # ROUND_CONV_EVEN is mm.cc's flag, reused rather than inventing a - # second spelling, and its polarity is mm.cc's too: absent means the - # core's power-up floor mode, even though this operator defaults the - # other way. It covers both conversions in the kernel -- the mmul - # and the epilogue's f32->bf16 store -- which must agree. + # mm.cc's flag and polarity, reused: absent means the core's + # power-up floor mode, though this operator defaults the other way. + # Covers both conversions in the kernel, which must agree. flags.append("-DROUND_CONV_EVEN") kernel_obj = KernelObjectArtifact( @@ -443,9 +383,8 @@ def get_kernel_artifacts(self): ) if self._link_file == self._kernel_object: return [kernel_obj] - # The tanh LUT's coefficient tables live in their own translation unit - # in mlir-aie's runtime lib, so on AIE2 the kernel object alone leaves - # tanh_lut_ab/tanh_lut_cd undefined at link time. See _link_file. + # The tanh LUT tables live in their own translation unit, so on AIE2 + # the kernel object alone leaves them undefined at link time. return [ KernelArchiveArtifact( self._link_file, @@ -456,16 +395,11 @@ def get_kernel_artifacts(self): def pack_B(self, B): """Reorder a row-major ``(K, N)`` weight matrix into consumption order. - Returns a flat uint8 tensor of bfp16ebs8 blocks on NPU2, where B is also - quantized, and a flat bf16 tensor on NPU1. Bound to the operator rather - than a static method because the layout depends on the resolved - ``tile_n`` and on the device; call ``op.pack_B(B)``. - - Packing all the way to consumption order is what makes both B hops - linear descriptors (design.py's b_recv_dims and b_send_dims are both - None), which in turn leaves the descriptor dimensions for a k slice - deep enough to halve the accumulator traffic. - See :mod:`iron.operators.flm.packing` for the layout itself. + Flat uint8 bfp16ebs8 blocks on NPU2, flat bf16 on NPU1. Bound to the + operator because the layout depends on the resolved ``tile_n`` and the + device. Packing to consumption order is what makes both B hops linear + descriptors, freeing the dimensions a deep k slice needs. See + :mod:`iron.operators.flm.packing`. """ return pack_b( B, @@ -485,12 +419,9 @@ def packed_B_size(self, K, N): def get_arg_spec(self): return [ AIERuntimeArgSpec("in", (self.M, self.K)), # A - # B arrives pre-packed by pack_B. On AIE2P it is also quantized to - # bfp16ebs8 -- 9 bytes per 8 values rather than bf16's 16 -- so it - # is declared in BYTES there, sizing the buffer from what pack_B - # actually returns; a (K, N) bf16 spec would over-allocate the - # largest buffer by 1.78x. On AIE2 B stays bf16 and the spec is the - # plain element count. + # On AIE2P B is quantized to bfp16ebs8, so it is declared in + # bytes and sized from what pack_B returns; a (K, N) bf16 spec + # would over-allocate by 1.78x. On AIE2 it is an element count. ( AIERuntimeArgSpec( "in", (self.packed_B_size(self.K, self.N),), dtype=np.uint8 diff --git a/iron/operators/flm/gemm/test.py b/iron/operators/flm/gemm/test.py index e11fc9608..77e8d0a25 100644 --- a/iron/operators/flm/gemm/test.py +++ b/iron/operators/flm/gemm/test.py @@ -47,17 +47,11 @@ def get_params(): if dev_name not in ("npu1", "npu2"): return [] - # The grid is 4 rows by as many columns as the device has, so the width of - # one full sweep -- N_TILE * COLS -- differs per device, and so do the N - # values that leave a trailing PARTIAL column-block. That trailing case is - # the interesting one: some columns compute the block while the rest only - # drain the A broadcast, and real transformer o/down projections always - # land there, since N = model dim is essentially never a multiple of the - # sweep width. - # - # At K = 512 there is a single k iteration, so tile_n defaults to 128 and a - # sweep is 1024 wide on NPU2 and 512 on NPU1; at K >= 1024 tile_n drops to - # 64, halving both. + # One full sweep is N_TILE * COLS wide, so both it and the N values that + # leave a trailing partial column-block differ per device. The trailing + # case is the interesting one: some columns compute the block while the + # rest only drain the A broadcast, and real o/down projections always land + # there. At K = 512 tile_n defaults to 128, halving at K >= 1024. # fmt: off if dev_name == "npu2": # M, K, N, epilogue, clamp, rounding @@ -83,15 +77,14 @@ def get_params(): ( 512, 1024, 2048, SILU, (-4.0, 4.0), CONV_EVEN), ( 256, 512, 1024, SILU, None, FLOOR), # K or N = 10240 at M > 256 overflows the shim BD's 20-bit - # mega_row iteration step, so that leg is issued as one transfer - # per mega_row, retired in windows. These are the real E4B FFN - # projections and were unsupported until that landed; they are - # the regression cover for it. M=2048 needs two windows, which is - # what exercises the windowing. + # mega_row iteration step, so that leg goes out as one transfer + # per mega_row against a bounded outstanding count. These are the + # real E4B FFN projections, unsupported until that landed, and + # M=2048 is what pushes past the bound. ( 1024, 10240, 2560, NONE, None, CONV_EVEN), # E4B down ( 1024, 2560, 10240, NONE, None, CONV_EVEN), # E4B gateup - ( 2048, 10240, 2560, NONE, None, CONV_EVEN), # A, 2 windows - ( 2048, 2560, 10240, NONE, None, CONV_EVEN), # C, 2 windows + ( 2048, 10240, 2560, NONE, None, CONV_EVEN), # E4B down, 2x + ( 2048, 2560, 10240, NONE, None, CONV_EVEN), # E4B gateup, 2x ] else: # npu1: _default_tile_n always returns 64 here, so with 4 columns # every sweep is N_TILE*COLS = 256 wide, not the 128*4=512 an @@ -130,18 +123,14 @@ def get_params(): def check_on_device(operator, golden_ref, K, rounding=CONV_EVEN): """Run ``operator`` against its golden reference and return run_test's result. - Bounds the error in ABSOLUTE terms as a fraction of the accumulated mass, - i.e. the expected size of the K reduction before cancellation, - K * mean|a| * mean|b|. A plain relative tolerance cannot work: with signed A - the K-sum cancels by ~sqrt(K), so |C| ends up far smaller than the mass - while the error tracks the mass, leaving near-zero outputs uncheckable. - - The fraction is per-architecture, because the two lower the same 8x8x8 mmul - onto very different arithmetic: NPU2 emulates it with bfp16, which drops - mantissa bits, while NPU1 has no bfp16 and lowers onto four native bf16 macs - accumulating in f32 -- exact up to the f32->bf16 store, so ~20x tighter. - floor truncates rather than rounding to nearest, so its bias accumulates - over the K reduction instead of cancelling and gets a looser bound on both. + Bounds the error absolutely, as a fraction of the accumulated mass + K * mean|a| * mean|b|. A relative tolerance cannot work: with signed A the + K-sum cancels by ~sqrt(K), so |C| ends up far smaller than the mass the + error tracks, leaving near-zero outputs uncheckable. + + The fraction is per-architecture, since NPU2 emulates the mmul with bfp16 + while NPU1 accumulates four native bf16 macs in f32 (~20x tighter). floor + truncates, so its bias accumulates and gets a looser bound on both. """ mass = ( K @@ -199,21 +188,12 @@ def test_gemm(M, K, N, epilogue, clamp, rounding, aie_context): assert not errors, "Test failed" -def test_gemm_split_leg_windowing(aie_context): - """K or N = 10240 at M > 256 overflows the shim BD's 20-bit mega_row - iteration step, so that leg is issued as one transfer per mega_row, with at - most SHIM_TASK_QUEUE of them outstanding. Two shim resources bound it and - NEITHER is modelled by the toolchain -- the BD ids (16/tile, freed without - a completion check) and the channel task queue (4 deep, pushed - unconditionally) -- so overrunning either is a silent device hang rather - than a diagnostic. - - The sequence retires the OLDEST transfer as it issues the next rather than - draining a whole window, so the live set is SHIM_TASK_QUEUE transfers plus - B and the unsplit leg for each of the two column-blocks a boundary spans: - 4 + 2 + 2 = 8 of 16 descriptors, and at most 4 outstanding per channel. - Assert that arithmetic here, since the numbers come from the hardware and a - future retune of SHIM_TASK_QUEUE could break it silently. +def test_gemm_split_leg_bounds(aie_context): + """K or N = 10240 overflows the shim BD's 20-bit mega_row step, so that leg + goes out one transfer per mega_row. Two unmodelled shim resources bound how + many may be live -- BD ids and the channel task queue -- and overrunning + either hangs silently. The live set is 4 + 2 + 2 = 8 of 16 descriptors; + assert that here, since retuning SHIM_TASK_QUEUE could break it silently. """ from aie.dialects.aie import get_target_model from iron.operators.flm.gemm.design import SHIM_TASK_QUEUE @@ -226,26 +206,18 @@ def test_gemm_split_leg_windowing(aie_context): "the split shapes will hang" ) - # The square case splits BOTH legs, which the real Gemma shapes never do - # (E4B's down-proj overflows on K and its gate/up on N, never both), so it - # is the only cover for the two-sided path. + # The square case splits both legs, which the real Gemma shapes never do + # (E4B's down overflows on K and its gate/up on N, never both), so it is + # the only cover for the two-sided path. GEMM(M=512, K=10240, N=10240, context=aie_context).compile() -def test_gemm_split_leg_windowing_runs(aie_context): +def test_gemm_split_leg_bounds_runs(aie_context): """Execute the two-sided split path, not just compile it. - test_gemm_split_leg_windowing above only compiles this shape: the failure - mode it guards against -- BD-id aliasing and shim task-queue overrun (see - that test's docstring) -- is a runtime device hang or silent corruption, - which compiling the MLIR cannot exercise. This dispatches the same shape on - hardware and checks the result. - - Regular rather than extensive despite being the largest shape here. What it - catches is a hang or silently wrong output, not a wrong number, and its - compile-only sibling is already regular, so leaving the executing half out - of the default run is the wrong side to err on. Costs ~8s against the - regular suite's ~13s. + The failure the sibling test guards against is a runtime hang or silent + corruption, which compiling cannot exercise. Regular rather than extensive + despite the size: ~8s against the suite's ~13s. """ M, K, N = 512, 10240, 10240 golden_ref = generate_golden_reference(M=M, K=K, N=N) @@ -259,16 +231,12 @@ def test_gemm_split_leg_windowing_runs(aie_context): def tile_option_params(): """Every (tile_n, tile_ma) the design accepts on this device. - The shape parameters above exercise only the DEFAULT tile geometry, because - __post_init__ resolves both knobs from the shape and the device. These cover - the knobs themselves, which change the blocked L1 layout: tile_n selects - CT_MAX_K and the B object width, tile_ma sets the mmul's rowA and the A - object height, and pack_B, the four stream-dimension lists and gather_dims - all key off them. A mismatch is silently wrong output rather than a build - error, so each combination has to actually run on hardware. - - The default tile_ma per tile_n stays in the regular suite; the overrides are - extensive, since each is its own kernel object and xclbin. + The shape parameters above only exercise the default geometry, since + __post_init__ resolves both knobs. These cover the knobs themselves, which + change the blocked L1 layout, and a mismatch is silently wrong output + rather than a build error, so each has to run on hardware. The defaults + stay in the regular suite; the overrides are extensive, since each is its + own kernel object and xclbin. """ dev = aie_utils.get_current_device() if dev is None or dev.resolve().name not in ("npu1", "npu2"): @@ -320,10 +288,9 @@ def test_gemm_tile_options(M, K, N, tile_n, tile_ma, aie_context): def test_artifact_stem_differs_from_generic_gemm(M, K, N, aie_context): """``flm.GEMM`` must never share an artifact stem with ``GEMM``. - Both classes are named ``GEMM``, and MLIROperator.name derives the stem - from the class name, while this repo's build cache keys on filename and - mtime rather than on source or flags -- so a shared stem would let the two - operators silently satisfy each other's builds in one build dir. + Both classes are named ``GEMM`` and MLIROperator.name derives the stem from + the class name, so with the cache keyed on filename the two operators would + silently satisfy each other's builds in one build dir. """ assert ( GEMM(M=M, K=K, N=N, context=aie_context).name @@ -334,18 +301,13 @@ def test_artifact_stem_differs_from_generic_gemm(M, K, N, aie_context): def test_one_xclbin_serves_every_shape(aie_context): """Several shapes back to back on one loaded xclbin. - This is what the runtime parameters are for, and the parametrised tests - above cannot cover it: each gets a fresh context, so the array is - reconfigured between cases and any state a dispatch leaves behind is - wiped. Here the shapes share one. - - They disagree on every parameter -- M, K, N, whether a column sits a block - out, and the activation -- and none of them may rebuild the xclbin. + The parametrised tests cannot cover this: each gets a fresh context, so + the array is reconfigured between cases. Here the shapes share one, they + disagree on every parameter, and none may rebuild the xclbin. """ - # Every shape here must resolve to the same m_chunk, because m_chunk - # shapes the core program and so the xclbin (see GEMM._config_tag). It - # buckets M by whether m_row_blocks is a multiple of it -- these are all - # even -- and excludes the K that would overflow the A descriptor's step. + # Every shape must resolve to the same m_chunk, which shapes the core + # program and so the xclbin. These are all even in m_row_blocks and exclude + # the K that would overflow the A descriptor's step. shapes = [ (512, 1536, 2048, "none"), (512, 1536, 256, "none"), # only 4 of 8 columns compute @@ -388,13 +350,9 @@ def test_one_xclbin_serves_every_shape(aie_context): def test_one_xclbin_serves_every_clamp_bound(aie_context): """Different clamp bounds back to back on one loaded xclbin. - The bounds are runtime parameters, so they must not rebuild anything; - only clamped-versus-unclamped is a build choice, because the clamped - instantiation costs program memory. See GEMM._clamp_capable. - - Deliberately separate from test_one_xclbin_serves_every_shape: that one - never clamps, so it cannot catch bounds leaking back into the - configuration, which is exactly what this asserts. + The bounds are runtime parameters, so they must not rebuild anything. + Separate from test_one_xclbin_serves_every_shape, which never clamps and so + cannot catch bounds leaking back into the configuration. """ M, K, N = 256, 512, 1024 bounds = [(-2.0, 2.0), (-4.0, 4.0), (-0.5, 0.5)] From 07368b5bf17e98e01de6de9bb9fcd58c6fd14d62 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 15 Sep 2026 15:03:21 -0600 Subject: [PATCH 11/14] flm_gemm: give the cores an explicit stack, which NPU1 needs The clamp-capable NPU1 build fails aiecc's measured_stack_sizes edge: error: stack_size is absent, so this core uses the device default of 1024 bytes, but it needs 1088 bytes which is every Phoenix failure in CI. The activation LUT path plus the epilogue's clamp vectors put it 64 bytes over the device default, and nothing had set stack_size, so the default applied. 2048 with the exact requirement named in the comment, since aiecc reports it if a change ever outgrows this. It also comes off the L1 budget rather than being left for aiecc to catch: _default_l1 was free to hand a buffer the bytes the stack needs, and relying on a build failure to notice is only tolerable while nothing has to fit in the gap. Verified not to move (tile_ma, b_depth) for any tile_n on either device, so the geometry and the measured performance are unchanged. Co-Authored-By: Claude --- iron/operators/flm/gemm/design.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index 70054c4c4..822061972 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -75,6 +75,12 @@ def compute_rows(dev): C_DEPTH = 2 # C fifo depth; also the core-body unroll B_DEPTH = 2 # B fifo depth; also the core-body unroll A_DEPTH = 2 +# L1 bytes reserved for the core's stack, which the buffer budget below must +# not hand out. The device default is 1024 and aiecc measures what a build +# actually needs: 1088 on NPU1, which is the activation LUT path plus the +# epilogue's clamp vectors, so the default fails the build outright. 2048 +# leaves headroom; aiecc names the exact requirement if a change outgrows it. +STACK_SIZE = 2048 # Row-blocks a core folds into one B fetch, cutting B's DDR reads by M_CHUNK # at the cost of that many L1 accumulators and forcing a_split. Off everywhere # for a contractual reason: it must divide m_row_blocks (M % 512 == 0) while @@ -196,9 +202,10 @@ def _default_l1(n_tile, ct_max_k, b_elem_bytes, budget, m_chunk=1): """Pick the largest working set that fits: (A-tile height, L1 B depth). Deeper B first, then the tallest A that still fits, since colA is worth - far more than B's L1 prefetch. No stack is reserved out of ``budget``; - aiecc fails the build if a core's measured requirement does not fit. + far more than B's L1 prefetch. ``budget`` is the whole local memory; the + stack comes off it here so callers can keep passing the raw size. """ + budget -= STACK_SIZE # m_chunk accumulators, since the core holds a B chunk across that many # row-blocks. The only term that scales with it. acc = m_chunk * M_TILE * n_tile * 4 @@ -220,6 +227,7 @@ def _b_depth_for(t_ma, n_tile, ct_max_k, b_elem_bytes, budget, m_chunk=1): ``_default_l1``'s depth is chosen with its own t_ma, which need not fit a caller-overridden one. Raise rather than reuse a depth that does not fit. """ + budget -= STACK_SIZE # Same terms as _default_l1; acc is the only one that scales with m_chunk. acc = m_chunk * M_TILE * n_tile * 4 cout = CT_OUT_LEN * 2 * C_DEPTH @@ -655,6 +663,7 @@ def sweep(group): my_cols[r][c], barriers[r][c], ], + stack_size=STACK_SIZE, ) ) From 5342f5f7f5153e5135ad5ba4dcd57841ea4e6edb Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 15 Sep 2026 15:05:31 -0600 Subject: [PATCH 12/14] deps: drop the release-note comment from requirements.txt Review feedback: the find-links line above stands on its own. Co-Authored-By: Claude --- requirements.txt | 4 ---- 1 file changed, 4 deletions(-) diff --git a/requirements.txt b/requirements.txt index b79a8960e..fdedee171 100755 --- a/requirements.txt +++ b/requirements.txt @@ -14,10 +14,6 @@ --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly --extra-index-url https://pypi.org/simple -# Tagged release wheels live under their own tag's asset page, not under -# latest-wheels-4 (which carries only the .dev builds), hence the extra -# find-links above. llvm-aie is not tagged in step with mlir-aie; this pin is -# the one mlir-aie v1.4.3 itself names in utils/peano-requirements.txt. mlir_aie==1.4.3 llvm-aie==22.0.0.2026090701+3e93bf7b From b02e51b9c91c6686145be2de040846dbef1fad08 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 15 Sep 2026 15:06:04 -0600 Subject: [PATCH 13/14] flm_gemm: retire split-leg capacity before issuing, not after emit_split appended a unit's TaskGroup and only then retired down to SHIM_TASK_QUEUE, but fill/drain pushes the DMA task immediately while TaskGroup.finish() is what emits the await. So each unit's transfers went onto the channel with four already outstanding -- a transient fifth against a queue documented as four deep, and one more than test_gemm_split_leg_bounds' 4 + 2 + 2 descriptor arithmetic assumes. Retire to SHIM_TASK_QUEUE - unit_cost first instead, so the push happens with room for it. The channel still never drains to empty, which is the property worth -12.4% on these shapes: three transfers stay in flight while the oldest is awaited. Reported by Copilot on #200. Build-verified on both devices at M=512 K=10240 N=10240, where both legs split -- a miscounted bound fails to close its task groups at build time. Not verified on hardware: this box cannot dispatch, so CI is the gate. Co-Authored-By: Claude --- iron/operators/flm/gemm/design.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index 822061972..4ac51a4df 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -847,13 +847,17 @@ def retire(limit): issue_a(mega_col, all_mb, tg_whole) for u in all_mb: + # Before issuing, not after: fill/drain pushes the task + # immediately while TaskGroup.finish() emits the await, so + # retiring afterwards would leave the queue transiently one + # over. Await down to where this unit's transfers fit. + retire(SHIM_TASK_QUEUE - unit_cost) tg_u = TaskGroup() if c_split: issue_c(mega_col, active_cols, [u], tg_u) if a_split: issue_a(mega_col, [u], tg_u, wait=True) pending.append((tg_u, unit_cost)) - retire(SHIM_TASK_QUEUE) # Not queue-counted: B and the unsplit leg ride channels the # units do not contend for. Still retired in order. From 2042939dd933eecd1d4eb9a960b5e482fccd00ff Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 15 Sep 2026 15:08:56 -0600 Subject: [PATCH 14/14] flm_gemm: clamp unconditionally, and name every runtime parameter Four review comments on #200, all about what belongs in the xclbin and what belongs in the instruction stream. MM_FUSED_CLAMP forked the build between clamped and unclamped callers, which reads as "a build per clamp" and is the odd one out: the activation is runtime-selected with every mode compiled in, so the clamp should be too. It now is. There is no unclamped instantiation to compile out -- an unclamped dispatch sends (-inf, +inf), and min(x, +inf) / max(x, -inf) leave every finite value bit-identical, so this costs no accuracy. Gone with it: the CLAMP template parameter, epilogue_dispatch, _clamp_capable, the cl axis in the config tag and the kernel object, and the clamp_enabled RTP word and kernel argument. The bounds are two always-sent words rather than a conditional trio, so an unclamped caller pays ~4 us of dispatch it did not before and a clamping one saves ~2. The README's -3.5%-at-M=256 figure was measured at four words and is annotated rather than restated, since it has not been re-measured. _config_tag folded into config_name: it had one caller besides config_name itself, and name now composes on config_name instead. name also carries the clamp bounds, as raw bit patterns. They are immediates in the runtime sequence and the build cache keys on filename and mtime, so without them a second operator at the same shape with different bounds is served the first one's instruction stream and silently clamps to the first one's values. _epilogue_mask ORs rather than sums, so two copies of a mode cannot carry into the neighbouring mode's bit, and __post_init__ now rejects an epilogue that epilogue_modes omits instead of letting it reach the kernel's default arm and apply no activation at all. Also trims the epilogue comment block, which review found excessive, and drops test.py's late import of get_target_model, already at module top. Build-verified on npu1 and npu2 across every tile_n, with and without a clamp, and with gelu for the LUT archive path. Not verified on hardware: this box ships a cp310 pyxrt that no mlir_aie 1.4.3 wheel matches, so CI is the gate on the numerics. Co-Authored-By: Claude --- aie_kernels/generic/mm_fused.cc | 66 ++++++++--------------- iron/operators/flm/gemm/README.md | 53 ++++++++++-------- iron/operators/flm/gemm/design.py | 54 +++++++++---------- iron/operators/flm/gemm/op.py | 89 ++++++++++++++++++++----------- iron/operators/flm/gemm/test.py | 19 ++++--- 5 files changed, 145 insertions(+), 136 deletions(-) diff --git a/aie_kernels/generic/mm_fused.cc b/aie_kernels/generic/mm_fused.cc index aae8b0be7..46a2173b0 100644 --- a/aie_kernels/generic/mm_fused.cc +++ b/aie_kernels/generic/mm_fused.cc @@ -30,11 +30,6 @@ #ifndef MM_FUSED_EPILOGUE_MODE_MASK #define MM_FUSED_EPILOGUE_MODE_MASK 0xF #endif -// Whether a clamp is compiled in at all: a capability, not a selection, since -// the clamped path costs program memory. The bounds are runtime. -#ifndef MM_FUSED_CLAMP -#define MM_FUSED_CLAMP 0 -#endif namespace { @@ -84,15 +79,16 @@ constexpr aie::rounding_mode round_mode = aie::rounding_mode::floor; #endif // One activation's inner loop. Templated so each mode compiles branch-free; // mm_fused_epilogue_chunk selects between them once per chunk. -template +// +// The clamp is unconditional. An unclamped caller sends (-inf, +inf), which +// leaves every finite value bit-identical, so there is no unclamped +// instantiation to compile and no clamped-versus-not fork in the build. +template static inline void epilogue_body(bfloat16 *__restrict y_out, const float *__restrict src, float clamp_min, float clamp_max) { - aie::vector lo, hi; - if constexpr (CLAMP) { - lo = aie::broadcast(clamp_min); - hi = aie::broadcast(clamp_max); - } + const aie::vector lo = aie::broadcast(clamp_min); + const aie::vector hi = aie::broadcast(clamp_max); AIE_LOOP_MAX_ITERATION_COUNT(CHUNK / V) for (int j = 0; j < CHUNK / V; j++) { @@ -106,8 +102,7 @@ epilogue_body(bfloat16 *__restrict y_out, const float *__restrict src, float cla f = silu_vec(f); else if constexpr (MODE == 3) f = sigmoid_vec(f); - if constexpr (CLAMP) - f = aie::max(aie::min(f, hi), lo); + f = aie::max(aie::min(f, hi), lo); aie::accum out; out.from_vector(f); // The assignment is the conversion: to_v16bfloat16 yields a raw @@ -116,24 +111,6 @@ epilogue_body(bfloat16 *__restrict y_out, const float *__restrict src, float cla aie::store_v(y_out + j * V, v); } } - -// Pick the clamped or unclamped instantiation. Only one exists unless -// MM_FUSED_CLAMP compiled the clamp in. -template -static inline void epilogue_dispatch(bfloat16 *__restrict y_out, - const float *__restrict src, - int32_t clamp_enabled, - float clamp_min, - float clamp_max) -{ -#if MM_FUSED_CLAMP - if (clamp_enabled) { - epilogue_body(y_out, src, clamp_min, clamp_max); - return; - } -#endif - epilogue_body(y_out, src, clamp_min, clamp_max); -} } // namespace extern "C" { @@ -162,21 +139,19 @@ void mm_fused_k_step(bfloat16 *a_buf, mm_fused_b_elem_t *b_buf, float *y_acc, in // Convert chunk (outer * C_DEPTH + half) of the f32 accumulator into a bf16 C // object the core body already acquired, applying an activation and clamp on -// the way out. +// the way out. Fusing them costs one more vector op per 16 elements instead of +// a separate pass over L1. // -// Fusing the activation is the point: the values are already in registers, so -// gelu/silu/sigmoid costs one more vector op per 16 elements rather than a -// separate pass over L1. The mode is runtime, tested once per chunk so the -// inner loop stays branch-free; the cost is program memory, since every mode -// in the mask is compiled in. The clamp splits the same way, its bounds -// arriving as raw int32 because npu_write_rtp only writes i32 words. The chunk -// index comes in two parts because the core body unrolls the drain. +// The mode is runtime, tested once per chunk so the inner loops stay +// branch-free; the cost is program memory, since every mode in the mask is +// compiled in. The bounds arrive as raw int32 because npu_write_rtp only +// writes i32 words, and the chunk index comes in two parts because the core +// body unrolls the drain. void mm_fused_epilogue_chunk(bfloat16 *y_out, float *y_acc, int32_t outer, int32_t half, int32_t mode, - int32_t clamp_enabled, int32_t clamp_min_bits, int32_t clamp_max_bits) { @@ -192,23 +167,24 @@ void mm_fused_epilogue_chunk(bfloat16 *y_out, switch (mode) { #if MM_FUSED_EPILOGUE_MODE_MASK & 2 case 1: - epilogue_dispatch<1>(y_out, src, clamp_enabled, clamp_min, clamp_max); + epilogue_body<1>(y_out, src, clamp_min, clamp_max); return; #endif #if MM_FUSED_EPILOGUE_MODE_MASK & 4 case 2: - epilogue_dispatch<2>(y_out, src, clamp_enabled, clamp_min, clamp_max); + epilogue_body<2>(y_out, src, clamp_min, clamp_max); return; #endif #if MM_FUSED_EPILOGUE_MODE_MASK & 8 case 3: - epilogue_dispatch<3>(y_out, src, clamp_enabled, clamp_min, clamp_max); + epilogue_body<3>(y_out, src, clamp_min, clamp_max); return; #endif // Mode 0 is always compiled, so a mode the mask leaves out yields an - // unactivated result rather than an unwritten buffer. + // unactivated result rather than an unwritten buffer. op.py rejects that + // combination up front; this is the backstop. default: - epilogue_dispatch<0>(y_out, src, clamp_enabled, clamp_min, clamp_max); + epilogue_body<0>(y_out, src, clamp_min, clamp_max); return; } } diff --git a/iron/operators/flm/gemm/README.md b/iron/operators/flm/gemm/README.md index 141f2dc46..28375662b 100644 --- a/iron/operators/flm/gemm/README.md +++ b/iron/operators/flm/gemm/README.md @@ -78,7 +78,7 @@ Two consequences of the native-vs-emulated split are worth knowing: ## Runtime parameters -**Four words** in an L1 buffer per core, written by the runtime sequence and +**Six words** in an L1 buffer per core, written by the runtime sequence and read by the core once its barrier opens: | word | value | @@ -87,20 +87,15 @@ read by the core once its barrier opens: | `m_row_blocks` | `M / 256` | | `k_iters` | `K / 512` | | epilogue | the `Epilogue` mode | +| `clamp_min` / `clamp_max` | the bounds, as raw `int32` bit patterns | A word is not free: each costs ~66 ns per core and the sequence writes `ROWS * COLS` = 32 of them, so **every word is ~2 us of dispatch latency** (measured by padding the buffer at a fixed core count). Against a ~107 us floor that is most of a short-prefill dispatch, so `rtp_layout()` in design.py sizes the buffer per configuration rather than sending words a build cannot -use. Three groups are therefore conditional or gone: - -* `clamp_enabled` / `clamp_min` / `clamp_max` are sent **only by a - clamp-capable build**. Where no clamped path is compiled in -- the default, - and every real projection -- they were written every dispatch and never - read. They stay raw `int32` bit patterns, because `npu_write_rtp` writes - i32 only; the kernel casts back with `__builtin_bit_cast`, since `memcpy` - leaves an unresolved external call rather than folding to a register move. +use. Two groups are therefore conditional or gone: + * `n_chunks` / `n_units` are sent **only when `m_chunk > 1`**. They are `m_row_blocks // M_CHUNK` and each other, so at the shipped `M_CHUNK = 1` the core just reads `m_row_blocks`. @@ -113,30 +108,44 @@ use. Three groups are therefore conditional or gone: n_work = (n_tiles - my_col + COLS - 1) // COLS n_drain = ((n_tiles + COLS - 1) // COLS) - n_work +The clamp bounds are the one group that is unconditional despite most callers +not clamping. The kernel has no unclamped instantiation to compile out: an +unclamped dispatch sends `(-inf, +inf)`, which leaves every finite value +bit-identical. Two always-sent words buy one xclbin for clamped and unclamped +callers alike, which is the whole point of the runtime parameters, and a +clamping caller now sends one word fewer than the old `clamp_enabled` trio did. +The bounds stay raw `int32` because `npu_write_rtp` writes i32 only; the kernel +casts back with `__builtin_bit_cast`, since `memcpy` leaves an unresolved +external call rather than folding to a register move. + The column index is per-tile **static data**, deliberately not a constant folded into the program: the 32 core programs today differ only in symbol names, and baking it into code would make them differ in instructions, foreclosing a future one-program xclbin. -Measured on the 30-shape suite: **-3.5% median at M=256** (best -11.3%, -E2B/kv), and within noise at M >= 1024 -- the saving is a constant ~12 us, so -it is a short-prefill and decode lever, not a prefill one. +Measured on the 30-shape suite at **four** words: **-3.5% median at M=256** +(best -11.3%, E2B/kv), and within noise at M >= 1024 -- the saving is a +constant ~12 us, so it is a short-prefill and decode lever, not a prefill one. +Making the clamp bounds unconditional put two words back, which the ~2 us per +word above prices at **~4 us of that ~12**; the shape of the result is +unchanged but the median has not been re-measured since. All columns are always built. One with no work for a shape gets `n_work = 0` and still drains its share of the A broadcast, because the memtile will not release an A object until every consumer has taken it. The two artifacts therefore carry different stems: the xclbin's `config_name` -covers tile_n, ct_max_k, tile_ma, the compiled activation set, whether a clamp exists, -rounding and the device, while `name` adds M, K, N and the activation. The -xclbin is built from a module emitted at a reference shape, whose runtime -sequence is discarded. - -Two things stay build-time, for the same reason -- each costs program memory: -which activations the epilogue can *select between* (`epilogue_modes`), and -whether a clamped path exists at all. Both land in the xclbin's name. Note the -asymmetry for clamp: *whether* to clamp is a build choice, but the *bounds* -are runtime, so `clamp=(-2, 2)` and `clamp=(-4, 4)` share one xclbin. +covers tile_n, ct_max_k, tile_ma, the compiled activation set, rounding and the +device, while `name` adds every runtime parameter -- M, K, N, the activation +and the clamp bounds. It has to: the sequence writes those as immediates and +the build cache keys on filename and mtime, so a stem that omits one serves the +first caller's instruction stream to the second. The xclbin is built from a +module emitted at a reference shape, whose runtime sequence is discarded. + +One thing stays build-time, because it costs program memory: which activations +the epilogue can *select between* (`epilogue_modes`). It lands in the xclbin's +name. The clamp does not -- every build compiles it, so `clamp=(-2, 2)`, +`clamp=(-4, 4)` and no clamp at all share one xclbin. The core releases its barrier straight after reading the parameters. `wait_for_value` emits `LockAction.Acquire`, which does not leave the lock diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index 4ac51a4df..01b140e15 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -85,7 +85,7 @@ def compute_rows(dev): # at the cost of that many L1 accumulators and forcing a_split. Off everywhere # for a contractual reason: it must divide m_row_blocks (M % 512 == 0) while # the overlay this replaces takes any multiple of 256, so a shape that cannot -# use it forks _config_tag. See README.md. +# use it forks config_name. See README.md. M_CHUNK_FOR_N = {16: 1, 32: 1, 64: 1, 128: 1} # How many column-blocks the runtime sequence keeps in flight. A block costs 3 # shim buffer descriptors on a column (A + B + C) of the 16 available, so the @@ -111,18 +111,25 @@ def mode(self) -> int: return list(Epilogue).index(self) -# The parameter buffer each core reads once its barrier opens. These four are +# The parameter buffer each core reads once its barrier opens. These six are # always present; conditional words follow at offsets rtp_layout() computes, # so a word a build cannot use is never allocated. +# +# The clamp bounds are unconditional even though most callers do not clamp, +# because the alternative is a second xclbin: the kernel's clamp is not +# compiled out, it is neutralised by sending (-inf, +inf). Two words is the +# price of that, and a clamping caller now pays one word less than it did. ( RTP_N_VAL, RTP_M_ROW_BLOCKS, RTP_K_ITERS, RTP_EPILOGUE, -) = range(4) + RTP_CLAMP_MIN, + RTP_CLAMP_MAX, +) = range(6) -def rtp_layout(clamp_capable, m_chunk): +def rtp_layout(m_chunk): """Slot index for each optional parameter, and the total word count. A word is not free: the sequence writes one per core, costing ~2 us of @@ -130,12 +137,7 @@ def rtp_layout(clamp_capable, m_chunk): rather than defaulted. """ slots = {} - n = 4 - if clamp_capable: - slots["clamp_enabled"] = n - slots["clamp_min"] = n + 1 - slots["clamp_max"] = n + 2 - n += 3 + n = 6 if m_chunk > 1: slots["n_chunks"] = n slots["n_units"] = n + 1 @@ -318,11 +320,11 @@ def gemm( epilogue = Epilogue(epilogue) # Clamp bounds ride the RTP buffer as raw int32 bit patterns, since - # npu_write_rtp writes i32 only. Whether a clamped path exists is still - # compile-time (op.py's -DMM_FUSED_CLAMP). - clamp_enabled = 1 if clamp is not None else 0 - rtp_slots, rtp_words = rtp_layout(clamp is not None, M_CHUNK) - clamp_lo, clamp_hi = clamp if clamp is not None else (0.0, 0.0) + # npu_write_rtp writes i32 only. No clamp means the identity bounds rather + # than a different build: min(x, +inf) and max(x, -inf) leave every finite + # value bit-identical, so an unclamped dispatch is numerically unchanged. + rtp_slots, rtp_words = rtp_layout(M_CHUNK) + clamp_lo, clamp_hi = clamp if clamp is not None else (-np.inf, np.inf) clamp_min_bits = int(np.float32(clamp_lo).view(np.int32)) clamp_max_bits = int(np.float32(clamp_hi).view(np.int32)) # A tile does a whole m x n block or nothing, so M and K must tile @@ -425,8 +427,8 @@ def unit_rows(u): epilogue_chunk = Kernel( EPILOGUE_SYMBOL, kernel_object, - # outer, half, mode, clamp_enabled, clamp_min_bits, clamp_max_bits - [ct_out_ty, ct_acc_ty] + [np.int32] * 6, + # outer, half, mode, clamp_min_bits, clamp_max_bits + [ct_out_ty, ct_acc_ty] + [np.int32] * 5, ) # --- Data movement ---------------------------------------------------- @@ -565,15 +567,10 @@ def core_fn(accs, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, my_col, barrier n_row_blocks = my_rtp[RTP_M_ROW_BLOCKS] n_k_iters = my_rtp[RTP_K_ITERS] epi_mode = my_rtp[RTP_EPILOGUE] - # Absent slots become compile-time constants rather than loads. The - # kernel's clamped path is compiled out when it is not capable, and at + clamp_min_bits = my_rtp[RTP_CLAMP_MIN] + clamp_max_bits = my_rtp[RTP_CLAMP_MAX] + # An absent slot becomes a compile-time constant rather than a load: at # M_CHUNK == 1 both chunk counts are just n_row_blocks. - if "clamp_enabled" in rtp_slots: - clamp_enabled = my_rtp[rtp_slots["clamp_enabled"]] - clamp_min_bits = my_rtp[rtp_slots["clamp_min"]] - clamp_max_bits = my_rtp[rtp_slots["clamp_max"]] - else: - clamp_enabled, clamp_min_bits, clamp_max_bits = 0, 0, 0 if "n_chunks" in rtp_slots: n_chunks = my_rtp[rtp_slots["n_chunks"]] n_units_rt = my_rtp[rtp_slots["n_units"]] @@ -615,7 +612,6 @@ def sweep(group): chunk, half, epi_mode, - clamp_enabled, clamp_min_bits, clamp_max_bits, ) @@ -752,11 +748,9 @@ def sequence(A, B, C, a_prods, b_prods, c_conses): rtps[r][c][RTP_M_ROW_BLOCKS] = m_row_blocks rtps[r][c][RTP_K_ITERS] = k_iters rtps[r][c][RTP_EPILOGUE] = epilogue.mode + rtps[r][c][RTP_CLAMP_MIN] = clamp_min_bits + rtps[r][c][RTP_CLAMP_MAX] = clamp_max_bits # Only what this configuration actually reads; see rtp_layout. - if "clamp_enabled" in rtp_slots: - rtps[r][c][rtp_slots["clamp_enabled"]] = clamp_enabled - rtps[r][c][rtp_slots["clamp_min"]] = clamp_min_bits - rtps[r][c][rtp_slots["clamp_max"]] = clamp_max_bits if "n_chunks" in rtp_slots: rtps[r][c][rtp_slots["n_chunks"]] = n_chunks rtps[r][c][rtp_slots["n_units"]] = n_units diff --git a/iron/operators/flm/gemm/op.py b/iron/operators/flm/gemm/op.py index f5e83bc34..b87063b75 100644 --- a/iron/operators/flm/gemm/op.py +++ b/iron/operators/flm/gemm/op.py @@ -67,7 +67,9 @@ class GEMM(MLIROperator): # should compile two. Unlike `epilogue`, this is part of the # configuration. epilogue_modes: tuple[Epilogue, ...] = tuple(Epilogue) - # Optional (min, max) applied after the activation. + # Optional (min, max) applied after the activation. The bounds are runtime + # parameters and the kernel always clamps, so this changes the instruction + # stream only -- clamped and unclamped callers share one xclbin. clamp: tuple[float, float] | None = None # n tile width. 64 halves the mmul's accumulator traffic per mac; 128 # halves A fetches instead. __post_init__ resolves None per device and @@ -145,6 +147,23 @@ def __post_init__(self): # the resolved fields still serialize into artifact names unchanged. self.epilogue = Epilogue(self.epilogue) self.rounding = Rounding(self.rounding) + # Deduplicated, since the mask ORs one bit per mode and a repeat would + # otherwise have to be tolerated by every consumer of the tuple. + self.epilogue_modes = tuple( + dict.fromkeys(Epilogue(m) for m in self.epilogue_modes) + ) + # A mode the mask leaves out reaches the kernel's default arm, which is + # NONE -- an unactivated result rather than a build or dispatch error. + # Refuse instead: this is the caller contradicting itself. + if ( + self.epilogue is not Epilogue.NONE + and self.epilogue not in self.epilogue_modes + ): + raise ValueError( + f"epilogue {self.epilogue} is not in epilogue_modes " + f"{tuple(str(m) for m in self.epilogue_modes)}, so it would not " + "be compiled in and the kernel would silently apply none" + ) if self.clamp is not None: lo, hi = self.clamp if lo > hi: @@ -155,25 +174,25 @@ def __post_init__(self): @property def _epilogue_mask(self) -> int: """Bitmask of the modes compiled into the epilogue. Mode 0 is always - present -- the kernel falls back to it.""" - return 1 | sum(1 << Epilogue(m).mode for m in self.epilogue_modes) - - @property - def _clamp_capable(self) -> int: - """Whether a clamped path is compiled in at all. + present -- the kernel falls back to it. - A capability, not a selection: the clamped instantiation costs - program memory. The bounds are runtime, so only clamped-versus-not - forks the build. + OR rather than sum: ``__post_init__`` deduplicates, but a sum would + make that a correctness requirement rather than tidiness, since two + copies of a mode carry into the neighbouring mode's bit. """ - return 1 if self.clamp is not None else 0 + mask = 1 + for m in self.epilogue_modes: + mask |= 1 << Epilogue(m).mode + return mask @property - def _config_tag(self) -> str: - """Everything that shapes the device configuration, and so the xclbin. + def config_name(self) -> str: + """Stem of the artifacts that do not depend on the shape. - M, K, N, the activation and the clamp bounds are absent: they are - runtime parameters. Whether a clamp exists at all does shape the build. + Everything here shapes the device configuration, and so the xclbin. M, + K, N, the activation and the clamp bounds are absent: they are runtime + parameters, so they reach the instruction stream instead -- see + ``name``. ``ck`` needs naming separately because retuning CT_MAX_K_FOR_N moves it while tn is unmoved, and tile_ma is caller-overridable. Omitting it @@ -181,27 +200,35 @@ def _config_tag(self) -> str: """ dev = aie_utils.get_current_device().resolve().name return ( - f"tn{self.tile_n}_ck{CT_MAX_K_FOR_N[self.tile_n]}" + f"FLM_GEMM_tn{self.tile_n}_ck{CT_MAX_K_FOR_N[self.tile_n]}" f"_ma{self.tile_ma}_mc{self.m_chunk}" - f"_em{self._epilogue_mask:x}_{self.rounding}" - f"_cl{self._clamp_capable}_{dev}" + f"_em{self._epilogue_mask:x}_{self.rounding}_{dev}" ) - @property - def config_name(self) -> str: - """Stem of the artifacts that do not depend on the shape.""" - return f"FLM_GEMM_{self._config_tag}" - @property def name(self) -> str: """Artifact stem for the instruction stream, which does depend on it. - Prefixed to disambiguate from ``iron.operators.GEMM``, which would - otherwise share a stem and satisfy this operator's cache lookups. + The configuration it runs on, then the runtime parameters on top. That + also inherits ``config_name``'s prefix, which disambiguates from + ``iron.operators.GEMM`` -- that class would otherwise share a stem and + satisfy this operator's cache lookups. + + Every runtime parameter has to appear, because the sequence writes them + as immediates and the build cache keys on filename and mtime: a stem + that omits one serves the first caller's instruction stream to the + second and silently applies the first caller's values. The clamp bounds + go in as raw bit patterns, so bounds that differ only below the printed + precision still get their own stem. """ - base = f"FLM_GEMM_M{self.M}_K{self.K}_N{self.N}_{self._config_tag}" + base = f"{self.config_name}_M{self.M}_K{self.K}_N{self.N}" if self.epilogue != Epilogue.NONE: base = f"{base}_epi{self.epilogue}" + if self.clamp is not None: + lo, hi = ( + int(np.float32(v).view(np.int32)) & 0xFFFFFFFF for v in self.clamp + ) + base = f"{base}_cl{lo:08x}{hi:08x}" return base @property @@ -232,7 +259,7 @@ def _kernel_object(self) -> str: f"mm_fused_{M_TILE}x{K_TILE}x{self.tile_n}" f"_ck{CT_MAX_K_FOR_N[self.tile_n]}" f"_r{R}t{T}_ma{self.tile_ma}_{self.rounding}" - f"_em{self._epilogue_mask:x}_cl{self._clamp_capable}.o" + f"_em{self._epilogue_mask:x}.o" ) @property @@ -303,13 +330,13 @@ def set_up_artifacts(self) -> None: kernels = self.get_kernel_artifacts() # Emitted at a reference shape and activation, so every shape sharing - # this configuration reuses it. Canonical bounds, not this instance's: - # the real ones reach only the discarded runtime sequence. + # this configuration reuses it. No clamp, not this instance's bounds: + # they reach only the discarded runtime sequence. config_mlir = self._mlir_artifact( f"{self.config_name}.mlir", *self._reference_shape, Epilogue.NONE, - (0.0, 0.0) if self._clamp_capable else None, + None, ) self.xclbin_artifact = XclbinArtifact( f"{self.config_name}.xclbin", @@ -356,8 +383,6 @@ def get_kernel_artifacts(self): f"-DMM_FUSED_OUT_CHUNK={CT_OUT_LEN}", f"-DMM_FUSED_C_DEPTH={C_DEPTH}", f"-DMM_FUSED_EPILOGUE_MODE_MASK={self._epilogue_mask}", - # Capability only -- the bounds are runtime. See _clamp_capable. - f"-DMM_FUSED_CLAMP={self._clamp_capable}", ] + arch_include if self._bfp16_b: flags += [ diff --git a/iron/operators/flm/gemm/test.py b/iron/operators/flm/gemm/test.py index 77e8d0a25..88a64306b 100644 --- a/iron/operators/flm/gemm/test.py +++ b/iron/operators/flm/gemm/test.py @@ -20,6 +20,7 @@ M_TILE, R, Rounding, + SHIM_TASK_QUEUE, _b_depth_for, _default_l1, ) @@ -195,9 +196,6 @@ def test_gemm_split_leg_bounds(aie_context): either hangs silently. The live set is 4 + 2 + 2 = 8 of 16 descriptors; assert that here, since retuning SHIM_TASK_QUEUE could break it silently. """ - from aie.dialects.aie import get_target_model - from iron.operators.flm.gemm.design import SHIM_TASK_QUEUE - dev = aie_utils.get_current_device() available = get_target_model(dev.resolve()).get_num_bds(0, 0) worst = SHIM_TASK_QUEUE + 2 + 2 @@ -373,9 +371,16 @@ def test_one_xclbin_serves_every_clamp_bound(aie_context): xclbin = stamp assert stamp == xclbin, f"clamp={clamp} rebuilt the xclbin" - # ...but an unclamped build is a different configuration, and must be: - # the clamped instantiation is compiled out entirely there. config_name - # rather than xclbin_artifact, which only exists once compile() has run. + # ...and neither does dropping the clamp: the kernel always clamps, and an + # unclamped caller neutralises it with (-inf, +inf) rather than compiling + # a second build. config_name rather than xclbin_artifact, which only + # exists once compile() has run. clamped = GEMM(M=M, K=K, N=N, clamp=bounds[0], context=aie_context) unclamped = GEMM(M=M, K=K, N=N, context=aie_context) - assert unclamped.config_name != clamped.config_name + assert unclamped.config_name == clamped.config_name + # The bounds do reach the instruction stream, though, so they must reach + # its stem or the build cache serves one caller's stream to another. + assert unclamped.name != clamped.name + assert ( + clamped.name != GEMM(M=M, K=K, N=N, clamp=bounds[1], context=aie_context).name + )