perf(ffn): roll the unrolled FFN inner-K accumulation into a hardware C_LOOP (sub-64 codegen −80%)#59
Merged
Merged
Conversation
…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).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
_emit_ffn_projection_chunkemitted the innerk_tile_countM_MMaccumulation fully unrolled (oneM_MM+ two address bumps per K-tile, baked as straight-line ASM). This is the live FFN path for the sub-64 regime:ffn(...)picksuse_loop_instructions = max_k_tiles <= mram_tile_capacity, which isFalsewhenever 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 hardwareC_LOOPreusing that register as the trip counter. The systolic accumulator sums across allk_tile_countM_MMs before theM_MM_WOflush, so the result is identical to the unrolled form.k_tile_count == 1keeps the straight-line body (no loop).This is the FFN analogue of the projection/attention rolling that
vram_sub_projection_asmalready does by default;_emit_ffn_projection_chunkwas the one remaining unrolled matmul site on the sub-64 decoder path.Verification
Numerically identical —
allclosematch 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).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-loopC_LOOP_START/ENDoverhead — i.e. it trades a little modeled latency for much smaller codegen, the same trade the existing default rolling makes.mlen ≥ 64is unaffected: those configs satisfymax_k_tiles <= mram_tile_capacityand already take theuse_loop_instructionsC_LOOP path.The roll is currently unconditional on this path; if the +2.1%
sim_latis undesirable on configs where host compile time is not a concern, it can be gated behind the existingunroll_loops/ATEN_OPS_UNROLLflag instead.