perf(kernels): edge blocks so matmul blocking covers any N; shared unroll skeleton (#103)

Two changes, one optimization and one refactor.

1. Edge blocks (bodies/matmul._blocked)

The blocked path used to fall back to the generic triple loop whenever N was
not a multiple of the block size.  That was invisible while the data points
were dense ladders, but once they became arbitrary integers it cost most of
the benefit: on a declared sample (range:4:64:8) six of eight points took the
fallback and scored 1.00x against it.

Now the blocked path is three segments, with mainI = N rounded down to mr and
mainJ likewise:

  main blocks   [0, mainI) x [0, mainJ)   full mr x nr
  right band    [0, mainI) x [mainJ, N)   one mr x 1 block per column
  bottom band   [mainI, N) x [0, N)       scalar triple loop

Measured on that sample:

  N     before      after
  4        520        536    0.97x   (was already blocked)
  13    20151       9101      2.21x
  21    80460      33434      2.41x
  30   228741      98059      2.33x
  38   459271     190905      2.41x
  47   861499     365786      2.36x
  55  1373241     572307      2.40x
  64   793767     793768     1.00x   (was already blocked)
  total 3817650 -> 2063896    1.85x

Against the -O0 reference the affected points went from about 3.4x to
5.4-9.0x, i.e. into the same band as the points that were already blocked.

A defect was found and fixed while doing this: all three segments test at the
bottom, so when a band is empty (mainJ == N, or mainI == N) the body ran once
anyway and wrote past the end.  N = 4, 52 and 64 segfaulted until each band
got an explicit skip.  Same class of mistake as the remainder tail loop in
5.1: the "remainder is zero" path has to be handled explicitly.

2. Shared unroll skeleton (loopgen.unrolled_loop)

add and reducesum each carried their own copy of the unroll-plus-remainder-tail
scaffolding (power-of-two check, round count, unrolled-segment end address,
tail loop, pointer advance).  It is now one function parameterised by the three
things that actually differed: the per-round pointer list, the single-element
body, and the exit label.

Verified equivalent rather than assumed: normalised output is identical for
all 18 generated artifacts (6 problems x unroll 1/8/32), the token sequences
match, and clang's .text for the changed pair is byte-identical.  The only
source-level difference was whitespace -- the original used one space after
the store mnemonic in the unrolled body and two in the tail, and the refactor
unifies them.

Also: the CLI now defaults -o to ./<problem>.s instead of printing to stdout,
which is what scattered 190 scratch .s files through /tmp; `-o -` still prints.

Docs updated with both findings (DEVELOPMENT.md 6.3, OPTIMIZATION.md 5 and a
new 5.1 on why "fall back the whole problem" is the most expensive kind of
conservative, and pipeline.py's table comments).

All six problems still pass 10/10 on the platform evaluator.

Co-authored-by: Claude <noreply@anthropic.com>

Author: jizhenjun

scratchv/backend/kernels/DEVELOPMENT.md

@@ -1725,30 +1725,83 @@ cost = 执行到的指令数 + 15 × (执行期间遇到的未命中)
 
 **在 N=64 上是 2.71×**，而且 `d_miss` 完全没变（773 → 773）——省下的全是指令。
 
-#### ⚠️ 但分块只在 `N % 4 == 0` 时生效
+#### 边角块：让分块对**任意 N** 生效
 
-看同一条采样上的逐点：
+**先看返工前的样子。** 最初的 `_blocked()` 开头有两条回退：
 
-| N | N%4 | 分块 cost | 通用 cost | 快多少 |
+```
+bltu a2, unit, .Lgeneric      # N 比粒度小 → 回退（这条留着）
+andi TMP, a2, unit-1
+bnez TMP, .Lgeneric           # N 不是粒度的倍数 → 整题回退（这条是问题）
+```
+
+第二条在"数据点都是 4 的倍数"时看不出来。换成任意整数后，同一条采样上：
+
+| N | N%4 | 返工前 | |
+|---:|---:|---:|---|
+| 4 | 0 | 520 | 走上分块 |
+| 13 | 1 | 20,151 | **1.00×** ← 整题退给通用 |
+| 21 | 1 | 80,460 | **1.00×** |
+| 30 | 2 | 228,741 | **1.00×** |
+| 38 | 2 | 459,271 | **1.00×** |
+| 47 | 3 | 861,499 | **1.00×** |
+| 55 | 3 | 1,373,241 | **1.00×** |
+| 64 | 0 | 793,767 | 走上分块 |
+
+**8 个点里 6 个一步没走到。**
+
+> **"数据点可以是任意整数"的第一个后果。** 当数据点全都规整时，分块几乎总是
+> 生效；一旦允许任意整数，它就成了少数派。**当初写"不是倍数就整题回退"时没想到
+> 这一点 —— 因为那时看得到那些规整的数据点。**
+
+**改法：把"整题回退"换成三段**（`mainI = N` 向下取整到 `mr`，`mainJ` 同理）：
+
+| 段 | 范围 | 用什么块 |
+|---|---|---|
+| **① 主块** | `i ∈ [0, mainI)`，`j ∈ [0, mainJ)` | 完整的 `mr×nr` |
+| **② 右带** | `i ∈ [0, mainI)`，`j ∈ [mainJ, N)` | 每列一个 `mr×1` 块 |
+| **③ 底带** | `i ∈ [mainI, N)`，`j ∈ [0, N)` | 标量三重循环 |
+
+**同一个采样上，返工后：**
+
+| N | N%4 | 返工前 | **返工后** | 快多少 |
 |---:|---:|---:|---:|---:|
-| 4 | 0 | 520 | 857 | 1.65× |
-| 13 | 1 | 20,151 | 20,116 | **1.00×** ← 通用回退 |
-| 21 | 1 | 80,460 | 80,425 | **1.00×** |
-| 30 | 2 | 228,741 | 228,721 | **1.00×** |
-| 38 | 2 | 459,271 | 459,251 | **1.00×** |
-| 47 | 3 | 861,499 | 861,479 | **1.00×** |
-| 55 | 3 | 1,373,241 | 1,373,221 | **1.00×** |
-| 64 | 0 | 793,767 | 2,154,214 | 2.71× |
-
-**8 个采样点里只有 2 个（4 和 64）走上分块**，其余 6 个是 `1.00×` ——
-**完全没走到**，因为 `_blocked()` 开头就把 `N % 4 != 0` 的整题退给了通用实现。
-
-**这不是 bug**（正确、且那些点仍比 `-O0` 快约 3.4×），但它是**一块明显的优化空间**：
-给分块路径加**边角块处理**（部分 tile），那 6 个点也能往 9× 靠。
-
-> **这正是"数据点可以是任意整数"的第一个后果。**
-> 当数据点全都规整时，分块几乎总是生效；一旦允许任意整数，它就成了少数派。
-> **当初写"`N % 4 != 0` 就整题回退"时没想到这一点 —— 因为那时看得到那些规整的数据点。**
+| 4 | 0 | 520 | 536 | 0.97× |
+| 13 | 1 | 20,151 | **9,101** | **2.21×** |
+| 21 | 1 | 80,460 | **33,434** | **2.41×** |
+| 30 | 2 | 228,741 | **98,059** | **2.33×** |
+| 38 | 2 | 459,271 | **190,905** | **2.41×** |
+| 47 | 3 | 861,499 | **365,786** | **2.36×** |
+| 55 | 3 | 1,373,241 | **572,307** | **2.40×** |
+| 64 | 0 | 793,767 | 793,768 | 1.00× |
+| | | **3,817,650** | **2,063,896** | **1.85×** |
+
+**非 4 倍数的点 2.2~2.4×；4 的倍数的点一点没变**（N=4 那个 0.97× 是多了两条守卫指令，
+在那个极小规模上摊销不开）。
+
+对 `-O0` 参考解的倍率也跟着上来了：那些点从 ~3.4× 提到 **5.4~9.0×**，
+和 4 的倍数的点落在同一档。
+
+#### ⚠️ 写边角块时踩的坑：循环是「先做后判」
+
+返工后的第一次跑：**非 4 倍数的点全过，N=4/52/64 段错误。**
+
+原因：这三段的循环都是**先执行后判断**（`do { 体 } while (J != N)`）。当边角为**空**时
+——比如 `mainJ == N`（右边没有余列）——循环体**照样跑一次**，于是从第 N 列写出去。
+
+**修法是进段之前先判空：**
+
+```asm
+    beq  MAINJ, a2, .Brs        # 右边没有余列 → 整段跳过
+    ...                         # ② 右带
+.Brs:
+    beq  MAINI, a2, .Lret       # 下边没有余行 → 直接返回
+    ...                         # ③ 底带
+```
+
+> **这不是分块特有的坑。** 任何"主循环 + 余数尾循环"的写法都有它：
+> §5.1 那个展开循环里，`beqz t2, .Ltail` 和 `beqz a2, .Lret` 就是同一件事——
+> **"余数为 0"这条路径必须显式处理，不能指望循环体自己不跑。**
 
 **分块多大**：受**寄存器数量**限制。
 

scratchv/backend/kernels/OPTIMIZATION.md

@@ -109,15 +109,30 @@ N=4 时 cost=466，其中指令 256 条、`15×miss = 210`（**占 45%**，其
 
 | # | 入口 | 现象/机制 | 怎么测 | 框架落点 | 状态 |
 |---|---|---|---|---|---|
+| **0** | **分块的边角块** ★已做 | 分块要求 `N % 4 == 0`，否则**整题退回通用**。数据点改成任意整数后，采样上 8 点里 **6 点一步没走到** | `probe.py` 看那些点是否从 `1.00x` 变成 ~2.4× | `bodies/matmul._blocked` 的三段 | ✅ **已完成**（采样合计 **1.85×**） |
 | **1** | **matmul N=64 的 L1 cache blocking** | §4.2 的容量门槛；可达 gap 的 ~79% | `measure.py` 看 `d_miss` 是否 3634 → ~1000 | `CacheBlocking` | `[P1]` 接口 |
 | **2** | **通用回退提速（GenericPlus）** | 通用版 9 条/元素 vs 特化 4.5 ⇒ **2.0×**；matmul 9 条/MAC vs 2.75 ⇒ **3.3×** | 强制 `ShapeKnowledge=Unknown` 跑对照 | `LoopFormLower` + `AddrStrengthReduce` | `[P1]` 接口 |
 | **3** | **小 N 的固定开销** | §4.3；N=4 单点 466 → 目标 ~440 | 把 256 条拆成 wrapper/分发/序言/主体分别量 | `DispatchLower` | `[P0]` 部分 |
 | **4** | **k 循环加「部分展开」档** | 现在只有「全展开」与「真循环」二选一；`i_miss` 在 N≥24 是 29，N=64 因转真循环降到 18 | 扫 unroll ∈ {2,4,8}，看 instr 与 i_miss 的**联合** cost | `LoopFormLower` | `[P1]` 接口 |
 | **5** | **per-kernel 起始对齐** | 现在只有段级 `RODATA_PAD`，每个 kernel 的起始地址被动地由前面所有 kernel 的长度决定 | 给每个 kernel 加独立 `.balign` 扫 | `SectionLayout` | `[P1]` 接口 |
 | **6** | **`RODATA_PAD` 重扫** | 现值是对 10 个已知 N 扫出来的；**换 wrapper / 工具链 / N 集合都必须重扫** | `padsweep.py` | `SectionLayout` | `[P1]` 接口 |
 | **7** | **fp32 的低指令地板**（内测比赛2） | f32 matmul = `fmul.s`+`fadd.s` = **2 条/MAC 且无需 `slli`**；Q16 是 `mulh`+`add`+摊薄的 `slli` | 同 #1 的方法 | `DtypePolicy` | `[P0]` 本阶段 |
+| **8** | **q16 也能分块**（内测比赛1） | 现在 q16 的 matmul 完全没有分块（整数组 23 个，装不下 (4,4) 的 25 个）。**但要注意**：`blocked_regs_needed()` 只数数据寄存器、不算机制那 13~17 个，所以 (2,2)/(3,3) 能过容量检查却会**撞车** | 先用 `_blocked()` 单独试 (2,2)，看生成的寄存器有没有重叠 | `blocked_regs_needed` + `_blocked` 的机制寄存器 | `[P1]` 未做 |
 | — | ~~指令调度~~ / ~~周期估算~~ | **成本口径不含周期**，收益恒为 0 | — | — | `[P2]` 不做 |
 
+### 5.1 入口 0 的经验：**"整题回退"是最贵的一种保守**
+
+写 `N % 4 != 0 → 整题回退` 时，代价只有"那些点慢一点"。**但"那些点"的数量取决于
+数据点怎么取** —— 全都是 4 的倍数时是 0 个，任意整数时是 70%。
+
+**判断这类取舍，不要问"这个约束合不合理"，要问"它排除掉的数据点占多少"。**
+
+改法（三段：主块 + 右带 + 底带）见 `DEVELOPMENT.md` §6.3。
+
+> 顺带一个**通用的坑**：三段循环都是"先做后判"，所以**边角为空时整段必须显式跳过**
+> （`beq MAINJ, a2, .Brs`）。第一版漏了这个 → N=4/52/64 段错误。
+> **"余数为 0"这条路径永远要显式处理** —— 展开循环的 `beqz` 守卫是同一个道理。
+
 **按「是否把某点从 <1.0 推到 ≥1.0」判据**：入口 1 修的是 N=64 单点（差 7.3%，最大单点）；
 入口 3+4 修的是其余 9 个点（各差 2~4%，是「面」）。
 

scratchv/backend/kernels/__main__.py

@@ -6,6 +6,7 @@
 """
 
 import argparse
+import os
 
 from scratchv.backend.kernels.bodies import PROBLEMS
 from scratchv.backend.kernels.pipeline import build_program
@@ -14,7 +15,8 @@
 def main(argv=None) -> int:
     ap = argparse.ArgumentParser(description='ScratchV 算子内核生成器')
     ap.add_argument('--problem', help='题名，见 --list')
-    ap.add_argument('-o', '--output', help='输出 .s 路径（默认打印到屏幕）')
+    ap.add_argument('-o', '--output', default=None,
+                    help='输出路径；省略则写到 ./<题名>.s；给 "-" 打到屏幕')
     ap.add_argument('--unroll', type=int, default=None,
                     help='展开因子（2 的幂）；省略则用实测表')
     ap.add_argument('--list', action='store_true', help='列出所有题名')
@@ -29,12 +31,16 @@ def main(argv=None) -> int:
         ap.error('要么给 --problem，要么给 --list')
 
     text = build_program(args.problem, args.unroll)
-    if args.output:
-        with open(args.output, 'w', encoding='utf-8') as f:
-            f.write(text)
-        print(f'已写出 {args.output}')
-    else:
+
+    # 省略 -o 时**落到当前目录**（与 gcc/clang 同一约定），而不是打到屏幕。
+    # 打屏幕那条路（`-o -`）留着，但必须是显式的 —— 否则产物会散在终端里没有人接。
+    if args.output == '-':
         print(text)
+        return 0
+    path = args.output or f'{args.problem}.s'
+    with open(path, 'w', encoding='utf-8') as f:
+        f.write(text)
+    print(f'已写出 {os.path.abspath(path)}')
     return 0
 
 

scratchv/backend/kernels/bodies/add.py

@@ -7,9 +7,11 @@
   unroll == 1  通用循环，每元素 9 条指令
   unroll > 1   展开 u 次 + 余数尾循环；偏移变立即数，循环开销摊到 u 个元素上
                每元素 ≈ 4 + 5/u 条指令。**不需要知道 N** —— 尾循环处理 N % u。
+
+骨架（幂校验、轮数、尾循环）在 `loopgen.unrolled_loop`，与 `reducesum` 共用。
 """
 
-from scratchv.backend.kernels.loopgen import prologue, epilogue, unrolled_body
+from scratchv.backend.kernels.loopgen import prologue, epilogue, unrolled_loop
 
 
 def build(target, dtype, unroll: int = 1) -> list[str]:
@@ -18,17 +20,22 @@ def build(target, dtype, unroll: int = 1) -> list[str]:
     ADD = dtype.add
     T1, T2 = dtype.tmp1, dtype.tmp2
 
+    # 单元素体：读 A、读 B、加、写 C。展开体与尾体共用这一段。
+    single = [
+        f'    {L}  {T1}, 0(a0)',
+        f'    {L}  {T2}, 0(t1)',
+        f'    {ADD} {T1}, {T1}, {T2}',
+        f'    {S}  {T1}, 0(a1)',
+    ]
+
     if unroll == 1:
         return [
             *prologue(),
             '    blez a2, .Lret',                # N <= 0 直接返回
             '    slli t0, a2, 2',                # t0 = N * 4
             '    add  t1, a0, t0',               # t1 = &B[0]
             '.Lloop:',
-            f'    {L}  {T1}, 0(a0)',             # 读 A[i]
-            f'    {L}  {T2}, 0(t1)',             # 读 B[i]
-            f'    {ADD} {T1}, {T1}, {T2}',       # 算
-            f'    {S}  {T1}, 0(a1)',             # 写 C[i]
+            *single,
             '    addi a0, a0, 4',                # 三个指针各前进一个元素
             '    addi t1, t1, 4',
             '    addi a1, a1, 4',
@@ -37,39 +44,25 @@ def build(target, dtype, unroll: int = 1) -> list[str]:
             *epilogue(),
         ]
 
-    if unroll & (unroll - 1):
-        raise ValueError(f'unroll 必须是 2 的幂，收到 {unroll}')
-
-    step = unroll * 4                # 每轮前进的字节数
-    sh = unroll.bit_length() - 1     # log2(unroll)
+    body: list[str] = []
+    for k in range(unroll):
+        off = k * 4
+        body += [
+            f'    {L}  {T1}, {off}(a0)',
+            f'    {L}  {T2}, {off}(t1)',
+            f'    {ADD} {T1}, {T1}, {T2}',
+            f'    {S}  {T1}, {off}(a1)',
+        ]
 
     return [
         *prologue(),
-        '    blez a2, .Lret',
-        '    slli t0, a2, 2',                # t0 = N * 4
-        '    add  t1, a0, t0',               # t1 = &B[0]
-        f'    srli t2, a2, {sh}',            # t2 = N / unroll（完整轮数）
-        '    beqz t2, .Ltail',               # N < unroll：直接走尾循环
-        f'    slli t3, t2, {sh + 2}',        # t3 = 轮数 * unroll * 4
-        '    add  t3, a0, t3',               # t3 = 展开段结束的地址
-        '.Lloop:',
-        *unrolled_body(dtype, unroll),       # 偏移全是立即数
-        f'    addi a0, a0, {step}',          # 三个指针每轮只前进一次
-        f'    addi t1, t1, {step}',
-        f'    addi a1, a1, {step}',
-        '    bne  a0, t3, .Lloop',
-        '.Ltail:',
-        f'    andi a2, a2, {unroll - 1}',    # 余数 = N % unroll
-        '    beqz a2, .Lret',
-        '.Ltail_loop:',
-        f'    {L}  {T1}, 0(a0)',
-        f'    {L}  {T2}, 0(t1)',
-        f'    {ADD} {T1}, {T1}, {T2}',
-        f'    {S}  {T1}, 0(a1)',
-        '    addi a0, a0, 4',
-        '    addi t1, t1, 4',
-        '    addi a1, a1, 4',
-        '    addi a2, a2, -1',
-        '    bnez a2, .Ltail_loop',
+        *unrolled_loop(
+            unroll=unroll,
+            prologue=['    slli t0, a2, 2',      # t0 = N * 4
+                      '    add  t1, a0, t0'],    # t1 = &B[0]
+            body=body,
+            tail=single,
+            ptrs=['a0', 't1', 'a1'],             # 三个指针每轮各前进 u 个元素
+        ),
         *epilogue(),
     ]