From a02f5606066b5eb18e7fb0286a2780f6a961294c Mon Sep 17 00:00:00 2001
From: booth-algo <kevinlauofficial01@gmail.com>
Date: Wed, 3 Jun 2026 14:29:54 +0100
Subject: [PATCH] perf(ffn): roll the inner K-accumulation of the unrolled FFN
 projection into a hardware C_LOOP
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

_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).
---
 asm_templates/ffn_asm.py | 26 ++++++++++++++++++--------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/asm_templates/ffn_asm.py b/asm_templates/ffn_asm.py
index 87491a6..04a2e56 100644
--- a/asm_templates/ffn_asm.py
+++ b/asm_templates/ffn_asm.py
@@ -592,14 +592,24 @@ def _emit_ffn_projection_chunk(
                 )
 
             lines.append(f"S_ADDI_INT gp{w_temp_register}, gp{w_actual_register}, 0 \n")
-            for _ in range(k_tile_count):
-                lines.append(f"M_MM 0, gp{w_temp_register}, gp{a_actual_register} \n")
-                lines.append(
-                    f"S_ADDI_INT gp{w_temp_register}, gp{w_temp_register}, {mlen * mlen} \n"
-                )
-                lines.append(
-                    f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * batch * seq_len} \n"
-                )
+            # Inner K-accumulation rolled into a hardware C_LOOP. The body is
+            # already loop-carried (w_temp += MLEN^2, a_actual += MLEN*B*S), and
+            # the weight HBM-offset register is dead here (only the per-MLEN-block
+            # prefetch above uses it, and it is reloaded each block), so it doubles
+            # 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; this collapses 3*k_tile_count emitted lines to ~5.
+            if k_tile_count > 1:
+                lines.append(f"C_LOOP_START gp{w_hbm_offset_register}, {k_tile_count} \n")
+            lines.append(f"M_MM 0, gp{w_temp_register}, gp{a_actual_register} \n")
+            lines.append(
+                f"S_ADDI_INT gp{w_temp_register}, gp{w_temp_register}, {mlen * mlen} \n"
+            )
+            lines.append(
+                f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * batch * seq_len} \n"
+            )
+            if k_tile_count > 1:
+                lines.append(f"C_LOOP_END gp{w_hbm_offset_register} \n")
             lines.append(f"M_MM_WO gp{intermediate_register}, gp0, 0 \n")
             lines.append(
                 f"S_ADDI_INT gp{intermediate_register}, gp{intermediate_register}, {blen * mlen} \n"