Python script for compare

validation/tools/compare_raptor_pimcomp.py

@@ -214,8 +214,17 @@ def prepare_pimcomp_model(model_path: Path, out_dir: Path) -> Path:
         if not equivalent:
             raise RuntimeError("Conv+BatchNormalization folding changed the model output")
         if any(node.op_type == "BatchNormalization" for node in model.graph.node):
-            raise RuntimeError("PIMCOMP model preparation did not eliminate BatchNormalization")
-        onnx.save(model, output_path)
+            import onnxruntime as ort
+
+            options = ort.SessionOptions()
+            options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
+            options.optimized_model_filepath = str(output_path)
+            ort.InferenceSession(str(model_path), options, providers=["CPUExecutionProvider"])
+            model = onnx.load(output_path)
+            if any(node.op_type == "BatchNormalization" for node in model.graph.node):
+                raise RuntimeError("PIMCOMP model preparation did not eliminate BatchNormalization")
+        else:
+            onnx.save(model, output_path)
     else:
         shutil.copy2(model_path, output_path)
     return output_path
@@ -271,9 +280,9 @@ def load_effective_hardware(args: argparse.Namespace) -> dict[str, int]:
 
 def write_pimsim_config(args: argparse.Namespace, out_dir: Path, hardware: dict[str, int]) -> Path:
     mesh_builder = load_mesh_builder()
-    example_config = REPO / "backend-simulators/pim/pimsim-nn/example/config/latency_config.json"
-    with open(example_config, "r", encoding="utf-8") as f:
+    with open(args.pimcomp_dir / "config.json", "r", encoding="utf-8") as f:
         config = json.load(f)
+    config["chip_config"]["core_config"].setdefault("rob_size", 1)
     config["chip_config"]["core_config"]["matrix_config"]["xbar_array_count"] = hardware["crossbar_count"]
     config["chip_config"]["core_config"]["matrix_config"]["xbar_size"] = [
         hardware["crossbar_size"],
@@ -491,10 +500,10 @@ def run_rust_validation(
     )
 
 
-def copy_pimcomp_outputs(args: argparse.Namespace, out_dir: Path):
+def copy_pimcomp_outputs(source_dir: Path, out_dir: Path):
     out_dir.mkdir(parents=True, exist_ok=True)
     for name in ("SimulationInfo.gz", "VerificationInfo.json", "MappingResult.txt"):
-        shutil.copy2(args.pimcomp_dir / "output" / name, out_dir / name)
+        shutil.copy2(source_dir / name, out_dir / name)
 
 
 def compile_pimcomp(
@@ -507,7 +516,7 @@ def compile_pimcomp(
     model_name = args.pimcomp_model_name or f"compare_{model_path.stem}"
     frontend_json = args.pimcomp_dir / "models/JSON" / f"{model_name}.json"
     frontend_cmd = [
-        "python3",
+        sys.executable,
         "frontend.py",
         "--model_path",
         str(model_path),
@@ -536,7 +545,7 @@ def compile_pimcomp(
         timeout_sec=args.timeout_seconds,
         steps=steps,
     )
-    copy_pimcomp_outputs(args, out_dir)
+    copy_pimcomp_outputs(args.pimcomp_dir / "output", out_dir)
     return out_dir / "VerificationInfo.json", out_dir / "SimulationInfo.gz"
 
 
@@ -798,7 +807,7 @@ def export_pimcomp_for_rust(
 def parse_pimsim_nn_report(output: str) -> dict[str, float | int | str]:
     patterns = {
         "output_count": r"output count:\s+([0-9]+)\s+samples",
-        "throughput": r"throughput:\s+([0-9.]+)\s+samples/s",
+        "throughput": r"throughput:\s+([0-9.eE+-]+)\s+samples/s",
         "average_latency_ms": r"average latency:\s+([0-9.eE+-]+)\s+ms",
         "latency_ms": r"latency:\s+([0-9.eE+-]+)\s+ms",
         "average_power_mw": r"average power:\s+([0-9.eE+-]+)\s+mW",
@@ -1192,6 +1201,11 @@ def main():
         type=Path,
         help="Reuse Raptor artifacts and results from an existing comparison_report.json.",
     )
+    parser.add_argument(
+        "--reuse-pimcomp-dir",
+        type=Path,
+        help="Reuse a directory containing PIMCOMP SimulationInfo.gz, VerificationInfo.json, and MappingResult.txt.",
+    )
     parser.add_argument("--skip-pimsim-nn", action="store_true")
     parser.add_argument("--verbose-raptor-compile", action="store_true")
     parser.add_argument("--raptor-extra-arg", action="append", default=[])
@@ -1365,17 +1379,31 @@ def main():
         out_dir / "pimcomp_model",
     )
 
-    compiled_pimcomp = try_stage(
-        failures,
-        "Compile PIMCOMP",
-        compile_pimcomp,
-        args,
-        pimcomp_model_path,
-        out_dir / "pimcomp",
-        steps,
-    ) if pimcomp_model_path is not None else None
-    if compiled_pimcomp is not None:
-        verification_info, simulation_info = compiled_pimcomp
+    if args.reuse_pimcomp_dir is not None:
+        reused_pimcomp_dir = args.reuse_pimcomp_dir.resolve()
+        copied_pimcomp = try_stage_success(
+            failures,
+            "Reuse PIMCOMP outputs",
+            copy_pimcomp_outputs,
+            reused_pimcomp_dir,
+            out_dir / "pimcomp",
+        )
+        if copied_pimcomp:
+            verification_info = out_dir / "pimcomp/VerificationInfo.json"
+            simulation_info = out_dir / "pimcomp/SimulationInfo.gz"
+            print(f"\n[Reuse PIMCOMP]\n  Directory: {reused_pimcomp_dir}")
+    else:
+        compiled_pimcomp = try_stage(
+            failures,
+            "Compile PIMCOMP",
+            compile_pimcomp,
+            args,
+            pimcomp_model_path,
+            out_dir / "pimcomp",
+            steps,
+        ) if pimcomp_model_path is not None else None
+        if compiled_pimcomp is not None:
+            verification_info, simulation_info = compiled_pimcomp
 
     if verification_info is not None and simulation_info is not None and model_io is not None:
         exported = try_stage(

validation/tools/run_pimcomp_paper_latency.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import re
+import shlex
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+
+REPO = Path(__file__).resolve().parents[2]
+SUITE = REPO / "validation/networks/pimcomp_models"
+PIMCOMP_SOURCE = REPO / "third_party/PIMCOMP-NN"
+COMPARE = REPO / "validation/tools/compare_raptor_pimcomp.py"
+MODELS = {
+    "vgg8": SUITE / "vgg8/vgg8-mnist-reconstructed.onnx",
+    "resnet18": SUITE / "resnet18/resnet18-v1-7.onnx",
+    "resnet34": SUITE / "resnet34/resnet34-v1-7.onnx",
+    "googlenet": SUITE / "googlenet/googlenet-12-no-softmax.onnx",
+}
+
+
+def run(command: list[str], *, dry_run: bool, check: bool = True) -> int:
+    print(f"$ {shlex.join(command)}", flush=True)
+    if dry_run:
+        return 0
+    return subprocess.run(command, cwd=REPO, check=check).returncode
+
+
+def prepare_pimcomp(work_dir: Path) -> None:
+    shutil.copytree(
+        PIMCOMP_SOURCE,
+        work_dir,
+        dirs_exist_ok=True,
+        ignore=shutil.ignore_patterns(".git", "build", "output"),
+    )
+    header = work_dir / "backend/GeneticAlgorithm.h"
+    source = header.read_text(encoding="utf-8")
+    if "int population_num = 200;" not in source:
+        raise RuntimeError("PIMCOMP GA population is not 200")
+    source, replacements = re.subn(
+        r"int max_iteration = \d+;",
+        "int max_iteration = 1000;",
+        source,
+    )
+    if replacements != 1:
+        raise RuntimeError("Could not set PIMCOMP GA max_iteration")
+    header.write_text(source, encoding="utf-8")
+    shutil.copy2(SUITE / "configs/arch-a.json", work_dir / "config.json")
+
+
+def comparison_command(model: Path, result_dir: Path, pimcomp_dir: Path, timeout: float) -> list[str]:
+    return [
+        sys.executable,
+        str(COMPARE),
+        "--model",
+        str(model),
+        "--out-dir",
+        str(result_dir),
+        "--pimcomp-dir",
+        str(pimcomp_dir),
+        "--core-count",
+        "168",
+        "--crossbar-count",
+        "96",
+        "--crossbar-size",
+        "128",
+        "--mesh-rows",
+        "12",
+        "--mesh-cols",
+        "14",
+        "--pimsim-mode",
+        "latency",
+        "--pimcomp-pipeline",
+        "element",
+        "--pimcomp-replication",
+        "GA",
+        "--timeout-seconds",
+        str(timeout),
+        "--fail-on-error",
+    ]
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Reproduce the serial Arch-A latency comparison from the PIMCOMP paper."
+    )
+    parser.add_argument("--out-dir", required=True, type=Path)
+    parser.add_argument("--models", nargs="+", choices=MODELS, default=list(MODELS))
+    parser.add_argument("--timeout-seconds", type=float, default=3600.0)
+    parser.add_argument(
+        "--resume",
+        action="store_true",
+        help="Keep the existing work tree and skip models with a completed JSON report.",
+    )
+    parser.add_argument("--dry-run", action="store_true", help="Print commands without modifying files.")
+    args = parser.parse_args()
+
+    out_dir = args.out_dir.resolve()
+    work_dir = out_dir / "pimcomp-ga1000"
+    if not args.dry_run and out_dir.exists() and any(out_dir.iterdir()) and not args.resume:
+        parser.error(f"{out_dir} is not empty; choose a fresh directory or pass --resume")
+
+    missing = [str(MODELS[name]) for name in args.models if not MODELS[name].exists()]
+    if missing:
+        parser.error(f"missing model(s): {', '.join(missing)}")
+
+    if args.dry_run:
+        print(f"# prepare isolated PIMCOMP GA build in {work_dir}")
+    else:
+        out_dir.mkdir(parents=True, exist_ok=True)
+        prepare_pimcomp(work_dir)
+
+    run(["cmake", "--build", str(REPO / "build_release")], dry_run=args.dry_run)
+    run(
+        ["cmake", "-S", str(work_dir), "-B", str(work_dir / "build")],
+        dry_run=args.dry_run,
+    )
+    run(
+        ["cmake", "--build", str(work_dir / "build"), "--target", "PIMCOMP-NN"],
+        dry_run=args.dry_run,
+    )
+
+    failed = []
+    for name in args.models:
+        result_dir = out_dir / name
+        if args.resume and (result_dir / "comparison_report.json").exists():
+            print(f"[{name}] completed report exists; skipping", flush=True)
+            continue
+        print(f"\n[{name}] Arch-A latency comparison", flush=True)
+        returncode = run(
+            comparison_command(MODELS[name], result_dir, work_dir, args.timeout_seconds),
+            dry_run=args.dry_run,
+            check=False,
+        )
+        if returncode:
+            failed.append(name)
+
+    if failed:
+        print(f"\nCompleted with failed comparisons: {', '.join(failed)}", file=sys.stderr)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

validation/tools/split_onnx_prefixes.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+
+import argparse
+from pathlib import Path
+
+import onnx
+
+
+def split_prefixes(model_path: Path, output_dir: Path, name: str) -> None:
+    model = onnx.shape_inference.infer_shapes(onnx.load(model_path))
+    initializer_names = {initializer.name for initializer in model.graph.initializer}
+    input_names = [value.name for value in model.graph.input if value.name not in initializer_names]
+    extractor = onnx.utils.Extractor(model)
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    for depth, node in enumerate(model.graph.node):
+        output_name = next(output for output in node.output if output)
+        prefix = extractor.extract_model(input_names, [output_name])
+        prefix.ir_version = max(prefix.ir_version, 4)
+        onnx.checker.check_model(prefix)
+
+        depth_name = f"depth_{depth:02d}"
+        depth_dir = output_dir / depth_name
+        depth_dir.mkdir(parents=True, exist_ok=True)
+        output_path = depth_dir / f"{name}_{depth_name}.onnx"
+        onnx.save(prefix, output_path)
+        print(f"{depth_name}: {node.op_type} -> {output_name} ({len(prefix.graph.node)} nodes)")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Split an ONNX graph into one ancestor prefix per node.")
+    parser.add_argument("model", type=Path)
+    parser.add_argument("output_dir", type=Path)
+    parser.add_argument("--name", required=True)
+    args = parser.parse_args()
+    split_prefixes(args.model, args.output_dir, args.name)
+
+
+if __name__ == "__main__":
+    main()