#!/usr/bin/env python3 from __future__ import annotations import argparse import csv import json import shlex import subprocess import sys from pathlib import Path from colorama import Fore, Style REPO = Path(__file__).resolve().parents[3] SUITE = REPO / "validation/networks/pimcomp_models" sys.path.insert(0, str(REPO / "validation")) from raptor_validation.pimsim_nn import parse_pimsim_nn_metrics # noqa: E402 from raptor_validation.validate_one import STAGE_COLORS # noqa: E402 PIMCOMP_SOURCE = REPO / "third_party/PIMCOMP-NN" PIMCOMP_CONFIG = REPO / "validation/pimsim_configs/pimcomp/arch-a/latency_config.json" COMPARE = Path(__file__).resolve().with_name("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-latency.onnx", } def result_dir(root: Path | None, name: str) -> Path: return root / name if root is not None else MODELS[name].parent def write_results_csv(root: Path | None) -> Path: output = (root or SUITE) / "results.csv" fields = ( "model", "raptor_latency_ms", "pimcomp_latency_ms", "raptor_energy_pj", "pimcomp_energy_pj", "faster_compiler", "speedup", ) rows = [] for name in MODELS: report_path = result_dir(root, name) / "pimcomp/comparison_report.json" if not report_path.exists(): continue report = json.loads(report_path.read_text(encoding="utf-8")) raptor = report.get("raptor_performance") or {} pimcomp = report.get("pimcomp_performance") or {} raptor_latency = raptor.get("latency_ms") pimcomp_latency = pimcomp.get("latency_ms") if raptor_latency is None or pimcomp_latency is None: continue raptor_energy = (raptor.get("average_energy_pj") or parse_pimsim_nn_metrics(raptor.get("raw_output", "")).get("average_energy_pj")) pimcomp_energy = (pimcomp.get("average_energy_pj") or parse_pimsim_nn_metrics(pimcomp.get("raw_output", "")).get("average_energy_pj")) faster = "raptor" if raptor_latency < pimcomp_latency else "pimcomp" rows.append({ "model": name, "raptor_latency_ms": f"{raptor_latency:.6f}", "pimcomp_latency_ms": f"{pimcomp_latency:.6f}", "raptor_energy_pj": "" if raptor_energy is None else f"{raptor_energy:.6f}", "pimcomp_energy_pj": "" if pimcomp_energy is None else f"{pimcomp_energy:.6f}", "faster_compiler": faster, "speedup": f"{max(raptor_latency, pimcomp_latency) / min(raptor_latency, pimcomp_latency):.2f}", }) with open(output, "w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=fields, lineterminator="\n") writer.writeheader() writer.writerows(rows) return output def print_stage(title: str, color: str) -> None: print("\n" + Style.BRIGHT + color + f"[{title}]" + Style.RESET_ALL, flush=True) def run(command: list[str], *, dry_run: bool, check: bool = True) -> int: print(f" cwd: {REPO}", flush=True) print(f" $ {shlex.join(command)}", flush=True) if dry_run: return 0 return subprocess.run(command, cwd=REPO, check=check).returncode def validate_pimcomp_source() -> None: header = PIMCOMP_SOURCE / "backend/GeneticAlgorithm.h" source = header.read_text(encoding="utf-8") 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, timeout: float) -> list[str]: return [ sys.executable, str(COMPARE), "--model", str(model), "--out-dir", str(result_dir), "--pimcomp-dir", str(PIMCOMP_SOURCE), "--pimcomp-config", str(PIMCOMP_CONFIG), "--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", 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="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() 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)}") validate_pimcomp_source() if out_dir is not None and not args.dry_run: out_dir.mkdir(parents=True, exist_ok=True) print(Style.BRIGHT + f"Found {len(args.models)} PIMCOMP model(s) to compare." + Style.RESET_ALL) print(f"Results root: {out_dir or SUITE}") print("=" * 72) print_stage("Build Raptor", STAGE_COLORS["Build Runner"]) run(["cmake", "--build", str(REPO / "build_release")], dry_run=args.dry_run) print_stage("Build PIMCOMP", STAGE_COLORS["Build Runner"]) run( ["cmake", "--build", str(PIMCOMP_SOURCE / "build"), "--target", "PIMCOMP-NN"], dry_run=args.dry_run, ) failed = [] for index, name in enumerate(args.models, start=1): model_result_dir = result_dir(out_dir, name) print( "\n" + Fore.CYAN + f"[{index}/{len(args.models)}]" + Style.RESET_ALL + f" {Style.BRIGHT}Comparing {name}{Style.RESET_ALL}", flush=True, ) if args.resume and (model_result_dir / "pimcomp/comparison_report.json").exists(): print( Fore.YELLOW + " Completed report exists; skipping" + Style.RESET_ALL, flush=True, ) continue returncode = run( comparison_command(MODELS[name], model_result_dir, args.timeout_seconds), dry_run=args.dry_run, check=False, ) if returncode: failed.append(name) if args.dry_run: return 1 if failed else 0 results_path = write_results_csv(out_dir) print_stage("Results", STAGE_COLORS["Compare Outputs"]) print(results_path.read_text(encoding="utf-8"), end="") print("\n" + Style.BRIGHT + Fore.CYAN + "Summary" + Style.RESET_ALL) print(Style.BRIGHT + f"Passed: {len(args.models) - len(failed)}" + Style.RESET_ALL) print(Style.BRIGHT + f"Failed: {len(failed)}" + Style.RESET_ALL) print(Style.BRIGHT + f"Results: {results_path}" + Style.RESET_ALL) if failed: print( Fore.RED + f"Failed comparisons: {', '.join(failed)}" + Style.RESET_ALL, file=sys.stderr, ) return 1 return 0 if __name__ == "__main__": raise SystemExit(main())