even faster on pimcomp models

validation/raptor_validation/validate_one.py

@@ -1,8 +1,10 @@
 import json
 import os
+import re
 import shutil
 import subprocess
 import sys
+import time
 import numpy as np
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -10,7 +12,7 @@ from colorama import Style, Fore
 from .gen_network_runner import gen_network_runner
 from .onnx_utils import gen_random_inputs, save_inputs_to_files, onnx_io, write_inputs_to_memory_bin, _ONNX_TO_NP
 from .raptor import compile_with_raptor
-from .pimsim_nn import export_raptor_latency_artifact, parse_pimsim_nn_metrics
+from .pimsim_nn import export_raptor_latency_artifact, parse_pimsim_nn_metrics, read_raptor_instruction_count
 from .subprocess_utils import run_command_with_reporter
 
 STAGE_TITLES = (
@@ -80,6 +82,35 @@ class ValidationResult:
     pimsim_power_mw: float | None = None
     pimsim_energy_pj: float | None = None
     pimsim_status: str = PIMSIM_SKIPPED
+    compile_time_s: float | None = None
+    host_memory_bytes: int | None = None
+    cores_memory_bytes: int | None = None
+    used_core_count: int | None = None
+    used_crossbar_count: int | None = None
+
+
+_MEMORY_UNITS = {"B": 1, "KB": 1 << 10, "MB": 1 << 20, "GB": 1 << 30}
+
+
+def collect_pim_resource_metrics(pim_dir):
+    pim_dir = Path(pim_dir)
+    report_path = pim_dir.parent / "reports" / "memory_report.txt"
+    report = report_path.read_text(encoding="utf-8") if report_path.exists() else ""
+
+    def memory_bytes(label):
+        match = re.search(rf"^\s*{re.escape(label)}:\s+([0-9.]+)\s+(B|KB|MB|GB)$", report, re.MULTILINE)
+        return round(float(match.group(1)) * _MEMORY_UNITS[match.group(2)]) if match else None
+
+    with open(pim_dir / "config.json", encoding="utf-8") as f:
+        config = json.load(f)
+    used_cores = sum(read_raptor_instruction_count(path) > 0 for path in pim_dir.glob("core_*.pim"))
+    used_crossbars = sum(sum(groups) for groups in config.get("array_group_map", {}).values())
+    return {
+        "host_memory_bytes": memory_bytes("Host memory"),
+        "cores_memory_bytes": memory_bytes("Local memory after reuse"),
+        "used_core_count": used_cores,
+        "used_crossbar_count": used_crossbars,
+    }
 
 
 class ProgressReporter:
@@ -416,8 +447,6 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
         pimsim_nn_build_dir = Path(pimsim_nn_build_dir).resolve()
         pimsim_config_path = Path(pimsim_config_path).resolve()
     compile_extra_args = list(raptor_extra_args or [])
-    if pimsim_enabled and "--pim-emit-json" not in compile_extra_args:
-        compile_extra_args.append("--pim-emit-json")
     owns_reporter = reporter is None
     reporter = reporter or ProgressReporter(model_total, stages_per_model=len(MODE_STAGE_TITLES[mode]), verbose=verbose)
 
@@ -435,6 +464,8 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
                  f" {Style.BRIGHT}Validating {network_onnx_path.name}{Style.RESET_ALL}")
     failed_with_exception = False
     pim_pass_timings = {}
+    compile_time_s = None
+    resource_metrics = {}
 
     try:
         stem = network_onnx_path.stem
@@ -443,6 +474,19 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
         runner_path = runner_build_dir / "runner"
         pim_output_base = raptor_dir / stem
 
+        def compile_pim():
+            nonlocal compile_time_s, resource_metrics
+            started = time.perf_counter()
+            timings = compile_with_raptor(
+                network_onnx_path, raptor_path, pim_output_base, crossbar_size,
+                crossbar_count, core_count=core_count,
+                raptor_extra_args=compile_extra_args, cwd=raptor_dir,
+                verbose=verbose, reporter=reporter,
+                timeout_sec=command_timeout_seconds)
+            compile_time_s = time.perf_counter() - started
+            resource_metrics = collect_pim_resource_metrics(raptor_dir / "pim")
+            return timings
+
         if mode != MODE_RUN_ONLY:
             print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compile ONNX")
             network_so_path, network_mlir_path = compile_onnx_network(
@@ -468,16 +512,14 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
 
             if mode == MODE_COMPILE_ONLY:
                 print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compile PIM")
-                pim_pass_timings = compile_with_raptor(
-                    network_onnx_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
-                    core_count=core_count,
-                    raptor_extra_args=compile_extra_args,
-                    cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds)
+                pim_pass_timings = compile_pim()
                 print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
                 reporter.advance()
                 reporter.record_result(True)
                 reporter.log(Style.BRIGHT + f"Result: {Fore.GREEN}PASS{Style.RESET_ALL}" + Style.RESET_ALL)
-                return ValidationResult(passed=True, pim_pass_timings=pim_pass_timings)
+                return ValidationResult(
+                    passed=True, pim_pass_timings=pim_pass_timings,
+                    compile_time_s=compile_time_s, **resource_metrics)
 
         if mode == MODE_RUN_ONLY:
             required_paths = [
@@ -489,6 +531,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
             missing = [f"{description} at {path}" for path, description in required_paths if not path.exists()]
             if missing:
                 raise FileNotFoundError("run-only mode requires existing artifacts:\n  " + "\n  ".join(missing))
+            resource_metrics = collect_pim_resource_metrics(raptor_dir / "pim")
 
         print_stage(reporter, model_index, model_total, network_onnx_path.name, "Generate Inputs")
         inputs_descriptor, outputs_descriptor = onnx_io(network_onnx_path)
@@ -508,11 +551,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
 
         if mode != MODE_RUN_ONLY:
             print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compile PIM")
-            pim_pass_timings = compile_with_raptor(
-                network_onnx_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
-                core_count=core_count,
-                raptor_extra_args=compile_extra_args,
-                cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds)
+            pim_pass_timings = compile_pim()
             print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
             reporter.advance()
 
@@ -587,6 +626,8 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
             pimsim_power_mw=pimsim_power_mw,
             pimsim_energy_pj=pimsim_energy_pj,
             pimsim_status=pimsim_status,
+            compile_time_s=compile_time_s,
+            **resource_metrics,
         )
     except Exception:
         failed_with_exception = True