fix: correct IR Conv2D and Reshape execution and tighten accuracy validation

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)

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,
+        )
 
 
 # ═══════════════════════════════════════════════════════════════════════════════