Skip to content

feat(ir): first-class ops for FWHT/WinogradConv/CSR-SpMM with C and machine-layer lowering - #99

Closed
FeelTheBeats wants to merge 5 commits into
ScratchV-Compiler:mainfrom
FeelTheBeats:sunxiaochuan02
Closed

FeelTheBeats wants to merge 5 commits into
ScratchV-Compiler:mainfrom
FeelTheBeats:sunxiaochuan02

Conversation

@FeelTheBeats

Copy link
Copy Markdown
Contributor

概述

把 FWHT / WinogradConv / CSR-SpMM 三个自定义张量算子从 standalone 生成器提升为
一等 IR 算子,并打通两条 lowering 路径:可移植 C 后端、平台汇编后端(经机器层
由 AsmEmitter 渲染)。主编译管线可直接产出尺寸无关的 .s。

改动分层

1. IR / 前端(ae4bf03)

  • 新增 FWHT / WINOGRAD_CONV / SPMM_CSR opcode、builder helper 与 verifier 签名
  • ir_numpy_ops 按 dtype 分派内核(INT32 → Q16.16,FLOAT32 → 原生 FP32)
  • ONNX 前端接入 org.scratchv 域白名单、shape 自动 wiring、可选 Winograd bias,
    以及编译期 F(2,3) kernel folding(复用 standalone 变换)

2. 可移植 C 后端(7669610)

  • 三个算子的 FP32 _infer/_emit:FWHT 蝶形 + 1/N 反变换缩放、CSR SpMM 的
    status-4 索引校验、Winograd F(2,3) tile 流水线(workspace_bytes 计入静态 scratch)

3. 平台汇编后端(d42976b / 0a2c043 / bb63617)

  • Phase 3:新增 --platform-asm / CompilerConfig.platform_asm,单算子图复用已验证的
    standalone FP32/rv32imf kernel,不支持的多算子图显式拒绝(d42976b)
  • 补齐机器层缺口:SLLI/SRLI/RET 终结符语义、FLW/FSW 内存操作数格式化(0a2c043)
  • platform_emit 重写为构造 MachineInstr kernel(FWHT / CSR SpMM / direct Conv)并由
    AsmEmitter 渲染,去除对 standalone 字符串生成器的依赖;新增 FMV_W_X opcode
    以产出 fmv.w.x(bb63617)

测试

  • tests/_op_ref.py 共享参考实现,IR / C / 汇编路径统一对拍
  • IR verifier、前端自定义算子 parse+run、定点算子、tensor-C、platform-emit 全量覆盖
  • platform-emit 数值测试在共享的 in-test RV32IMF 解释器中真实执行产物

依赖与说明

  • 无新增运行时依赖
  • 多算子图暂不支持 --platform-asm(显式报错,后续可扩展)

Add FWHT/WINOGRAD_CONV/SPMM_CSR opcodes, builder helpers and verifier
signatures, plus dtype-dispatched kernels (Q16.16 on INT32, native FP32 on
FLOAT32) in ir_numpy_ops. Wire the ONNX frontend with the org.scratchv domain
whitelist, self-contained shape wiring, optional Winograd bias and compile-time
F(2,3) kernel folding that reuses the standalone transform.

Tests: IR kernel/verifier coverage and frontend parse+run vs the standalone
references (tests/_op_ref.py).
Add FP32 _infer/_emit for the three operators: FWHT butterfly with 1/N inverse
scaling, CSR SpMM with status-4 index checks, and the Winograd F(2,3) tile
pipeline with tracked static scratch (workspace_bytes now accounts for it).

Tests execute host-compiled C against the shared references, cover static
rejection, and run a CompilerDriver end-to-end path.
…e 3)

Add --platform-asm / CompilerConfig.platform_asm, which reuses the verified
standalone FP32/rv32imf kernels for single-operator graphs (Fwht, Conv/
WinogradConv, SpmmCsr) instead of lowering IR through the scalar allocator/
ABI-frame path. Unsupported graphs are rejected explicitly.

Tests cover the cnn_entry contract, size independence, multi-op rejection, the
CLI flag wiring, and CompilerDriver end-to-end emission.
Fill machine-layer gaps needed by tensor-kernel lowering: immediate shift
opcodes, a RET terminator with explicit semantics, and AsmEmitter formatting
for float load/store memory operands.
Rewrite platform_emit to build MachineInstr kernels (FWHT / CSR SpMM / direct
Conv) and render them with AsmEmitter, dropping the dependency on the standalone
string generators. Add a distinct FMV_W_X opcode so float zeroing emits the
standalone/interpreter mnemonic fmv.w.x. Numeric tests execute the emitted
listings in the shared in-test RV32IMF interpreter.
@github-actions

github-actions Bot commented Oct 6, 2026

Copy link
Copy Markdown

🤖 AI Code Review

共审查 10 个变更文件
⚠️ 另有 8 个文件超过上限(最多 10 个)未审查

📁 scratchv/analysis/ir_verifier.py

🔴 Bug: cout gate silently bypasses bias check — Line ~417: if cout is missing, non-int, 0, or negative, the entire bias-shape validation is skipped. A malformed/absent cout attr should be its own error, not a silent pass.
Suggestion: error out when cout is absent or <= 0 (assuming cout is required), and validate bias shape against a positive cout.

🟡 T type tag on FWHT/WINOGRAD_CONV may be decorative — OpcodeSpec gets "T" for these, but the added custom validation only handles SPMM_CSR. If there's no generic "T == all-operand-same-dtype" checker elsewhere in the verifier, the tag is misleading and FWHT/WINOGRAD operand type consistency is unenforced.
Suggestion: confirm a generic path exists; otherwise add an explicit dtype-equality check for these opcodes.

🟡 count not shown in diff — WINOGRAD_CONV branch uses count but it's not defined in the visible context. If count isn't guaranteed to be len(inst.operands) at this point in the function, the count == 3 bias-selection logic is fragile.
Suggestion: use len(inst.operands) directly or add a local alias near the top of the block.

🟡 SPMM_CSR missing shape sanity — No check that values.shape == col.shape (or that rowptr.shape[0] matches expected rows). If shape validation is not done upstream, malformed CSR passes here silently.
Suggestion: at minimum assert tuple(values.shape) == tuple(col.shape).

🟡 Boolean guard idiom — isinstance(cout, int) and not isinstance(cout, bool) works (bool subclasses int), but type(cout) is int is shorter and unambiguous about exact typing.

💭 Error messages — "SPMM_CSR values and B must share a dtype" would be more actionable if it included actual dtypes (e.g., values=%s vs b=%s), matching the diagnostic style typical in the rest of the verifier.

💭 Comment — The header comment explaining the FWHT/WINOGRAD/SPMM_CSR split is good; consider also noting that SPMM_CSR shape rules are NOT enforced here (if that's the intended scope).


📁 scratchv/backend/asm_emit.py

🔴 Missing guard parity with other load/store paths — This block special-cases FLW/FSW for memory-address formatting, but the diff doesn't show how LW/SW/LBU/SB are formatted elsewhere. If those go through the default path without _memory_address, you'll get inconsistent output (e.g. lw x5, x6 vs flw f0, (x6)). Either extract a shared _format_load_store() helper covering all load/store ops, or confirm the others are handled and remove the duplication.

🟡 DRY: hardcoded op strings — OP_NAMES was just extended with FLW/FSW, but the f-strings below hardcode "flw"/"fsw". If someone renames an entry in OP_NAMES, these silently drift. Use:

name = OP_NAMES[instr.op]
return f"  {name} {_fmt_op(instr.dst)}, {address}{_comment_suffix(instr)}"

Then collapse the two if blocks into one if instr.op in (MachineOp.FLW, MachineOp.FSW) and instr.dst and instr.src1:.

🟡 Semantic naming concern for FSW — instr.dst is being read as the FP source register (fs), since FSW's semantics are fsw fs, offset(rs1). That's fine internally, but it means the field name lies about its meaning for this op. Worth a comment at the branch, or a helper like _fp_source(instr) if FSW/FSUBS etc. share this quirk elsewhere.

💭 Consistency check for new mappings — Verify SLLI/SRLI/RET/FMV_W_X are actually emitted by the codegen pass somewhere; adding them to OP_NAMES is harmless if unused, but it also means unhandled-op tests should expect these to appear in the "known ops" set. If there's an __init__.pyi or test enumerating valid machine ops, update in the same commit.

💭 Ordering — Confirm this block runs before any generic "single-register op" formatter that might otherwise catch FLW/FSW and emit flw f0, x6 (no parens/offset). Placing it after the CALL/J block at line 176+ would be safer if those share code paths; currently positioned after other special-cases which looks right, just double-check.


📁 scratchv/backend/machine_semantics.py

Code Review

🟡 RET: implicit x0 destination not documented — The semantics omit defs (correct, since jalr x0 writes to the zero register which is ignored), but the comment only mentions ra. Consider noting why defs is intentionally empty:

# ret -> jalr x0, 0(ra); x0 is the zero register, so no defs.

This helps future maintainers not add defs=(0,) by mistake.


No other issues. The diff is correct and consistent:

  • SLLI/SRLI — correctly mirror SRAI semantics (same register/imm layout, just different shift direction handled elsewhere).
  • FMV_W_X — reuses _DEF_USE exactly like the existing FMV_S_X. Clean.
  • RET is_pseudo — correctly omitted; comment confirms the assembler emits it verbatim, so no lowering pass needed. The implicit_uses pattern matches JAL's treatment of ra.
  • n_phys=0 (default) for RET — correct, since ra is tracked via implicit_uses (hardwired, non-allocatable).

📁 scratchv/backend/machine_types.py

🟡 Completeness risk — These new enum values (SLLI, SRLI, RET, FMV_W_X) presumably need corresponding handling in:

  • Instruction encoding/decoding logic
  • Assembler/disassembler tables
  • VM execution / step-through logic

If any of those use exhaustive match/switch on MachineOp, this diff alone will cause a compile error or silent fallthrough. Confirm the rest of the codebase handles these.

🟡 Semantic clarity for FMV_W_X vs FMV_S_X — In standard RISC-V, fmv.w.x and fmv.s.x are functionally identical (both move an integer to a 32-bit float register). Are they distinct in this VM's encoding? If they're true aliases, consider adding a comment noting they map to the same encoding/behavior, or unify them to avoid divergence in execution logic.

💭 Nit — RET comment says jalr x0, 0(ra) which is correct. Minor consistency: existing pseudo-instructions (MAX) document operands as rd, rs1, rs2. You could match that style: # pseudo: jalr zero, 0(ra).


📁 scratchv/backend/platform_emit.py

🔴 Bug: Conv output store indexing produces NCHW, not NHWC — _lower_conv store block: a3 = ((b*Cout + oc)*H + oh)*W + ow. For NHWC (as documented, and consistent with the input indexing a few lines above which is ((b*H + ih)*W + iw)*Cin + c), the correct form is ((b*H + oh)*W + ow)*Cout + oc. As written, the kernel writes into a plane-per-channel layout, so any NHWC consumer reads garbage. Fix the store indexing or fix the docstring; both cannot be true.

🔴 Bug: OpCode.WINOGRAD_CONV → _lower_conv — Winograd and direct convolution are different algorithms. Silently emitting the direct-conv kernel for a Winograd op will produce wrong numerics (or at best the wrong transform). If this is intentional as a placeholder, raise PlatformEmitError instead; if it's a TODO, gate it behind a comment that explicitly calls it out.

🟡 Callee-saved register ABI violated — All three lowerings clobber s0–s11 freely (_lower_spmm uses all of them, _lower_conv uses s0–s6, s7–s11, _lower_fwht uses s2–s8) with no save/restore. Under the standard RISC-V calling convention s-regs are callee-saved, so if cnn_entry is called from C the caller's s-regs get corrupted. Either emit prologue/epilogue push/pops for the used s-regs or update the module docstring to state that s0–s11 are not preserved (and make the callers account for it).

🟡 Scratch regs reuse argument regs — _lower_conv uses a3/a4 as scratch in the innermost loop. Legal (caller-saved) but misleading given the ABI comment only names a0/a1/a2; consider using t* to make intent clear and to leave a* slots available for the caller.

🟡 render_platform string-sniffing is brittle — if ".size cnn_entry" not in text keys off a literal substring, which would silently misfire if AsmEmitter ever emits it inside a comment or with different spacing. Prefer pushing .size emission into AsmEmitter (or a flag) rather than post-processing the rendered string.

🟡 generate_platform_asm cross-function scan — Collects matching opcodes across all functions and blocks, then emits one kernel keyed off the first match. A program with one CONV in one function and a stray matching opcode elsewhere still passes. Consider scoping to a single entry function/block to fail loudly on malformed graphs.

💭 Layout comment nit — Docstring for _lower_spmm says header is [M,K,N,nnz] but the code never loads K (offset 4) — fine, just don't advertise a field you don't consume, or note it's reserved.

💭 Loop-carried s-reg pressure — _lower_spmm holds lo/hi/p/j/acc etc. in s-regs across the whole triple-nested loop; consider moving loop counters into t-regs to free s-regs for future expansion and to align with the "register allocation is not involved" claim (which is currently overstated — you are allocating, by hand).


📁 scratchv/backend/tensor_c_codegen.py

🔴 Bug: Winograd scratch is static local, not in arena — _emit line ~517:

self.body.append("static float sv_wg_d[16], sv_wg_v[16], sv_wg_acc[16];")
self.extra_workspace += 3 * 16 * 4

Two problems:

  1. static means the scratch lives in BSS, not in sv_arena — so extra_workspace is being added to a "workspace budget" the arena doesn't actually occupy. Either move these into the arena (via the normal allocator, so high grows) or stop counting them against max_workspace_bytes.
  2. static locals make scratchv_run non-reentrant. Two threads (or two nested/call-through invocations) will race on sv_wg_d/v/acc. Every other op in this file is allocation-free or arena-backed for a reason — Winograd breaks that invariant.

🔴 Bug: _sum_terms negative-coefficient formatting — line ~42:

expression += (" - " + text) if expression else ("-" + text)

Correct for the first term ("-v"), but "-1.0 * v" would be safer for mixed-coefficient cases. As written it's fine for the {0,±1} contract in the docstring, but the docstring isn't enforced anywhere — a future caller passing coef=2 silently emits +2*text but coef=-2 emits -2*text with no space, which is valid C but unreadable. Suggest: emit sign + " " + (coef>1 ? f"{coef}*" : "") + text.

🔴 Bug: Winograd output boundary uses output height, but tiles iterate ceil(H/2) — line ~555:

if ({oh}<{height}ULL && {ow}<{width}ULL) {dst}[...] = {expression};

Guard is correct, but the write index is ((n*cout+oc)*H+oh)*W+ow — if H or W is odd, the last tile's second output is guarded correctly, good. However, tiles_h = (height + 1) // 2 for H=1 gives tiles_h=1, so we iterate ti=0 and try to write row 0 (valid) and row 1 (guarded). But we still load input at gi=-1,0,1,2 — all bounds-checked, OK. Just flagging: for degenerate H=1 or W=1 inputs the algorithm wastes a tile but stays correct.

🟡 Missing bias shape validation — Winograd accepts 3rd operand as bias but never checks bias.shape[0] == cout. A shape mismatch here causes silent OOB read of sv_wg_acc output — well, of bias[oc] where oc < cout, so it's an OOB read if bias is shorter than cout. Add:

if len(xs) == 3 and xs[2].shape != (cout,):
    raise TensorCCodegenError("Winograd bias must be [cout]")

🟡 Winograd transform correctness unverifiable from this file — the comment says values match _WINO_F23_BT/_WINO_F23_AT in another file. Please either move the canonical copy here and let the other file import it, or add a unit test that cross-checks the two. Two hand-written 4×4 matrices with a magic-comment link is the exact pattern that rots.

🟡 SPMM_CSR validation split across _infer and runtime — _infer checks values.shape == col.shape but not values.shape[0] == rowptr[-1] (i.e., nnz matches total nonzeros). This is caught at runtime via e > nnz (return 4), but a mismatch surfaces only when the program is called, not at codegen time. Consider validating in _infer when rowptr is a constant, otherwise leave as-is.

🟡 FWHT inverse precision — dst[i] /= float(n): fine for FP32, but for n=2^20 the divisor is exact; for larger powers of two it starts to round. Not a bug, just noting the contract.

🟡 extra_workspace grows monotonically inside _emit but _winograd_declared is only set on first Winograd op — correct, but the check workspace_total > max_workspace_bytes runs after body construction. If it raises, all the _emit work was wasted. Minor: consider validating the total earlier (e.g., count declared Winograd ops in a prepass), or accept the wasted work as negligible.

💭 Style: SPMM_CSR inner loop line ~520 has no spaces in for(size_t j=0;j<...;++j) — inconsistent with the rest of the file (for (size_t j=0; j<...; ++j)).

💭 Test coverage: no tests in this diff. FWHT inverse-vs-forward round-trip, SPMM_CSR against a dense reference, and Winograd against a naive 3×3 conv would all catch regressions in the transforms and the emitted C.


📁 scratchv/compiler.py

🟡 **Suggestion: Undocumented precedence between `platform_asm` and `backend`** — Lines 504-507: The new flag silently wins over any `backend` setting. If a user sets `backend = "ir"` and `platform_asm = True`, there's no warning or validation. Consider either mutual-exclusivity validation in `CompilerConfig.__post_init__` or a comment on the field documenting that `platform_asm` takes priority.

🟡 **Suggestion: No error handling or output validation** — Line 507: `generate_platform_asm(program)` could raise or return `None`/empty string, which would propagate silently through the compile pipeline. At minimum consider a guard:
```python
result = generate_platform_asm(program)
if not result:
    raise CompilerError("platform_asm emit produced no output")
return result

💭 Nit: Stale docstring — The method docstring says "instruction selection + regalloc + emit", but the platform_asm path skips all three. A brief note like "(or delegates to platform emit)" would keep it accurate.

💭 Nit: No path-logging — Compilers benefit from knowing which codegen path ran. Consider a logger.debug("using platform_asm emit") for debuggability during development.


---

## 📁 `scratchv/frontend/onnx_parser.py`



🔴 **Bug: `self.initializers` likely undefined / wrong lookup** — `_prepare_winograd_kernels` uses `self.initializers.get(w_name)` and `if u_name not in self.initializers`, but `__init__` initializes `_value_map`, `_constant_cache`, `_producers` — not `initializers`. If `initializers` isn't populated elsewhere, `weight` falls through to re-load from protobuf (duplicate work already done in the initializer loop) and the dedup `u_name not in self.initializers` never triggers, recomputing U for repeated weights. Verify the attribute name — you likely want `self._value_map` or a dedicated set.

🔴 **Bug: Winograd ignores stride/padding/dilution/group attributes** — `_handle_winogradconv` hardcodes `result.shape = (batch, cout, height, width)` and passes only `cout/cin/bias` to the builder. Any ONNX WinogradConv with non-unit stride, asymmetric padding, dilations, or groups will silently produce wrong shapes/results. Either raise `ONNXParseError` for unsupported attrs or forward them.

🔴 **Bug: `inputs[0..3]` unguarded in `_handle_spmmcsr`** — if a producer is missing or an input name was empty, indexing raises `IndexError` instead of a parse error. Validate `len(inputs) == 4` with a descriptive message.

🟡 **`_handle_fwht` breaks on dynamic shapes** — `math.prod(x.shape)` will raise `TypeError` if any dim is `None`. Add an explicit check before `prod`.

🟡 **In-function imports duplicate module-level ones** — `onnx` is already imported at module scope (used by `onnx.load`, `onnx.numpy_helper`). The local `import onnx` is redundant; `numpy` likely is too. Remove or hoist.

🟡 **Per-channel Python loop for U folding is slow** — nested `for oc`/`for ic` calling `.tolist()` per 3×3 slice is fine for small kernels but scales poorly with large `cout*cin`. Consider vectorizing `winograd_f23_kernel_transform` if the standalone version accepts arrays.

🟡 **Folded U stored as constant but handler uses `_value_map`** — `_bind_constant` produces a Value; the handler then checks `len(u.shape) != 4`. Ensure `_bind_constant` actually registers under the exact `u_name` string (not wrapped/mangled), otherwise `_value_map[u_name]` KeyError.

💭 **`dict[str, str]` annotation requires Python ≥3.9** — fine if that's the floor, but note it in the module docstring if not.

💭 **Duplicate domain check** — domains validated in `parse()` *and* `_translate_node`. The second check is dead code after the first; either trust the pre-validation and drop it, or keep only the node-level one.

---

## 📁 `scratchv/ir/builder.py`



🔴 **Bug: `winograd_conv` — `cout=0, cin=0` 哨兵值有歧义** — 调用方若忘记传 `cout/cin`,下游拿到的是 `0` 还是"未知"?无法区分。建议用 `Optional[int] = None`,emit 前校验非 None,或改为从 `x.shape`/`u.shape` 推导。

🟡 **`fwht` — `direction` 是无约束字符串** — 拼错(如 `"foward"`)会静默生成无效 IR。建议用 `Literal["forward", "inverse"]` 或在 emit 前 assert 合法值。

🟡 **三个新方法均无输入校验** — `spmm_csr` 未检查 `values.shape[0] == rowptr.shape[0] - 1`,`winograd_conv` 未检查 `bias` 的 rank/shape 与 `cout` 一致。对比同文件的 `gemm`,缺少 builder 层的防御会导致错误延迟到后端执行才暴露。建议至少加 `assert` 做基本 shape 一致性检查。

🟡 **`winograd_conv` — 可变 operand 列表构建可读性差** — `[x, u] + ([bias] if bias is not None else [])` 在多处 builder 里模式重复。考虑提取一个 `_emit_with_optional` helper,或统一偏置参数到 kwargs 中处理。

---

## 📁 `scratchv/ir/types.py`



🟡 **`FWHT` / `SPMM_CSR` in `is_nn()` — Semantic mismatch** — Lines 82-84. These are linear-algebra ops (Hadamard transform, sparse matrix multiply), not neural-network primitives. If `is_nn()` is used to gate NN-specific optimizations or dispatch, putting them here will incorrectly apply NN-only passes to non-NN IR nodes. `WINOGRAD_CONV` is debatable (conv optimization), but the other two are clearly Linalg.

Suggestion: Introduce an `is_linalg()` (or `is_transform()`) category, or at minimum don't conflate these into `is_nn()`.

🟡 **No corresponding IR node class or lowering path visible** — Adding opcodes without matching `Value`/`Node` subclasses or compiler passes means these are dead enum entries until wired up. Is there a companion PR with the lowering, or are these placeholders? If placeholders, consider a brief `TODO` referencing the issue.

💭 **Comment style** — "Standalone-verified" is unclear to future readers. Something like `# Verified in standalone kernels, promoted to IR` would be self-documenting.

---


---
### ⚠️ 未审查的文件
  - scratchv/main.py
  - scratchv/verification/ir_numpy_ops.py
  - tests/_op_ref.py
  - tests/test_frontend_custom_ops.py
  - tests/test_ir_fixedpoint_ops.py
  - tests/test_ir_verifier.py
  - tests/test_platform_emit.py
  - tests/test_tensor_c_codegen.py

@jizhenjun jizhenjun closed this Oct 6, 2026
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