Skip to content

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

Merged
jizhenjun merged 1 commit into
mainfrom
dev/scratch
Oct 7, 2026
Merged

jizhenjun merged 1 commit into
mainfrom
dev/scratch

Conversation

@jizhenjun

Copy link
Copy Markdown
Collaborator

…roll 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.

  1. 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.

…roll 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 ./<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>
@github-actions

github-actions Bot commented Oct 7, 2026

Copy link
Copy Markdown

🤖 AI Code Review

共审查 7 个变更文件

📁 scratchv/backend/kernels/DEVELOPMENT.md

🔴 Bug risk in asm snippet — The guard beq MAINJ, a2, .Brs skips the right strip entirely when mainJ == N. But if mainI is also 0 (N < tile), neither guard fires and you fall through into ③ with uninitialized MAINI. Add a top-level beq MAINI, a2, .Lret or confirm the caller guarantees mainI ≥ 1 when N ≥ unit.

🟡 Self-contained table — "那些点从 3.4× 提到 5.49.0×" is asserted but not backed by numbers in the table. A reader must hunt for -O0 baselines elsewhere. Consider adding a "vs -O0" column or a footnote with the raw numbers so this section stands alone.

🟡 Cross-reference fragility — §5.1 is a positional anchor. If sections get inserted/renumbered this silently breaks. Prefer a named anchor (e.g. [§tail-loop]) or a relative reference like "the expanded-loop guard described above."

💭 Column naming — "返工前" in the second table means "cost under old code" but reads like "cost before the current table." The old table's "分块 cost / 通用 cost" was more precise. Consider "通用回退 cost" to match.


📁 scratchv/backend/kernels/OPTIMIZATION.md

🟡 Entry #8 状态定为 P1 不合理 — 寄存器撞车 = 错误结果,不是性能问题。blocked_regs_needed() 不数机制寄存器意味着当前 q16 的 (2,2)/(3,3) 路径会产生静默的数据损坏。建议标为 [P0],并在 §5.1 之后加一段类似 "如果 blocked_regs_needed 高估了可用容量,所有依赖它的分块决策都是不可信的"。

🟡 底部判据段落后未同步更新 — Entry 0 标为已完成,但紧接着的段落仍说"入口 3+4 修的是其余 9 个点"。Entry 0 已把 6 个采样点从 1.00× 拉上来,"其余"的基数和描述都已过时。

💭 §5.1 的核心教训可以更前置 — "不要问约束合不合理,要问它排除掉的数据点占多少"是通用设计原则,埋在 §5.1 里太深。考虑在表格前的引言段加一句引用,或在文档顶部设一个"踩坑原则"小节。

💭 边界为空的 beq 守卫值得独立成条 — 当前它在 blockquote 里作为附注,但它本质上是 Entry 0 的第二个 bug(第一版段错误),和"整题回退"一样是实际踩过的坑。建议改成独立的 bullet 或与 Entry 0 的修复说明并列,避免读者忽略。


📁 scratchv/backend/kernels/__main__.py

🔴 Breaking change: default output channel — Behavior flips silently: previously omitting -o printed to stdout, now it writes ./<problem>.s. Any downstream consumer (Makefile, CI step, | diff, command -v-style wrappers, tests capturing stdout) that relied on the old contract will break with no error. The rationale in the comment is sound, but this should be a flagged migration — add a CHANGELOG entry and, ideally, a one-release deprecation window (print to stdout with a warning, or gate the new behavior behind an env var / flag).

🟡 Unvalidated filename from args.problem — Line path = args.output or f'{args.problem}.s'. If PROBLEMS keys are ever user-supplied (e.g. via a future --custom flag or import from YAML), a key like ../foo or a/b would escape the cwd. Today the keys are presumably hardcoded constants so this is safe by construction — but it's load-bearing on that assumption. Consider os.path.basename(args.problem) or asserting args.problem matches ^[A-Za-z0-9_.-]+$.

🟡 No error handling around file write — open(path, 'w') will raise on PermissionError, IsADirectoryError (if <problem>.s is a dir), or FileNotFoundError (if user passes -o /no/such/dir/x.s). As a CLI, ap.error(f'无法写入 {path}: {e}') with a clean exit code is friendlier than a traceback, and it's cheap to add:

try:
    with open(path, 'w', encoding='utf-8') as f:
        f.write(text)
except OSError as e:
    ap.error(f'无法写入 {path}: {e}')

🟡 Missing test for the -o - branch — The return 0 after print(text) is load-bearing (without it control falls through to the file-write path and overwrites). Worth a unit test that asserts -o - writes nothing to disk and stdout equals text, so a future refactor doesn't accidentally drop the early return.

💭 default=None is redundant — argparse already defaults to None; being explicit is fine but adds noise. Minor.

💭 Comment block is good — Clearly explains why the convention matches gcc/clang and why implicit stdout was rejected. Exactly the kind of comment that survives the next refactor.


📁 scratchv/backend/kernels/bodies/add.py

🔴 Bug: 缺失 blez a2, .Lret — 旧的 unroll > 1 路径在 *prologue() 之后有 blez a2, .Lret 作为 N ≤ 0 的早期退出。新代码在传入 unrolled_loop 时省略了这一行。若 N 为负,srli(逻辑右移)会把负值当无符号数处理,产生巨大的轮数,导致越界读写。

  • 如果 unrolled_loop 内部补了这行则没问题;若没有,需在 *unrolled_loop(...) 之前加回:
    '    blez a2, .Lret',

🟡 Suggestion: 2 的幂校验被移除 — 旧代码有 if unroll & (unroll - 1): raise ValueError(...),现已删除。如果 unrolled_loop 没有做同样的校验,传入 unroll=3 会生成错误的 srli/slli 移位量(bit_length()-1 = 1,但实际应为 log2(3) 不存在),产生未定义行为。

  • 确认 unrolled_loop 中有等价校验;若没有,补回来或加断言。

💭 Nit: single 和 body 内每条指令的模式完全一致(load T1, load T2, add, store),可以考虑用同一个生成函数以 off 为参数,进一步减少重复。当前规模下不影响可读性,仅供参考。


📁 scratchv/backend/kernels/bodies/reducesum.py

Code Review: reducesum.py refactoring to unrolled_loop


🔴 Removed power-of-2 validation — The explicit guard if unroll & (unroll - 1): raise ValueError(...) is gone. If unrolled_loop doesn't validate this internally, a non-power-of-2 unroll (e.g. 3, 6) will silently generate incorrect code — the srli shift and andi mask in the tail are only correct for powers of 2. Either ensure unrolled_loop validates this, or keep the check here as a cheap defense.

🔴 Missing N ≤ 0 early exit on the unrolled path — The old code had blez a2, .Lret before zero_acc and the loop setup. That guard is no longer present in the caller. If unrolled_loop doesn't emit it internally, when a2 ≤ 0:

  • srli t2, a2, sh on a negative a2 yields a huge positive value (RISC-V logical shift)
  • The loop condition bne a0, t3, .Lloop then iterates into unmapped memory

Verify that unrolled_loop generates blez a2, .Lret before the accumulator init. If it does, consider passing the early-exit label as a parameter for clarity.

🟡 Contract with unrolled_loop is implicit — This refactoring is correct only if unrolled_loop handles: (1) the blez a2, .Lret guard, (2) power-of-2 validation, and (3) computes sh internally. None of that is visible in this diff. Add a one-line docstring on build stating the contract, or pull those invariants into the call site to keep the caller self-contained:

if unroll & (unroll - 1):
    raise ValueError(f'unroll 必须是 2 的幂,收到 {unroll}')

🟡 single is used for both unroll == 1 body and unrolled_loop tail — but they have subtly different semantics — In the unroll == 1 loop, single is the entire loop body (load + add), with pointer/counter increments following separately. In unrolled_loop's tail, single is just the inner instructions, and unrolled_loop presumably appends addi a0, a0, 4 + addi a2, a2, -1 + branch. This is correct today, but the reuse obscures that single plays two roles. A brief comment like # 2-instruction inner body; outer loop adds pointer/counter would help future readers.

💭 exit_label='.Lout' — The comment explains why (# 余数为 0 也要去写出累加器), which is good. Just double-check that unrolled_loop uses .Lout consistently as the "tail remainder = 0" jump target and not as a generic exit — if unrolled_loop is later reused by a kernel without a store step, the name could mislead.


📁 scratchv/backend/kernels/loopgen.py

🔴 Bug: Hardcoded .Lret skips prologue for N≤0 — First generated line 'blez a2, .Lret' ignores the exit_label parameter. For reducesum with N=0, this skips the prologue (accumulator init) AND jumps to .Lret instead of .Lout, so the output pointer gets nothing written — corrupt/garbage result. Fix: move blez after *prologue and use the parameter:

*prologue,
f'    blez a2, {exit_label}',   # N<=0: prologue already ran, now skip to exit
f'    srli t2, a2, {sh}',
...

🟡 Suggestion: Document a2 must not be in ptrs — a2 is used internally as the N count and remainder. If a caller passes a2 in ptrs, the addi advances corrupt the remainder mid-tail-loop, causing an infinite loop or wrong element count. The docstring already warns about t2/t3 but should add a2 to that list.

# Add to docstring:
#   寄存器占用:t2(轮数)、t3(结束地址)、a2(计数/余数)。调用方不要用这三个,
#   且 ptrs 里也不要出现 a2。

📁 scratchv/backend/kernels/pipeline.py

纯注释变更,无逻辑风险。

🟡 **性能数字需标注来源** — "快 2.2~2.4×"、"采样合计 1.85×":
    若为实测值,建议补充基准配置(芯片/编译器/数据规模),否则读者
    无法复现或判断适用性。若为预估,应标注为 projected。

🟡 **文档引用一致性** — 新增引用 "DEVELOPMENT.md §6.3":
    确认该章节实际描述了边角块方案(主块 + 右带 + 底带),而非仍指向
    旧版 §6.3 的 N%4 约束内容。注释与文档若不同步,比注释本身更容易
    腐烂。

@FeelTheBeats FeelTheBeats left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

单算子性能优化与边界bug修复

@jizhenjun
jizhenjun merged commit 149ae47 into main Oct 7, 2026
8 checks passed
@jizhenjun
jizhenjun deleted the dev/scratch branch October 7, 2026 14:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants