From d0a7066d1adf97a8b58da605ab3aa211e1b21166 Mon Sep 17 00:00:00 2001 From: Seven Gao <799889633@qq.com> Date: Sat, 3 Oct 2026 23:37:40 +0800 Subject: [PATCH 1/5] feat(standalone): support INT32 initializers in MemoryPlan (F2) CSR sparse ops (SpmmCsr) need col_indices/row_ptr laid out verbatim as one 32-bit word per element, never through the Q16.16 float conversion. - add ONNX_INT32 constant - parse int32_data (field 5) and raw_data, restoring negative values - shape fallback for INT32 tensors - layout_weights INT32 branch packs raw int32 words Tests: tests/test_standalone_int32_weights.py (4 passed). Full-suite failed set unchanged vs baseline (226 pre-existing llvm-mca/env). --- .../standalone/onnx_to_riscv_standalone.py | 31 +++++ tests/test_standalone_int32_weights.py | 115 ++++++++++++++++++ 2 files changed, 146 insertions(+) create mode 100644 tests/test_standalone_int32_weights.py diff --git a/scratchv/standalone/onnx_to_riscv_standalone.py b/scratchv/standalone/onnx_to_riscv_standalone.py index 960b12ee..bd9188b0 100644 --- a/scratchv/standalone/onnx_to_riscv_standalone.py +++ b/scratchv/standalone/onnx_to_riscv_standalone.py @@ -154,6 +154,7 @@ def _safe_decode(b: object) -> str: # ONNX data type constants ONNX_FLOAT = 1 +ONNX_INT32 = 6 ONNX_INT64 = 7 @@ -394,6 +395,25 @@ def _parse_tensor(self, data: bytes) -> TensorInfo: arr.append(struct.pack("= 0x80000000: + u -= 0x100000000 + out.append(struct.pack(" bytes: # Convert to Q16.16 if tensor.data_type == ONNX_FLOAT: q16_vals = float32_to_q16(tensor.data) + elif tensor.data_type == ONNX_INT32: + # INT32 tensors (e.g. CSR col_indices / row_ptr) are stored + # verbatim as one 32-bit word per element — do NOT run them + # through the Q16.16 conversion. + n = len(tensor.data) // 4 + int_vals = struct.unpack(f"<{n}i", tensor.data) + q16_vals = [v & 0xFFFFFFFF for v in int_vals] elif tensor.data_type == ONNX_INT64: # INT64 tensors (like Reshape target shapes) — keep as-is q16_vals = [] diff --git a/tests/test_standalone_int32_weights.py b/tests/test_standalone_int32_weights.py new file mode 100644 index 00000000..514fe417 --- /dev/null +++ b/tests/test_standalone_int32_weights.py @@ -0,0 +1,115 @@ +"""F2: INT32 initializer support in the standalone ONNX parser / memory plan. + +CSR sparse operators (SpmmCsr) need integer index arrays (`col_indices`, +`row_ptr`) laid out verbatim as one 32-bit word per element — they must NOT be +run through the Q16.16 float conversion used for weights. +""" + +import struct + +import numpy as np +import onnx +from onnx import TensorProto, helper, numpy_helper +import pytest + +from scratchv.standalone.onnx_to_riscv_standalone import ( + MemoryPlan, + ONNXModel, + ONNX_INT32, +) + + +def _relu_model(initializers): + """Minimal valid graph carrying the given initializers.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [4]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [4]) + node = helper.make_node("Relu", ["x"], ["y"]) + graph = helper.make_graph([node], "int32_test", [x], [y], initializers) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + + +def _roundtrip(model, tmp_path): + path = tmp_path / "model.onnx" + onnx.save(model, str(path)) + return ONNXModel.from_file(str(path)) + + +def test_parser_reads_int32_raw_data(tmp_path): + """numpy_helper stores raw_data; parser must report dtype INT32 and shape.""" + values = np.array([10, -20, 30, -40], dtype=np.int32) + init = numpy_helper.from_array(values, name="idx") + model = _relu_model([init]) + + parsed = _roundtrip(model, tmp_path) + tensor = parsed.initializers["idx"] + + assert tensor.data_type == ONNX_INT32 + assert tensor.shape == (4,) + assert struct.unpack("<4i", tensor.data) == (10, -20, 30, -40) + + +def test_parser_reads_int32_typed_field(tmp_path): + """helper.make_tensor(raw=False) stores int32_data (field 5), incl. negatives.""" + init = helper.make_tensor( + "idx", TensorProto.INT32, [4], [100, -200, 300, -400] + ) + model = _relu_model([init]) + + parsed = _roundtrip(model, tmp_path) + tensor = parsed.initializers["idx"] + + assert tensor.data_type == ONNX_INT32 + assert tensor.shape == (4,) + assert struct.unpack("<4i", tensor.data) == (100, -200, 300, -400) + + +def test_layout_weights_int32_is_verbatim_one_word_per_element(): + """INT32 tensors must be laid out as raw int32 words, not Q16.16.""" + model = ONNXModel() + from scratchv.standalone.onnx_to_riscv_standalone import TensorInfo + + t = TensorInfo() + t.name = "row_ptr" + t.shape = (3,) + t.data_type = ONNX_INT32 + t.data = struct.pack("<3i", 0, 2, 3) + model.initializers["row_ptr"] = t + + mem = MemoryPlan() + blob = mem.layout_weights(model.initializers) + + assert mem.get_weight_offset("row_ptr") == 0 + assert blob == struct.pack("<3i", 0, 2, 3) + # A Q16.16 conversion would inflate 2 -> 131072; ensure it did not happen. + assert struct.unpack("<3i", blob) == (0, 2, 3) + + +def test_layout_weights_mixes_float_and_int32(): + """Float weights still become Q16.16; INT32 arrays stay verbatim.""" + model = ONNXModel() + from scratchv.standalone.onnx_to_riscv_standalone import TensorInfo + + w = TensorInfo() + w.name = "w" + w.shape = (2,) + w.data_type = 1 # ONNX_FLOAT + w.data = struct.pack("<2f", 0.5, -0.25) + + idx = TensorInfo() + idx.name = "idx" + idx.shape = (2,) + idx.data_type = ONNX_INT32 + idx.data = struct.pack("<2i", 7, -7) + model.initializers["w"] = w + model.initializers["idx"] = idx + + mem = MemoryPlan() + blob = mem.layout_weights(model.initializers) + + w_off = mem.get_weight_offset("w") + i_off = mem.get_weight_offset("idx") + assert w_off == 0 + assert i_off == 8 # 2 float words + + assert struct.unpack("<2i", blob[w_off : w_off + 8]) == (0x8000, -0x4000) + assert struct.unpack("<2i", blob[i_off : i_off + 8]) == (7, -7) From f6f697266dfc4b5a58ecb0df90de5eeb58f551dd Mon Sep 17 00:00:00 2001 From: Seven Gao <799889633@qq.com> Date: Sat, 3 Oct 2026 23:38:41 +0800 Subject: [PATCH 2/5] fix(standalone): resolve Spike tools via env/PATH, drop hardcoded path (F1) spike_sim.py hardcoded /home/kinsomwang/... which is unusable elsewhere. Now: SCRATCHV_SPIKE / SCRATCHV_SPIKE_DASM / SCRATCHV_SPIKE_LOG_PARSER override, else PATH lookup, else bare command name. Tests: tests/test_standalone_spike_paths.py (4 passed). --- scratchv/standalone/spike_sim.py | 16 ++++++++++---- tests/test_standalone_spike_paths.py | 33 ++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 4 deletions(-) create mode 100644 tests/test_standalone_spike_paths.py diff --git a/scratchv/standalone/spike_sim.py b/scratchv/standalone/spike_sim.py index 12e972f5..f07f6b6a 100644 --- a/scratchv/standalone/spike_sim.py +++ b/scratchv/standalone/spike_sim.py @@ -12,7 +12,7 @@ - Instruction trace samples - Spike execution time -Spike binary: /home/kinsomwang/workspace/coralnpu-spike-rv32/bin/spike +Spike binary: $SCRATCHV_SPIKE, else `spike` resolved on PATH (no hardcoded path) ScratchV binary: output.bin Usage: @@ -25,6 +25,7 @@ import argparse import os +import shutil import struct import subprocess import sys @@ -34,9 +35,16 @@ from dataclasses import dataclass, field # ── Paths ────────────────────────────────────────────────────────────────── -SPIKE = "/home/kinsomwang/workspace/coralnpu-spike-rv32/bin/spike" -SPIKE_DASM = "/home/kinsomwang/workspace/coralnpu-spike-rv32/bin/spike-dasm" -SPIKE_LOG_PARSER = "/home/kinsomwang/workspace/coralnpu-spike-rv32/bin/spike-log-parser" +# Resolve Spike tools portably: explicit env override first, then PATH. +# No hardcoded machine-specific paths. +def _resolve_tool(name: str, env_var: str) -> str: + """Resolve a Spike tool path: env override → PATH lookup → bare name.""" + return os.environ.get(env_var) or shutil.which(name) or name + + +SPIKE = _resolve_tool("spike", "SCRATCHV_SPIKE") +SPIKE_DASM = _resolve_tool("spike-dasm", "SCRATCHV_SPIKE_DASM") +SPIKE_LOG_PARSER = _resolve_tool("spike-log-parser", "SCRATCHV_SPIKE_LOG_PARSER") # ── Constants ────────────────────────────────────────────────────────────── ELF_BASE = 0x80000000 # RISC-V DRAM base (Spike default) diff --git a/tests/test_standalone_spike_paths.py b/tests/test_standalone_spike_paths.py new file mode 100644 index 00000000..4bf7df3f --- /dev/null +++ b/tests/test_standalone_spike_paths.py @@ -0,0 +1,33 @@ +"""F1: Spike tool path resolution must be portable (no hardcoded paths).""" + +from scratchv.standalone.spike_sim import _resolve_tool + + +def test_env_override_wins(monkeypatch): + monkeypatch.setenv("SCRATCHV_SPIKE", "/opt/riscv/bin/spike") + assert _resolve_tool("spike", "SCRATCHV_SPIKE") == "/opt/riscv/bin/spike" + + +def test_falls_back_to_path_lookup(monkeypatch): + monkeypatch.delenv("SCRATCHV_SPIKE", raising=False) + monkeypatch.setattr( + "scratchv.standalone.spike_sim.shutil.which", + lambda name: f"/usr/bin/{name}", + ) + assert _resolve_tool("spike", "SCRATCHV_SPIKE") == "/usr/bin/spike" + + +def test_falls_back_to_bare_name(monkeypatch): + monkeypatch.delenv("SCRATCHV_SPIKE", raising=False) + monkeypatch.setattr("scratchv.standalone.spike_sim.shutil.which", lambda name: None) + assert _resolve_tool("spike", "SCRATCHV_SPIKE") == "spike" + + +def test_no_hardcoded_foreign_path_remains(): + import inspect + + from scratchv.standalone import spike_sim + + source = inspect.getsource(spike_sim) + assert "kinsomwang" not in source + assert "/home/" not in source From 2e3baa7738b8a3a1594366df34c7d15e61f95e01 Mon Sep 17 00:00:00 2001 From: Seven Gao <799889633@qq.com> Date: Mon, 5 Oct 2026 00:41:52 +0800 Subject: [PATCH 3/5] build: require Python >=3.11 and add toolchain/runtime requirements Pin the interpreter to 3.11+ and add optional requirements for the RISC-V toolchain and simulation runtime. --- pyproject.toml | 5 +---- requirements/ci.txt | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 4 deletions(-) create mode 100644 requirements/ci.txt diff --git a/pyproject.toml b/pyproject.toml index a588db51..7a792698 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ description = "A compiler from ONNX models to RISC-V assembly and LLVM IR" readme = "README.md" license = {text = "MIT"} keywords = ["compiler", "risc-v", "onnx", "llvm", "machine-learning"] -requires-python = ">=3.8" +requires-python = ">=3.11" dependencies = [ "onnx>=1.14", "numpy>=1.24", @@ -21,9 +21,6 @@ classifiers = [ "Intended Audience :: Education", "License :: OSI Approved :: MIT License", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Topic :: Software Development :: Compilers", diff --git a/requirements/ci.txt b/requirements/ci.txt new file mode 100644 index 00000000..e99f7b59 --- /dev/null +++ b/requirements/ci.txt @@ -0,0 +1,15 @@ +# CI / verification environment +# Target Python: 3.11 - 3.12 (the bundled Python 3.8 .venv is broken; do not use) +# CI and local runs use this file. Install: +# python3.11 -m venv .venv311 && source .venv311/bin/activate +# pip install -r requirements/ci.txt +# pip install -e ".[all]" # optional: tinyfive / llvmlite / onnxruntime + +numpy>=1.26,<3 +onnx>=1.14,<2 +protobuf>=4.21 +pytest>=7 + +# Optional (accelerators / independent verification): uncomment to use +# onnxruntime +# llvmlite From bbd0cb39183daa75b577a54fbb98788ad68abe88 Mon Sep 17 00:00:00 2001 From: Seven Gao <799889633@qq.com> Date: Mon, 5 Oct 2026 00:41:52 +0800 Subject: [PATCH 4/5] feat(standalone): add transform, convolution and sparse-matmul operators - Conv/Gemm: accept models without a bias input - Fwht: scalar fast Walsh-Hadamard transform with forward/inverse directions, boundary handling and impulse-property regression - SpmmCsr: CSR sparse matrix-multiply operator with Q16.16 reference test - WinogradConv: F(2,3) kernel transform and convolution code generation - Conv: infer kernel size from the weight shape (fixes non-3x3 kernels) - Conv: support batch (N>1) via compile-time unrolling --- .../standalone/onnx_to_riscv_standalone.py | 1301 ++++++++++++----- tests/test_standalone_conv_batch.py | 106 ++ tests/test_standalone_fwht.py | 88 ++ tests/test_standalone_optional_bias.py | 95 ++ tests/test_standalone_spmm.py | 115 ++ tests/test_standalone_winograd.py | 109 ++ tests/test_winograd_kernel_transform.py | 63 + 7 files changed, 1543 insertions(+), 334 deletions(-) create mode 100644 tests/test_standalone_conv_batch.py create mode 100644 tests/test_standalone_fwht.py create mode 100644 tests/test_standalone_optional_bias.py create mode 100644 tests/test_standalone_spmm.py create mode 100644 tests/test_standalone_winograd.py create mode 100644 tests/test_winograd_kernel_transform.py diff --git a/scratchv/standalone/onnx_to_riscv_standalone.py b/scratchv/standalone/onnx_to_riscv_standalone.py index bd9188b0..79143eb4 100644 --- a/scratchv/standalone/onnx_to_riscv_standalone.py +++ b/scratchv/standalone/onnx_to_riscv_standalone.py @@ -553,6 +553,13 @@ def _infer_node_shape(self, node: NodeInfo) -> None: self._set_output_shape(node, input_shapes[0]) elif op == "Reshape": self._infer_reshape_shape(node, input_shapes) + elif op == "Fwht": + if input_shapes and input_shapes[0]: + self._set_output_shape(node, input_shapes[0]) + elif op == "SpmmCsr": + self._infer_spmm_shape(node, input_shapes) + elif op == "WinogradConv": + self._infer_conv_shape(node, input_shapes) def _infer_conv_shape( self, node: NodeInfo, shapes: list[tuple[int, ...]] @@ -565,11 +572,18 @@ def _infer_conv_shape( out_channels = w_shape[0] if len(w_shape) >= 1 else 1 - kernel = node.attrs.get("kernel_shape", [3, 3]) - if isinstance(kernel, list) and len(kernel) >= 2: - kh, kw = int(kernel[0]), int(kernel[1]) + # Prefer the actual weight kernel dims over the kernel_shape attr / + # default 3, otherwise non-3x3 kernels (e.g. 5x5) are mis-sized. + if len(w_shape) >= 4 and w_shape[2] > 0 and w_shape[3] > 0: + kh, kw = int(w_shape[2]), int(w_shape[3]) + elif len(w_shape) >= 3 and w_shape[2] > 0: + kh = kw = int(w_shape[2]) else: - kh = kw = 3 + kernel = node.attrs.get("kernel_shape", [3, 3]) + if isinstance(kernel, list) and len(kernel) >= 2: + kh, kw = int(kernel[0]), int(kernel[1]) + else: + kh = kw = 3 stride = node.attrs.get("strides", [1, 1]) if isinstance(stride, list) and len(stride) >= 2: @@ -687,6 +701,22 @@ def _infer_reshape_shape( if len(shapes) >= 2 and shapes[1]: self._set_output_shape(node, shapes[1]) + def _infer_spmm_shape( + self, node: NodeInfo, shapes: list[tuple[int, ...]] + ) -> None: + """Infer SpmmCsr output (M, N) from row_ptr length and dense B shape. + + Inputs: (values, col_indices, row_ptr, B); M = len(row_ptr) - 1, + N = B.shape[1] (or 1 for a 1-D B / SpMV). + """ + if len(shapes) < 4 or not shapes[3]: + return + row_shape = shapes[2] + m = (row_shape[0] - 1) if row_shape else 0 + b_shape = shapes[3] + n = b_shape[1] if len(b_shape) >= 2 else 1 + self._set_output_shape(node, (m, n)) + def _set_output_shape( self, node: NodeInfo, shape: tuple[int, ...] ) -> None: @@ -1091,6 +1121,64 @@ def int64_to_shape(raw_bytes: bytes) -> tuple[int, ...]: return tuple(int(v) for v in vals) +# ═══════════════════════════════════════════════════════════════════════════ +# Winograd F(2,3) transforms (compile-time helpers, pure Python) +# ═══════════════════════════════════════════════════════════════════════════ +# 2D minimal filtering: Y = A^T [ (G g G^T) ⊙ (B^T d B) ] A +# F(2,3): input tile 4x4, kernel 3x3, output 2x2. +# G contains 1/2 entries; B^T and A^T are integer {0, ±1}. + +_WINO_F23_BT = ((1, 0, -1, 0), (0, 1, 1, 0), (0, -1, 1, 0), (0, 1, 0, -1)) +_WINO_F23_G = ((1.0, 0.0, 0.0), (0.5, 0.5, 0.5), + (0.5, -0.5, 0.5), (0.0, 0.0, 1.0)) +_WINO_F23_AT = ((1.0, 1.0, 1.0, 0.0), (0.0, 1.0, -1.0, -1.0)) + + +def winograd_f23_kernel_transform(g): + """Return U = G g G^T (4x4) for a 3x3 kernel g (nested sequences).""" + P = [[sum(_WINO_F23_G[i][k] * g[k][j] for k in range(3)) for j in range(3)] + for i in range(4)] + return [[sum(P[i][k] * _WINO_F23_G[j][k] for k in range(3)) for j in range(4)] + for i in range(4)] + + +def _prepare_winograd_kernels(model) -> None: + """Fold WinogradConv 3x3 kernels into the Winograd domain at compile time. + + For each WinogradConv node with FLOAT weight [Cout, Cin, 3, 3], compute + U = [Cout, Cin, 4, 4] and register it as a new FLOAT initializer. The node + stores the tensor name in ``attrs["_winograd_u"]`` for the code generator. + """ + for node in model.nodes: + if node.op_type != "WinogradConv" or len(node.inputs) < 2: + continue + w = model.initializers.get(node.inputs[1]) + if w is None or not w.data or len(w.shape) != 4: + continue + cout, cin, kh, kw = w.shape + if (kh, kw) != (3, 3): + raise ValueError("WinogradConv currently supports 3x3 kernels only") + + floats = struct.unpack(f"<{cout * cin * 9}f", w.data) + u_bytes = bytearray() + for oc in range(cout): + for ic in range(cin): + base = (oc * cin + ic) * 9 + g = [[floats[base + r * 3 + c] for c in range(3)] for r in range(3)] + u = winograd_f23_kernel_transform(g) + for r in range(4): + for c in range(4): + u_bytes += struct.pack(" None: """ x_name = node.inputs[0] w_name = node.inputs[1] - b_name = node.inputs[2] + # Conv bias is optional; a missing third input means "no bias" (zero). + b_name = node.inputs[2] if len(node.inputs) > 2 and node.inputs[2] else None out_name = node.outputs[0] x_shape = self.model.get_shape(x_name) @@ -1761,7 +1850,8 @@ def _gen_conv(self, node: NodeInfo) -> None: # Load base addresses self._get_workspace_addr(x_name, self.S2) # s2 = input base self._get_weight_addr(w_name, self.S3) # s3 = weight base - self._get_weight_addr(b_name, self.S4) # s4 = bias base + if b_name is not None: + self._get_weight_addr(b_name, self.S4) # s4 = bias base self._get_workspace_addr(out_name, self.S5) # s5 = output base # ── Preload loop-invariant constants ────────────────────────── @@ -1813,336 +1903,366 @@ def _gen_conv(self, node: NodeInfo) -> None: K_K = K * K W_minus_K_times_4 = (W - K) * 4 - # ── oc loop ──────────────────────────────────────────────────── - self.emit.emit(rv_addi(OC_REG, _R_ZERO, 0), f"oc=0 (C_out={C_out})") - L("_conv_oc_loop") - - # ── oh loop ──────────────────────────────────────────────────── - self.emit.emit(rv_addi(OH_REG, _R_ZERO, 0), f"oh=0 (H_out={H_out})") - L("_conv_oh_loop") - - # ih_base = oh * stride_h (precompute for this oh row) - self.emit.emit(rv_mul(IH_BASE_REG, OH_REG, STRIDE_H_REG), - "ih_base = oh * stride_h") - - # ── ow loop ──────────────────────────────────────────────────── - self.emit.emit(rv_addi(OW_REG, _R_ZERO, 0), f"ow=0 (W_out={W_out})") - L("_conv_ow_loop") - - # iw_base = ow * stride_w (precompute for this column) - self.emit.emit(rv_mul(IW_BASE_REG, OW_REG, STRIDE_W_REG), - "iw_base = ow * stride_w") - - # Each output element starts a new reduction, including its bias. - self.emit.emit(rv_slli(TMP_REG, OC_REG, 2), "offset = oc*4") - self.emit.emit(rv_add(TMP_REG, S4_BIAS_BASE, TMP_REG), "+ bias_base") - self.emit.emit(rv_lw(ACC_REG, TMP_REG, 0), "acc = bias[oc]") - - if no_pad and K == 3: - # ══════════════════════════════════════════════════════════════ - # Unrolled K=3 no-pad fast path: fully unroll kw+kh loops - # + pointer-increment ic iteration (no address recalculation). - # ~7 instr/MAC vs ~12 before. 2026-06-08 - # ══════════════════════════════════════════════════════════════ - # ic_advance = bytes to advance IN_PTR to next ic's same spatial pos - # K*K loads plus K-1 row skips move (K-1)*W + K elements. - # The next channel's window starts H*W elements from this one. - ic_advance_bytes = (H * W - (K - 1) * W - K) * 4 - IC_ADV_REG = KH_REG # s10 freed: no kh counter needed - ROW_ADV_REG = KW_REG # s11 freed: no kw counter needed - - self.emit.emit_li32(IC_ADV_REG, ic_advance_bytes, - f"ic_adv={ic_advance_bytes}") - self.emit.emit_li32(ROW_ADV_REG, W_minus_K_times_4, - f"row_adv={W_minus_K_times_4}") - - # ic = 0 - self.emit.emit(rv_addi(IC_REG, _R_ZERO, 0), - f"ic=0 (C_in={C_in})") - - # ── Compute starting pointers for ic=0 (once, outside loop) ─ - # IN_PTR = input_base + (ih_base * W + iw_base) * 4 - self.emit.emit(rv_mul(T6_REG, IH_BASE_REG, W_REG), - "inoff = ih_base * W") - self.emit.emit(rv_add(T6_REG, T6_REG, IW_BASE_REG), - "+ iw_base") - self.emit.emit(rv_slli(T6_REG, T6_REG, 2), - "* 4") - self.emit.emit(rv_add(IN_PTR, S2_IN_BASE, T6_REG), - "in_ptr = &input[0, ih_base, iw_base]") - # WT_PTR = weight_base + oc * C_in*K*K * 4 - self.emit.emit_li32(T6_REG, C_in_K_K, f"oc_wt_step={C_in_K_K}") - self.emit.emit(rv_mul(T6_REG, OC_REG, T6_REG), "wt_oc_base") - self.emit.emit(rv_slli(T6_REG, T6_REG, 2), - f"wtoff = oc * {C_in_K_K} * 4") - self.emit.emit(rv_add(WT_PTR, S3_W_BASE, T6_REG), - "wt_ptr = &weight[oc, 0, 0, 0]") - - L("_conv_ic_loop") - - # ── Fully unrolled 3×3 MAC block (9 MACs, no loop control) ─ - # Row 0 (kh=0): MAC 0,1,2 then row advance - for _ in range(3): - self.emit.emit(rv_lw(VAL_REG, IN_PTR, 0), - "load input") - self.emit.emit(rv_lw(T6_REG, WT_PTR, 0), - "load weight") - self.emit.emit(rv_mul(T6_REG, VAL_REG, T6_REG), - "Q16.16 mul") - self.emit.emit(rv_srai(T6_REG, T6_REG, 16), - ">>16") - self.emit.emit(rv_add(ACC_REG, ACC_REG, T6_REG), - "acc += prod") - self.emit.emit(rv_addi(IN_PTR, IN_PTR, 4), - "in_ptr++") - self.emit.emit(rv_addi(WT_PTR, WT_PTR, 4), - "wt_ptr++") - # Row advance: in_ptr += (W-K)*4 - self.emit.emit(rv_add(IN_PTR, IN_PTR, ROW_ADV_REG), - "in_ptr += row_adv to next row") - - # Row 1 (kh=1): MAC 3,4,5 then row advance - for _ in range(3): - self.emit.emit(rv_lw(VAL_REG, IN_PTR, 0), - "load input") - self.emit.emit(rv_lw(T6_REG, WT_PTR, 0), - "load weight") - self.emit.emit(rv_mul(T6_REG, VAL_REG, T6_REG), - "Q16.16 mul") - self.emit.emit(rv_srai(T6_REG, T6_REG, 16), - ">>16") - self.emit.emit(rv_add(ACC_REG, ACC_REG, T6_REG), - "acc += prod") - self.emit.emit(rv_addi(IN_PTR, IN_PTR, 4), - "in_ptr++") - self.emit.emit(rv_addi(WT_PTR, WT_PTR, 4), - "wt_ptr++") - # Row advance - self.emit.emit(rv_add(IN_PTR, IN_PTR, ROW_ADV_REG), - "in_ptr += row_adv to next row") - - # Row 2 (kh=2): MAC 6,7,8 (last row, no row advance after) - for _ in range(3): - self.emit.emit(rv_lw(VAL_REG, IN_PTR, 0), - "load input") - self.emit.emit(rv_lw(T6_REG, WT_PTR, 0), - "load weight") - self.emit.emit(rv_mul(T6_REG, VAL_REG, T6_REG), - "Q16.16 mul") - self.emit.emit(rv_srai(T6_REG, T6_REG, 16), - ">>16") - self.emit.emit(rv_add(ACC_REG, ACC_REG, T6_REG), - "acc += prod") - self.emit.emit(rv_addi(IN_PTR, IN_PTR, 4), - "in_ptr++") - self.emit.emit(rv_addi(WT_PTR, WT_PTR, 4), - "wt_ptr++") - - # ── Advance IN_PTR to next ic's same spatial position ───── - # WT_PTR already at next ic (weights contiguous across ic) - self.emit.emit(rv_add(IN_PTR, IN_PTR, IC_ADV_REG), - "in_ptr += ic_advance") - # ic++ and loop - self.emit.emit(rv_addi(IC_REG, IC_REG, 1), "ic++") - self.emit.emit(rv_slt(COND_REG, IC_REG, C_IN_REG), "ic < C_in?") - self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, - "_conv_ic_loop", "loop ic") - - else: - # ══════════════════════════════════════════════════════════════ - # General loop path (K ≠ 3 or padded) - # ══════════════════════════════════════════════════════════════ - # ── ic loop ──────────────────────────────────────────────── - self.emit.emit(rv_addi(IC_REG, _R_ZERO, 0), f"ic=0 (C_in={C_in})") - L("_conv_ic_loop") - - # Step 1: IN_PTR = input_base + ic * H*W * 4 - self.emit.emit(rv_mul(VAL_REG, IC_REG, HW_REG), - "ic_offset = ic * H*W") - self.emit.emit(rv_slli(VAL_REG, VAL_REG, 2), - "ic_byte_off = ic_offset * 4") - self.emit.emit(rv_add(IN_PTR, S2_IN_BASE, VAL_REG), - "in_ptr = input_base + ic*H*W*4") - - # Step 2: WT_PTR = weight_base + oc * C_in*K*K * 4 - self.emit.emit_li32(VAL_REG, C_in_K_K, f"oc_wt_step={C_in_K_K}") - self.emit.emit(rv_mul(VAL_REG, OC_REG, VAL_REG), "wt_oc_base") - self.emit.emit(rv_slli(VAL_REG, VAL_REG, 2), - f"wt_oc_byte = {C_in_K_K}*4") - self.emit.emit(rv_add(WT_PTR, S3_W_BASE, VAL_REG), - "wt_ptr = weight_base + oc_wt_byte") - - # Step 3: IN_PTR += (ih_base * W + iw_base) * 4 - self.emit.emit(rv_mul(VAL_REG, IH_BASE_REG, W_REG), - "ih_off = ih_base * W") - self.emit.emit(rv_add(VAL_REG, VAL_REG, IW_BASE_REG), - "row_off = ih*W + iw_base") - self.emit.emit(rv_slli(VAL_REG, VAL_REG, 2), - "row_byte = row_off * 4") - self.emit.emit(rv_add(IN_PTR, IN_PTR, VAL_REG), - "in_ptr = &input[ic, ih_base, iw_base]") - - # Step 4: WT_PTR += ic * K*K * 4 - self.emit.emit_li32(T6_REG, K_K, f"K*K = {K_K}") - self.emit.emit(rv_mul(VAL_REG, IC_REG, T6_REG), - f"ic_wt_off = ic * {K_K}") - self.emit.emit(rv_slli(VAL_REG, VAL_REG, 2), - "ic_wt_byte = ic_wt_off * 4") - self.emit.emit(rv_add(WT_PTR, WT_PTR, VAL_REG), - "wt_ptr += ic * K*K * 4") - - # ── kh loop ──────────────────────────────────────────────── - self.emit.emit(rv_addi(KH_REG, _R_ZERO, 0), f"kh=0 (K={K})") - L("_conv_kh_loop") - - # ── kw loop (innermost) ──────────────────────────────────── - self.emit.emit(rv_addi(KW_REG, _R_ZERO, 0), f"kw=0 (K={K})") - L("_conv_kw_loop") - - if no_pad: - # No-padding fast path: pointer-walking - self.emit.emit(rv_lw(VAL_REG, IN_PTR, 0), - "load input[ic,ih,iw]") - self.emit.emit(rv_lw(T6_REG, WT_PTR, 0), - "load weight[oc,ic,kh,kw]") - self.emit.emit(rv_mul(T6_REG, VAL_REG, T6_REG), - "Q16.16 mul") - self.emit.emit(rv_srai(T6_REG, T6_REG, 16), - ">> 16") - self.emit.emit(rv_add(ACC_REG, ACC_REG, T6_REG), - "acc += product") - self.emit.emit(rv_addi(IN_PTR, IN_PTR, 4), - "in_ptr += 4") - self.emit.emit(rv_addi(WT_PTR, WT_PTR, 4), - "wt_ptr += 4") - self.emit.emit(rv_addi(KW_REG, KW_REG, 1), "kw++") - self.emit.emit(rv_slt(COND_REG, KW_REG, K_REG), "kw < K?") + # Batch (N) is handled by emitting the whole computation N times with a + # compile-time slice offset, rather than a runtime loop: N is known at + # codegen time and is small (<= a few). N==1 emits exactly the old code. + in_batch_stride = C_in * H * W + out_batch_stride = C_out * H_out * W_out + + for _n in range(N): + # Unique label prefix per batch copy (labels repeat across copies). + previous_prefix = self.emit.label_prefix + self.emit.label_prefix = f"{previous_prefix}n{_n}_" + if _n > 0: + # Advance both working bases by one sample stride. Bases carry + # over between copies, so add an incremental stride (not n*). + self.emit.emit_li32(TMP_REG, in_batch_stride, + f"batch n={_n} in stride") + self.emit.emit(rv_slli(TMP_REG, TMP_REG, 2), "*4") + self.emit.emit(rv_add(S2_IN_BASE, S2_IN_BASE, TMP_REG), + "input base += slice") + self.emit.emit_li32(TMP_REG, out_batch_stride, + f"batch n={_n} out stride") + self.emit.emit(rv_slli(TMP_REG, TMP_REG, 2), "*4") + self.emit.emit(rv_add(S5_OUT_BASE, S5_OUT_BASE, TMP_REG), + "output base += slice") + + # ── oc loop ──────────────────────────────────────────────── + self.emit.emit(rv_addi(OC_REG, _R_ZERO, 0), f"oc=0 (C_out={C_out})") + L("_conv_oc_loop") + + # ── oh loop ──────────────────────────────────────────────────── + self.emit.emit(rv_addi(OH_REG, _R_ZERO, 0), f"oh=0 (H_out={H_out})") + L("_conv_oh_loop") + + # ih_base = oh * stride_h (precompute for this oh row) + self.emit.emit(rv_mul(IH_BASE_REG, OH_REG, STRIDE_H_REG), + "ih_base = oh * stride_h") + + # ── ow loop ──────────────────────────────────────────────────── + self.emit.emit(rv_addi(OW_REG, _R_ZERO, 0), f"ow=0 (W_out={W_out})") + L("_conv_ow_loop") + + # iw_base = ow * stride_w (precompute for this column) + self.emit.emit(rv_mul(IW_BASE_REG, OW_REG, STRIDE_W_REG), + "iw_base = ow * stride_w") + + # Each output element starts a new reduction, including its bias. + if b_name is not None: + self.emit.emit(rv_slli(TMP_REG, OC_REG, 2), "offset = oc*4") + self.emit.emit(rv_add(TMP_REG, S4_BIAS_BASE, TMP_REG), "+ bias_base") + self.emit.emit(rv_lw(ACC_REG, TMP_REG, 0), "acc = bias[oc]") + else: + self.emit.emit(rv_addi(ACC_REG, _R_ZERO, 0), "acc = 0 (no bias)") + + if no_pad and K == 3: + # ══════════════════════════════════════════════════════════════ + # Unrolled K=3 no-pad fast path: fully unroll kw+kh loops + # + pointer-increment ic iteration (no address recalculation). + # ~7 instr/MAC vs ~12 before. 2026-06-08 + # ══════════════════════════════════════════════════════════════ + # ic_advance = bytes to advance IN_PTR to next ic's same spatial pos + # K*K loads plus K-1 row skips move (K-1)*W + K elements. + # The next channel's window starts H*W elements from this one. + ic_advance_bytes = (H * W - (K - 1) * W - K) * 4 + IC_ADV_REG = KH_REG # s10 freed: no kh counter needed + ROW_ADV_REG = KW_REG # s11 freed: no kw counter needed + + self.emit.emit_li32(IC_ADV_REG, ic_advance_bytes, + f"ic_adv={ic_advance_bytes}") + self.emit.emit_li32(ROW_ADV_REG, W_minus_K_times_4, + f"row_adv={W_minus_K_times_4}") + + # ic = 0 + self.emit.emit(rv_addi(IC_REG, _R_ZERO, 0), + f"ic=0 (C_in={C_in})") + + # ── Compute starting pointers for ic=0 (once, outside loop) ─ + # IN_PTR = input_base + (ih_base * W + iw_base) * 4 + self.emit.emit(rv_mul(T6_REG, IH_BASE_REG, W_REG), + "inoff = ih_base * W") + self.emit.emit(rv_add(T6_REG, T6_REG, IW_BASE_REG), + "+ iw_base") + self.emit.emit(rv_slli(T6_REG, T6_REG, 2), + "* 4") + self.emit.emit(rv_add(IN_PTR, S2_IN_BASE, T6_REG), + "in_ptr = &input[0, ih_base, iw_base]") + # WT_PTR = weight_base + oc * C_in*K*K * 4 + self.emit.emit_li32(T6_REG, C_in_K_K, f"oc_wt_step={C_in_K_K}") + self.emit.emit(rv_mul(T6_REG, OC_REG, T6_REG), "wt_oc_base") + self.emit.emit(rv_slli(T6_REG, T6_REG, 2), + f"wtoff = oc * {C_in_K_K} * 4") + self.emit.emit(rv_add(WT_PTR, S3_W_BASE, T6_REG), + "wt_ptr = &weight[oc, 0, 0, 0]") + + L("_conv_ic_loop") + + # ── Fully unrolled 3×3 MAC block (9 MACs, no loop control) ─ + # Row 0 (kh=0): MAC 0,1,2 then row advance + for _ in range(3): + self.emit.emit(rv_lw(VAL_REG, IN_PTR, 0), + "load input") + self.emit.emit(rv_lw(T6_REG, WT_PTR, 0), + "load weight") + self.emit.emit(rv_mul(T6_REG, VAL_REG, T6_REG), + "Q16.16 mul") + self.emit.emit(rv_srai(T6_REG, T6_REG, 16), + ">>16") + self.emit.emit(rv_add(ACC_REG, ACC_REG, T6_REG), + "acc += prod") + self.emit.emit(rv_addi(IN_PTR, IN_PTR, 4), + "in_ptr++") + self.emit.emit(rv_addi(WT_PTR, WT_PTR, 4), + "wt_ptr++") + # Row advance: in_ptr += (W-K)*4 + self.emit.emit(rv_add(IN_PTR, IN_PTR, ROW_ADV_REG), + "in_ptr += row_adv to next row") + + # Row 1 (kh=1): MAC 3,4,5 then row advance + for _ in range(3): + self.emit.emit(rv_lw(VAL_REG, IN_PTR, 0), + "load input") + self.emit.emit(rv_lw(T6_REG, WT_PTR, 0), + "load weight") + self.emit.emit(rv_mul(T6_REG, VAL_REG, T6_REG), + "Q16.16 mul") + self.emit.emit(rv_srai(T6_REG, T6_REG, 16), + ">>16") + self.emit.emit(rv_add(ACC_REG, ACC_REG, T6_REG), + "acc += prod") + self.emit.emit(rv_addi(IN_PTR, IN_PTR, 4), + "in_ptr++") + self.emit.emit(rv_addi(WT_PTR, WT_PTR, 4), + "wt_ptr++") + # Row advance + self.emit.emit(rv_add(IN_PTR, IN_PTR, ROW_ADV_REG), + "in_ptr += row_adv to next row") + + # Row 2 (kh=2): MAC 6,7,8 (last row, no row advance after) + for _ in range(3): + self.emit.emit(rv_lw(VAL_REG, IN_PTR, 0), + "load input") + self.emit.emit(rv_lw(T6_REG, WT_PTR, 0), + "load weight") + self.emit.emit(rv_mul(T6_REG, VAL_REG, T6_REG), + "Q16.16 mul") + self.emit.emit(rv_srai(T6_REG, T6_REG, 16), + ">>16") + self.emit.emit(rv_add(ACC_REG, ACC_REG, T6_REG), + "acc += prod") + self.emit.emit(rv_addi(IN_PTR, IN_PTR, 4), + "in_ptr++") + self.emit.emit(rv_addi(WT_PTR, WT_PTR, 4), + "wt_ptr++") + + # ── Advance IN_PTR to next ic's same spatial position ───── + # WT_PTR already at next ic (weights contiguous across ic) + self.emit.emit(rv_add(IN_PTR, IN_PTR, IC_ADV_REG), + "in_ptr += ic_advance") + # ic++ and loop + self.emit.emit(rv_addi(IC_REG, IC_REG, 1), "ic++") + self.emit.emit(rv_slt(COND_REG, IC_REG, C_IN_REG), "ic < C_in?") self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, - "_conv_kw_loop", "loop kw") + "_conv_ic_loop", "loop ic") + else: - # Padding path: check coordinates before any input load. - self.emit.emit(rv_mul(TMP_REG, OH_REG, STRIDE_H_REG), - "tmp = oh * stride_h") - self.emit.emit(rv_add(TMP_REG, TMP_REG, KH_REG), "tmp += kh") - self.emit.emit(rv_addi(TMP_REG, TMP_REG, -ph), - f"tmp -= {ph}") - - skip_label = f"_conv_skip_{len(self.emit.labels)}" - self.emit.emit(rv_slti(COND_REG, TMP_REG, 0), "ih < 0?") - self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, skip_label, - "skip if ih < 0") - self.emit.emit_li(T6_REG, H, f"H = {H}") - self.emit.emit(rv_slt(COND_REG, TMP_REG, T6_REG), "ih < H?") - self.emit.emit_branch(rv_beq, COND_REG, _R_ZERO, skip_label, - "skip if ih >= H") - - self.emit.emit(rv_mul(VAL_REG, OW_REG, STRIDE_W_REG), - "val = ow * stride_w") - self.emit.emit(rv_add(VAL_REG, VAL_REG, KW_REG), "val += kw") - self.emit.emit(rv_addi(VAL_REG, VAL_REG, -pw), - f"val -= {pw}") - self.emit.emit(rv_slti(COND_REG, VAL_REG, 0), "iw < 0?") - self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, skip_label, - "skip if iw < 0") - - self.emit.emit_li(T6_REG, W, f"W = {W}") - self.emit.emit(rv_slt(COND_REG, VAL_REG, T6_REG), "iw < W?") - self.emit.emit_branch(rv_beq, COND_REG, _R_ZERO, skip_label, - "skip if iw >= W") - - self.emit.emit(rv_mul(T6_REG, IC_REG, HW_REG), - "addr = ic*H*W") - self.emit.emit(rv_mul(COND_REG, TMP_REG, W_REG), - "tmp = ih*W") - self.emit.emit(rv_add(T6_REG, T6_REG, COND_REG), - "addr += ih*W") - self.emit.emit(rv_add(T6_REG, T6_REG, VAL_REG), - "addr += iw") - self.emit.emit(rv_slli(T6_REG, T6_REG, 2), "addr *= 4") - self.emit.emit(rv_add(T6_REG, S2_IN_BASE, T6_REG), - "addr += input_base") - self.emit.emit(rv_lw(VAL_REG, T6_REG, 0), - "load input[ic,ih,iw]") - - # WT_PTR already points at weight[oc, ic, 0, 0]. - self.emit.emit(rv_mul(T6_REG, KH_REG, K_REG), "addr = kh*K") - self.emit.emit(rv_add(T6_REG, T6_REG, KW_REG), - "addr += kw") - self.emit.emit(rv_slli(T6_REG, T6_REG, 2), "addr *= 4") - self.emit.emit(rv_add(T6_REG, WT_PTR, T6_REG), - "addr += weight[oc,ic] base") - self.emit.emit(rv_lw(COND_REG, T6_REG, 0), - "load weight[oc,ic,kh,kw]") - - self.emit.emit(rv_mul(T6_REG, VAL_REG, COND_REG), - "Q16.16 mul") - self.emit.emit(rv_srai(T6_REG, T6_REG, 16), - ">> 16") - self.emit.emit(rv_add(ACC_REG, ACC_REG, T6_REG), - "acc += product") - - self.emit.label(skip_label) - - self.emit.emit(rv_addi(KW_REG, KW_REG, 1), "kw++") - self.emit.emit(rv_slt(COND_REG, KW_REG, K_REG), "kw < K?") + # ══════════════════════════════════════════════════════════════ + # General loop path (K ≠ 3 or padded) + # ══════════════════════════════════════════════════════════════ + # ── ic loop ──────────────────────────────────────────────── + self.emit.emit(rv_addi(IC_REG, _R_ZERO, 0), f"ic=0 (C_in={C_in})") + L("_conv_ic_loop") + + # Step 1: IN_PTR = input_base + ic * H*W * 4 + self.emit.emit(rv_mul(VAL_REG, IC_REG, HW_REG), + "ic_offset = ic * H*W") + self.emit.emit(rv_slli(VAL_REG, VAL_REG, 2), + "ic_byte_off = ic_offset * 4") + self.emit.emit(rv_add(IN_PTR, S2_IN_BASE, VAL_REG), + "in_ptr = input_base + ic*H*W*4") + + # Step 2: WT_PTR = weight_base + oc * C_in*K*K * 4 + self.emit.emit_li32(VAL_REG, C_in_K_K, f"oc_wt_step={C_in_K_K}") + self.emit.emit(rv_mul(VAL_REG, OC_REG, VAL_REG), "wt_oc_base") + self.emit.emit(rv_slli(VAL_REG, VAL_REG, 2), + f"wt_oc_byte = {C_in_K_K}*4") + self.emit.emit(rv_add(WT_PTR, S3_W_BASE, VAL_REG), + "wt_ptr = weight_base + oc_wt_byte") + + # Step 3: IN_PTR += (ih_base * W + iw_base) * 4 + self.emit.emit(rv_mul(VAL_REG, IH_BASE_REG, W_REG), + "ih_off = ih_base * W") + self.emit.emit(rv_add(VAL_REG, VAL_REG, IW_BASE_REG), + "row_off = ih*W + iw_base") + self.emit.emit(rv_slli(VAL_REG, VAL_REG, 2), + "row_byte = row_off * 4") + self.emit.emit(rv_add(IN_PTR, IN_PTR, VAL_REG), + "in_ptr = &input[ic, ih_base, iw_base]") + + # Step 4: WT_PTR += ic * K*K * 4 + self.emit.emit_li32(T6_REG, K_K, f"K*K = {K_K}") + self.emit.emit(rv_mul(VAL_REG, IC_REG, T6_REG), + f"ic_wt_off = ic * {K_K}") + self.emit.emit(rv_slli(VAL_REG, VAL_REG, 2), + "ic_wt_byte = ic_wt_off * 4") + self.emit.emit(rv_add(WT_PTR, WT_PTR, VAL_REG), + "wt_ptr += ic * K*K * 4") + + # ── kh loop ──────────────────────────────────────────────── + self.emit.emit(rv_addi(KH_REG, _R_ZERO, 0), f"kh=0 (K={K})") + L("_conv_kh_loop") + + # ── kw loop (innermost) ──────────────────────────────────── + self.emit.emit(rv_addi(KW_REG, _R_ZERO, 0), f"kw=0 (K={K})") + L("_conv_kw_loop") + + if no_pad: + # No-padding fast path: pointer-walking + self.emit.emit(rv_lw(VAL_REG, IN_PTR, 0), + "load input[ic,ih,iw]") + self.emit.emit(rv_lw(T6_REG, WT_PTR, 0), + "load weight[oc,ic,kh,kw]") + self.emit.emit(rv_mul(T6_REG, VAL_REG, T6_REG), + "Q16.16 mul") + self.emit.emit(rv_srai(T6_REG, T6_REG, 16), + ">> 16") + self.emit.emit(rv_add(ACC_REG, ACC_REG, T6_REG), + "acc += product") + self.emit.emit(rv_addi(IN_PTR, IN_PTR, 4), + "in_ptr += 4") + self.emit.emit(rv_addi(WT_PTR, WT_PTR, 4), + "wt_ptr += 4") + self.emit.emit(rv_addi(KW_REG, KW_REG, 1), "kw++") + self.emit.emit(rv_slt(COND_REG, KW_REG, K_REG), "kw < K?") + self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, + "_conv_kw_loop", "loop kw") + else: + # Padding path: check coordinates before any input load. + self.emit.emit(rv_mul(TMP_REG, OH_REG, STRIDE_H_REG), + "tmp = oh * stride_h") + self.emit.emit(rv_add(TMP_REG, TMP_REG, KH_REG), "tmp += kh") + self.emit.emit(rv_addi(TMP_REG, TMP_REG, -ph), + f"tmp -= {ph}") + + skip_label = f"_conv_skip_{len(self.emit.labels)}" + self.emit.emit(rv_slti(COND_REG, TMP_REG, 0), "ih < 0?") + self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, skip_label, + "skip if ih < 0") + self.emit.emit_li(T6_REG, H, f"H = {H}") + self.emit.emit(rv_slt(COND_REG, TMP_REG, T6_REG), "ih < H?") + self.emit.emit_branch(rv_beq, COND_REG, _R_ZERO, skip_label, + "skip if ih >= H") + + self.emit.emit(rv_mul(VAL_REG, OW_REG, STRIDE_W_REG), + "val = ow * stride_w") + self.emit.emit(rv_add(VAL_REG, VAL_REG, KW_REG), "val += kw") + self.emit.emit(rv_addi(VAL_REG, VAL_REG, -pw), + f"val -= {pw}") + self.emit.emit(rv_slti(COND_REG, VAL_REG, 0), "iw < 0?") + self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, skip_label, + "skip if iw < 0") + + self.emit.emit_li(T6_REG, W, f"W = {W}") + self.emit.emit(rv_slt(COND_REG, VAL_REG, T6_REG), "iw < W?") + self.emit.emit_branch(rv_beq, COND_REG, _R_ZERO, skip_label, + "skip if iw >= W") + + self.emit.emit(rv_mul(T6_REG, IC_REG, HW_REG), + "addr = ic*H*W") + self.emit.emit(rv_mul(COND_REG, TMP_REG, W_REG), + "tmp = ih*W") + self.emit.emit(rv_add(T6_REG, T6_REG, COND_REG), + "addr += ih*W") + self.emit.emit(rv_add(T6_REG, T6_REG, VAL_REG), + "addr += iw") + self.emit.emit(rv_slli(T6_REG, T6_REG, 2), "addr *= 4") + self.emit.emit(rv_add(T6_REG, S2_IN_BASE, T6_REG), + "addr += input_base") + self.emit.emit(rv_lw(VAL_REG, T6_REG, 0), + "load input[ic,ih,iw]") + + # WT_PTR already points at weight[oc, ic, 0, 0]. + self.emit.emit(rv_mul(T6_REG, KH_REG, K_REG), "addr = kh*K") + self.emit.emit(rv_add(T6_REG, T6_REG, KW_REG), + "addr += kw") + self.emit.emit(rv_slli(T6_REG, T6_REG, 2), "addr *= 4") + self.emit.emit(rv_add(T6_REG, WT_PTR, T6_REG), + "addr += weight[oc,ic] base") + self.emit.emit(rv_lw(COND_REG, T6_REG, 0), + "load weight[oc,ic,kh,kw]") + + self.emit.emit(rv_mul(T6_REG, VAL_REG, COND_REG), + "Q16.16 mul") + self.emit.emit(rv_srai(T6_REG, T6_REG, 16), + ">> 16") + self.emit.emit(rv_add(ACC_REG, ACC_REG, T6_REG), + "acc += product") + + self.emit.label(skip_label) + + self.emit.emit(rv_addi(KW_REG, KW_REG, 1), "kw++") + self.emit.emit(rv_slt(COND_REG, KW_REG, K_REG), "kw < K?") + self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, + "_conv_kw_loop", "loop kw") + + # ── After kw loop: advance for next kh row ───────────────── + if no_pad: + self.emit.emit(rv_addi(IN_PTR, IN_PTR, W_minus_K_times_4), + f"in_ptr += (W({W})-K({K}))*4 = {W_minus_K_times_4}") + + # ── Increment kh ──────────────────────────────────────────── + self.emit.emit(rv_addi(KH_REG, KH_REG, 1), "kh++") + self.emit.emit(rv_slt(COND_REG, KH_REG, K_REG), "kh < K?") self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, - "_conv_kw_loop", "loop kw") + "_conv_kh_loop", "loop kh") - # ── After kw loop: advance for next kh row ───────────────── - if no_pad: - self.emit.emit(rv_addi(IN_PTR, IN_PTR, W_minus_K_times_4), - f"in_ptr += (W({W})-K({K}))*4 = {W_minus_K_times_4}") + # ── Increment ic ──────────────────────────────────────────── + self.emit.emit(rv_addi(IC_REG, IC_REG, 1), "ic++") + self.emit.emit(rv_slt(COND_REG, IC_REG, C_IN_REG), "ic < C_in?") + self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, + "_conv_ic_loop", "loop ic") + + # ── Store output[oc, oh, ow] = acc ───────────────────────────── + # out_offset = oc*H_out*W_out + oh*W_out + ow + self.emit.emit_li(TMP_REG, H_out_W_out, f"H_out*W_out={H_out_W_out}") + self.emit.emit(rv_mul(TMP_REG, OC_REG, TMP_REG), + "addr = oc * H_out_W_out") + self.emit.emit_li(T6_REG, W_out, f"W_out={W_out}") + self.emit.emit(rv_mul(T6_REG, OH_REG, T6_REG), + "tmp = oh * W_out") + self.emit.emit(rv_add(TMP_REG, TMP_REG, T6_REG), + "addr += oh * W_out") + self.emit.emit(rv_add(TMP_REG, TMP_REG, OW_REG), + "addr += ow") + self.emit.emit(rv_slli(TMP_REG, TMP_REG, 2), "addr *= 4") + self.emit.emit(rv_add(TMP_REG, S5_OUT_BASE, TMP_REG), + "+ output_base") + self.emit.emit(rv_sw(TMP_REG, ACC_REG, 0), + "store output[oc,oh,ow]") + + # ── Increment ow ──────────────────────────────────────────────── + self.emit.emit(rv_addi(OW_REG, OW_REG, 1), "ow++") + self.emit.emit_li(T6_REG, W_out, f"W_out={W_out}") + self.emit.emit(rv_slt(COND_REG, OW_REG, T6_REG), "ow < W_out?") + self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, + "_conv_ow_loop", "loop ow") - # ── Increment kh ──────────────────────────────────────────── - self.emit.emit(rv_addi(KH_REG, KH_REG, 1), "kh++") - self.emit.emit(rv_slt(COND_REG, KH_REG, K_REG), "kh < K?") + # ── Increment oh ──────────────────────────────────────────────── + self.emit.emit(rv_addi(OH_REG, OH_REG, 1), "oh++") + self.emit.emit_li(T6_REG, H_out, f"H_out={H_out}") + self.emit.emit(rv_slt(COND_REG, OH_REG, T6_REG), "oh < H_out?") self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, - "_conv_kh_loop", "loop kh") + "_conv_oh_loop", "loop oh") - # ── Increment ic ──────────────────────────────────────────── - self.emit.emit(rv_addi(IC_REG, IC_REG, 1), "ic++") - self.emit.emit(rv_slt(COND_REG, IC_REG, C_IN_REG), "ic < C_in?") + # ── Increment oc ──────────────────────────────────────────────── + self.emit.emit(rv_addi(OC_REG, OC_REG, 1), "oc++") + self.emit.emit_li(T6_REG, C_out, f"C_out={C_out}") + self.emit.emit(rv_slt(COND_REG, OC_REG, T6_REG), "oc < C_out?") self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, - "_conv_ic_loop", "loop ic") - - # ── Store output[oc, oh, ow] = acc ───────────────────────────── - # out_offset = oc*H_out*W_out + oh*W_out + ow - self.emit.emit_li(TMP_REG, H_out_W_out, f"H_out*W_out={H_out_W_out}") - self.emit.emit(rv_mul(TMP_REG, OC_REG, TMP_REG), - "addr = oc * H_out_W_out") - self.emit.emit_li(T6_REG, W_out, f"W_out={W_out}") - self.emit.emit(rv_mul(T6_REG, OH_REG, T6_REG), - "tmp = oh * W_out") - self.emit.emit(rv_add(TMP_REG, TMP_REG, T6_REG), - "addr += oh * W_out") - self.emit.emit(rv_add(TMP_REG, TMP_REG, OW_REG), - "addr += ow") - self.emit.emit(rv_slli(TMP_REG, TMP_REG, 2), "addr *= 4") - self.emit.emit(rv_add(TMP_REG, S5_OUT_BASE, TMP_REG), - "+ output_base") - self.emit.emit(rv_sw(TMP_REG, ACC_REG, 0), - "store output[oc,oh,ow]") - - # ── Increment ow ──────────────────────────────────────────────── - self.emit.emit(rv_addi(OW_REG, OW_REG, 1), "ow++") - self.emit.emit_li(T6_REG, W_out, f"W_out={W_out}") - self.emit.emit(rv_slt(COND_REG, OW_REG, T6_REG), "ow < W_out?") - self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, - "_conv_ow_loop", "loop ow") - - # ── Increment oh ──────────────────────────────────────────────── - self.emit.emit(rv_addi(OH_REG, OH_REG, 1), "oh++") - self.emit.emit_li(T6_REG, H_out, f"H_out={H_out}") - self.emit.emit(rv_slt(COND_REG, OH_REG, T6_REG), "oh < H_out?") - self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, - "_conv_oh_loop", "loop oh") - - # ── Increment oc ──────────────────────────────────────────────── - self.emit.emit(rv_addi(OC_REG, OC_REG, 1), "oc++") - self.emit.emit_li(T6_REG, C_out, f"C_out={C_out}") - self.emit.emit(rv_slt(COND_REG, OC_REG, T6_REG), "oc < C_out?") - self.emit.emit_branch(rv_bne, COND_REG, _R_ZERO, - "_conv_oc_loop", "loop oc") + "_conv_oc_loop", "loop oc") + + # restore label prefix after this batch copy + self.emit.label_prefix = previous_prefix # ── ReLU ─────────────────────────────────────────────────────────── @@ -2352,7 +2472,8 @@ def _gen_gemm(self, node: NodeInfo) -> None: """ a_name = node.inputs[0] w_name = node.inputs[1] - b_name = node.inputs[2] + # Gemm bias (C) is optional; a missing third input means "no bias". + b_name = node.inputs[2] if len(node.inputs) > 2 and node.inputs[2] else None out_name = node.outputs[0] a_shape = self.model.get_shape(a_name) @@ -2376,7 +2497,8 @@ def _gen_gemm(self, node: NodeInfo) -> None: self._get_workspace_addr(a_name, self.S2) self._get_weight_addr(w_name, self.S3) - self._get_weight_addr(b_name, self.S4) + if b_name is not None: + self._get_weight_addr(b_name, self.S4) self._get_workspace_addr(out_name, self.S5) # Nested loops: for i in range(M), for j in range(N) @@ -2399,10 +2521,13 @@ def _gen_gemm(self, node: NodeInfo) -> None: self.emit.emit(rv_addi(j_reg, _R_ZERO, 0), f"j=0 (N={N})") L("_gemm_j_loop") - # Load bias[j] → acc - self.emit.emit(rv_slli(addr_reg, j_reg, 2), "addr = j*4") - self.emit.emit(rv_add(addr_reg, self.S4, addr_reg), "addr += bias_base") - self.emit.emit(rv_lw(acc_reg, addr_reg, 0), "acc = bias[j]") + # Load bias[j] → acc (or zero when Gemm has no C input) + if b_name is not None: + self.emit.emit(rv_slli(addr_reg, j_reg, 2), "addr = j*4") + self.emit.emit(rv_add(addr_reg, self.S4, addr_reg), "addr += bias_base") + self.emit.emit(rv_lw(acc_reg, addr_reg, 0), "acc = bias[j]") + else: + self.emit.emit(rv_addi(acc_reg, _R_ZERO, 0), "acc = 0 (no bias)") # k = 0 self.emit.emit(rv_addi(k_reg, _R_ZERO, 0), f"k=0 (K={K})") @@ -2586,6 +2711,513 @@ def _gen_reshape(self, node: NodeInfo) -> None: self.emit.emit(rv_slt(cond_reg, i_reg, self.T6), "i < num_el?") self.emit.emit_branch(rv_bne, cond_reg, _R_ZERO, "_reshape_loop", "loop") + # ── FWHT ─────────────────────────────────────────────────────────── + + def _gen_fwht(self, node: NodeInfo) -> None: + """Scalar iterative FWHT over the flattened input (N a power of two). + + Algorithm (standard in-place butterfly on Q16.16 values): + for length in 1, 2, 4, ..., N/2: + for i in range(0, N, 2*length): + for j in range(length): + u = a[i+j]; v = a[i+j+length] + a[i+j] = u+v; a[i+j+length] = u-v + + The input is first copied to the output workspace, then transformed + in place there (input and output are distinct buffers). + """ + x_name = node.inputs[0] + out_name = node.outputs[0] + shape = self.model.get_shape(x_name) + + num_el = 1 + for d in shape: + num_el *= d + + if num_el <= 0 or (num_el & (num_el - 1)) != 0: + raise ValueError( + f"Fwht requires a power-of-two element count, got {num_el}" + ) + + direction = node.attrs.get("direction", "forward") + if direction not in ("forward", "inverse"): + raise ValueError( + f"Fwht direction must be 'forward' or 'inverse', got {direction!r}" + ) + # Inverse = forward Hadamard (H is symmetric, H^2 = N*I) then scale 1/N. + # N is a power of two, so 1/N is an arithmetic right shift by log2(N). + inv_shift = num_el.bit_length() - 1 + + self.mem.alloc_workspace(out_name, num_el) + + self._get_workspace_addr(x_name, self.S2) + self._get_workspace_addr(out_name, self.S5) + + L = self.emit.label + + # ── Copy input → output workspace ────────────────────────────── + i_reg = self.S6 + addr_src = self.T3 + addr_dst = self.T5 + val_reg = self.T1 + cond_reg = self.T4 + + self.emit.emit(rv_addi(i_reg, _R_ZERO, 0), f"fwht copy ({num_el} el)") + L("_fwht_copy_loop") + self.emit.emit(rv_slli(addr_src, i_reg, 2), "") + self.emit.emit(rv_add(addr_src, self.S2, addr_src), "") + self.emit.emit(rv_lw(val_reg, addr_src, 0), "load src[i]") + self.emit.emit(rv_slli(addr_dst, i_reg, 2), "") + self.emit.emit(rv_add(addr_dst, self.S5, addr_dst), "") + self.emit.emit(rv_sw(addr_dst, val_reg, 0), "store dst[i]") + self.emit.emit(rv_addi(i_reg, i_reg, 1), "i++") + self.emit.emit_li32(self.T6, num_el, f"n={num_el}") + self.emit.emit(rv_slt(cond_reg, i_reg, self.T6), "i < n?") + self.emit.emit_branch(rv_bne, cond_reg, _R_ZERO, "_fwht_copy_loop", "loop") + + # ── Butterflies (in place on output workspace) ───────────────── + len_reg = self.S6 # current butterfly length + i_reg = self.S7 # block start index + j_reg = self.S8 # offset inside block + idx_reg = self.S9 # i + j + step_reg = self.S10 # 2 * length + n_reg = self.S11 # N (loop-invariant) + + addr_a = self.T3 + addr_b = self.T5 + u_reg = self.T0 + v_reg = self.T1 + cond_reg = self.T4 + tmp_reg = self.T6 + + self.emit.emit_li32(n_reg, num_el, f"N={num_el}") + self.emit.emit(rv_addi(len_reg, _R_ZERO, 1), "length = 1") + + L("_fwht_len_loop") + self.emit.emit(rv_slt(cond_reg, len_reg, n_reg), "length < N?") + self.emit.emit_branch(rv_beq, cond_reg, _R_ZERO, "_fwht_done", "exit") + self.emit.emit(rv_slli(step_reg, len_reg, 1), "step = 2*length") + + self.emit.emit(rv_addi(i_reg, _R_ZERO, 0), "i = 0") + L("_fwht_i_loop") + self.emit.emit(rv_addi(j_reg, _R_ZERO, 0), "j = 0") + L("_fwht_j_loop") + + # addr_a = out + (i+j)*4 ; addr_b = out + (i+j+length)*4 + self.emit.emit(rv_add(idx_reg, i_reg, j_reg), "idx = i+j") + self.emit.emit(rv_slli(addr_a, idx_reg, 2), "idx*4") + self.emit.emit(rv_add(addr_a, self.S5, addr_a), "addr_a") + self.emit.emit(rv_add(tmp_reg, idx_reg, len_reg), "idx+length") + self.emit.emit(rv_slli(tmp_reg, tmp_reg, 2), "(idx+length)*4") + self.emit.emit(rv_add(addr_b, self.S5, tmp_reg), "addr_b") + + self.emit.emit(rv_lw(u_reg, addr_a, 0), "u = a[idx]") + self.emit.emit(rv_lw(v_reg, addr_b, 0), "v = a[idx+length]") + self.emit.emit(rv_add(tmp_reg, u_reg, v_reg), "u+v") + self.emit.emit(rv_sw(addr_a, tmp_reg, 0), "a[idx] = u+v") + self.emit.emit(rv_sub(tmp_reg, u_reg, v_reg), "u-v") + self.emit.emit(rv_sw(addr_b, tmp_reg, 0), "a[idx+length] = u-v") + + # j++ ; j < length? + self.emit.emit(rv_addi(j_reg, j_reg, 1), "j++") + self.emit.emit(rv_slt(cond_reg, j_reg, len_reg), "j < length?") + self.emit.emit_branch(rv_bne, cond_reg, _R_ZERO, "_fwht_j_loop", "loop j") + + # i += step ; i < N? + self.emit.emit(rv_add(i_reg, i_reg, step_reg), "i += step") + self.emit.emit(rv_slt(cond_reg, i_reg, n_reg), "i < N?") + self.emit.emit_branch(rv_bne, cond_reg, _R_ZERO, "_fwht_i_loop", "loop i") + + # length *= 2 ; loop + self.emit.emit(rv_slli(len_reg, len_reg, 1), "length *= 2") + self.emit.emit_jump(rv_j, "_fwht_len_loop", "loop length") + + L("_fwht_done") + + # ── Inverse scaling: a[i] = a[i] >> log2(N) ──────────────────── + if direction == "inverse" and inv_shift > 0: + i_reg = self.S6 + addr = self.T3 + val = self.T1 + cond = self.T4 + self.emit.emit(rv_addi(i_reg, _R_ZERO, 0), "inv-scale i=0") + L("_fwht_scale_loop") + self.emit.emit(rv_slli(addr, i_reg, 2), "") + self.emit.emit(rv_add(addr, self.S5, addr), "") + self.emit.emit(rv_lw(val, addr, 0), "load a[i]") + self.emit.emit(rv_srai(val, val, inv_shift), f"a[i] >>= {inv_shift}") + self.emit.emit(rv_sw(addr, val, 0), "store a[i]") + self.emit.emit(rv_addi(i_reg, i_reg, 1), "i++") + self.emit.emit_li32(self.T6, num_el, f"n={num_el}") + self.emit.emit(rv_slt(cond, i_reg, self.T6), "i < n?") + self.emit.emit_branch(rv_bne, cond, _R_ZERO, "_fwht_scale_loop", "loop") + + # ── Sparse matrix multiply (CSR) ─────────────────────────────────── + + def _gen_spmmcsr(self, node: NodeInfo) -> None: + """CSR sparse matrix × dense matrix (SpMM; N=1 degenerates to SpMV). + + Inputs: (values[q16, nnz], col_indices[int32, nnz], + row_ptr[int32, M+1], B[q16, K, N]) -> C[q16, M, N] + + for i in range(M): + for n in range(N): C[i,n] = 0 + for j in row_ptr[i] .. row_ptr[i+1]-1: + a = values[j]; k = col_indices[j] + for n in range(N): + C[i,n] += (a * B[k,n]) >> 16 + """ + v_name = node.inputs[0] + col_name = node.inputs[1] + row_name = node.inputs[2] + b_name = node.inputs[3] + out_name = node.outputs[0] + + row_shape = self.model.get_shape(row_name) + M = (row_shape[0] - 1) if row_shape else 0 + b_shape = self.model.get_shape(b_name) + N = b_shape[1] if len(b_shape) >= 2 else 1 + + self.mem.alloc_workspace(out_name, M * N) + + base_v = self.S2 + base_col = self.S3 + base_row = self.S4 + base_b = self.S5 + base_c = self.S6 + + i_reg = self.S7 + j_reg = self.S8 + jend_reg = self.S9 + n_reg = self.S10 + ncol_reg = self.S11 + + acc = self.T0 + a_reg = self.T1 + k_reg = self.T2 + addr = self.T3 + cond = self.T4 + tmp = self.T6 + + self._get_weight_addr(v_name, base_v) + self._get_weight_addr(col_name, base_col) + self._get_weight_addr(row_name, base_row) + self._get_workspace_addr(b_name, base_b) + self._get_workspace_addr(out_name, base_c) + + self.emit.emit_li(ncol_reg, N, f"N={N}") + + L = self.emit.label + + # for i in range(M) — condition at top so M=0 is a no-op + self.emit.emit(rv_addi(i_reg, _R_ZERO, 0), f"i=0 (M={M})") + L("_spmm_i_loop") + self.emit.emit_li32(tmp, M, f"M={M}") + self.emit.emit(rv_slt(cond, i_reg, tmp), "i < M?") + self.emit.emit_branch(rv_beq, cond, _R_ZERO, "_spmm_done", "exit") + + # j = row_ptr[i]; j_end = row_ptr[i+1] + self.emit.emit(rv_slli(addr, i_reg, 2), "i*4") + self.emit.emit(rv_add(addr, base_row, addr), "&row_ptr[i]") + self.emit.emit(rv_lw(j_reg, addr, 0), "j = row_ptr[i]") + self.emit.emit(rv_lw(jend_reg, addr, 4), "j_end = row_ptr[i+1]") + + # Zero the C row: for n in range(N) + self.emit.emit(rv_addi(n_reg, _R_ZERO, 0), "n=0 (zero)") + L("_spmm_zero_loop") + self.emit.emit(rv_mul(tmp, i_reg, ncol_reg), "i*N") + self.emit.emit(rv_add(tmp, tmp, n_reg), "+ n") + self.emit.emit(rv_slli(tmp, tmp, 2), "*4") + self.emit.emit(rv_add(tmp, base_c, tmp), "&C[i,n]") + self.emit.emit(rv_sw(tmp, _R_ZERO, 0), "C[i,n] = 0") + self.emit.emit(rv_addi(n_reg, n_reg, 1), "n++") + self.emit.emit(rv_slt(cond, n_reg, ncol_reg), "n < N?") + self.emit.emit_branch(rv_bne, cond, _R_ZERO, "_spmm_zero_loop", "loop") + + # for j in [j, j_end) + L("_spmm_j_loop") + self.emit.emit(rv_slt(cond, j_reg, jend_reg), "j < j_end?") + self.emit.emit_branch(rv_beq, cond, _R_ZERO, "_spmm_j_done", "exit") + + # a = values[j] + self.emit.emit(rv_slli(addr, j_reg, 2), "j*4") + self.emit.emit(rv_add(addr, base_v, addr), "&values[j]") + self.emit.emit(rv_lw(a_reg, addr, 0), "a = values[j]") + # k = col_indices[j] (indirect load -> addresses B's row) + self.emit.emit(rv_slli(addr, j_reg, 2), "j*4") + self.emit.emit(rv_add(addr, base_col, addr), "&col[j]") + self.emit.emit(rv_lw(k_reg, addr, 0), "k = col[j]") + + # for n in range(N) + self.emit.emit(rv_addi(n_reg, _R_ZERO, 0), "n=0") + L("_spmm_n_loop") + # B[k,n] + self.emit.emit(rv_mul(tmp, k_reg, ncol_reg), "k*N") + self.emit.emit(rv_add(tmp, tmp, n_reg), "+ n") + self.emit.emit(rv_slli(tmp, tmp, 2), "*4") + self.emit.emit(rv_add(tmp, base_b, tmp), "&B[k,n]") + self.emit.emit(rv_lw(tmp, tmp, 0), "load B[k,n]") + self.emit.emit(rv_mul(tmp, a_reg, tmp), "a*B (low32)") + self.emit.emit(rv_srai(tmp, tmp, 16), ">>16") + # C[i,n] += prod + self.emit.emit(rv_mul(addr, i_reg, ncol_reg), "i*N") + self.emit.emit(rv_add(addr, addr, n_reg), "+ n") + self.emit.emit(rv_slli(addr, addr, 2), "*4") + self.emit.emit(rv_add(addr, base_c, addr), "&C[i,n]") + self.emit.emit(rv_lw(acc, addr, 0), "old C") + self.emit.emit(rv_add(acc, acc, tmp), "+= prod") + self.emit.emit(rv_sw(addr, acc, 0), "store C[i,n]") + # n++ + self.emit.emit(rv_addi(n_reg, n_reg, 1), "n++") + self.emit.emit(rv_slt(cond, n_reg, ncol_reg), "n < N?") + self.emit.emit_branch(rv_bne, cond, _R_ZERO, "_spmm_n_loop", "loop") + + # j++ + self.emit.emit(rv_addi(j_reg, j_reg, 1), "j++") + self.emit.emit_jump(rv_j, "_spmm_j_loop", "loop j") + + L("_spmm_j_done") + # i++ + self.emit.emit(rv_addi(i_reg, i_reg, 1), "i++") + self.emit.emit_jump(rv_j, "_spmm_i_loop", "loop i") + + L("_spmm_done") + + # ── Winograd F(2,3) convolution ───────────────────────────────────── + + def _gen_winogradconv(self, node: NodeInfo) -> None: + """2D Winograd F(2,3) convolution (stride=1, SAME padding). + + Y_tile = A^T [ sum_ic (G g G^T) ⊙ (B^T d B) ] A + Kernel transform U = G g G^T is folded at compile time (2.2) into the + initializer named in ``attrs["_winograd_u"]``. + """ + x_name = node.inputs[0] + w_name = node.inputs[1] + b_name = node.inputs[2] if len(node.inputs) > 2 and node.inputs[2] else None + out_name = node.outputs[0] + + x_shape = self.model.get_shape(x_name) + w_shape = self.model.get_shape(w_name) + N = x_shape[0] if len(x_shape) > 0 else 1 + Cin = x_shape[1] if len(x_shape) > 1 else 1 + H = x_shape[2] if len(x_shape) > 2 else 1 + W = x_shape[3] if len(x_shape) > 3 else 1 + Cout = w_shape[0] if len(w_shape) > 0 else 1 + KH = w_shape[2] if len(w_shape) > 2 else 3 + + attrs = node.attrs + pads = list(attrs.get("pads", [1, 1, 1, 1]))[:4] + if pads != [1, 1, 1, 1]: + raise ValueError("WinogradConv requires pads=[1,1,1,1] (SAME, stride 1)") + if list(attrs.get("strides", [1, 1]))[:2] != [1, 1]: + raise ValueError("WinogradConv requires strides=[1,1]") + if KH != 3: + raise ValueError("WinogradConv supports 3x3 kernels only") + u_name = attrs.get("_winograd_u") + if not u_name: + raise ValueError("WinogradConv missing precomputed kernel (2.2)") + + TH = (H + 1) // 2 + TW = (W + 1) // 2 + Hout, Wout = H, W + + d_name = f"{out_name}__d" + v_name = f"{out_name}__v" + acc_name = f"{out_name}__acc" + for nm in (d_name, v_name, acc_name): + self.mem.alloc_workspace(nm, 16) + self.mem.alloc_workspace(out_name, N * Cout * Hout * Wout) + + base_in = self.S2 + base_u = self.S3 + base_out = self.S4 + base_d = self.S5 + base_v = self.S6 + base_acc = self.S7 + base_bias = self.S8 + + n_reg = self.S9 + oc_reg = self.S10 + ti_reg = self.S11 + tj_reg = self.T4 + ic_reg = self.T5 + + T0, T1, T2, T3, T6 = self.T0, self.T1, self.T2, self.T3, self.T6 + + self._get_workspace_addr(x_name, base_in) + self._get_weight_addr(u_name, base_u) + self._get_workspace_addr(out_name, base_out) + self._get_workspace_addr(d_name, base_d) + self._get_workspace_addr(v_name, base_v) + self._get_workspace_addr(acc_name, base_acc) + if b_name is not None: + self._get_weight_addr(b_name, base_bias) + + BT = _WINO_F23_BT + Bm = [[BT[b][a] for b in range(4)] for a in range(4)] # B[a][b] = BT[b][a] + AT = _WINO_F23_AT + A = [[AT[b][a] for b in range(2)] for a in range(4)] # A[a][b] = AT[b][a] + + E = self.emit.emit + Br = self.emit.emit_branch + J = self.emit.emit_jump + L = self.emit.label + li = self.emit.emit_li + + # for n + E(rv_addi(n_reg, _R_ZERO, 0), f"n=0 (N={N})") + L("_wg_n_loop") + E(rv_addi(oc_reg, _R_ZERO, 0), f"oc=0 (Cout={Cout})") + L("_wg_oc_loop") + E(rv_addi(ti_reg, _R_ZERO, 0), f"ti=0 (TH={TH})") + L("_wg_ti_loop") + E(rv_addi(tj_reg, _R_ZERO, 0), f"tj=0 (TW={TW})") + L("_wg_tj_loop") + + # zero the 4x4 Winograd-domain accumulator + for p in range(16): + E(rv_sw(base_acc, _R_ZERO, 4 * p), "" if p else "acc[:] = 0") + + E(rv_addi(ic_reg, _R_ZERO, 0), f"ic=0 (Cin={Cin})") + L("_wg_ic_loop") + + # ── load padded 4x4 input tile into dbuf ────────────────────── + for r in range(4): + for c in range(4): + E(rv_slli(T0, ti_reg, 1), "") + E(rv_addi(T0, T0, r - 1), "") # gi = 2*ti-1+r + E(rv_slli(T1, tj_reg, 1), "") + E(rv_addi(T1, T1, c - 1), "") # gj = 2*tj-1+c + E(rv_addi(T2, _R_ZERO, 0), "") # default 0 + skip = f"_wg_dskip_{r}_{c}" + Br(rv_blt, T0, _R_ZERO, skip, "") # gi < 0 + Br(rv_blt, T1, _R_ZERO, skip, "") # gj < 0 + li(T6, H) + E(rv_slt(T3, T0, T6), "") + Br(rv_beq, T3, _R_ZERO, skip, "") # gi >= H + li(T6, W) + E(rv_slt(T3, T1, T6), "") + Br(rv_beq, T3, _R_ZERO, skip, "") # gj >= W + li(T3, Cin) + E(rv_mul(T3, n_reg, T3), "") + E(rv_add(T3, T3, ic_reg), "") + li(T6, H) + E(rv_mul(T3, T3, T6), "") + E(rv_add(T3, T3, T0), "") + li(T6, W) + E(rv_mul(T3, T3, T6), "") + E(rv_add(T3, T3, T1), "") + E(rv_slli(T3, T3, 2), "") + E(rv_add(T3, base_in, T3), "") + E(rv_lw(T2, T3, 0), "") + L(skip) + E(rv_sw(base_d, T2, 4 * (r * 4 + c)), "") + + # ── input transform: vbuf = B^T dbuf B (signed sums of dbuf) ─── + for i in range(4): + for j in range(4): + E(rv_addi(T2, _R_ZERO, 0), "") + for a in range(4): + for b in range(4): + coef = BT[i][a] * Bm[b][j] + if coef == 0: + continue + E(rv_lw(T0, base_d, 4 * (a * 4 + b)), "") + E(rv_add(T2, T2, T0) if coef > 0 else rv_sub(T2, T2, T0), "") + E(rv_sw(base_v, T2, 4 * (i * 4 + j)), "") + + # ── Winograd-domain MAC: acc[p] += (U[oc,ic,p] * V[p]) >> 16 ─── + li(T0, Cin) + E(rv_mul(T0, oc_reg, T0), "") + E(rv_add(T0, T0, ic_reg), "") + E(rv_slli(T0, T0, 6), "((oc*Cin+ic)*16)*4") + E(rv_add(T0, base_u, T0), "&U[oc,ic,0]") + for p in range(16): + E(rv_lw(T1, T0, 4 * p), "") + E(rv_lw(T2, base_v, 4 * p), "") + # Full-precision Q16.16 multiply: (U*V) >> 16 using MULH+MUL. + # Winograd intermediates (V, U) can exceed 2^15, so the low-32 + # MUL + SRAI used elsewhere would wrap; the 64-bit product >>16 + # stays in range and matches the ideal reference. + E(rv_mul(T6, T1, T2), "mul lo") + E(rv_mulh(T1, T1, T2), "mul hi") + E(rv_srli(T6, T6, 16), "logical >>16 (low word can be negative)") + E(rv_slli(T1, T1, 16), "") + E(rv_or(T1, T1, T6), "(U*V)>>16") + E(rv_lw(T6, base_acc, 4 * p), "") + E(rv_add(T6, T6, T1), "") + E(rv_sw(base_acc, T6, 4 * p), "") + + E(rv_addi(ic_reg, ic_reg, 1), "ic++") + li(T6, Cin) + E(rv_slt(T3, ic_reg, T6), "") + Br(rv_bne, T3, _R_ZERO, "_wg_ic_loop", "loop ic") + + # ── output transform Y = A^T acc A, + bias, store in-bounds ──── + for oy in range(2): + for ox in range(2): + E(rv_addi(T2, _R_ZERO, 0), "") + for a in range(4): + for b in range(4): + coef = AT[oy][a] * AT[ox][b] + if coef == 0: + continue + E(rv_lw(T0, base_acc, 4 * (a * 4 + b)), "") + E(rv_add(T2, T2, T0) if coef > 0 else rv_sub(T2, T2, T0), "") + if b_name is not None: + E(rv_slli(T3, oc_reg, 2), "") + E(rv_add(T3, base_bias, T3), "") + E(rv_lw(T0, T3, 0), "bias[oc]") + E(rv_add(T2, T2, T0), "") + # bounds: oh = 2*ti+oy < H, ow = 2*tj+ox < W + oskip = f"_wg_oskip_{oy}_{ox}" + E(rv_slli(T0, ti_reg, 1), "") + E(rv_addi(T0, T0, oy), "oh") + li(T6, H) + E(rv_slt(T3, T0, T6), "") + Br(rv_beq, T3, _R_ZERO, oskip, "") + E(rv_slli(T1, tj_reg, 1), "") + E(rv_addi(T1, T1, ox), "ow") + li(T6, W) + E(rv_slt(T3, T1, T6), "") + Br(rv_beq, T3, _R_ZERO, oskip, "") + li(T3, Cout) + E(rv_mul(T3, n_reg, T3), "") + E(rv_add(T3, T3, oc_reg), "") + li(T6, H) + E(rv_mul(T3, T3, T6), "") + E(rv_add(T3, T3, T0), "") + li(T6, W) + E(rv_mul(T3, T3, T6), "") + E(rv_add(T3, T3, T1), "") + E(rv_slli(T3, T3, 2), "") + E(rv_add(T3, base_out, T3), "") + E(rv_sw(T3, T2, 0), "") + L(oskip) + + E(rv_addi(tj_reg, tj_reg, 1), "tj++") + li(T6, TW) + E(rv_slt(T3, tj_reg, T6), "") + Br(rv_bne, T3, _R_ZERO, "_wg_tj_loop", "loop tj") + + E(rv_addi(ti_reg, ti_reg, 1), "ti++") + li(T6, TH) + E(rv_slt(T3, ti_reg, T6), "") + Br(rv_bne, T3, _R_ZERO, "_wg_ti_loop", "loop ti") + + E(rv_addi(oc_reg, oc_reg, 1), "oc++") + li(T6, Cout) + E(rv_slt(T3, oc_reg, T6), "") + Br(rv_bne, T3, _R_ZERO, "_wg_oc_loop", "loop oc") + + E(rv_addi(n_reg, n_reg, 1), "n++") + li(T6, N) + E(rv_slt(T3, n_reg, T6), "") + Br(rv_bne, T3, _R_ZERO, "_wg_n_loop", "loop n") + + L("_wg_done") + # ═══════════════════════════════════════════════════════════════════════════ # Part 8: Main Pipeline # ═══════════════════════════════════════════════════════════════════════════ @@ -2793,6 +3425,7 @@ def convert_onnx_to_riscv( # ── Step 2: Layout weights and plan memory ───────────────────────── print("\n[2/5] Converting weights to Q16.16 fixed-point and planning memory...") + _prepare_winograd_kernels(model) # fold WinogradConv 3x3 kernels (if any) memory = MemoryPlan() weight_data = memory.layout_weights(model.initializers) print(f" Weight data: {memory.data_size:,} bytes ({memory.data_size/1024/1024:.1f} MB)") diff --git a/tests/test_standalone_conv_batch.py b/tests/test_standalone_conv_batch.py new file mode 100644 index 00000000..9a5e9c52 --- /dev/null +++ b/tests/test_standalone_conv_batch.py @@ -0,0 +1,106 @@ +"""Conv batch (N>1) regression (task A). + +_gen_conv emits the computation once per sample (compile-time unrolled) with +an incremental input/output slice offset. Verify N=1/2/3/4 against an +independent Q16.16 direct-conv reference, including K=5 and odd sizes. +""" + +import contextlib +import io +import struct +import tempfile + +import numpy as np +import onnx +from onnx import TensorProto, helper +import pytest + +from scratchv.standalone.onnx_to_riscv_standalone import convert_onnx_to_riscv +from scratchv.standalone.benchmark import PROFILES, RV32EmulatorFast + +Q = 65536 + + +def _q16(a): + return np.trunc(np.asarray(a, dtype=np.float64) * Q).astype(np.int64) + + +def _wrap32(v): + v &= 0xFFFFFFFF + return v - 0x100000000 if v >= 0x80000000 else v + + +def _ref_conv_direct(x_q, w_q, pads): + N, C, H, W = x_q.shape + Cout, _, KH, KW = w_q.shape + ph, pw = pads + xp = np.pad(x_q, ((0, 0), (0, 0), (ph, ph), (pw, pw))) + out = np.zeros((N, Cout, H, W), dtype=np.int64) + for n in range(N): + for oc in range(Cout): + for oh in range(H): + for ow in range(W): + acc = 0 + for ic in range(C): + for kh in range(KH): + for kw in range(KW): + prod = _wrap32(int(xp[n, ic, oh + kh, ow + kw]) + * int(w_q[oc, ic, kh, kw])) >> 16 + acc = _wrap32(acc + prod) + out[n, oc, oh, ow] = acc + return out + + +def _model(weight, n, cin, h, w, pad): + cout, _, k, _ = weight.shape + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [n, cin, h, w]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [n, cout, h, w]) + wt = helper.make_tensor("w", TensorProto.FLOAT, list(weight.shape), + weight.ravel().tolist()) + node = helper.make_node("Conv", ["x", "w"], ["y"], + pads=[pad, pad, pad, pad]) + graph = helper.make_graph([node], "c", [x], [y], [wt]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + + +def _run(model, x_q): + with tempfile.TemporaryDirectory() as td: + path = f"{td}/m.onnx" + onnx.save(model, path) + meta = {} + with contextlib.redirect_stdout(io.StringIO()): + rc = convert_onnx_to_riscv(path, f"{td}/m.bin", f"{td}/m.s", + benchmark=False, metadata=meta) + assert rc == 0 + with open(f"{td}/m.bin", "rb") as f: + binary = f.read() + flat = [int(v) for v in np.asarray(x_q).reshape(-1)] + in_bytes = struct.pack(f"<{len(flat)}i", *flat) + emu = RV32EmulatorFast(mem_size_mb=128) + emu.load_unified_binary(binary, int(meta["code_bytes"])) + emu.regs[10] = 0x04000000 + emu.regs[11] = 0x05000000 + for i, byte in enumerate(in_bytes): + emu.mem[0x04000000 + i] = byte + emu.run(max_instr=100_000_000, uarch=PROFILES["basic"]) + out_el = int(meta["output_elements"]) + return np.array([_wrap32(emu.read_mem_i32(0x05000000 + 4 * i)) + for i in range(out_el)], dtype=np.int64) + + +@pytest.mark.parametrize("n,cin,cout,h,w,k", [ + (1, 3, 4, 8, 8, 3), + (2, 3, 5, 9, 9, 3), + (3, 3, 4, 10, 10, 5), + (4, 4, 8, 8, 8, 3), +]) +def test_conv_batch_matches_reference(n, cin, cout, h, w, k): + rng = np.random.default_rng(n) + weight = rng.normal(0, 0.3, (cout, cin, k, k)).astype(np.float32) + x_q = _q16(rng.uniform(-0.25, 0.25, (n, cin, h, w))) + pad = k // 2 + + got = _run(_model(weight, n, cin, h, w, pad), x_q) + ref = _ref_conv_direct(x_q, _q16(weight), [pad, pad]).reshape(-1) + assert got.size == ref.size + assert int(np.abs(got - ref).max()) == 0 diff --git a/tests/test_standalone_fwht.py b/tests/test_standalone_fwht.py new file mode 100644 index 00000000..c046eafc --- /dev/null +++ b/tests/test_standalone_fwht.py @@ -0,0 +1,88 @@ +"""Regression for the scalar Fwht operator (task 1.3). + +Uses an analytically-known property instead of a duplicated reference: +FWHT of the impulse [1, 0, 0, ...] is the all-ones vector (first Hadamard +column), and the inverse additionally scales by 1/N. +""" + +import contextlib +import io +import struct +import tempfile + +import numpy as np +import onnx +from onnx import TensorProto, helper +import pytest + +from scratchv.standalone.onnx_to_riscv_standalone import convert_onnx_to_riscv +from scratchv.standalone.benchmark import PROFILES, RV32EmulatorFast + +Q = 65536 + + +def _fwht_model(n, direction="forward"): + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, n]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, n]) + node = helper.make_node("Fwht", ["X"], ["Y"], domain="org.scratchv", + direction=direction) + graph = helper.make_graph([node], "fwht", [x], [y]) + return helper.make_model( + graph, + opset_imports=[helper.make_opsetid("org.scratchv", 1), + helper.make_opsetid("", 13)], + ) + + +def _signed(v): + v &= 0xFFFFFFFF + return v - 0x100000000 if v >= 0x80000000 else v + + +def _run(model, input_q16): + with tempfile.TemporaryDirectory() as td: + path = f"{td}/m.onnx" + onnx.save(model, path) + meta = {} + with contextlib.redirect_stdout(io.StringIO()): + rc = convert_onnx_to_riscv(path, f"{td}/m.bin", f"{td}/m.s", + benchmark=False, metadata=meta) + assert rc == 0 + with open(f"{td}/m.bin", "rb") as f: + binary = f.read() + + in_bytes = struct.pack(f"<{len(input_q16)}i", *input_q16) + emu = RV32EmulatorFast(mem_size_mb=128) + emu.load_unified_binary(binary, int(meta["code_bytes"])) + emu.regs[10] = 0x04000000 + emu.regs[11] = 0x05000000 + for i, b in enumerate(in_bytes): + emu.mem[0x04000000 + i] = b + emu.run(max_instr=5_000_000, uarch=PROFILES["basic"]) + out_el = int(meta["output_elements"]) + return [_signed(emu.read_mem_i32(0x05000000 + 4 * i)) for i in range(out_el)] + + +@pytest.mark.parametrize("n", [2, 16, 64, 256]) +def test_forward_impulse_is_all_ones(n): + impulse = [Q] + [0] * (n - 1) + assert _run(_fwht_model(n, "forward"), impulse) == [Q] * n + + +@pytest.mark.parametrize("n", [16, 64, 256]) +def test_inverse_impulse_is_scaled_ones(n): + impulse = [Q] + [0] * (n - 1) + expected = Q >> (n.bit_length() - 1) # 1/N in Q16.16 + assert _run(_fwht_model(n, "inverse"), impulse) == [expected] * n + + +def test_non_power_of_two_raises(): + with pytest.raises(Exception) as exc: + _run(_fwht_model(48), [0] * 48) + assert "power-of-two" in str(exc.value) + + +def test_unknown_direction_raises(): + with pytest.raises(Exception) as exc: + _run(_fwht_model(16, "sideways"), [0] * 16) + assert "direction" in str(exc.value) diff --git a/tests/test_standalone_optional_bias.py b/tests/test_standalone_optional_bias.py new file mode 100644 index 00000000..9909e1c3 --- /dev/null +++ b/tests/test_standalone_optional_bias.py @@ -0,0 +1,95 @@ +"""Regression: Conv/Gemm must accept a missing bias input (treated as zero). + +Before the fix, `_gen_conv` / `_gen_gemm` hard-indexed `node.inputs[2]`, so a +bias-less model raised IndexError at codegen time. The no-bias output must be +identical to the same model with an explicit all-zero bias. +""" + +import contextlib +import io +import struct +import tempfile + +import numpy as np +import onnx +from onnx import TensorProto, helper +import pytest + +from scratchv.standalone.onnx_to_riscv_standalone import convert_onnx_to_riscv +from scratchv.standalone.benchmark import PROFILES, RV32EmulatorFast + +Q = 65536 + + +def _conv_model(bias): + rng = np.random.default_rng(0) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 8, 8]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 8, 8, 8]) + wt = helper.make_tensor("w", TensorProto.FLOAT, [8, 3, 3, 3], + rng.normal(0, 0.3, size=8 * 3 * 9).astype(np.float32).ravel().tolist()) + inputs = ["x", "w"] + inits = [wt] + if bias is not None: + bs = helper.make_tensor("b", TensorProto.FLOAT, [8], bias) + inputs.append("b") + inits.append(bs) + node = helper.make_node("Conv", inputs, ["y"], pads=[1, 1, 1, 1]) + g = helper.make_graph([node], "conv", [x], [y], inits) + return helper.make_model(g, opset_imports=[helper.make_opsetid("", 13)]) + + +def _gemm_model(bias): + rng = np.random.default_rng(1) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 8]) + wt = helper.make_tensor("w", TensorProto.FLOAT, [8, 4], + rng.normal(0, 0.3, size=32).astype(np.float32).ravel().tolist()) + inputs = ["x", "w"] + inits = [wt] + if bias is not None: + bs = helper.make_tensor("b", TensorProto.FLOAT, [8], bias) + inputs.append("b") + inits.append(bs) + node = helper.make_node("Gemm", inputs, ["y"], transB=1) + g = helper.make_graph([node], "gemm", [x], [y], inits) + return helper.make_model(g, opset_imports=[helper.make_opsetid("", 13)]) + + +def _run(model, in_el): + with tempfile.TemporaryDirectory() as td: + path = f"{td}/m.onnx" + onnx.save(model, path) + meta = {} + with contextlib.redirect_stdout(io.StringIO()): + rc = convert_onnx_to_riscv(path, f"{td}/m.bin", f"{td}/m.s", + benchmark=False, metadata=meta) + assert rc == 0 + with open(f"{td}/m.bin", "rb") as f: + binary = f.read() + + rng = np.random.default_rng(7) + q = np.trunc(rng.uniform(-0.25, 0.25, size=in_el) * Q).astype(np.int64) + input_bytes = struct.pack(f"<{in_el}i", *q.tolist()) + + emu = RV32EmulatorFast(mem_size_mb=128) + emu.load_unified_binary(binary, int(meta["code_bytes"])) + emu.regs[10] = 0x04000000 + emu.regs[11] = 0x05000000 + for i, b in enumerate(input_bytes): + emu.mem[0x04000000 + i] = b + emu.run(max_instr=5_000_000, uarch=PROFILES["basic"]) + out_el = int(meta["output_elements"]) + out = [emu.read_mem_i32(0x05000000 + 4 * i) for i in range(out_el)] + return [v - 0x100000000 if v >= 0x80000000 else v for v in out] + + +def test_conv_without_bias_matches_zero_bias(): + no_bias = _run(_conv_model(None), 1 * 3 * 8 * 8) + zero_bias = _run(_conv_model([0.0] * 8), 1 * 3 * 8 * 8) + assert no_bias == zero_bias + + +def test_gemm_without_bias_matches_zero_bias(): + no_bias = _run(_gemm_model(None), 4) + zero_bias = _run(_gemm_model([0.0] * 8), 4) + assert no_bias == zero_bias diff --git a/tests/test_standalone_spmm.py b/tests/test_standalone_spmm.py new file mode 100644 index 00000000..6527fc67 --- /dev/null +++ b/tests/test_standalone_spmm.py @@ -0,0 +1,115 @@ +"""Regression for the CSR SpmmCsr operator (task 3.4). + +Compiles a small CSR sparse × dense model, runs it on the RV32 emulator and +compares against an independent Q16.16 reference. Covers SpMV (N=1), SpMM +(N>1), a non-square matrix and an all-zero (empty) row. +""" + +import contextlib +import io +import struct +import tempfile + +import numpy as np +import onnx +from onnx import TensorProto, helper +import pytest + +from scratchv.standalone.onnx_to_riscv_standalone import convert_onnx_to_riscv +from scratchv.standalone.benchmark import PROFILES, RV32EmulatorFast + +Q = 65536 + + +def _q16(a): + return np.trunc(np.asarray(a, dtype=np.float64) * Q).astype(np.int64) + + +def _wrap32(v): + v &= 0xFFFFFFFF + return v - 0x100000000 if v >= 0x80000000 else v + + +def _ref_csr(values_q, col, row, b_q, m, n): + out = [[0] * n for _ in range(m)] + for i in range(m): + for j in range(int(row[i]), int(row[i + 1])): + a = int(values_q[j]) + k = int(col[j]) + for c in range(n): + prod = _wrap32(a * int(b_q[k][c])) >> 16 + out[i][c] = _wrap32(out[i][c] + prod) + return out + + +def _spmm_model(values, col, row, m, k, n): + x = helper.make_tensor_value_info("B", TensorProto.FLOAT, [k, n]) + y = helper.make_tensor_value_info("C", TensorProto.FLOAT, [m, n]) + nnz = len(values) + vals = helper.make_tensor("values", TensorProto.FLOAT, [nnz], + [float(v) for v in values]) + coli = helper.make_tensor("col", TensorProto.INT32, [nnz], [int(c) for c in col]) + rowi = helper.make_tensor("rowptr", TensorProto.INT32, [m + 1], [int(r) for r in row]) + node = helper.make_node("SpmmCsr", ["values", "col", "rowptr", "B"], ["C"]) + graph = helper.make_graph([node], "spmm", [x], [y], [vals, coli, rowi]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + + +def _run(model, b_q): + with tempfile.TemporaryDirectory() as td: + path = f"{td}/m.onnx" + onnx.save(model, path) + meta = {} + with contextlib.redirect_stdout(io.StringIO()): + rc = convert_onnx_to_riscv(path, f"{td}/m.bin", f"{td}/m.s", + benchmark=False, metadata=meta) + assert rc == 0, "compile failed" + with open(f"{td}/m.bin", "rb") as f: + binary = f.read() + + flat = [int(v) for v in np.asarray(b_q).reshape(-1)] + in_bytes = struct.pack(f"<{len(flat)}i", *flat) + emu = RV32EmulatorFast(mem_size_mb=128) + emu.load_unified_binary(binary, int(meta["code_bytes"])) + emu.regs[10] = 0x04000000 + emu.regs[11] = 0x05000000 + for i, byte in enumerate(in_bytes): + emu.mem[0x04000000 + i] = byte + emu.run(max_instr=50_000_000, uarch=PROFILES["basic"]) + out_el = int(meta["output_elements"]) + return [_wrap32(emu.read_mem_i32(0x05000000 + 4 * i)) for i in range(out_el)] + + +# A = [[v,v,0,0],[0,0,0,0],[0,v,0,0],[0,0,0,v]] (row 1 empty) +_ROW = [0, 2, 2, 3, 4] +_COL = [0, 2, 1, 3] +_VALUES = [0.25, -0.5, 0.75, 0.125] + + +def _case(m, k, n): + rng = np.random.default_rng(3) + b = _q16(rng.uniform(-0.4, 0.4, (k, n))) + model = _spmm_model(_VALUES, _COL, _ROW, m, k, n) + got = _run(model, b) + ref = _ref_csr(_q16(_VALUES), _COL, _ROW, b, m, n) + return got, [v for row in ref for v in row] + + +def test_spmv_n1(): + got, ref = _case(4, 4, 1) + assert got == ref + + +def test_spmm_n2_with_empty_row(): + got, ref = _case(4, 4, 2) + assert got == ref + + +def test_non_square(): + # uses only the first 4 rows/cols; extend B to k=6 to exercise non-square + rng = np.random.default_rng(5) + b = _q16(rng.uniform(-0.4, 0.4, (6, 3))) + model = _spmm_model(_VALUES, _COL, _ROW, 4, 6, 3) + got = _run(model, b) + ref = _ref_csr(_q16(_VALUES), _COL, _ROW, b, 4, 3) + assert got == [v for row in ref for v in row] diff --git a/tests/test_standalone_winograd.py b/tests/test_standalone_winograd.py new file mode 100644 index 00000000..1f320ebb --- /dev/null +++ b/tests/test_standalone_winograd.py @@ -0,0 +1,109 @@ +"""Winograd F(2,3) convolution regression (tasks 2.3 / 2.5). + +Verifies the generated WinogradConv against an independent direct-convolution +Q16.16 reference (stride 1, SAME). Winograd trades a little precision for +fewer multiplies, so the tolerance scales with the reduction length Cin. +""" + +import contextlib +import io +import struct +import tempfile + +import numpy as np +import onnx +from onnx import TensorProto, helper +import pytest + +from scratchv.standalone.onnx_to_riscv_standalone import convert_onnx_to_riscv +from scratchv.standalone.benchmark import PROFILES, RV32EmulatorFast + +Q = 65536 + + +def _q16(a): + return np.trunc(np.asarray(a, dtype=np.float64) * Q).astype(np.int64) + + +def _wrap32(v): + v &= 0xFFFFFFFF + return v - 0x100000000 if v >= 0x80000000 else v + + +def _ref_conv_direct(x_q, w_q, pads): + """Q16.16 direct convolution (stride 1, arbitrary symmetric pads).""" + N, C, H, W = x_q.shape + Cout = w_q.shape[0] + ph, pw = pads[0], pads[1] + xp = np.pad(x_q, ((0, 0), (0, 0), (ph, ph), (pw, pw))) + out = np.zeros((N, Cout, H, W), dtype=np.int64) + for n in range(N): + for oc in range(Cout): + for oh in range(H): + for ow in range(W): + acc = 0 + for ic in range(C): + for kh in range(3): + for kw in range(3): + prod = _wrap32(int(xp[n, ic, oh + kh, ow + kw]) + * int(w_q[oc, ic, kh, kw])) >> 16 + acc = _wrap32(acc + prod) + out[n, oc, oh, ow] = acc + return out + + +def _model(weight_f32, cin, h, w): + cout = weight_f32.shape[0] + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, cin, h, w]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, cout, h, w]) + wt = helper.make_tensor("w", TensorProto.FLOAT, list(weight_f32.shape), + weight_f32.ravel().tolist()) + node = helper.make_node("WinogradConv", ["x", "w"], ["y"], + domain="org.scratchv", pads=[1, 1, 1, 1]) + graph = helper.make_graph([node], "wino", [x], [y], [wt]) + return helper.make_model( + graph, + opset_imports=[helper.make_opsetid("org.scratchv", 1), + helper.make_opsetid("", 13)], + ) + + +def _run(model, x_q): + with tempfile.TemporaryDirectory() as td: + path = f"{td}/m.onnx" + onnx.save(model, path) + meta = {} + with contextlib.redirect_stdout(io.StringIO()): + rc = convert_onnx_to_riscv(path, f"{td}/m.bin", f"{td}/m.s", + benchmark=False, metadata=meta) + assert rc == 0, "compile failed" + with open(f"{td}/m.bin", "rb") as f: + binary = f.read() + + flat = [int(v) for v in np.asarray(x_q).reshape(-1)] + in_bytes = struct.pack(f"<{len(flat)}i", *flat) + emu = RV32EmulatorFast(mem_size_mb=128) + emu.load_unified_binary(binary, int(meta["code_bytes"])) + emu.regs[10] = 0x04000000 + emu.regs[11] = 0x05000000 + for i, byte in enumerate(in_bytes): + emu.mem[0x04000000 + i] = byte + emu.run(max_instr=100_000_000, uarch=PROFILES["basic"]) + out_el = int(meta["output_elements"]) + return np.array([_wrap32(emu.read_mem_i32(0x05000000 + 4 * i)) + for i in range(out_el)], dtype=np.int64) + + +@pytest.mark.parametrize("cin,cout,h,w", [(3, 4, 8, 8), (4, 2, 7, 7), (8, 3, 16, 16)]) +def test_winograd_matches_direct_within_tolerance(cin, cout, h, w): + rng = np.random.default_rng(0) + weight = rng.normal(0, 0.3, (cout, cin, 3, 3)).astype(np.float32) + x_q = _q16(rng.uniform(-0.25, 0.25, (1, cin, h, w))) + + got = _run(_model(weight, cin, h, w), x_q) + ref = _ref_conv_direct(x_q, _q16(weight), [1, 1]).reshape(-1) + + assert got.size == ref.size + max_lsb = int(np.abs(got - ref).max()) + tol = 16 * cin + 8 # Winograd Q16 truncation grows with the reduction + assert max_lsb <= tol, f"max_lsb={max_lsb} > tol={tol}" diff --git a/tests/test_winograd_kernel_transform.py b/tests/test_winograd_kernel_transform.py new file mode 100644 index 00000000..3f8583e0 --- /dev/null +++ b/tests/test_winograd_kernel_transform.py @@ -0,0 +1,63 @@ +"""Winograd F(2,3) compile-time kernel transform (task 2.2).""" + +import struct +import tempfile + +import numpy as np +import onnx +from onnx import TensorProto, helper +import pytest + +from scratchv.standalone.onnx_to_riscv_standalone import ( + ONNXModel, + _prepare_winograd_kernels, + winograd_f23_kernel_transform, +) + +G = np.array([[1.0, 0.0, 0.0], + [0.5, 0.5, 0.5], + [0.5, -0.5, 0.5], + [0.0, 0.0, 1.0]]) + + +def test_kernel_transform_matches_numpy(): + rng = np.random.default_rng(0) + g = rng.normal(0, 0.5, (3, 3)) + got = np.array(winograd_f23_kernel_transform(g.tolist())) + want = G @ g @ G.T + assert np.allclose(got, want, atol=1e-6) + + +def _winograd_model(weight): + cout, cin = weight.shape[0], weight.shape[1] + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, cin, 8, 8]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, cout, 8, 8]) + wt = helper.make_tensor("w", TensorProto.FLOAT, list(weight.shape), + weight.astype(np.float32).ravel().tolist()) + node = helper.make_node("WinogradConv", ["x", "w"], ["y"], pads=[1, 1, 1, 1]) + graph = helper.make_graph([node], "wino", [x], [y], [wt]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + + +def test_prepare_winograd_kernels_registers_u(): + rng = np.random.default_rng(1) + weight = rng.normal(0, 0.3, (2, 3, 3, 3)).astype(np.float32) + with tempfile.TemporaryDirectory() as td: + path = f"{td}/m.onnx" + onnx.save(_winograd_model(weight), path) + model = ONNXModel.from_file(path) + _prepare_winograd_kernels(model) + + assert "w__wino23" in model.initializers + u = model.initializers["w__wino23"] + assert u.shape == (2, 3, 4, 4) + assert u.data_type == onnx.TensorProto.FLOAT + + got = np.array(struct.unpack(f"<{2 * 3 * 16}f", u.data)).reshape(2, 3, 4, 4) + want = np.stack([[G @ weight[oc, ic] @ G.T for ic in range(3)] + for oc in range(2)]) + assert np.allclose(got, want, atol=1e-5) + + # node records the transformed tensor name + node = model.nodes[0] + assert node.attrs.get("_winograd_u") == "w__wino23" From a22f869044f886a5bfebd1825b422fabd6acb83b Mon Sep 17 00:00:00 2001 From: Seven Gao <799889633@qq.com> Date: Mon, 5 Oct 2026 00:41:52 +0800 Subject: [PATCH 5/5] feat(standalone): platform-standard assembly backend with runtime-sized kernels - --platform-asm: clang-assemblable listing (global cnn_entry, label branches, inline .word data, la gp) - size-independent kernels for Fwht / Conv2D / SpmmCsr that read all dimensions from the argument block at runtime - tests: one kernel per operator verified end-to-end over many shapes, plus fast regression cases --- .../standalone/onnx_to_riscv_standalone.py | 359 +++++++++++++++++- tests/test_platform_asm.py | 113 ++++++ tests/test_platform_kernels.py | 283 ++++++++++++++ tests/test_platform_stage2.py | 257 +++++++++++++ 4 files changed, 1011 insertions(+), 1 deletion(-) create mode 100644 tests/test_platform_asm.py create mode 100644 tests/test_platform_kernels.py create mode 100644 tests/test_platform_stage2.py diff --git a/scratchv/standalone/onnx_to_riscv_standalone.py b/scratchv/standalone/onnx_to_riscv_standalone.py index 79143eb4..121a58ad 100644 --- a/scratchv/standalone/onnx_to_riscv_standalone.py +++ b/scratchv/standalone/onnx_to_riscv_standalone.py @@ -3223,6 +3223,354 @@ def _gen_winogradconv(self, node: NodeInfo) -> None: # ═══════════════════════════════════════════════════════════════════════════ +def _emit_platform_fwht() -> str: + """Size-independent FWHT kernel (platform ABI: a0=in, a1=out, a2=N).""" + lines = [ + "# ScratchV platform kernel: FWHT (size-independent, a2 = N)", + " .option norelax", + " .text", + " .globl cnn_entry", + " .type cnn_entry, @function", + "cnn_entry:", + " mv t0, a0 # src = in", + " mv t1, a1 # dst = out", + " mv t2, a2 # N", + " mv t3, zero # i", + "_p_fwht_copy:", + " lw t4, 0(t0)", + " sw t4, 0(t1)", + " addi t0, t0, 4", + " addi t1, t1, 4", + " addi t3, t3, 1", + " blt t3, t2, _p_fwht_copy", + " li t5, 1 # length = 1", + "_p_fwht_len:", + " bge t5, t2, _p_fwht_done", + " slli t6, t5, 1 # step = 2*length", + " li s2, 0 # i = block start", + "_p_fwht_i:", + " bge s2, t2, _p_fwht_next_len", + " li s3, 0 # j = 0", + "_p_fwht_j:", + " bge s3, t5, _p_fwht_next_block", + " add s4, s2, s3 # idx = i + j", + " slli s5, s4, 2", + " add s6, a1, s5 # &a[idx]", + " add s7, s4, t5 # idx + length", + " slli s7, s7, 2", + " add s8, a1, s7 # &a[idx+length]", + " lw s9, 0(s6) # u", + " lw s10, 0(s8) # v", + " add s11, s9, s10 # u+v", + " sw s11, 0(s6)", + " sub s11, s9, s10 # u-v", + " sw s11, 0(s8)", + " addi s3, s3, 1", + " j _p_fwht_j", + "_p_fwht_next_block:", + " add s2, s2, t6", + " j _p_fwht_i", + "_p_fwht_next_len:", + " slli t5, t5, 1", + " j _p_fwht_len", + "_p_fwht_done:", + " ret", + " .size cnn_entry, . - cnn_entry", + ] + return "\n".join(lines) + "\n" + + +def _emit_platform_conv() -> str: + """Size-independent direct Conv2D kernel (platform ABI). + + a0 = input (NCHW), a1 = output (NCHW), a2 = param block + (N,C_in,H,W,C_out,K,pad,stride,H_out,W_out,w_ptr,b_ptr). + Semantics per docs/输入规范.md section 3.5. + """ + lines = [ + "# ScratchV platform kernel: direct Conv2D (size-independent)", + " .option norelax", + " .text", + " .globl cnn_entry", + " .type cnn_entry, @function", + "cnn_entry:", + " lw s0, 0(a2) # N", + " lw s1, 4(a2) # C_in", + " lw s2, 8(a2) # H", + " lw s3, 12(a2) # W", + " lw s4, 16(a2) # C_out", + " lw s5, 20(a2) # K", + " lw s6, 24(a2) # pad", + " lw s7, 32(a2) # H_out", + " lw s8, 36(a2) # W_out", + " lw s9, 40(a2) # w_ptr", + " lw s10, 44(a2) # b_ptr", + " mv s11, a0 # in base", + " mv a7, a1 # out base", + " li t0, 0 # n = 0", + "_p_conv_n:", + " bge t0, s0, _p_conv_done", + " li t1, 0 # oc = 0", + "_p_conv_oc:", + " bge t1, s4, _p_conv_next_n", + " li t2, 0 # oh = 0", + "_p_conv_oh:", + " bge t2, s7, _p_conv_next_oc", + " li t3, 0 # ow = 0", + "_p_conv_ow:", + " bge t3, s8, _p_conv_next_oh", + " li t6, 0 # acc = 0", + " beq s10, zero, _p_conv_nobias", + " slli a3, t1, 2", + " add a3, s10, a3", + " lw t6, 0(a3) # acc = bias[oc]", + "_p_conv_nobias:", + " li t4, 0 # ic = 0", + "_p_conv_ic:", + " bge t4, s1, _p_conv_store", + " li t5, 0 # kh = 0", + "_p_conv_kh:", + " bge t5, s5, _p_conv_next_ic", + " li a6, 0 # kw = 0", + "_p_conv_kw:", + " bge a6, s5, _p_conv_next_kh", + " add a3, t2, t5", + " sub a3, a3, s6 # ih", + " add a4, t3, a6", + " sub a4, a4, s6 # iw", + " blt a3, zero, _p_conv_kw_skip", + " bge a3, s2, _p_conv_kw_skip", + " blt a4, zero, _p_conv_kw_skip", + " bge a4, s3, _p_conv_kw_skip", + " mul a5, t0, s1", + " add a5, a5, t4", + " mul a5, a5, s2", + " add a5, a5, a3", + " mul a5, a5, s3", + " add a5, a5, a4", + " slli a5, a5, 2", + " add a5, s11, a5", + " lw a5, 0(a5) # in value", + " mul a4, t1, s1", + " add a4, a4, t4", + " mul a4, a4, s5", + " add a4, a4, t5", + " mul a4, a4, s5", + " add a4, a4, a6", + " slli a4, a4, 2", + " add a4, s9, a4", + " lw a4, 0(a4) # weight value", + " mul a5, a5, a4", + " srai a5, a5, 16", + " add t6, t6, a5 # acc += prod", + "_p_conv_kw_skip:", + " addi a6, a6, 1", + " j _p_conv_kw", + "_p_conv_next_kh:", + " addi t5, t5, 1", + " j _p_conv_kh", + "_p_conv_next_ic:", + " addi t4, t4, 1", + " j _p_conv_ic", + "_p_conv_store:", + " mul a3, t0, s4", + " add a3, a3, t1", + " mul a3, a3, s7", + " add a3, a3, t2", + " mul a3, a3, s8", + " add a3, a3, t3", + " slli a3, a3, 2", + " add a3, a7, a3", + " sw t6, 0(a3)", + " addi t3, t3, 1", + " j _p_conv_ow", + "_p_conv_next_oh:", + " addi t2, t2, 1", + " j _p_conv_oh", + "_p_conv_next_oc:", + " addi t1, t1, 1", + " j _p_conv_oc", + "_p_conv_next_n:", + " addi t0, t0, 1", + " j _p_conv_n", + "_p_conv_done:", + " ret", + " .size cnn_entry, . - cnn_entry", + ] + return "\n".join(lines) + "\n" + + +def _emit_platform_spmm() -> str: + """Size-independent CSR SpMM kernel (platform ABI). + + a0 = CSR block (values_ptr, col_ptr, row_ptr), a1 = out (M*N), a2 = params + (M, K, N, nnz, B_ptr). Semantics per docs/输入规范.md section 4.5. + """ + lines = [ + "# ScratchV platform kernel: CSR SpMM (size-independent)", + " .option norelax", + " .text", + " .globl cnn_entry", + " .type cnn_entry, @function", + "cnn_entry:", + " lw t2, 0(a0) # values_ptr", + " lw t3, 4(a0) # col_ptr", + " lw t4, 8(a0) # row_ptr", + " lw t5, 0(a2) # M", + " lw t6, 8(a2) # N", + " lw s2, 16(a2) # B_ptr", + " mv a0, t2 # values", + " mv a2, t3 # col", + " mv a3, t4 # row_ptr", + " mv a4, t5 # M", + " mv a5, t6 # N", + " mv a6, s2 # B", + " li s3, 0 # i = 0", + "_p_spmm_i:", + " bge s3, a4, _p_spmm_done", + " slli t0, s3, 2", + " add t0, a3, t0 # &row_ptr[i]", + " lw s5, 0(t0) # j = row_ptr[i]", + " lw s6, 4(t0) # j_end = row_ptr[i+1]", + " li s4, 0 # n = 0", + "_p_spmm_zero:", + " mul t0, s3, a5", + " add t0, t0, s4", + " slli t0, t0, 2", + " add t0, a1, t0", + " sw zero, 0(t0)", + " addi s4, s4, 1", + " blt s4, a5, _p_spmm_zero", + "_p_spmm_j:", + " bge s5, s6, _p_spmm_next_i", + " slli t1, s5, 2", + " add t1, a0, t1", + " lw t1, 0(t1) # a = values[j]", + " slli t0, s5, 2", + " add t0, a2, t0", + " lw t0, 0(t0) # k = col[j]", + " li s4, 0 # n = 0", + "_p_spmm_n:", + " mul t2, t0, a5", + " add t2, t2, s4", + " slli t2, t2, 2", + " add t2, a6, t2", + " lw t2, 0(t2) # B[k,n]", + " mul t2, t1, t2", + " srai t2, t2, 16", + " mul t3, s3, a5", + " add t3, t3, s4", + " slli t3, t3, 2", + " add t3, a1, t3", + " lw t4, 0(t3)", + " add t4, t4, t2", + " sw t4, 0(t3)", + " addi s4, s4, 1", + " blt s4, a5, _p_spmm_n", + " addi s5, s5, 1", + " j _p_spmm_j", + "_p_spmm_next_i:", + " addi s3, s3, 1", + " j _p_spmm_i", + "_p_spmm_done:", + " ret", + " .size cnn_entry, . - cnn_entry", + ] + return "\n".join(lines) + "\n" + + +def emit_platform_asm(generator, weight_data: bytes, model=None) -> str: + """Render a platform-standard, clang-assemblable listing (``--platform-asm``). + + If *model* is a single size-independent-supported problem node (currently + FWHT), emit the dedicated runtime-sized kernel; otherwise emit the + fixed-shape standard listing (B-1). + """ + if model is not None and len(model.nodes) == 1 and model.nodes[0].op_type == "Fwht": + return _emit_platform_fwht() + if model is not None and len(model.nodes) == 1 and model.nodes[0].op_type == "SpmmCsr": + return _emit_platform_spmm() + if model is not None and len(model.nodes) == 1 and model.nodes[0].op_type == "Conv": + return _emit_platform_conv() + return _emit_platform_listing(generator, weight_data) + + +def _emit_platform_listing(generator, weight_data: bytes) -> str: + """Fixed-shape platform-standard listing (B-1). + + Differences from the default debug listing: + - global entry symbol ``cnn_entry`` (not ``_start``); + - branch/jump targets are real labels (not numeric offsets); + - the ``_init_data_base`` placeholder pair becomes ``la gp, __data_start``; + - weights/constants are emitted inline as ``.word`` in a ``.data`` section + (no ``.incbin`` / ``.include``), preceded by ``__data_start``. + + Platform contract: a0 = input base, a1 = output base, a2 = size scalar; + sp is already set by the caller. + """ + import re as _re + + emit = generator.emit + dis = emit.disassemble(symbolic=True).splitlines() + + # Prefer the emitter's own label names; else a generated .Lpc_. + idx_to_name: dict[int, str] = {} + for name, idx in emit.labels.items(): + if idx not in idx_to_name: + idx_to_name[idx] = name + + def label_for(idx: int) -> str: + if idx == 0: + return "cnn_entry" + return idx_to_name.get(idx, f".Lpc_{idx}") + + out: list[str] = [] + for line in dis: + stripped = line.strip() + if stripped.endswith(":") and stripped[:-1].startswith(".Lsv_"): + idx = int(stripped[len(".Lsv_"):-1]) + out.append(f"{label_for(idx)}:") + else: + out.append(_re.sub(r"\.Lsv_(\d+)", lambda m: label_for(int(m.group(1))), line)) + + # Replace the auipc/addi placeholder pair with a standard la. + fixed: list[str] = [] + skip_next = False + for line in out: + if skip_next: + skip_next = False + continue + if "auipc gp, 0x0" in line: + indent = line[: len(line) - len(line.lstrip())] + fixed.append(f"{indent}la gp, __data_start # data section base") + skip_next = True # drop the following addi/mv placeholder + continue + fixed.append(line) + + header = [ + "# Generated by ScratchV standalone --platform-asm", + "# Platform contract: a0=input base, a1=output base, a2=size; sp preset.", + "# Assemble: clang -march=rv32im -mabi=ilp32 -nostdlib -c this.s", + "", + " .option norelax", + " .text", + " .globl cnn_entry", + " .type cnn_entry, @function", + ] + tail = [" .size cnn_entry, . - cnn_entry", ""] + + data_lines = ["", " .data", " .balign 4", "__data_start:"] + if weight_data: + words = struct.unpack(f"<{len(weight_data) // 4}I", weight_data) + for i in range(0, len(words), 8): + chunk = ", ".join(f"0x{w:08x}" for w in words[i:i + 8]) + data_lines.append(f" .word {chunk}") + else: + data_lines.append(" .word 0") + + return "\n".join(header + fixed + tail + data_lines) + "\n" + + def patch_gp_data_base( generator: CNNRISCVGenerator, code_bytes: bytes, @@ -3395,6 +3743,7 @@ def convert_onnx_to_riscv( metadata: dict | None = None, llvm_mca: str | None = None, symbolic_asm: bool = False, + platform_asm: bool = False, ) -> int: """Full pipeline: ONNX model → RISC-V RV32IM binary. @@ -3590,7 +3939,10 @@ def count_listing_instructions(asm_text: str) -> int: print(f" Binary: {output_bin} ({len(binary):,} bytes)") # Disassembly for verification - asm_text = generator.emit.disassemble(symbolic=schedule or symbolic_asm) + if platform_asm: + asm_text = emit_platform_asm(generator, weight_data, model) + else: + asm_text = generator.emit.disassemble(symbolic=schedule or symbolic_asm) asm_path = output_asm or output_bin.replace(".bin", ".s") with open(asm_path, "w") as f: f.write(asm_text) @@ -3777,6 +4129,10 @@ def main() -> int: help="Schedule physical-register instructions before writing the binary") parser.add_argument("--symbolic-asm", action="store_true", help="Emit reassemblable symbolic targets for scheduling analysis") + parser.add_argument("--platform-asm", action="store_true", + help="Emit platform-standard assembly: global entry cnn_entry, " + "label branch targets, inline .word weights, la gp (clang-" + "assemblable, -march=rv32im -nostdlib)") parser.add_argument( "--uarch", default="basic", choices=["single", "fast", "basic", "slow"], help="Microarchitecture profile for cycle-accurate emulation: " @@ -3818,6 +4174,7 @@ def main() -> int: const_merge=args.const_merge, schedule=args.schedule, symbolic_asm=args.symbolic_asm, + platform_asm=args.platform_asm, ) diff --git a/tests/test_platform_asm.py b/tests/test_platform_asm.py new file mode 100644 index 00000000..039ba3c7 --- /dev/null +++ b/tests/test_platform_asm.py @@ -0,0 +1,113 @@ +"""--platform-asm output shape (platform contract; no RISC-V toolchain needed).""" + +import contextlib +import io +import tempfile + +import onnx +from onnx import TensorProto, helper +import pytest + +from scratchv.standalone.onnx_to_riscv_standalone import ( + ONNXModel, + CNNRISCVGenerator, + MemoryPlan, + emit_platform_asm, +) + + +def _conv_model(): + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 8, 8]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 8, 8, 8]) + import numpy as np + rng = np.random.default_rng(0) + w = helper.make_tensor("w", TensorProto.FLOAT, [8, 3, 3, 3], + rng.normal(0, 0.3, size=8 * 3 * 9).astype(np.float32).ravel().tolist()) + node = helper.make_node("Conv", ["x", "w"], ["y"], pads=[1, 1, 1, 1]) + graph = helper.make_graph([node], "c", [x], [y], [w]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + + +def _build_platform_asm(tmp_path): + path = tmp_path / "m.onnx" + onnx.save(_conv_model(), str(path)) + model = ONNXModel.from_file(str(path)) + mem = MemoryPlan() + for vi in model.inputs: + n = 1 + for d in model.get_shape(vi.name): + n *= d + mem.alloc_workspace(vi.name, n) + weight_data = mem.layout_weights(model.initializers) + gen = CNNRISCVGenerator(model, mem) + with contextlib.redirect_stdout(io.StringIO()): + gen.generate() + return emit_platform_asm(gen, weight_data) + + +def _fwht_model(): + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 64]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 64]) + node = helper.make_node("Fwht", ["X"], ["Y"], domain="org.scratchv") + graph = helper.make_graph([node], "fwht", [x], [y]) + return helper.make_model( + graph, + opset_imports=[helper.make_opsetid("org.scratchv", 1), + helper.make_opsetid("", 13)], + ) + + +def test_platform_asm_is_size_independent_for_fwht(tmp_path): + from scratchv.standalone.onnx_to_riscv_standalone import emit_platform_asm + + path = tmp_path / "fwht.onnx" + onnx.save(_fwht_model(), str(path)) + model = ONNXModel.from_file(str(path)) + mem = MemoryPlan() + for vi in model.inputs: + n = 1 + for d in model.get_shape(vi.name): + n *= d + mem.alloc_workspace(vi.name, n) + weight_data = mem.layout_weights(model.initializers) + gen = CNNRISCVGenerator(model, mem) + with contextlib.redirect_stdout(io.StringIO()): + gen.generate() + + asm = emit_platform_asm(gen, weight_data, model) + assert ".globl cnn_entry" in asm + # runtime-sized: bounds come from a2, not a compile-time constant + assert "mv t2, a2" in asm + assert "blt t3, t2, _p_fwht_copy" in asm # copy uses runtime N + # no hard-coded 64 anywhere as a loop bound + assert "li t2, 64" not in asm + + +def test_platform_asm_contract(tmp_path): + asm = _build_platform_asm(tmp_path) + + # entry symbol is global cnn_entry, and old _start is gone + assert ".globl cnn_entry" in asm + assert "\ncnn_entry:" in asm + assert not any(l.strip() == "_start:" for l in asm.splitlines()) + + # no forbidden directives + assert ".incbin" not in asm + assert ".include" not in asm + + # no unresolved placeholder / symbolic temp labels leak + assert "auipc gp, 0x0" not in asm + assert ".Lsv_" not in asm + assert "la gp, __data_start" in asm + + # branch/jump targets use labels, not numeric offsets + import re + for line in asm.splitlines(): + code = line.split("#", 1)[0] + if re.match(r"\s*(bne|beq|blt|bge|j|jal)\b", code): + assert not re.search(r",\s*-?\d+\s*$", code), f"numeric target: {line}" + + # inline data section with .word, no incbin + assert "__data_start:" in asm + assert ".word" in asm + assert asm.count(".word") > 10 diff --git a/tests/test_platform_kernels.py b/tests/test_platform_kernels.py new file mode 100644 index 00000000..796dd70b --- /dev/null +++ b/tests/test_platform_kernels.py @@ -0,0 +1,283 @@ +"""Size-independent platform kernels: FWHT and CSR SpMM. + +The generated ``.s`` is executed by a tiny in-test RV32IM interpreter (only the +instruction subset the kernels use), so the kernels' logic is verified without +a RISC-V toolchain. +""" + +import io +import contextlib +import re + +import numpy as np +import pytest + + +# ── Minimal RV32IM interpreter (subset used by the kernels) ──────────────── +_RN = {"zero": 0, "ra": 1, "sp": 2, "gp": 3, "tp": 4, "t0": 5, "t1": 6, + "t2": 7, "s0": 8, "s1": 9, "a0": 10, "a1": 11, "a2": 12, "a3": 13, + "a4": 14, "a5": 15, "a6": 16, "a7": 17, "s2": 18, "s3": 19, "s4": 20, + "s5": 21, "s6": 22, "s7": 23, "s8": 24, "s9": 25, "s10": 26, "s11": 27, + "t3": 28, "t4": 29, "t5": 30, "t6": 31} + + +def _parse(asm): + prog, labels = [], {} + for raw in asm.splitlines(): + line = raw.split("#", 1)[0].strip() + if not line or line.startswith("."): + continue + if line.endswith(":"): + labels[line[:-1]] = len(prog) + continue + prog.append(line) + return prog, labels + + +def _s32(v): + v &= 0xFFFFFFFF + return v - 0x100000000 if v >= 0x80000000 else v + + +def _run(prog, labels, regs, mem, max_steps=10_000_000): + pc = steps = 0 + while pc < len(prog) and steps < max_steps: + steps += 1 + p = prog[pc].replace(",", " ").split() + op, r = p[0], p[1:] + nxt = pc + 1 + g = lambda i: regs[i] + st = lambda i, v: regs.__setitem__(i, v & 0xFFFFFFFF) + + if op == "li": + st(_RN[r[0]], int(r[1])) + elif op == "mv": + st(_RN[r[0]], g(_RN[r[1]])) + elif op in ("lw", "sw"): + off = int(r[1][:r[1].index("(")]) + bs = r[1][r[1].index("(") + 1:-1] + addr = (g(_RN[bs]) + off) & 0xFFFFFFFF + if op == "lw": + st(_RN[r[0]], int.from_bytes(mem[addr:addr + 4], "little")) + else: + mem[addr:addr + 4] = int(g(_RN[r[0]])).to_bytes(4, "little") + elif op == "add": + st(_RN[r[0]], g(_RN[r[1]]) + g(_RN[r[2]])) + elif op == "sub": + st(_RN[r[0]], g(_RN[r[1]]) - g(_RN[r[2]])) + elif op == "mul": + st(_RN[r[0]], g(_RN[r[1]]) * g(_RN[r[2]])) + elif op == "srai": + st(_RN[r[0]], _s32(g(_RN[r[1]])) >> int(r[2])) + elif op == "slli": + st(_RN[r[0]], g(_RN[r[1]]) << int(r[2])) + elif op == "addi": + st(_RN[r[0]], g(_RN[r[1]]) + int(r[2])) + elif op in ("blt", "bge", "beq", "bne"): + a, b = _s32(g(_RN[r[0]])), _s32(g(_RN[r[1]])) + cond = {"blt": a < b, "bge": a >= b, + "beq": a == b, "bne": a != b}[op] + nxt = labels[r[2]] if cond else nxt + elif op == "j": + nxt = labels[r[0]] + elif op == "ret": + return steps + else: + raise AssertionError(f"unknown op {op}") + pc = nxt + return steps + + +def _asm_for(model_builder, tmp_path): + import onnx + from scratchv.standalone.onnx_to_riscv_standalone import ( + ONNXModel, CNNRISCVGenerator, MemoryPlan, emit_platform_asm, + ) + path = tmp_path / "m.onnx" + onnx.save(model_builder(), str(path)) + model = ONNXModel.from_file(str(path)) + mem = MemoryPlan() + for vi in model.inputs: + n = 1 + for d in model.get_shape(vi.name): + n *= d + mem.alloc_workspace(vi.name, n) + wd = mem.layout_weights(model.initializers) + gen = CNNRISCVGenerator(model, mem) + with contextlib.redirect_stdout(io.StringIO()): + gen.generate() + return emit_platform_asm(gen, wd, model) + + +def _fwht_model(): + from onnx import TensorProto, helper + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 64]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 64]) + node = helper.make_node("Fwht", ["X"], ["Y"], domain="org.scratchv") + return helper.make_model( + helper.make_graph([node], "f", [x], [y]), + opset_imports=[helper.make_opsetid("org.scratchv", 1), + helper.make_opsetid("", 13)]) + + +def _spmm_model(): + from onnx import TensorProto, helper + # empty weight arrays are fine; the platform kernel reads runtime memory + x = helper.make_tensor_value_info("B", TensorProto.FLOAT, [4, 2]) + y = helper.make_tensor_value_info("C", TensorProto.FLOAT, [4, 2]) + vals = helper.make_tensor("values", TensorProto.FLOAT, [4], [0.0] * 4) + col = helper.make_tensor("col", TensorProto.INT32, [4], [0, 1, 2, 3]) + row = helper.make_tensor("rowptr", TensorProto.INT32, [5], [0, 1, 2, 3, 4]) + node = helper.make_node("SpmmCsr", ["values", "col", "rowptr", "B"], ["C"]) + return helper.make_model( + helper.make_graph([node], "s", [x], [y], [vals, col, row]), + opset_imports=[helper.make_opsetid("", 13)]) + + +@pytest.mark.parametrize("n", [8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]) +def test_platform_fwht_kernel_matches_reference(tmp_path, n): + asm = _asm_for(_fwht_model, tmp_path) + prog, labels = _parse(asm) + + Q = 65536 + rng = np.random.default_rng(0) + x = [int(np.trunc(v * Q)) for v in rng.uniform(-0.4, 0.4, n)] + mem = bytearray(0x200000) + inp, out = 0x10000, 0x80000 + for i, v in enumerate(x): + mem[inp + 4 * i:inp + 4 * i + 4] = int(v & 0xFFFFFFFF).to_bytes(4, "little") + regs = [0] * 32 + regs[10], regs[11], regs[12] = inp, out, n + _run(prog, labels, regs, mem) + + got = [_s32(int.from_bytes(mem[out + 4 * i:out + 4 * i + 4], "little")) + for i in range(n)] + a = [float(v) / Q for v in x] + length = 1 + while length < n: + for i in range(0, n, 2 * length): + for j in range(length): + u, v = a[i + j], a[i + j + length] + a[i + j], a[i + j + length] = u + v, u - v + length <<= 1 + ref = [int(np.trunc(v * Q)) for v in a] + assert got == ref + + +def test_platform_spmm_kernel_matches_reference(tmp_path): + asm = _asm_for(_spmm_model, tmp_path) + prog, labels = _parse(asm) + + Q = 65536 + M, K, N = 4, 4, 2 + values = [0.25, -0.5, 0.75, 0.125] + col = [0, 2, 1, 3] + row = [0, 2, 2, 3, 4] + rng = np.random.default_rng(1) + B = rng.uniform(-0.4, 0.4, (K, N)) + vq = [int(np.trunc(v * Q)) for v in values] + Bq = [[int(np.trunc(x * Q)) for x in r] for r in B] + + mem = bytearray(0x40000) + w = lambda a, v: mem.__setitem__(slice(a, a + 4), + int(v & 0xFFFFFFFF).to_bytes(4, "little")) + r32 = lambda a: _s32(int.from_bytes(mem[a:a + 4], "little")) + block, va, ca, ra, ba, out, par = (0x1000, 0x10000, 0x10100, 0x10200, + 0x10300, 0x10400, 0x11000) + w(block, va); w(block + 4, ca); w(block + 8, ra) + for i, v in enumerate(vq): + w(va + 4 * i, v) + for i, v in enumerate(col): + w(ca + 4 * i, v) + for i, v in enumerate(row): + w(ra + 4 * i, v) + for kk in range(K): + for nn in range(N): + w(ba + 4 * (kk * N + nn), Bq[kk][nn]) + w(par + 0, M); w(par + 4, K); w(par + 8, N); w(par + 12, len(vq)); w(par + 16, ba) + + regs = [0] * 32 + regs[10], regs[11], regs[12] = block, out, par + _run(prog, labels, regs, mem) + + C = [[0] * N for _ in range(M)] + for i in range(M): + for j in range(row[i], row[i + 1]): + a, kk = vq[j], col[j] + for nn in range(N): + prod = _s32((a * Bq[kk][nn]) & 0xFFFFFFFF) >> 16 + C[i][nn] = _s32(C[i][nn] + prod) + got = [r32(out + 4 * (i * N + nn)) for i in range(M) for nn in range(N)] + ref = [C[i][nn] for i in range(M) for nn in range(N)] + assert got == ref + + +def _conv_model(cin=3, h=8, w=8, k=3, cout=8): + from onnx import TensorProto, helper + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, cin, h, w]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, cout, h, w]) + wt = helper.make_tensor("w", TensorProto.FLOAT, [cout, cin, k, k], + [0.0] * (cout * cin * k * k)) + b = helper.make_tensor("b", TensorProto.FLOAT, [cout], [0.0] * cout) + pad = k // 2 + node = helper.make_node("Conv", ["x", "w", "b"], ["y"], + pads=[pad, pad, pad, pad]) + return helper.make_model( + helper.make_graph([node], "c", [x], [y], [wt, b]), + opset_imports=[helper.make_opsetid("", 13)]) + + +@pytest.mark.parametrize("cin,h,w,k,cout", [(3, 8, 8, 3, 8), (3, 16, 16, 5, 8)]) +def test_platform_conv_kernel_matches_reference(tmp_path, cin, h, w, k, cout): + asm = _asm_for(lambda: _conv_model(cin, h, w, k, cout), tmp_path) + prog, labels = _parse(asm) + + Q = 65536 + rng = np.random.default_rng(0) + Nn, Cin, H, W = 1, cin, h, w + Cout, K, pad = cout, k, k // 2 + Hout = Wout = (H + 2 * pad - K) + 1 + xin = rng.uniform(-0.4, 0.4, (Nn, Cin, H, W)) + wt = rng.uniform(-0.4, 0.4, (Cout, Cin, K, K)) + bias = rng.uniform(-0.4, 0.4, (Cout,)) + xq = np.trunc(xin * Q).astype(np.int64) + wq = np.trunc(wt * Q).astype(np.int64) + bq = np.trunc(bias * Q).astype(np.int64) + + mem = bytearray(0x800000) + w = lambda a, v: mem.__setitem__(slice(a, a + 4), + int(v & 0xFFFFFFFF).to_bytes(4, "little")) + r32 = lambda a: _s32(int.from_bytes(mem[a:a + 4], "little")) + inp, out, wp, bp, par = 0x10000, 0x80000, 0x100000, 0x120000, 0x140000 + + def put(base, arr): + for i, v in enumerate(np.asarray(arr).reshape(-1)): + w(base + 4 * i, int(v)) + put(inp, xq) + put(wp, wq) + put(bp, bq) + w(par + 0, Nn); w(par + 4, Cin); w(par + 8, H); w(par + 12, W) + w(par + 16, Cout); w(par + 20, K); w(par + 24, pad); w(par + 28, 1) + w(par + 32, Hout); w(par + 36, Wout); w(par + 40, wp); w(par + 44, bp) + + regs = [0] * 32 + regs[10], regs[11], regs[12] = inp, out, par + _run(prog, labels, regs, mem) + + got = [r32(out + 4 * i) for i in range(Nn * Cout * Hout * Wout)] + # reference: per-product arithmetic shift then 32-bit accumulate (matches kernel) + xp = np.zeros((Cin, H + 2 * pad, W + 2 * pad), dtype=np.int64) + xp[:, pad:pad + H, pad:pad + W] = xq[0] + ref = [] + for oc in range(Cout): + for oh in range(Hout): + for ow in range(Wout): + acc = int(bq[oc]) + for ic in range(Cin): + for kh in range(K): + for kw in range(K): + p = _s32((int(xp[ic, oh + kh, ow + kw]) + * int(wq[oc, ic, kh, kw])) & 0xFFFFFFFF) >> 16 + acc = _s32(acc + p) + ref.append(acc) + assert got == ref diff --git a/tests/test_platform_stage2.py b/tests/test_platform_stage2.py new file mode 100644 index 00000000..edda2ca7 --- /dev/null +++ b/tests/test_platform_stage2.py @@ -0,0 +1,257 @@ +"""Full Stage-2 (30 data point) verification of the size-independent +``--platform-asm`` kernels for the three contest problems. + +Each problem's 10 data-point shapes (from ``stage2数据点.md`` v1.2) are +generated on the fly, the emitted ``.s`` is executed by the tiny RV32IM +interpreter from ``test_platform_kernels.py``, and the result is compared +against a NumPy reference. This is the end-to-end check that one +size-independent ``.s`` per problem covers every data point. + +Skipped by default: the largest SpMM points take tens of seconds each under +the pure-Python interpreter (~3 min total). Run explicitly with:: + + SCRATCHV_PLATFORM_STAGE2=1 PYTHONPATH=. python3.11 \ + -m pytest tests/test_platform_stage2.py -q +""" + +import contextlib +import importlib.util +import io +import os +import pathlib + +import numpy as np +import pytest + +pytestmark = pytest.mark.skipif( + os.environ.get("SCRATCHV_PLATFORM_STAGE2") != "1", + reason="set SCRATCHV_PLATFORM_STAGE2=1 to run the full 30-point " + "platform-kernel check (~3 min)", +) + +_here = pathlib.Path(__file__).resolve().parent +_spec = importlib.util.spec_from_file_location( + "_tpk", _here / "test_platform_kernels.py") +_tpk = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_tpk) + +_s32 = _tpk._s32 +Q = 65536 +_MASK = 0xFFFFFFFF + +# ── Stage-2 data-point specs (stage2数据点.md v1.2) ───────────────────────── +FWHT = [8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096] +CONV = [(1, 3, 16, 16, 3, 8), (1, 8, 16, 16, 3, 16), (1, 8, 8, 8, 3, 16), + (1, 8, 12, 12, 3, 16), (1, 6, 12, 12, 3, 12), (1, 3, 16, 16, 5, 8), + (1, 4, 12, 12, 5, 8), (1, 8, 8, 8, 3, 8), (1, 8, 6, 6, 3, 16), + (1, 4, 8, 8, 3, 8)] +SPMM = [(512, 512, 1, 0.90), (1024, 1024, 1, 0.95), (256, 512, 4, 0.85), + (1024, 1024, 1, 0.98), (1280, 128, 16, 0.99), (512, 1024, 1, 0.90), + (256, 256, 4, 0.80), (256, 2048, 1, 0.95), (512, 512, 2, 0.70), + (256, 256, 8, 0.50)] + + +def _q(a): + return np.trunc(np.asarray(a, np.float64) * Q).astype(np.int64) + + +def _emit(model, tmp_path): + import onnx + from scratchv.standalone.onnx_to_riscv_standalone import ( + ONNXModel, CNNRISCVGenerator, MemoryPlan, emit_platform_asm) + path = tmp_path / "m.onnx" + onnx.save(model, str(path)) + m = ONNXModel.from_file(str(path)) + mem = MemoryPlan() + for vi in m.inputs: + n = 1 + for d in m.get_shape(vi.name): + n *= d + mem.alloc_workspace(vi.name, n) + wd = mem.layout_weights(m.initializers) + gen = CNNRISCVGenerator(m, mem) + with contextlib.redirect_stdout(io.StringIO()): + gen.generate() + return emit_platform_asm(gen, wd, m) + + +def _run(asm, regs_init, mem, max_steps=400_000_000): + prog, labels = _tpk._parse(asm) + regs = [0] * 32 + for i, v in regs_init.items(): + regs[i] = v + _tpk._run(prog, labels, regs, mem, max_steps=max_steps) + + +# ── builders ────────────────────────────────────────────────────────────── +def _fwht_model(n): + from onnx import TensorProto, helper + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, n]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, n]) + node = helper.make_node("Fwht", ["X"], ["Y"], domain="org.scratchv") + return helper.make_model( + helper.make_graph([node], "f", [x], [y]), + opset_imports=[helper.make_opsetid("org.scratchv", 1), + helper.make_opsetid("", 13)]) + + +def _conv_model(N, C, H, W, K, Cout): + from onnx import TensorProto, helper + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [N, C, H, W]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [N, Cout, H, W]) + rng = np.random.default_rng(abs(hash((N, C, H, W, K, Cout))) % 2**32) + wt = helper.make_tensor( + "w", TensorProto.FLOAT, [Cout, C, K, K], + rng.normal(0, 0.3, size=Cout * C * K * K).astype(np.float32).ravel().tolist()) + bs = helper.make_tensor( + "b", TensorProto.FLOAT, [Cout], + rng.normal(0, 0.3, size=Cout).astype(np.float32).ravel().tolist()) + pad = K // 2 + node = helper.make_node("Conv", ["x", "w", "b"], ["y"], + pads=[pad, pad, pad, pad]) + return helper.make_model( + helper.make_graph([node], "c", [x], [y], [wt, bs]), + opset_imports=[helper.make_opsetid("", 13)]) + + +def _spmm_model(M, K, N, dens, seed): + from onnx import TensorProto, helper + rng = np.random.default_rng(seed) + mask = rng.random((M, K)) < dens + A = (mask * rng.uniform(-0.4, 0.4, (M, K))).astype(np.float32) + vals, col, row = [], [], [0] + for i in range(M): + for j in np.nonzero(A[i])[0]: + col.append(int(j)); vals.append(float(A[i, j])) + row.append(len(col)) + B = rng.uniform(-0.4, 0.4, (K, N)).astype(np.float32) + x = helper.make_tensor_value_info("B", TensorProto.FLOAT, [K, N]) + y = helper.make_tensor_value_info("C", TensorProto.FLOAT, [M, N]) + v = helper.make_tensor("values", TensorProto.FLOAT, [len(vals)], vals) + c = helper.make_tensor("col", TensorProto.INT32, [len(col)], col) + r = helper.make_tensor("rowptr", TensorProto.INT32, [M + 1], row) + node = helper.make_node("SpmmCsr", ["values", "col", "rowptr", "B"], ["C"]) + return helper.make_model( + helper.make_graph([node], "s", [x], [y], [v, c, r]), + opset_imports=[helper.make_opsetid("", 13)]) + + +# ── reference helpers (per-product >>16 then accumulate, matching kernel) ── +def _fwht_ref(x): + a = list(x) + n = len(a) + length = 1 + while length < n: + for i in range(0, n, 2 * length): + for j in range(length): + u, v = a[i + j], a[i + j + length] + a[i + j], a[i + j + length] = _s32(u + v), _s32(u - v) + length <<= 1 + return a + + +def _conv_ref(xq, wq, bq, pad, K): + C, H, W = xq.shape + Cout, _, _, _ = wq.shape + Hout, Wout = H, W + xp = np.zeros((C, H + 2 * pad, W + 2 * pad), np.int64) + xp[:, pad:pad + H, pad:pad + W] = xq + acc = np.zeros((Cout, Hout, Wout), np.int64) + for ic in range(C): + for kh in range(K): + for kw in range(K): + patch = xp[ic, kh:kh + Hout, kw:kw + Wout] + acc += (patch[None, :, :] * wq[:, ic, kh, kw][:, None, None]) >> 16 + acc += bq[:, None, None] + return acc.reshape(-1) + + +def _spmm_ref(M, K, N, vq, col, row, Bq): + rowid = np.repeat(np.arange(M), np.diff(row)) + contrib = (vq[:, None] * Bq[col]) >> 16 + C = np.zeros((M, N), np.int64) + np.add.at(C, rowid, contrib) + return C.reshape(-1) + + +# ── FWHT ────────────────────────────────────────────────────────────────── +@pytest.mark.parametrize("n", FWHT, ids=[f"n{n}" for n in FWHT]) +def test_stage2_fwht(tmp_path, n): + asm = _emit(_fwht_model(n), tmp_path) + rng = np.random.default_rng(0) + x = [int(np.trunc(v * Q)) for v in rng.uniform(-0.4, 0.4, n)] + mem = bytearray(0x200000) + inp, out = 0x10000, 0x80000 + for i, v in enumerate(x): + mem[inp + 4 * i:inp + 4 * i + 4] = int(v & _MASK).to_bytes(4, "little") + _run(asm, {10: inp, 11: out, 12: n}, mem) + got = [_s32(int.from_bytes(mem[out + 4 * i:out + 4 * i + 4], "little")) + for i in range(n)] + assert got == _fwht_ref(x) + + +# ── Conv ────────────────────────────────────────────────────────────────── +@pytest.mark.parametrize("cfg", CONV, ids=[f"c{i}" for i in range(1, 11)]) +def test_stage2_conv(tmp_path, cfg): + N, C, H, W, K, Cout = cfg + model = _conv_model(*cfg) + asm = _emit(model, tmp_path) + from onnx import numpy_helper + inits = {t.name: numpy_helper.to_array(t) for t in model.graph.initializer} + wq = _q(inits["w"]) + bq = _q(inits["b"]) + pad = K // 2 + rng = np.random.default_rng(0) + xq = _q(rng.uniform(-0.4, 0.4, (N, C, H, W))) + mem = bytearray(0x200000) + w = lambda a, v: mem.__setitem__(slice(a, a + 4), int(v & _MASK).to_bytes(4, "little")) + r32 = lambda a: _s32(int.from_bytes(mem[a:a + 4], "little")) + inp, out, wp, bp, par = 0x10000, 0x80000, 0x100000, 0x120000, 0x140000 + for i, v in enumerate(xq.reshape(-1)): + w(inp + 4 * i, int(v)) + for i, v in enumerate(wq.reshape(-1)): + w(wp + 4 * i, int(v)) + for i, v in enumerate(bq): + w(bp + 4 * i, int(v)) + for off, val in enumerate([N, C, H, W, Cout, K, pad, 1, H, W, wp, bp]): + w(par + 4 * off, int(val)) + _run(asm, {10: inp, 11: out, 12: par}, mem) + got = [r32(out + 4 * i) for i in range(N * Cout * H * W)] + ref = _conv_ref(xq[0], wq, bq, pad, K) + assert got == list(ref) + + +# ── SpMM ────────────────────────────────────────────────────────────────── +@pytest.mark.parametrize("cfg", SPMM, ids=[f"s{i}" for i in range(1, 11)]) +def test_stage2_spmm(tmp_path, cfg): + M, K, N, dens = cfg + model = _spmm_model(M, K, N, dens, SPMM.index(cfg) + 1) + asm = _emit(model, tmp_path) + from onnx import numpy_helper + d = {t.name: numpy_helper.to_array(t) for t in model.graph.initializer} + vq = _q(d["values"]) + col = d["col"] + row = d["rowptr"] + rng = np.random.default_rng(1) + Bq = _q(rng.uniform(-0.4, 0.4, (K, N))) + mem = bytearray(0x4000000) + w = lambda a, v: mem.__setitem__(slice(a, a + 4), int(v & _MASK).to_bytes(4, "little")) + r32 = lambda a: _s32(int.from_bytes(mem[a:a + 4], "little")) + block, va, ca, ra, ba, out, par = (0x1000, 0x100000, 0x600000, 0xB00000, + 0x1000000, 0x1100000, 0x1200000) + w(block, va); w(block + 4, ca); w(block + 8, ra) + for i, v in enumerate(vq): + w(va + 4 * i, int(v)) + for i, v in enumerate(col): + w(ca + 4 * i, int(v)) + for i, v in enumerate(row): + w(ra + 4 * i, int(v)) + for kk in range(K): + for nn in range(N): + w(ba + 4 * (kk * N + nn), int(Bq[kk][nn])) + for off, val in enumerate([M, K, N, len(vq), ba]): + w(par + 4 * off, int(val)) + _run(asm, {10: block, 11: out, 12: par}, mem) + got = [r32(out + 4 * (i * N + nn)) for i in range(M) for nn in range(N)] + ref = _spmm_ref(M, K, N, vq, col, row, Bq) + assert got == list(ref)