ATen codegen fixes: matmul column-stride, FFN weight row-align, single-pass softmax#63
Merged
Conversation
vram_sub_projection (non-transposed M_MM) strode the weight-column sub-tile base by blen (sub-row). The RTL matrix SRAM read is row-granular (raddr >> log2(MLEN)), so successive BLEN-column output sub-tiles collapsed onto the same MRAM rows and the output columns replicated ([a,b,c,d]x4) -> ATen linear = 27% on RTL. Stride by blen*mlen (one full row per sub-tile, matching the projection_asm template and the transposed M_TMM path). ATen linear 27% -> 95.17% on RTL. Prefetch is block-granular (mlen*mlen), so no prefetch change is needed. Note: the transactional emulator's M_MM weight model is element-granular (matrix_machine.rs asserts mat_offset < mlen) and rejects this row-aligned stride; that emulator model is infaithful to the RTL and is a separate fix.
_emit_ffn_projection_chunk advanced the weight-read pointer by blen between sub-column tiles. The RTL matrix SRAM read is row-granular (raddr >> log2(MLEN)), so +blen collapsed sub-columns 1..mlen/blen-1 onto MRAM row 0 (output cols 4-15 = 0). Stride the weight pointer by blen*mlen (row-aligned, matching projection_asm); the output-write pointer stays at +blen. FFN UP projection 7.6% -> 94.3% on RTL. Same class as the vram_sub_projection mat_col_stride fix.
Rework softmax_plena to compute each row in one vector pass, keeping the reduction results in FP registers and feeding them directly to the next vector op (V_RED_MAX f2 -> V_SUB_VF ..f2 ; V_RED_SUM f3 -> S_RECI_FP f3 -> V_MUL_VF ..f3). Drops the online/streaming path that relied on S_MAX_FP and round-tripped max/sum through FP SRAM, which tripped a reduction-result capture hazard (softmax ATen 51.95% -> 100%). Attention's online-softmax is unchanged.
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
1 participant
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.
Three ATen
PlenaCompilercodegen fixes, verified end-to-end on the PLENA RTL simulator.mat_col_stridein the sub-projection lowering (RTL linear ATen 27% → ~95%).softmax_plenato a per-row single pass, keeping max/sum in FP registers and feeding them directly to the vector consumers (no FP-SRAM round-trip, noS_MAX_FP). Fixes the reduction-result capture hazard the batched/online path hit (softmax ATen 51.95% → 100%). Attention's online-softmax is untouched.Verified on RTL: softmax ATen 100%, ffn ATen 91.6%, linear ATen ~96% (= template), rms/scratchpad 100%.