From 0fad210034f19d30d211c1db3a13af4ea53d6584 Mon Sep 17 00:00:00 2001 From: jizhenjun <1569778661@qq.com> Date: Wed, 7 Oct 2026 22:29:07 +0800 Subject: [PATCH] perf(kernels): edge blocks so matmul blocking covers any N; shared unroll skeleton 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 ./.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 --- scratchv/backend/kernels/DEVELOPMENT.md | 95 +++++++++--- scratchv/backend/kernels/OPTIMIZATION.md | 15 ++ scratchv/backend/kernels/__main__.py | 18 ++- scratchv/backend/kernels/bodies/add.py | 65 ++++---- scratchv/backend/kernels/bodies/matmul.py | 147 ++++++++++++++----- scratchv/backend/kernels/bodies/reducesum.py | 50 +++---- scratchv/backend/kernels/loopgen.py | 48 ++++++ scratchv/backend/kernels/pipeline.py | 6 +- 8 files changed, 313 insertions(+), 131 deletions(-) diff --git a/scratchv/backend/kernels/DEVELOPMENT.md b/scratchv/backend/kernels/DEVELOPMENT.md index 7dc80013..10b61fe7 100644 --- a/scratchv/backend/kernels/DEVELOPMENT.md +++ b/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"这条路径必须显式处理，不能指望循环体自己不跑。** **分块多大**：受**寄存器数量**限制。 diff --git a/scratchv/backend/kernels/OPTIMIZATION.md b/scratchv/backend/kernels/OPTIMIZATION.md index 55f89ed6..fb3271e6 100644 --- a/scratchv/backend/kernels/OPTIMIZATION.md +++ b/scratchv/backend/kernels/OPTIMIZATION.md @@ -109,6 +109,7 @@ 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]` 部分 | @@ -116,8 +117,22 @@ N=4 时 cost=466，其中指令 256 条、`15×miss = 210`（**占 45%**，其 | **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%，是「面」）。 diff --git a/scratchv/backend/kernels/__main__.py b/scratchv/backend/kernels/__main__.py index 27906a11..ed808a40 100644 --- a/scratchv/backend/kernels/__main__.py +++ b/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 diff --git a/scratchv/backend/kernels/bodies/add.py b/scratchv/backend/kernels/bodies/add.py index f40213e9..367c89de 100644 --- a/scratchv/backend/kernels/bodies/add.py +++ b/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,6 +20,14 @@ 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(), @@ -25,10 +35,7 @@ def build(target, dtype, unroll: int = 1) -> list[str]: ' 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(), ] diff --git a/scratchv/backend/kernels/bodies/matmul.py b/scratchv/backend/kernels/bodies/matmul.py index 62252e74..1d63f726 100644 --- a/scratchv/backend/kernels/bodies/matmul.py +++ b/scratchv/backend/kernels/bodies/matmul.py @@ -97,48 +97,46 @@ def _blocked(target, dtype, mr: int, nr: int) -> list[str]: a_cur = ['s4', 's5', 's6', 's7'][:mr] # A 的 mr 个游标 b_cur, b_base, c_base = 's8', 's9', 's10' I, J, K, TMP, ADDR = 't1', 't2', 't3', 't4', 'a3' + MAINI, MAINJ = 'a4', 't5' # 完整块的行/列边界 + + def zero(reg: str) -> str: + return f' fmv.w.x {reg}, zero' if dtype.bank == 'fp' else f' li {reg}, 0' + + def reset_rows() -> list[str]: + """把 A 的行基址与 C 的块基址复位到第 0 行。②③ 段要靠它回到起点。""" + lines = [f' mv {a_base[0]}, a0'] + for r in range(1, mr): + lines.append(f' add {a_base[r]}, {a_base[r - 1]}, {STRIDE}') + lines.append(f' mv {c_base}, a1') + return lines unit = max(mr, nr) # 两者都是 2 的幂，取大者即最小公倍数的掩码位数 out = [ ' blez a2, .Lret', f' li {TMP}, {unit}', - f' bltu a2, {TMP}, .Lgeneric', # N < 分块粒度：走通用实现 - f' andi {TMP}, a2, {unit - 1}', - f' bnez {TMP}, .Lgeneric', # N 不是粒度的倍数：走通用实现 + f' bltu a2, {TMP}, .Lgeneric', # N < 分块粒度：太小，不值得分块 '', f' slli {STRIDE}, a2, 2', # STRIDE = N*4 f' mul {TMP}, a2, a2', f' slli {TMP}, {TMP}, 2', # N*N*4 f' add {b_base}, a0, {TMP}', # &B[0][0] + # 完整块边界：向下取整到 mr / nr。用 `andi x, x, -4` 取 N 的低位清零。 + f' andi {MAINI}, a2, {-mr}', + f' andi {MAINJ}, a2, {-nr}', ] - # A 的 mr 个行基址，彼此相隔一个行步长 - out.append(f' mv {a_base[0]}, a0') - for r in range(1, mr): - out.append(f' add {a_base[r]}, {a_base[r - 1]}, {STRIDE}') - out += [ - f' mv {c_base}, a1', # C 的块基址 - f' li {I}, 0', - '.Bli:', - f' li {J}, 0', - '.Blj:', - ] - # 每个 j：A 游标复位到行基址，B 游标指向 B[0][j] + + # ── ① 主块：i ∈ [0, mainI)，j ∈ [0, mainJ) ────────────────────────────── + out += reset_rows() + out += [f' li {I}, 0', '.Bli:', f' li {J}, 0', '.Blj:'] for r in range(mr): out.append(f' mv {a_cur[r]}, {a_base[r]}') out += [ f' slli {TMP}, {J}, 2', f' add {b_cur}, {b_base}, {TMP}', # &B[0][j] ] - # 累加器清零（x0 恒为 0，直接搬） for reg in acc: - if dtype.bank == 'fp': - out.append(f' fmv.w.x {reg}, zero') - else: - out.append(f' li {reg}, 0') - out += [ - f' li {K}, 0', - '.Blk:', - ] + out.append(zero(reg)) + out += [f' li {K}, 0', '.Blk:'] for r in range(mr): out.append(f' {L} {av[r]}, 0({a_cur[r]})') for c in range(nr): @@ -152,10 +150,9 @@ def _blocked(target, dtype, mr: int, nr: int) -> list[str]: f' add {b_cur}, {b_cur}, {STRIDE}', f' addi {K}, {K}, 1', f' bne {K}, a2, .Blk', + f' slli {TMP}, {J}, 2', + f' add {ADDR}, {c_base}, {TMP}', # &C[i][j] ] - # 写回 C 的 mr × nr 块：一个地址寄存器逐行加行步长 - out.append(f' slli {TMP}, {J}, 2') - out.append(f' add {ADDR}, {c_base}, {TMP}') for r in range(mr): for c in range(nr): out.append(f' {S} {acc[r * nr + c]}, {c * 4}({ADDR})') @@ -163,18 +160,98 @@ def _blocked(target, dtype, mr: int, nr: int) -> list[str]: out.append(f' add {ADDR}, {ADDR}, {STRIDE}') out += [ f' addi {J}, {J}, {nr}', - f' bne {J}, a2, .Blj', - f' slli {TMP}, {STRIDE}, 2', # mr == 4 时：前进 4 行 + f' bne {J}, {MAINJ}, .Blj', # j 到 mainJ（不是 N） + f' slli {TMP}, {STRIDE}, {mr.bit_length() - 1}', # mr 行 = mr * STRIDE + ] + for reg in a_base: + out.append(f' add {reg}, {reg}, {TMP}') + out += [ + f' add {c_base}, {c_base}, {TMP}', + f' addi {I}, {I}, {mr}', + f' bne {I}, {MAINI}, .Bli', # i 到 mainI（不是 N） + ] + + # ── ② 右带：i ∈ [0, mainI)，j ∈ [mainJ, N) —— 每列一个 mr×1 块 ────────── + # ⚠️ 循环是「先做后判」，所以 mainJ == N（右边没有余列）时必须**整段跳过**， + # 否则会从第 N 列写出去。这一条在 N 是 4 的倍数时必然触发。 + out.append(f' beq {MAINJ}, a2, .Brs') + out += reset_rows() + out += [f' li {I}, 0', '.Bri:', f' mv {J}, {MAINJ}', '.Brj:'] + for r in range(mr): + out.append(f' mv {a_cur[r]}, {a_base[r]}') + out += [ + f' slli {TMP}, {J}, 2', + f' add {b_cur}, {b_base}, {TMP}', + ] + for r in range(mr): + out.append(zero(acc[r])) + out += [f' li {K}, 0', '.Brk:'] + for r in range(mr): + out.append(f' {L} {av[r]}, 0({a_cur[r]})') + out.append(f' {L} {bv[0]}, 0({b_cur})') + for r in range(mr): + out += mac_instrs(dtype, acc[r], av[r], bv[0], prod) + for r in range(mr): + out.append(f' addi {a_cur[r]}, {a_cur[r]}, 4') + out += [ + f' add {b_cur}, {b_cur}, {STRIDE}', + f' addi {K}, {K}, 1', + f' bne {K}, a2, .Brk', + f' slli {TMP}, {J}, 2', + f' add {ADDR}, {c_base}, {TMP}', # &C[i][j]，这一列 + ] + for r in range(mr): + out.append(f' {S} {acc[r]}, 0({ADDR})') + if r != mr - 1: + out.append(f' add {ADDR}, {ADDR}, {STRIDE}') + out += [ + f' addi {J}, {J}, 1', # 一列一列走 + f' bne {J}, a2, .Brj', # j 到 N + f' slli {TMP}, {STRIDE}, {mr.bit_length() - 1}', ] - if mr != 4: - raise ValueError('分块前进的行数目前只实现了 mr == 4') for reg in a_base: out.append(f' add {reg}, {reg}, {TMP}') out += [ f' add {c_base}, {c_base}, {TMP}', f' addi {I}, {I}, {mr}', - f' bne {I}, a2, .Bli', - ' j .Lret', # 分块跑完就出去，别落进通用实现 + f' bne {I}, {MAINI}, .Bri', # i 到 mainI + '.Brs:', + ] + + # ── ③ 底带：i ∈ [mainI, N)，j ∈ [0, N) —— 标量 ───────────────────────── + # 同理：mainI == N（下边没有余行）时整段跳过。 + out.append(f' beq {MAINI}, a2, .Lret') + out += [ + f' mul {TMP}, {MAINI}, {STRIDE}', # mainI 行处 + f' add {a_base[0]}, a0, {TMP}', # &A[mainI][0] + f' add {c_base}, a1, {TMP}', # &C[mainI][0] + f' mv {I}, {MAINI}', + '.Bbi:', + f' li {J}, 0', + '.Bbj:', + zero(acc[0]), + f' mv {ADDR}, {a_base[0]}', # &A[i][0] + f' slli {TMP}, {J}, 2', + f' add {b_cur}, {b_base}, {TMP}', # &B[0][j] + f' li {K}, 0', + '.Bbk:', + f' {L} {av[0]}, 0({ADDR})', + f' {L} {bv[0]}, 0({b_cur})', + *mac_instrs(dtype, acc[0], av[0], bv[0], prod), + f' addi {ADDR}, {ADDR}, 4', + f' add {b_cur}, {b_cur}, {STRIDE}', + f' addi {K}, {K}, 1', + f' bne {K}, a2, .Bbk', + f' slli {TMP}, {J}, 2', + f' add {TMP}, {c_base}, {TMP}', + f' {S} {acc[0]}, 0({TMP})', + f' addi {J}, {J}, 1', + f' bne {J}, a2, .Bbj', + f' add {a_base[0]}, {a_base[0]}, {STRIDE}', + f' add {c_base}, {c_base}, {STRIDE}', + f' addi {I}, {I}, 1', + f' bne {I}, a2, .Bbi', + ' j .Lret', # 三段跑完就出去，别落进通用实现 '', ] return out @@ -192,6 +269,6 @@ def build(target, dtype, unroll: int = 1, blocking: tuple[int, int] = (0, 0)) -> return [*head, *_generic(target, dtype)] mr, nr = blocking - if mr != 4: - raise ValueError('目前只实现了 mr == 4') + if mr > 4: + raise ValueError('分块的行数上限是 4（A 的行基址/游标各写死了 4 个）') return [*head, '', *_blocked(target, dtype, mr, nr), *_generic(target, dtype)] diff --git a/scratchv/backend/kernels/bodies/reducesum.py b/scratchv/backend/kernels/bodies/reducesum.py index 06611e9b..e9fec57f 100644 --- a/scratchv/backend/kernels/bodies/reducesum.py +++ b/scratchv/backend/kernels/bodies/reducesum.py @@ -6,9 +6,13 @@ 两个档： unroll == 1 通用循环，每元素 5 条指令 unroll > 1 展开 u 次 + 余数尾循环。**不需要知道 N**。 + +骨架与 `add` 共用（`loopgen.unrolled_loop`）。与 `add` 的区别只有三处，全部从 +参数进去：只前进一个指针（`ptrs=['a0']`）、累加器清零、以及余数为 0 时跳到 +`.Lout` 而不是 `.Lret`（累加器要在那里写出去）。 """ -from scratchv.backend.kernels.loopgen import prologue, epilogue +from scratchv.backend.kernels.loopgen import prologue, epilogue, unrolled_loop from scratchv.backend.kernels.dtypes import zero_acc @@ -18,14 +22,18 @@ def build(target, dtype, unroll: int = 1) -> list[str]: ADD = dtype.add ACC, T1 = dtype.acc, dtype.tmp1 + single = [ + f' {L} {T1}, 0(a0)', + f' {ADD} {ACC}, {ACC}, {T1}', + ] + if unroll == 1: return [ *prologue(), ' blez a2, .Lret', *zero_acc(dtype, 't6'), # t6 在这道题里没别的用途，安全 '.Lloop:', - f' {L} {T1}, 0(a0)', - f' {ADD} {ACC}, {ACC}, {T1}', + *single, ' addi a0, a0, 4', ' addi a2, a2, -1', ' bnez a2, .Lloop', @@ -33,41 +41,23 @@ 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 - - body = [] + body: list[str] = [] for k in range(unroll): body += [ f' {L} {T1}, {k * 4}(a0)', f' {ADD} {ACC}, {ACC}, {T1}', ] - tail = [f' {L} {T1}, 0(a0)', f' {ADD} {ACC}, {ACC}, {T1}'] - return [ *prologue(), - ' blez a2, .Lret', - *zero_acc(dtype, 't6'), - f' srli t2, a2, {sh}', # t2 = N / unroll - ' beqz t2, .Ltail', - f' slli t3, t2, {sh + 2}', - ' add t3, a0, t3', # 展开段结束地址 - '.Lloop:', - *body, - f' addi a0, a0, {step}', - ' bne a0, t3, .Lloop', - '.Ltail:', - f' andi a2, a2, {unroll - 1}', - ' beqz a2, .Lout', - '.Ltail_loop:', - *tail, - ' addi a0, a0, 4', - ' addi a2, a2, -1', - ' bnez a2, .Ltail_loop', + *unrolled_loop( + unroll=unroll, + prologue=zero_acc(dtype, 't6'), + body=body, + tail=single, + ptrs=['a0'], # 只有一个输入流要前进 + exit_label='.Lout', # 余数为 0 也要去写出累加器 + ), '.Lout:', f' {S} {ACC}, 0(a1)', *epilogue(), diff --git a/scratchv/backend/kernels/loopgen.py b/scratchv/backend/kernels/loopgen.py index d97cdc2d..c6646ff9 100644 --- a/scratchv/backend/kernels/loopgen.py +++ b/scratchv/backend/kernels/loopgen.py @@ -36,3 +36,51 @@ def unrolled_body(dtype, unroll: int) -> list[str]: f' {dtype.store} {dtype.tmp1}, {off}(a1)', ] return out + + +def unrolled_loop(*, unroll: int, prologue: list[str], body: list[str], + tail: list[str], ptrs: list[str], + exit_label: str = '.Lret') -> list[str]: + """把一段【单元素循环体】套成「展开 u 轮 + 余数尾循环」——**对任意 N 正确**。 + + 这是 `add` 与 `reducesum` 共用的骨架：幂校验、轮数计算、展开段结束地址、 + 余数尾循环。两题的差别只有三处，都从参数进来： + + prologue 循环前的一次性准备（算 B 基址、累加器清零…）。 + 出来后 `a0` 必须指向第一个待处理元素。 + body 展开体，一**轮** u 个元素；偏移用立即数 `0/4/…/(u-1)*4`，从 `a0` 起算。 + tail 单体体，同样从 `a0` 起算。**它是"任意 N 正确"的关键** + —— 比赛的数据点若都是 u 的倍数，这段永远不执行。 + ptrs 需要每轮前进的指针寄存器，如 `['a0','t1','a1']`。 + **`a0` 必须在里面**（循环哨兵用它）。尾循环里这些指针按单元素前进。 + exit_label 余数为 0 时跳到哪。`add` 跳 `.Lret`（不需要收尾）； + `reducesum` 跳 `.Lout`（要在那里把累加器写出去）。 + + 寄存器占用：`t2`（轮数）、`t3`（结束地址）。调用方不要用这两个。 + """ + if unroll < 1 or unroll & (unroll - 1): + raise ValueError(f'unroll 必须是 2 的幂，收到 {unroll}') + if 'a0' not in ptrs: + raise ValueError("a0 必须在内 ptrs 里——循环哨兵要用它") + step = unroll * 4 + sh = unroll.bit_length() - 1 # log2(unroll) + return [ + ' blez a2, .Lret', + *prologue, + 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', # 展开段结束的地址 + '.Lloop:', + *body, + *[f' addi {p}, {p}, {step}' for p in ptrs], + ' bne a0, t3, .Lloop', + '.Ltail:', + f' andi a2, a2, {unroll - 1}', # 余数 = N % unroll + f' beqz a2, {exit_label}', + '.Ltail_loop:', + *tail, + *[f' addi {p}, {p}, 4' for p in ptrs], + ' addi a2, a2, -1', + ' bnez a2, .Ltail_loop', + ] diff --git a/scratchv/backend/kernels/pipeline.py b/scratchv/backend/kernels/pipeline.py index e03e3b59..869e2cd0 100644 --- a/scratchv/backend/kernels/pipeline.py +++ b/scratchv/backend/kernels/pipeline.py @@ -44,9 +44,9 @@ # 【采样】range:4:64:8 → N = {4, 13, 21, 30, 38, 47, 55, 64} # 【目标】总 cost 最小。 # -# ⚠️ 分块只在 `N % 4 == 0` 时生效，其余整题退回通用实现。在那条采样上 -# **8 个点只有 2 个（4 和 64）真正走分块**，其余 6 个是 1.00x（一步没走到）。 -# 数据点允许任意整数之后，这条约束的代价被放大 —— 见 §6.3。 +# ⚠️ 分块**对任意 N 生效**（边角块：主块 + 右带 + 底带，见 DEVELOPMENT.md §6.3）。 +# 早期版本要求 `N % 4 == 0`，否则整题退回通用实现 —— 在那条采样上 8 个点只有 +# 2 个走得到分块。改成边角后，非 4 倍数的点快 2.2~2.4×，采样合计 1.85×。 # # 容量约束：mr*nr + mr + nr + 1 个**数据**寄存器，从 TargetDesc 对应银行取。 # f32：数据在浮点组（32 个）→ (4,4) 需要 25 个，放得下