Raptor/validation/tools/pimcomp/run_pimcomp_paper_latency.py

#!/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())