diff --git a/scratchv/simulator/rv32_emulator.py b/scratchv/simulator/rv32_emulator.py index 98b9bc42..d2f20c5e 100644 --- a/scratchv/simulator/rv32_emulator.py +++ b/scratchv/simulator/rv32_emulator.py @@ -482,18 +482,16 @@ def _op_load_const(self, instr): def _op_reshape(self, instr): ops = self._get_operands(instr) - result = ops[0] - # Reshape to flatten: typically from (N,C,H,W) to (N, C*H*W) - # If the second operand is an initializer with shape info - if len(ops) > 1 and hasattr(ops[1], 'shape') and ops[1].size > 0: - target_shape = tuple(int(v) for v in ops[1].flatten() if v > 0) - if len(target_shape) > 0: - try: - result = ops[0].reshape(target_shape) - except (ValueError, RuntimeError): - # Fallback: flatten to 2D - result = ops[0].reshape(ops[0].shape[0], -1) - if result is ops[0]: + target_shape = tuple(int(v) for v in instr.attrs.get("shape", ())) + if target_shape: + # ONNX Reshape uses 0 to copy the corresponding input dimension. + resolved_shape = tuple( + ops[0].shape[index] if dim == 0 else dim + for index, dim in enumerate(target_shape) + ) + result = ops[0].reshape(resolved_shape) + else: + # Preserve the historical fallback for IR without shape metadata. result = ops[0].reshape(ops[0].shape[0], -1) if instr.dest: self._values[instr.dest.name] = result @@ -542,8 +540,8 @@ def _conv2d_numpy(x: np.ndarray, w: np.ndarray, b: np.ndarray, for j in range(K): patch = x[:, :, i:i + H_out * stride:stride, j:j + W_out * stride:stride] - cols[(i * K + j) * C_in:(i * K + j + 1) * C_in, :] = \ - patch.reshape(N * C_in, H_out * W_out) + cols[i * K + j::K * K, :] = \ + patch.transpose(1, 0, 2, 3).reshape(C_in, -1) # Weight: (C_out, C_in, K, K) → (C_out, C_in*K*K) w_mat = w.reshape(C_out, -1) diff --git a/tests/test_cnn_pipeline.py b/tests/test_cnn_pipeline.py index 6dd1b4ed..c6717654 100644 --- a/tests/test_cnn_pipeline.py +++ b/tests/test_cnn_pipeline.py @@ -1,12 +1,12 @@ -"""End-to-end test: ONNX CNN model → RISC-V binary → data verification. +"""CNN pipeline tests, including explicit IR semantic verification. Full pipeline test: 1. Parse cnn.onnx → ScratchV IR 2. IR → RISC-V assembly (instruction selection + register alloc + emit) 3. RISC-V assembly → 32-bit machine code (RV32IM encoder) - 4. IR trace execution (numpy) → compiled output + 4. IR trace execution (numpy) → IR semantic output 5. ONNX Runtime reference inference → expected output - 6. MSE / MAE comparison + 6. Strict shape and numerical comparison """ from __future__ import annotations @@ -18,6 +18,8 @@ # ── Path to project root ──────────────────────────────────────────────────── PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) MODEL_PATH = os.path.join(PROJECT_ROOT, "models", "graph", "cnn.onnx") +IR_RTOL = 1e-4 +IR_ATOL = 1e-5 pytestmark = pytest.mark.skipif( not os.path.exists(MODEL_PATH), @@ -41,6 +43,11 @@ def onnx_model(): shape = [d.dim_value for d in inp.type.tensor_type.shape.dim] inputs.append({"name": inp.name, "shape": shape}) + outputs = [] + for output in graph.output: + shape = [d.dim_value for d in output.type.tensor_type.shape.dim] + outputs.append({"name": output.name, "shape": shape}) + nodes = [] for node in graph.node: nodes.append({ @@ -57,6 +64,7 @@ def onnx_model(): return { "graph_name": graph.name, "inputs": inputs, + "outputs": outputs, "nodes": nodes, "num_nodes": len(nodes), "initializers": initializers, @@ -102,8 +110,8 @@ def test_input(): @pytest.fixture(scope="module") -def scratchv_output(ir_program, onnx_model, test_input): - """Step 4: Execute IR via trace executor → compiled output.""" +def ir_output(ir_program, onnx_model, test_input): + """Step 4: Execute IR via trace executor → IR semantic output.""" from scratchv.simulator.rv32_emulator import IRTraceExecutor executor = IRTraceExecutor(ir_program, onnx_model["initializers"]) return executor.run(test_input) @@ -292,31 +300,31 @@ def test_binary_deterministic(self, riscv_assembly): class TestExecution: """Verify IR trace execution produces valid output.""" - def test_output_is_array(self, scratchv_output): + def test_output_is_array(self, ir_output): """Output should be a numpy array.""" - assert isinstance(scratchv_output, np.ndarray) + assert isinstance(ir_output, np.ndarray) - def test_output_is_finite(self, scratchv_output): + def test_output_is_finite(self, ir_output): """Output should not contain NaN or Inf.""" - assert np.all(np.isfinite(scratchv_output)) + assert np.all(np.isfinite(ir_output)) - def test_output_range_reasonable(self, scratchv_output): + def test_output_range_reasonable(self, ir_output): """Sigmoid output should be in [0, 1].""" - out = np.asarray(scratchv_output, dtype=np.float32).flatten() + out = np.asarray(ir_output, dtype=np.float32).flatten() for v in out: assert 0.0 <= v <= 1.0, f"Output {v} outside sigmoid range [0,1]" - def test_output_shape(self, scratchv_output): + def test_output_shape(self, ir_output): """Final output should be scalar-like (sigmoid).""" - assert scratchv_output.size == 1 + assert ir_output.size == 1 # ═══════════════════════════════════════════════════════════════════════════════ -# Test 6: Verification (MSE / MAE) +# Test 6: IR semantic verification against ONNX Runtime # ═══════════════════════════════════════════════════════════════════════════════ class TestVerification: - """Compare ScratchV output against ONNX Runtime reference.""" + """Compare IR trace output against the named ONNX Runtime output.""" @pytest.mark.skipif( not __import__("importlib").util.find_spec("onnxruntime"), @@ -331,40 +339,36 @@ def test_onnxruntime_available(self, onnxrt_output): not __import__("importlib").util.find_spec("onnxruntime"), reason="onnxruntime not installed", ) - def test_mse_below_threshold(self, scratchv_output, onnxrt_output): - """MSE should be finite.""" - ref = list(onnxrt_output.values())[0] - comp = np.asarray(scratchv_output, dtype=np.float32).flatten() - ref = np.asarray(ref, dtype=np.float32).flatten() - - mse = float(np.mean((ref[:len(comp)] - comp) ** 2)) - assert np.isfinite(mse) - assert mse < 1.0 # should be well below 1.0 for sigmoid output - - @pytest.mark.skipif( - not __import__("importlib").util.find_spec("onnxruntime"), - reason="onnxruntime not installed", - ) - def test_mae_below_threshold(self, scratchv_output, onnxrt_output): - """MAE should be finite.""" - ref = list(onnxrt_output.values())[0] - comp = np.asarray(scratchv_output, dtype=np.float32).flatten() - ref = np.asarray(ref, dtype=np.float32).flatten() - - mae = float(np.mean(np.abs(ref[:len(comp)] - comp))) - assert np.isfinite(mae) - assert mae < 1.0 + def test_output_names_match_model(self, onnx_model, onnxrt_output): + """Runtime outputs should match the names declared by the model.""" + expected_names = {output["name"] for output in onnx_model["outputs"]} + assert set(onnxrt_output) == expected_names @pytest.mark.skipif( not __import__("importlib").util.find_spec("onnxruntime"), reason="onnxruntime not installed", ) - def test_output_not_identical(self, scratchv_output, onnxrt_output): - """Outputs should differ (different implementations).""" - ref = list(onnxrt_output.values())[0] - comp = np.asarray(scratchv_output, dtype=np.float32).flatten() - ref = np.asarray(ref, dtype=np.float32).flatten() - assert not np.allclose(ref[:len(comp)], comp, rtol=1e-6) + def test_ir_output_matches_onnxruntime( + self, ir_output, onnx_model, onnxrt_output): + """IR execution should match the corresponding runtime output.""" + output_names = [output["name"] for output in onnx_model["outputs"]] + assert len(output_names) == 1, ( + "IRTraceExecutor returns one final value; the model must expose " + "exactly one output for this comparison" + ) + + output_name = output_names[0] + assert output_name in onnxrt_output + + actual = np.asarray(ir_output, dtype=np.float32) + expected = np.asarray(onnxrt_output[output_name], dtype=np.float32) + assert actual.shape == expected.shape + np.testing.assert_allclose( + actual, + expected, + rtol=IR_RTOL, + atol=IR_ATOL, + ) # ═══════════════════════════════════════════════════════════════════════════════