simplify pimcomp compare workflow

validation/tools/run_pimcomp_paper_latency.py

@@ -2,17 +2,18 @@
 from __future__ import annotations
 
 import argparse
-import re
 import shlex
-import shutil
 import subprocess
 import sys
 from pathlib import Path
 
+from colorama import Fore, Style
+
 
 REPO = Path(__file__).resolve().parents[2]
 SUITE = REPO / "validation/networks/pimcomp_models"
 PIMCOMP_SOURCE = REPO / "third_party/PIMCOMP-NN"
+PIMCOMP_CONFIG = REPO / "validation/pimsim_configs/pimcomp/arch-a/latency_config.json"
 COMPARE = REPO / "validation/tools/compare_raptor_pimcomp.py"
 MODELS = {
     "vgg8": SUITE / "vgg8/vgg8-mnist-reconstructed.onnx",
@@ -23,38 +24,21 @@ MODELS = {
 
 
 def run(command: list[str], *, dry_run: bool, check: bool = True) -> int:
-    print(f"$ {shlex.join(command)}", flush=True)
+    print(Fore.CYAN + "$ " + Style.RESET_ALL + 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"
+def validate_pimcomp_source() -> None:
+    header = PIMCOMP_SOURCE / "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(
-        REPO / "validation/pimsim_configs/pimcomp/arch-a/latency_config.json",
-        work_dir / "config.json",
-    )
+    for setting in ("int population_num = 200;", "int max_iteration = 1000;"):
+        if setting not in source:
+            raise RuntimeError(f"PIMCOMP paper setting is missing: {setting}")
 
 
-def comparison_command(model: Path, result_dir: Path, pimcomp_dir: Path, timeout: float) -> list[str]:
+def comparison_command(model: Path, result_dir: Path, timeout: float) -> list[str]:
     return [
         sys.executable,
         str(COMPARE),
@@ -63,7 +47,9 @@ def comparison_command(model: Path, result_dir: Path, pimcomp_dir: Path, timeout
         "--out-dir",
         str(result_dir),
         "--pimcomp-dir",
-        str(pimcomp_dir),
+        str(PIMCOMP_SOURCE),
+        "--pimcomp-config",
+        str(PIMCOMP_CONFIG),
         "--core-count",
         "168",
         "--crossbar-count",
@@ -82,6 +68,7 @@ def comparison_command(model: Path, result_dir: Path, pimcomp_dir: Path, timeout
         "GA",
         "--timeout-seconds",
         str(timeout),
+        "--fail-on-error",
     ]
 
 
@@ -89,51 +76,53 @@ 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(
+        "--out-dir",
+        type=Path,
+        help="Result root (default: artifacts beside each model under validation/).",
+    )
     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.",
+        help="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")
+    out_dir = args.out_dir.resolve() if args.out_dir is not None else None
 
     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:
+    validate_pimcomp_source()
+    if out_dir is not None and not args.dry_run:
         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"],
+        ["cmake", "--build", str(PIMCOMP_SOURCE / "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)
+        result_dir = out_dir / name if out_dir is not None else MODELS[name].parent
+        if args.resume and (result_dir / "pimcomp/comparison_report.json").exists():
+            print(
+                Fore.YELLOW + f"[{name}] completed report exists; skipping" + Style.RESET_ALL,
+                flush=True,
+            )
             continue
-        print(f"\n[{name}] Arch-A latency comparison", flush=True)
+        print(
+            "\n" + Fore.CYAN + f"[{name}]" + Style.RESET_ALL
+            + f" {Style.BRIGHT}Arch-A latency comparison{Style.RESET_ALL}",
+            flush=True,
+        )
         returncode = run(
-            comparison_command(MODELS[name], result_dir, work_dir, args.timeout_seconds),
+            comparison_command(MODELS[name], result_dir, args.timeout_seconds),
             dry_run=args.dry_run,
             check=False,
         )
@@ -141,8 +130,17 @@ def main() -> int:
             failed.append(name)
 
     if failed:
-        print(f"\nCompleted with failed comparisons: {', '.join(failed)}", file=sys.stderr)
+        print(
+            "\n" + Style.BRIGHT + Fore.RED + "Result: FAIL" + Style.RESET_ALL,
+            file=sys.stderr,
+        )
+        print(
+            Fore.RED + f"Failed comparisons: {', '.join(failed)}" + Style.RESET_ALL,
+            file=sys.stderr,
+        )
         return 1
+    if not args.dry_run:
+        print("\n" + Style.BRIGHT + f"Result: {Fore.GREEN}PASS" + Style.RESET_ALL)
     return 0