Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jizhenjun merged 1 commit into from
Oct 7, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 74 additions & 21 deletions scratchv/backend/kernels/DEVELOPMENT.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"这条路径必须显式处理，不能指望循环体自己不跑。**

**分块多大**：受**寄存器数量**限制。

Expand Down
15 changes: 15 additions & 0 deletions scratchv/backend/kernels/OPTIMIZATION.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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%，是「面」）。

Expand Down
18 changes: 12 additions & 6 deletions scratchv/backend/kernels/__main__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
"""

import argparse
import os

from scratchv.backend.kernels.bodies import PROBLEMS
from scratchv.backend.kernels.pipeline import build_program
Expand All @@ -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='列出所有题名')
Expand All @@ -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


Expand Down
65 changes: 29 additions & 36 deletions scratchv/backend/kernels/bodies/add.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand All @@ -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',
Expand All @@ -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(),
]
Loading
Loading