Skip to content

perf(ffn): roll the unrolled FFN inner-K accumulation into a hardware C_LOOP (sub-64 codegen −80%)#59

Merged
booth-algo merged 1 commit into
mainfrom
perf/ffn-roll
Jun 3, 2026
Merged

perf(ffn): roll the unrolled FFN inner-K accumulation into a hardware C_LOOP (sub-64 codegen −80%)#59
booth-algo merged 1 commit into
mainfrom
perf/ffn-roll

Conversation

@booth-algo

Copy link
Copy Markdown
Collaborator

_emit_ffn_projection_chunk emitted the inner k_tile_count M_MM accumulation fully unrolled (one M_MM + two address bumps per K-tile, baked as straight-line ASM). This is the live FFN path for the sub-64 regime: ffn(...) picks use_loop_instructions = max_k_tiles <= mram_tile_capacity, which is False whenever the K dimension exceeds MRAM tile capacity — exactly the small-mlen/small-capacity Artix-7 case (e.g. inter_dim=1536, mram_tile_capacity=4). So the FFN up/gate/down projections were emitted fully unrolled (with K-split chunking) precisely where it hurts most, and they dominate the decoder program there (~80% of static ISA lines).

The inner accumulation body is already in loop-carried form (w_temp += MLEN², a_actual += MLEN·B·S), and the weight HBM-offset register is dead inside the activation-column loop (only the per-MLEN-block prefetch reads it, and it is reloaded at each block boundary), so it rolls cleanly into a hardware C_LOOP reusing that register as the trip counter. The systolic accumulator sums across all k_tile_count M_MMs before the M_MM_WO flush, so the result is identical to the unrolled form. k_tile_count == 1 keeps the straight-line body (no loop).

This is the FFN analogue of the projection/attention rolling that vram_sub_projection_asm already does by default; _emit_ffn_projection_chunk was the one remaining unrolled matmul site on the sub-64 decoder path.

Verification

Numerically identical — allclose match rate unchanged at 98.52% PASS on both sub-64 decoder configs; vision is untouched (it uses a separate MLP emitter, _emit_vision_mlp_block, and its ISA is byte-for-byte unchanged).

config ISA lines before after Δ allclose
native_32x32x4 decoder 1,130,045 207,677 −81.6% PASS (98.52%)
native_16x16x4 decoder 2,405,087 498,911 −79.3% PASS
native_32x32x4 vision 569,058 569,058 0% (untouched) PASS

This is a code-size / host-compile-time win (the multi-million-line sub-64 decoder compiles shrink ~5×), at a small +2.1% modeled sim_lat (native_32x32x4: 16.716 → 17.063ms) from per-loop C_LOOP_START/END overhead — i.e. it trades a little modeled latency for much smaller codegen, the same trade the existing default rolling makes. mlen ≥ 64 is unaffected: those configs satisfy max_k_tiles <= mram_tile_capacity and already take the use_loop_instructions C_LOOP path.

The roll is currently unconditional on this path; if the +2.1% sim_lat is undesirable on configs where host compile time is not a concern, it can be gated behind the existing unroll_loops / ATEN_OPS_UNROLL flag instead.

…on into a hardware C_LOOP

_emit_ffn_projection_chunk (the live sub-64 FFN path: default unrolled, used when max_k_tiles > mram_tile_capacity, covers up/gate/down single-pass + K-split) emitted the inner k_tile_count M_MM accumulation fully unrolled. The body is already loop-carried (w_temp += MLEN^2, a_actual += MLEN*B*S) and the weight HBM-offset register is dead inside the act_col loop (only the per-MLEN-block prefetch reads it, reloaded each block), so it rolls cleanly into a C_LOOP using that register as the trip counter. The systolic accumulator sums across all k_tile_count M_MMs before the M_MM_WO flush, identical to the unrolled form.

Verified numerically identical (allclose 98.52% PASS, unchanged) on native_32x32x4 and native_16x16x4 decoders; vision is untouched (separate MLP emitter). Static ISA lines: native_32x32x4 decoder 1,130,045 -> 207,677 (-81.6%); native_16x16x4 decoder 2,405,087 -> 498,911 (-79.3%). sim_lat +2.1% (native_32x32x4: 16.716 -> 17.063ms) from per-loop C_LOOP overhead — a code-size/host-compile-time win, not a latency win. mlen>=64 is unaffected (uses the C_LOOP-based use_loop_instructions path).
@booth-algo booth-algo merged commit e6f4955 into main Jun 3, 2026
3 checks passed
@booth-algo booth-algo deleted the perf/ffn-roll branch June 3, 2026 15:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant