Python script for compare

2026-07-24 12:47:37 +02:00
parent e7611be8e1
commit f3a4e19f7c
14 changed files with 684 additions and 21 deletions
@@ -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())