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 个，放得下