more complete pimcomp comparison scripts
Validate Operations / validate-operations (push) Waiting to run

update pimsim-nn submodule
This commit is contained in:
NiccoloN
2026-08-06 21:49:54 +02:00
parent 4acd3b0c81
commit 2e76164aed
20 changed files with 1626 additions and 548 deletions
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,755 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import csv
import json
import os
import shlex
import shutil
import subprocess
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from contextlib import nullcontext, redirect_stderr, redirect_stdout
from dataclasses import dataclass
from pathlib import Path
from tempfile import TemporaryDirectory
from colorama import Fore, Style
REPO = Path(__file__).resolve().parents[5]
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_CONFIGS = REPO / "validation/pimsim_configs/pimcomp"
COMPARE = Path(__file__).resolve().with_name("compare_raptor_pimcomp.py")
ARCHES = tuple(sorted(path.name for path in PIMCOMP_CONFIGS.iterdir() if path.is_dir()))
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-pimsim-nn.onnx",
"yolo11n": SUITE / "yolo11n/yolo11n-pimsim-nn.onnx",
}
COMPARISONS = (
("latency", 1, "element"),
("throughput", 2, "batch"),
("throughput", 4, "batch"),
("throughput", 8, "batch"),
)
@dataclass(frozen=True)
class ComparisonSpec:
label: str
model: Path
output_dir: Path
common_dir: Path
config: Path
mode: str
pipeline: int
pimcomp_pipeline: str
shared_pimcomp_dir: Path
anchor: bool
def model_dir(root: Path | None, name: str) -> Path:
return root / name if root is not None else MODELS[name].parent
def result_dir(root: Path | None, name: str, arch: str, mode: str, pipeline: int) -> Path:
base = model_dir(root, name)
suffix = "latency" if mode == "latency" else f"throughput/pipeline{pipeline}"
return base / arch / suffix
def clean_artifacts(root: Path | None, models: list[str], arches: list[str]) -> int:
removed = 0
for name in models:
base = model_dir(root, name)
for arch in arches:
for path in (base / arch / "latency", base / arch / "throughput"):
if path.is_dir() and not path.is_symlink():
shutil.rmtree(path)
removed += 1
common = base / "common"
if common.is_dir() and not common.is_symlink():
shutil.rmtree(common)
removed += 1
for path in (
(root or SUITE) / "results.csv",
(root or SUITE) / "results_latency.csv",
(root or SUITE) / "results_throughput.csv",
):
if path.is_file() or path.is_symlink():
path.unlink()
removed += 1
return removed
def write_results_csv(
root: Path | None,
arch: str,
models: list[str],
comparisons: tuple[tuple[str, int, str], ...] = COMPARISONS,
) -> Path:
output = (root or SUITE) / "results.csv"
fields = (
"model",
"arch",
"mode",
"raptor_pipeline",
"pimcomp_pipeline",
"raptor_functional_validation",
"pimcomp_functional_validation",
"raptor_throughput_samples_s",
"pimcomp_throughput_samples_s",
"raptor_latency_ms",
"pimcomp_latency_ms",
"raptor_power_mw",
"pimcomp_power_mw",
"raptor_energy_pj",
"pimcomp_energy_pj",
"better_compiler",
"speedup",
)
selected = {
(name, arch, mode, str(pipeline))
for name in models
for mode, pipeline, _ in comparisons
}
rows = []
if output.exists():
with open(output, newline="", encoding="utf-8") as f:
reader = csv.DictReader(f)
if reader.fieldnames and "model" in reader.fieldnames:
for row in reader:
key = (
row.get("model"),
row.get("arch"),
row.get("mode"),
row.get("raptor_pipeline"),
)
if key in selected:
continue
row.pop("status", None)
rows.append({field: row.get(field, "NA") for field in fields})
for name in models:
for comparison_mode, pipeline, pimcomp_pipeline in comparisons:
report_path = result_dir(root, name, arch, comparison_mode, pipeline) / "pimcomp/comparison_report.json"
row = {
"model": name,
"arch": arch,
"mode": comparison_mode,
"raptor_pipeline": pipeline,
"pimcomp_pipeline": pimcomp_pipeline,
"raptor_functional_validation": "NA",
"pimcomp_functional_validation": "NA",
"raptor_throughput_samples_s": "NA",
"pimcomp_throughput_samples_s": "NA",
"raptor_latency_ms": "NA",
"pimcomp_latency_ms": "NA",
"raptor_power_mw": "NA",
"pimcomp_power_mw": "NA",
"raptor_energy_pj": "NA",
"pimcomp_energy_pj": "NA",
"better_compiler": "NA",
"speedup": "NA",
}
if report_path.exists():
report = json.loads(report_path.read_text(encoding="utf-8"))
raptor_values = performance_values(report.get("raptor_performance") or {})
pimcomp_values = performance_values(report.get("pimcomp_performance") or {})
raptor_metric = raptor_values["throughput"] if comparison_mode == "throughput" else raptor_values["latency"]
pimcomp_metric = pimcomp_values["throughput"] if comparison_mode == "throughput" else pimcomp_values["latency"]
row.update(
raptor_functional_validation=functional_validation_status(
report.get("raptor_validation")
),
pimcomp_functional_validation=functional_validation_status(
report.get("pimcomp_validation")
),
raptor_throughput_samples_s=format_value(raptor_values["throughput"]),
pimcomp_throughput_samples_s=format_value(pimcomp_values["throughput"]),
raptor_latency_ms=format_value(raptor_values["latency"]),
pimcomp_latency_ms=format_value(pimcomp_values["latency"]),
raptor_power_mw=format_value(raptor_values["power"]),
pimcomp_power_mw=format_value(pimcomp_values["power"]),
raptor_energy_pj=format_value(raptor_values["energy"]),
pimcomp_energy_pj=format_value(pimcomp_values["energy"]),
)
if raptor_metric is not None and pimcomp_metric is not None:
row.update(
better_compiler=comparison_winner(comparison_mode, raptor_metric, pimcomp_metric),
speedup=f"{max(raptor_metric, pimcomp_metric) / min(raptor_metric, pimcomp_metric):.2f}",
)
rows.append(row)
rows.sort(key=lambda row: (list(MODELS).index(row["model"]), row["mode"], int(row["raptor_pipeline"])))
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 performance_values(performance: dict) -> dict[str, float | None]:
parsed = parse_pimsim_nn_metrics(performance.get("raw_output", ""))
return {
"throughput": performance.get("throughput") or parsed.get("throughput"),
"latency": (
performance.get("latency_ms")
or performance.get("average_latency_ms")
or parsed.get("latency_ms")
or parsed.get("average_latency_ms")
),
"power": performance.get("average_power_mw") or parsed.get("average_power_mw"),
"energy": performance.get("average_energy_pj") or parsed.get("average_energy_pj"),
}
def comparison_passed(report: dict) -> bool:
if report.get("failures"):
return False
for key in ("raptor_validation", "pimcomp_validation"):
result = report.get(key) or {}
if result.get("status") != "done" or not result.get("passed"):
return False
for key in ("raptor_performance", "pimcomp_performance"):
performance = report.get(key) or {}
if performance.get("error") or performance.get("skipped"):
return False
return True
def functional_validation_status(result: dict | None) -> str:
if result is None:
return "NA"
return "PASS" if result.get("status") == "done" and result.get("passed") else "FAIL"
def comparison_winner(mode: str, raptor: float, pimcomp: float) -> str:
if raptor == pimcomp:
return "tie"
if mode == "throughput":
return "raptor" if raptor > pimcomp else "pimcomp"
return "raptor" if raptor < pimcomp else "pimcomp"
def format_value(value: float | None) -> str:
return "NA" if value is None else f"{value:.6f}"
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 run_comparison_job(job: tuple[str, list[str], Path | None]) -> tuple[str, int, str | None]:
label, command, log_path = job
sys.stdout.flush()
sys.stderr.flush()
if log_path is None:
return label, run(command, dry_run=False, check=False), None
saved_stdout = os.dup(1)
saved_stderr = os.dup(2)
try:
with open(log_path, "w", encoding="utf-8", buffering=1) as log:
os.dup2(log.fileno(), 1)
os.dup2(log.fileno(), 2)
with redirect_stdout(log), redirect_stderr(log):
returncode = run(command, dry_run=False, check=False)
finally:
os.dup2(saved_stdout, 1)
os.dup2(saved_stderr, 2)
os.close(saved_stdout)
os.close(saved_stderr)
return label, returncode, str(log_path)
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,
common_dir: Path,
config: Path,
mode: str,
pipeline: int,
pimcomp_pipeline: str,
pimsim_time_ms: int,
timeout: float,
reuse_raptor_report: Path | None = None,
reuse_pimcomp_dir: Path | None = None,
) -> list[str]:
time_args = ["--pimsim-time-ms", str(pimsim_time_ms)] if mode == "throughput" else []
reuse_args = []
if reuse_raptor_report is not None:
reuse_args.extend(["--reuse-raptor-report", str(reuse_raptor_report)])
if reuse_pimcomp_dir is not None:
reuse_args.extend(["--reuse-pimcomp-dir", str(reuse_pimcomp_dir)])
return [
sys.executable,
str(COMPARE),
"--model",
str(model),
"--out-dir",
str(result_dir),
"--common-dir",
str(common_dir),
"--pimcomp-dir",
str(PIMCOMP_SOURCE),
"--pimcomp-config",
str(config),
"--pimsim-mode",
mode,
*time_args,
"--pimcomp-pipeline",
pimcomp_pipeline,
"--pimcomp-replication",
"GA",
f"--raptor-extra-arg=--pipeline={pipeline}",
"--timeout-seconds",
str(timeout),
"--fail-on-error",
*reuse_args,
]
def prepare_common_command(model: Path, common_dir: Path, timeout: float) -> list[str]:
return [
sys.executable,
str(COMPARE),
"--model",
str(model),
"--out-dir",
str(common_dir),
"--common-dir",
str(common_dir),
"--prepare-common",
"--timeout-seconds",
str(timeout),
]
def comparison_command_for(
spec: ComparisonSpec,
args: argparse.Namespace,
*,
reuse_shared_pimcomp: bool,
) -> list[str]:
reuse_pimcomp_dir = None
if reuse_shared_pimcomp:
reuse_pimcomp_dir = spec.shared_pimcomp_dir
elif args.only == "raptor":
reuse_pimcomp_dir = spec.output_dir / "pimcomp/output"
return comparison_command(
spec.model,
spec.output_dir,
spec.common_dir,
spec.config,
spec.mode,
spec.pipeline,
spec.pimcomp_pipeline,
args.pimsim_time_ms,
args.timeout_seconds,
reuse_raptor_report=(
spec.output_dir / "pimcomp/comparison_report.json"
if args.only == "pimcomp"
else None
),
reuse_pimcomp_dir=reuse_pimcomp_dir,
)
def pimcomp_artifact_ready(output_dir: Path) -> bool:
return all(
(output_dir / name).is_file()
for name in ("SimulationInfo.gz", "VerificationInfo.json", "MappingResult.txt")
) and (output_dir.parent / "comparison_report.json").is_file()
def run_comparison_jobs(
comparison_jobs: list[tuple[str, list[str], Path | None]],
max_workers: int,
log_dir: Path | None,
log_offset: int,
) -> tuple[list[str], int]:
if not comparison_jobs:
return [], log_offset
print_directly = min(max_workers, len(comparison_jobs)) == 1
jobs = [
(
label,
command,
None if print_directly or log_dir is None else log_dir / f"{log_offset + index}.log",
)
for index, (label, command, _) in enumerate(comparison_jobs)
]
with (nullcontext(None) if print_directly else ProcessPoolExecutor(max_workers=max_workers)) as executor:
completed = (
map(run_comparison_job, jobs)
if print_directly
else (
future.result()
for future in as_completed(executor.submit(run_comparison_job, job) for job in jobs)
)
)
failed = []
for label, returncode, result_log in completed:
if result_log:
output = Path(result_log).read_text(encoding="utf-8", errors="replace")
if output:
print(output, end="" if output.endswith("\n") else "\n")
if returncode:
failed.append(label)
return failed, log_offset + len(jobs)
def ensure_throughput_config(arch: str, sim_time_ms: int) -> Path:
source = PIMCOMP_CONFIGS / arch / "throughput_config_1000ms.json"
target = PIMCOMP_CONFIGS / arch / f"throughput_config_{sim_time_ms}ms.json"
if target.exists():
return target
if sim_time_ms == 1000:
return source
config = json.loads(source.read_text(encoding="utf-8"))
config["sim_config"]["sim_time"] = sim_time_ms
target.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")
return target
def config_path(arch: str, mode: str, sim_time_ms: int, *, write: bool) -> Path:
path = (
PIMCOMP_CONFIGS / arch / "latency_config.json"
if mode == "latency"
else PIMCOMP_CONFIGS / arch / f"throughput_config_{sim_time_ms}ms.json"
)
if mode == "throughput" and write:
path = ensure_throughput_config(arch, sim_time_ms)
if not path.exists() and (write or mode == "latency"):
raise ValueError(f"{arch} has no {mode} config: {path}")
return path
def core_count(config: Path) -> int:
if not config.exists() and config.name.startswith("throughput_config_"):
config = config.with_name("throughput_config_1000ms.json")
with open(config, encoding="utf-8") as f:
return int(json.load(f)["chip_config"]["core_cnt"])
def main() -> int:
parser = argparse.ArgumentParser(
description="Compare supported PIMCOMP models with Raptor latency and throughput schedules."
)
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(
"--arch",
nargs="+",
choices=ARCHES,
default=list(ARCHES),
help="PIM architectures to run (default: all architectures).",
)
parser.add_argument(
"--mode",
nargs="+",
choices=("latency", "throughput"),
default=["latency", "throughput"],
help="Pimsim modes to run (default: both).",
)
parser.add_argument(
"--only",
choices=("raptor", "pimcomp"),
help="Re-run only this compiler and reuse the other side's existing artifacts.",
)
parser.add_argument(
"--pipeline",
type=int,
choices=(1, 2, 4, 8),
help="Run only this Raptor pipeline within the selected mode.",
)
parser.add_argument(
"--pimsim-time-ms",
type=int,
default=100,
help="throughput pimsim-nn horizon in ms (default: 100).",
)
parser.add_argument(
"--timeout-seconds",
type=float,
default=0.0,
help="Per-stage timeout in seconds; 0 means no timeout (default).",
)
parser.add_argument(
"-j",
"--jobs",
type=int,
default=os.cpu_count() or 1,
help="Number of comparisons to run in parallel (default: all available CPUs).",
)
parser.add_argument(
"--clean",
action="store_true",
help="Remove generated comparison artifacts and result summaries, then exit.",
)
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
args.arch = list(dict.fromkeys(args.arch))
args.mode = list(dict.fromkeys(args.mode))
if args.clean:
print(f"Removed {clean_artifacts(out_dir, args.models, args.arch)} comparison artifact path(s).")
return 0
if args.jobs < 1:
parser.error("--jobs must be at least 1")
if args.pimsim_time_ms <= 0:
parser.error("--pimsim-time-ms must be positive")
if args.timeout_seconds < 0:
parser.error("--timeout-seconds must be non-negative")
comparisons_by_arch: dict[str, tuple[tuple[str, int, str], ...]] = {}
configs_by_arch: dict[str, dict[str, Path]] = {}
for arch in args.arch:
comparisons = tuple(
comparison for comparison in COMPARISONS
if comparison[0] in args.mode
and (args.pipeline is None or comparison[1] == args.pipeline)
)
if not comparisons:
parser.error("--pipeline does not belong to the selected --mode")
configs = {
mode: config_path(arch, mode, args.pimsim_time_ms, write=not args.dry_run)
for mode, _, _ in comparisons
}
unsupported = [
pipeline for mode, pipeline, _ in comparisons
if mode == "throughput" and core_count(configs[mode]) % pipeline
]
if unsupported:
if args.pipeline is not None:
parser.error(
f"{arch} has {core_count(configs['throughput'])} cores; "
f"throughput pipelines must divide that count (invalid: {unsupported})"
)
print(
Fore.YELLOW
+ f"Skipping unsupported {arch} throughput pipeline(s): {unsupported}"
+ Style.RESET_ALL
)
comparisons = tuple(comparison for comparison in comparisons if comparison[1] not in unsupported)
comparisons_by_arch[arch] = comparisons
configs_by_arch[arch] = configs
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.only is not None:
missing_reuse = []
for arch in args.arch:
for name in args.models:
for mode, pipeline, _ in comparisons_by_arch[arch]:
comparison_dir = result_dir(out_dir, name, arch, mode, pipeline)
required = (
comparison_dir / "pimcomp/comparison_report.json"
if args.only == "pimcomp"
else comparison_dir / "pimcomp/output/SimulationInfo.gz"
)
if not required.exists():
missing_reuse.append(str(required))
elif args.only == "raptor":
report = comparison_dir / "pimcomp/comparison_report.json"
if not report.exists() or json.loads(report.read_text(encoding="utf-8")).get(
"pimcomp_model_source"
) != "original_onnx":
missing_reuse.append(f"{report} (not generated from the original ONNX model)")
if missing_reuse:
parser.error(
f"--only {args.only} needs existing comparison artifacts: "
+ ", ".join(missing_reuse)
)
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"Architectures: {', '.join(args.arch)}")
print(f"Modes: {', '.join(args.mode)}")
print(f"Throughput pimsim time: {args.pimsim_time_ms} ms")
print(f"Max parallel jobs: {args.jobs}")
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)
if args.only != "raptor":
print_stage("Build PIMCOMP", STAGE_COLORS["Build Runner"])
run(
["cmake", "--build", str(PIMCOMP_SOURCE / "build"), "--target", "PIMCOMP-NN"],
dry_run=args.dry_run,
)
print_stage("Prepare shared artifacts", STAGE_COLORS["Build Runner"])
for name in args.models:
run(
prepare_common_command(
MODELS[name],
model_dir(out_dir, name) / "common",
args.timeout_seconds,
),
dry_run=args.dry_run,
)
failed = []
comparison_specs: list[ComparisonSpec] = []
shared_pimcomp_by_group: dict[tuple[str, str, str], Path] = {}
for arch in args.arch:
comparisons = comparisons_by_arch[arch]
configs = configs_by_arch[arch]
for index, name in enumerate(args.models, start=1):
for mode, pipeline, pimcomp_pipeline in comparisons:
model_result_dir = result_dir(out_dir, name, arch, mode, pipeline)
label = f"{arch}/{name}/{mode}/pipeline{pipeline}"
group = (arch, name, mode)
shared_pimcomp_dir = shared_pimcomp_by_group.get(group)
anchor = shared_pimcomp_dir is None
if anchor:
shared_pimcomp_dir = model_result_dir / "pimcomp/output"
shared_pimcomp_by_group[group] = shared_pimcomp_dir
print(
"\n" + Fore.CYAN + f"[{arch} {index}/{len(args.models)}]" + Style.RESET_ALL
+ f" {Style.BRIGHT}Comparing {name} ({mode}, pipeline={pipeline}){Style.RESET_ALL}",
flush=True,
)
comparison_specs.append(
ComparisonSpec(
label=label,
model=MODELS[name],
output_dir=model_result_dir,
common_dir=model_dir(out_dir, name) / "common",
config=configs[mode],
mode=mode,
pipeline=pipeline,
pimcomp_pipeline=pimcomp_pipeline,
shared_pimcomp_dir=shared_pimcomp_dir,
anchor=anchor,
)
)
if args.dry_run:
for spec in comparison_specs:
command = comparison_command_for(
spec,
args,
reuse_shared_pimcomp=not spec.anchor,
)
if run(command, dry_run=True, check=False):
failed.append(spec.label)
elif comparison_specs:
print(f"Running {len(comparison_specs)} comparison(s)")
anchor_jobs = [
(
spec.label,
comparison_command_for(spec, args, reuse_shared_pimcomp=False),
None,
)
for spec in comparison_specs
if spec.anchor
]
dependent_specs = [spec for spec in comparison_specs if not spec.anchor]
print_directly = min(args.jobs, len(comparison_specs)) == 1
with (nullcontext(None) if print_directly else TemporaryDirectory(prefix="raptor-pimcomp-")) as log_dir:
anchor_failed, log_offset = run_comparison_jobs(
anchor_jobs,
args.jobs,
Path(log_dir) if log_dir else None,
0,
)
failed.extend(anchor_failed)
dependent_jobs = [
(
spec.label,
comparison_command_for(
spec,
args,
reuse_shared_pimcomp=pimcomp_artifact_ready(spec.shared_pimcomp_dir),
),
None,
)
for spec in dependent_specs
]
dependent_failed, _ = run_comparison_jobs(
dependent_jobs,
args.jobs,
Path(log_dir) if log_dir else None,
log_offset,
)
failed.extend(dependent_failed)
if args.dry_run:
return 1 if failed else 0
results_path = None
for arch in args.arch:
results_path = write_results_csv(
out_dir,
arch,
args.models,
comparisons_by_arch[arch],
)
assert results_path is not None
print_stage(results_path.name, STAGE_COLORS["Compare Outputs"])
print(results_path.read_text(encoding="utf-8"), end="")
for arch in args.arch:
for name in args.models:
for mode, pipeline, _ in comparisons_by_arch[arch]:
label = f"{arch}/{name}/{mode}/pipeline{pipeline}"
report_path = result_dir(out_dir, name, arch, mode, pipeline) / "pimcomp/comparison_report.json"
if not report_path.exists() or not comparison_passed(
json.loads(report_path.read_text(encoding="utf-8"))
):
if label not in failed:
failed.append(label)
print("\n" + Style.BRIGHT + Fore.CYAN + "Summary" + Style.RESET_ALL)
total_jobs = sum(len(args.models) * len(comparisons) for comparisons in comparisons_by_arch.values())
print(Style.BRIGHT + f"Passed: {total_jobs - 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())
@@ -0,0 +1,49 @@
# PIMCOMP batch correctness reproduction
PIMCOMP's batch scheduler currently emits an incomplete standalone program for
models containing post operations. The generated `VerificationInfo.json` uses a
negative `source_address` to identify the preceding node, but the batch
verifier resolves that address by copying the provider tensor directly from
ONNX Runtime's `intermediate_result` table. It does not read the simulated
global-memory buffer.
The corresponding post-operation scheduling call is commented out in
[`BatchPipelineSchedule.cpp`](../../../../../third_party/PIMCOMP-NN/backend/BatchPipelineSchedule.cpp#L1535).
Consequently, Pool, ReLU, Concat, and Reshape provider tensors can be consumed
by later `LD` instructions without any generated `ST` producer. The comparison
exporter initializes those host-side intermediate buffers to zero, which makes
the missing computation visible in Rust functional validation.
For the checked-in GoogLeNet throughput/pipeline2 artifact, 39 provider tensors
are referenced by batch loads. Nineteen have generated stores; twenty are
never written. Preloading the provider tensors with the same ONNX Runtime
intermediates used by PIMCOMP's verifier makes the exported program pass. This
reproduces the verifier's input contract; it does not repair PIMCOMP's batch
schedule.
Run the default reproduction from the repository root:
```bash
.venv/bin/python validation/tools/pim/pimcomp/correctness/run_prefill_experiment.py
```
The launcher accepts an alternate comparison directory, model, and work
work directory, and shared reference-artifact directory:
```bash
.venv/bin/python validation/tools/pim/pimcomp/correctness/run_prefill_experiment.py \
validation/networks/pimcomp_models/googlenet/arch-a/throughput/pipeline2 \
validation/networks/pimcomp_models/googlenet/googlenet-12-pimsim-nn.onnx \
/tmp/pimcomp-prefill-googlenet \
validation/networks/pimcomp_models/googlenet/common
```
It runs the exported artifact once with its original memory image and once
with [`prefill_batch_memory.py`](prefill_batch_memory.py), then compares both
outputs with the recorded native reference. The expected GoogLeNet result is a
baseline maximum difference near `6.70705` and a prefilled maximum difference
near `4.05e-6`.
The issue is in PIMCOMP batch scheduling/validation semantics, not in the Rust
simulator's vector-length interpretation. Vector lengths remain element counts
as specified by the reference ISA.
@@ -0,0 +1,171 @@
#!/usr/bin/env python3
"""Populate the host-side intermediate buffers expected by PIMCOMP batch mode."""
from __future__ import annotations
import argparse
import json
import shutil
from pathlib import Path
import numpy as np
import onnx
import onnxruntime as ort
def sanitize_output_name(name: str) -> str:
return "".join(char if char.isalnum() or char in "_.-" else "_" for char in name[:255])
def flatten_reference(value: np.ndarray) -> np.ndarray:
if value.ndim == 4:
value = value.transpose(0, 2, 3, 1)
elif value.ndim == 2:
value = value.transpose()
else:
raise ValueError(f"PIMCOMP batch verification only flattens 2D/4D tensors, got {value.shape}")
return value.astype(np.float32, copy=False).reshape(-1)
def provider_indices(verification_info: dict) -> list[int]:
indices = set()
for core in verification_info["instruction"]["core_list"]:
for instruction in core or []:
if (
instruction["operation"] == "LD"
and instruction.get("stage") in ("INPUT", "POST")
and instruction["node_index"] != 1
):
source_address = instruction["source_address"]
if source_address >= 0:
raise ValueError(f"Expected a negative provider address, got {source_address}")
indices.add(-source_address)
return sorted(indices)
def output_base(
model: onnx.ModelProto,
config: dict,
node_list: list[dict],
max_output: int,
) -> tuple[int, list[dict]]:
if len(model.graph.output) != len(config["outputs_addresses"]):
raise ValueError("Exported output addresses do not match the ONNX graph outputs")
nodes_by_name = {node["name"]: node for node in node_list}
bases = []
outputs = []
for index, graph_output in enumerate(model.graph.output):
node = nodes_by_name[graph_output.name]
address = config["outputs_addresses"][index]
base = int(address - node["new_node_index"] * max_output * 4)
bases.append(base)
outputs.append({"name": graph_output.name, "address": address})
if len(set(bases)) != 1:
raise ValueError(f"Exported output addresses use different memory bases: {bases}")
return bases[0], outputs
def prefill_batch_memory(
model_path: Path,
comparison_dir: Path,
input_path: Path,
output_memory: Path,
metadata_path: Path,
) -> None:
comparison_dir = comparison_dir.resolve()
model = onnx.load(model_path)
verification_info = json.loads(
(comparison_dir / "pimcomp/output/VerificationInfo.json").read_text(encoding="utf-8")
)
config = json.loads((comparison_dir / "pimcomp/exported/config.json").read_text(encoding="utf-8"))
node_list = verification_info["node_list"]
max_output = max(int(np.prod(node["output_dim"], dtype=int)) for node in node_list)
base, outputs = output_base(model, config, node_list, max_output)
providers = provider_indices(verification_info)
runtime_model = onnx.ModelProto()
runtime_model.CopyFrom(model)
existing_outputs = {output.name for output in runtime_model.graph.output}
for index in providers:
name = node_list[index]["name"]
if name not in existing_outputs:
runtime_model.graph.output.append(onnx.ValueInfoProto(name=name))
session = ort.InferenceSession(runtime_model.SerializeToString(), providers=["CPUExecutionProvider"])
session_inputs = session.get_inputs()
if len(session_inputs) != 1:
raise ValueError("PIMCOMP export currently requires exactly one runtime input tensor")
input_meta = session_inputs[0]
input_tensor = np.loadtxt(input_path, delimiter=",", dtype=np.float32).reshape(input_meta.shape)
provider_names = [node_list[index]["name"] for index in providers]
provider_values = session.run(provider_names, {input_meta.name: input_tensor})
output_memory = output_memory.resolve()
metadata_path = metadata_path.resolve()
if output_memory == (comparison_dir / "pimcomp/exported/memory.bin").resolve():
raise ValueError("Refusing to overwrite the original exported memory.bin")
output_memory.parent.mkdir(parents=True, exist_ok=True)
shutil.copy2(comparison_dir / "pimcomp/exported/memory.bin", output_memory)
nodes_by_index = {index: node for index, node in enumerate(node_list)}
with output_memory.open("r+b") as memory:
memory_size = output_memory.stat().st_size
for index, value in zip(providers, provider_values):
flattened = flatten_reference(value)
node = nodes_by_index[index]
if node["operation"] in ("OP_CONV", "OP_FC") and node.get("with_act") == 1:
flattened = np.maximum(flattened, 0)
address = base + index * max_output * 4
end = address + flattened.nbytes
if end > memory_size:
raise ValueError(f"Provider {index} exceeds exported memory: {end} > {memory_size}")
memory.seek(address)
memory.write(flattened.tobytes())
output_metadata = {output["name"]: output for output in outputs}
for index, output_meta in enumerate(session.get_outputs()):
if output_meta.name not in output_metadata:
continue
output = output_metadata[output_meta.name]
output["dump"] = f"{output['address']},{int(np.prod(output_meta.shape, dtype=int)) * 4}"
output["reference"] = str(
comparison_dir / f"outputs/output{index}_{sanitize_output_name(output_meta.name)}.csv"
)
metadata_path.parent.mkdir(parents=True, exist_ok=True)
metadata_path.write_text(
json.dumps(
{
"memory": str(output_memory),
"output_base": base,
"provider_indices": providers,
"outputs": outputs,
},
indent=2,
)
+ "\n",
encoding="utf-8",
)
print(f"Prefilled {len(providers)} provider tensor(s) into {output_memory}")
print(f"Metadata: {metadata_path}")
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", type=Path, required=True)
parser.add_argument("--comparison-dir", type=Path, required=True)
parser.add_argument("--input", type=Path, required=True)
parser.add_argument("--output-memory", type=Path, required=True)
parser.add_argument("--metadata", type=Path, required=True)
args = parser.parse_args()
prefill_batch_memory(
args.model,
args.comparison_dir,
args.input,
args.output_memory,
args.metadata,
)
return 0
if __name__ == "__main__":
raise SystemExit(main())
@@ -0,0 +1,127 @@
#!/usr/bin/env python3
"""Reproduce the PIMCOMP batch prefill correctness experiment."""
from __future__ import annotations
import argparse
import json
import subprocess
import sys
import tempfile
from pathlib import Path
import numpy as np
from prefill_batch_memory import prefill_batch_memory
SCRIPT_DIR = Path(__file__).resolve().parent
REPO_ROOT = SCRIPT_DIR.parents[4]
DEFAULT_COMPARISON_DIR = REPO_ROOT / "validation/networks/pimcomp_models/googlenet/arch-a/throughput/pipeline2"
DEFAULT_MODEL = REPO_ROOT / "validation/networks/pimcomp_models/googlenet/googlenet-12-pimsim-nn.onnx"
DEFAULT_COMMON_DIR = REPO_ROOT / "validation/networks/pimcomp_models/googlenet/common"
SIMULATOR_MANIFEST = REPO_ROOT / "backend-simulators/pim/pim-simulator/Cargo.toml"
def run_simulator(
comparison_dir: Path,
memory: Path,
output: Path,
dump: str,
) -> None:
subprocess.run(
[
"cargo",
"run",
"--manifest-path",
str(SIMULATOR_MANIFEST),
"--no-default-features",
"--release",
"--bin",
"pim-simulator",
"--",
"-f",
str(comparison_dir / "pimcomp/exported"),
"--memory",
str(memory),
"-o",
str(output),
"-d",
dump,
],
cwd=REPO_ROOT,
check=True,
)
def compare_outputs(
baseline_path: Path,
prefilled_path: Path,
reference_path: Path,
work_dir: Path,
) -> None:
baseline = np.fromfile(baseline_path, dtype=np.float32)
prefilled = np.fromfile(prefilled_path, dtype=np.float32)
reference = np.loadtxt(reference_path, delimiter=",", dtype=np.float32).reshape(-1)
baseline_diff = np.max(np.abs(baseline - reference))
prefilled_diff = np.max(np.abs(prefilled - reference))
print(f"work directory: {work_dir}")
print(f"baseline max diff: {baseline_diff:.9g}")
print(f"prefilled max diff: {prefilled_diff:.9g}")
if np.allclose(baseline, reference, atol=1e-4, rtol=1e-3):
raise RuntimeError("baseline unexpectedly passes")
if not np.allclose(prefilled, reference, atol=1e-4, rtol=1e-3):
raise RuntimeError("prefilled run does not pass")
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("comparison_dir", nargs="?", type=Path, default=DEFAULT_COMPARISON_DIR)
parser.add_argument("model", nargs="?", type=Path, default=DEFAULT_MODEL)
parser.add_argument("work_dir", nargs="?", type=Path)
parser.add_argument("common_dir", nargs="?", type=Path, default=DEFAULT_COMMON_DIR)
args = parser.parse_args()
comparison_dir = args.comparison_dir.resolve()
model = args.model.resolve()
common_dir = args.common_dir.resolve()
if args.work_dir is None:
work_dir = Path(tempfile.mkdtemp(prefix="pimcomp-prefill."))
else:
work_dir = args.work_dir.resolve()
work_dir.mkdir(parents=True, exist_ok=True)
input_path = common_dir / "inputs/in0.csv"
if not input_path.is_file():
input_path = comparison_dir / "inputs/in0.csv"
prefilled_memory = work_dir / "prefilled_memory.bin"
metadata_path = work_dir / "metadata.json"
baseline_output = work_dir / "baseline.out.bin"
prefilled_output = work_dir / "prefilled.out.bin"
prefill_batch_memory(
model,
comparison_dir,
input_path,
prefilled_memory,
metadata_path,
)
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
output = metadata["outputs"][0]
dump = output["dump"]
reference = Path(output["reference"])
run_simulator(
comparison_dir,
comparison_dir / "pimcomp/exported/memory.bin",
baseline_output,
dump,
)
run_simulator(comparison_dir, prefilled_memory, prefilled_output, dump)
compare_outputs(baseline_output, prefilled_output, reference, work_dir)
return 0
if __name__ == "__main__":
raise SystemExit(main())
@@ -0,0 +1,438 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import json
import re
import struct
import sys
from collections import Counter, namedtuple
from pathlib import Path
Event = namedtuple("Event", "op peer size ordinal instruction")
Program = namedtuple("Program", "events operations starts_inactive ends_with_send")
Transfer = namedtuple(
"Transfer",
"sender receiver size sender_ordinal receiver_ordinal sender_instruction receiver_instruction "
"sender_last_instruction receiver_last_instruction count",
)
CORE_FILE = re.compile(r"core_(\d+)\.(json|pim)$")
HEADER = struct.Struct("<4sII")
RECORD = struct.Struct("<BBBBiiii")
OPCODE_NAMES = (
"nop", "sldi", "sld", "sadd", "ssub", "smul", "saddi", "smuli", "setbw",
"mvmul", "vvadd", "vvsub", "vvmul", "vvdmul", "vvmax", "vvsll", "vvsra",
"vavg", "vrelu", "vtanh", "vsigm", "vsoftmax", "vmv", "vrsu", "vrsl",
"ld", "st", "lldi", "lmv", "send", "recv", "wait", "sync",
)
IGNORED_OPS = {"nop", "sldi", "lldi", "setbw"}
MEMORY_OPS = {"sld", "ld", "st", "vmv", "vrsu", "vrsl", "lmv"}
DISPLAY_ORDER = (
"ld", "lmv", "st", "sld", "vmv", "vrsu", "vrsl",
"vvmul", "mvmul", "vvdmul", "smul", "smuli",
"vvadd", "vvsub", "vvmax", "vavg", "sadd", "saddi", "ssub",
"vvsll", "vvsra", "vrelu", "vtanh", "vsigm", "vsoftmax", "wait", "sync",
)
DISPLAY_RANK = {op: rank for rank, op in enumerate(DISPLAY_ORDER)}
ARROW_HEIGHT = 2.0
def summarize_operations(operations: tuple[str, ...] | list[str]) -> str:
counts = Counter(op for op in operations if op not in IGNORED_OPS and op not in ("send", "recv"))
ordered = sorted(counts, key=lambda op: (DISPLAY_RANK.get(op, len(DISPLAY_RANK)), op))
groups = ([op for op in ordered if op in MEMORY_OPS], [op for op in ordered if op not in MEMORY_OPS])
rows = [
[(counts[op], op) for op in group[index : index + 2]]
for group in groups
for index in range(0, len(group), 2)
]
if not rows:
return ""
count_width = max(len(str(count)) for row in rows for count, _ in row)
op_width = max(len(op) for row in rows for _, op in row)
return '""' + "\\n".join(
" ".join(f"{count:>{count_width}}x {op:<{op_width}}" for count, op in row).rstrip()
for row in rows
) + '""'
def read_json(path: Path) -> Program:
instructions = json.loads(path.read_text(encoding="utf-8"))
if not isinstance(instructions, list):
raise ValueError(f"{path}: expected a JSON instruction array")
operations = tuple(
str(instruction.get("op", "unknown")) if isinstance(instruction, dict) else "unknown"
for instruction in instructions
)
events = []
for index, instruction in enumerate(instructions):
if not isinstance(instruction, dict):
continue
op = instruction.get("op")
if op not in ("send", "recv"):
continue
try:
peer = int(instruction["core"])
size = int(instruction["size"])
except (KeyError, TypeError, ValueError) as error:
raise ValueError(f"{path}: invalid communication instruction {instruction!r}") from error
events.append(Event(op, peer, size, len(events), index))
starts_inactive = bool(
events
and events[0].op == "recv"
and all(
isinstance(instruction, dict) and instruction.get("op") == "sldi"
for instruction in instructions[:events[0].instruction]
)
)
last_op = instructions[-1].get("op") if instructions and isinstance(instructions[-1], dict) else None
return Program(events, operations, starts_inactive, last_op == "send")
def read_binary(path: Path) -> Program:
data = path.read_bytes()
if len(data) < HEADER.size:
raise ValueError(f"{path}: binary core file is too small")
magic, version, count = HEADER.unpack_from(data)
expected_size = HEADER.size + count * RECORD.size
if magic != b"PIMB":
raise ValueError(f"{path}: invalid PIM binary magic")
if version != 1:
raise ValueError(f"{path}: unsupported PIM binary version {version}")
if len(data) != expected_size:
raise ValueError(f"{path}: expected {expected_size} bytes, found {len(data)}")
events = []
operations = []
last_opcode = None
for index in range(count):
opcode, _, _, _, peer, _, _, size = RECORD.unpack_from(data, HEADER.size + index * RECORD.size)
last_opcode = opcode
operations.append(OPCODE_NAMES[opcode] if opcode < len(OPCODE_NAMES) else f"opcode_{opcode}")
if opcode in (29, 30):
events.append(Event("send" if opcode == 29 else "recv", peer, size, len(events), index))
starts_inactive = bool(
events
and events[0].op == "recv"
and all(data[HEADER.size + index * RECORD.size] == 1 for index in range(events[0].instruction))
)
return Program(events, tuple(operations), starts_inactive, last_opcode == 29)
def read_programs(directory: Path, artifact_format: str) -> dict[int, Program]:
candidates: dict[str, dict[int, Path]] = {"json": {}, "pim": {}}
for path in directory.iterdir():
match = CORE_FILE.fullmatch(path.name)
if match:
candidates[match.group(2)][int(match.group(1))] = path
if artifact_format == "auto":
artifact_format = "json" if candidates["json"] else "pim"
files = candidates[artifact_format]
if not files:
raise ValueError(f"{directory}: no core_*.{artifact_format} files found")
reader = read_json if artifact_format == "json" else read_binary
return {core: reader(path) for core, path in files.items()}
def match_transfers(programs: dict[int, Program]) -> list[Transfer]:
positions = {core: 0 for core in programs}
transfers = []
while True:
made_progress = False
for core in sorted(programs):
events = programs[core].events
position = positions[core]
if position == len(events):
continue
event = events[position]
if event.peer not in programs:
sender, receiver = (core, event.peer) if event.op == "send" else (event.peer, core)
transfers.append(
Transfer(
sender,
receiver,
event.size,
event.ordinal if event.op == "send" else None,
event.ordinal if event.op == "recv" else None,
event.instruction if event.op == "send" else None,
event.instruction if event.op == "recv" else None,
event.instruction if event.op == "send" else None,
event.instruction if event.op == "recv" else None,
1,
)
)
positions[core] += 1
made_progress = True
continue
peer_events = programs[event.peer].events
peer_position = positions[event.peer]
if peer_position == len(peer_events):
continue
peer_event = peer_events[peer_position]
if (
event.peer == core
or peer_event.peer != core
or peer_event.op == event.op
or peer_event.size != event.size
):
continue
send = event if event.op == "send" else peer_event
receive = peer_event if event.op == "send" else event
sender, receiver = (core, event.peer) if event.op == "send" else (event.peer, core)
transfers.append(
Transfer(
sender,
receiver,
event.size,
send.ordinal,
receive.ordinal,
send.instruction,
receive.instruction,
send.instruction,
receive.instruction,
1,
)
)
positions[core] += 1
positions[event.peer] += 1
made_progress = True
if made_progress:
continue
remaining = {
core: programs[core].events[position]
for core, position in positions.items()
if position < len(programs[core].events)
}
if not remaining:
return transfers
details = ", ".join(
f"core {core}: {event.op} {event.peer} ({event.size} B)"
for core, event in sorted(remaining.items())
)
raise ValueError(f"communication streams cannot be matched at {details}")
def visible_transfers(
transfers: list[Transfer],
selected: set[int],
programs: dict[int, Program],
) -> list[Transfer]:
visible = [transfer for transfer in transfers if transfer.sender in selected or transfer.receiver in selected]
grouped = []
for transfer in visible:
previous = grouped[-1] if grouped else None
can_group = (
previous is not None
and transfer.sender_ordinal is not None
and previous.sender == transfer.sender
and previous.receiver == transfer.receiver
and previous.sender_ordinal + previous.count == transfer.sender_ordinal
and (
transfer.receiver_ordinal is None
or previous.receiver_ordinal + previous.count == transfer.receiver_ordinal
)
)
if can_group:
sender_gap = programs[transfer.sender].operations[
previous.sender_last_instruction + 1 : transfer.sender_instruction
]
receiver_gap = (
programs[transfer.receiver].operations[
previous.receiver_last_instruction + 1 : transfer.receiver_instruction
]
if transfer.receiver in programs
else ()
)
can_group = not summarize_operations(sender_gap) and not summarize_operations(receiver_gap)
if can_group:
grouped[-1] = previous._replace(
size=previous.size + transfer.size,
sender_last_instruction=transfer.sender_last_instruction,
receiver_last_instruction=transfer.receiver_last_instruction,
count=previous.count + transfer.count,
)
else:
grouped.append(transfer)
return grouped
def collect_operation_notes(
cores: list[int],
transfers: list[Transfer],
programs: dict[int, Program],
) -> dict[int, list[tuple[int, str]]]:
notes: dict[int, list[tuple[int, str]]] = {}
last_instructions = {core: -1 for core in cores}
anchors = {core: -1 for core in cores}
for transfer_index, transfer in enumerate(transfers):
endpoints = {
transfer.sender: (transfer.sender_instruction, transfer.sender_last_instruction),
transfer.receiver: (transfer.receiver_instruction, transfer.receiver_last_instruction),
}
for core, (instruction, last_instruction) in endpoints.items():
if core not in last_instructions or instruction is None:
continue
summary = summarize_operations(
programs[core].operations[last_instructions[core] + 1 : instruction]
)
if summary:
notes.setdefault(anchors[core], []).append((core, summary))
last_instructions[core] = last_instruction
anchors[core] = transfer_index
for core in cores:
summary = summarize_operations(programs[core].operations[last_instructions[core] + 1 :])
if summary:
notes.setdefault(anchors[core], []).append((core, summary))
return notes
def note_height(summary: str) -> float:
return 2.5 + 1.6 * (summary.count("\\n") + 1)
def parallel_note_order(
notes: list[tuple[int, str]],
cores: list[int],
) -> list[tuple[int, str]]:
positions = {core: position for position, core in enumerate(cores)}
packed: list[list[tuple[int, str]]] = []
for note in sorted(notes, key=lambda item: positions[item[0]]):
row = next(
(
row
for row in packed
if positions[note[0]] - positions[row[-1][0]] > 1
),
None,
)
if row is None:
row = []
packed.append(row)
row.append(note)
return [note for row in packed for note in row]
def render_text(cores: list[int], transfers: list[Transfer], programs: dict[int, Program]) -> str:
aliases = {core: f"C{core}" for core in cores}
positions = {core: position for position, core in enumerate(cores)}
lines = [f'participant "Core {core}" as {aliases[core]}' for core in cores]
first_receives = {
core: next((event.ordinal for event in programs[core].events if event.op == "recv"), None)
for core in cores
}
active = {core for core in cores if not programs[core].starts_inactive}
lines.extend(f"activate {aliases[core]}" for core in cores if core in active)
notes = collect_operation_notes(cores, transfers, programs)
availability = [0.0] * (2 * len(cores) + 1)
cursor = 0.0
def emit(line: str, left: int, right: int, height: float) -> None:
nonlocal cursor
start = max(availability[left : right + 1])
gap = start - cursor
if abs(gap) >= 0.05:
value = f"{gap:.1f}".rstrip("0").rstrip(".")
lines.append(f"space {value}")
lines.append(line)
cursor = start + height
availability[left : right + 1] = [cursor] * (right - left + 1)
def emit_notes(anchor: int) -> None:
for core, summary in parallel_note_order(notes.get(anchor, []), cores):
center = 2 * positions[core] + 1
emit(f"note over {aliases[core]}:{summary}", center - 1, center + 1, note_height(summary))
emit_notes(-1)
for transfer_index, transfer in enumerate(transfers):
label = f"{transfer.size} B"
if transfer.count > 1:
label = f"{transfer.count} sends, {label}"
if transfer.sender in aliases and transfer.receiver in aliases:
sender = 2 * positions[transfer.sender] + 1
receiver = 2 * positions[transfer.receiver] + 1
line = f"{aliases[transfer.sender]}->(1){aliases[transfer.receiver]}:{label}"
left, right = sorted((sender, receiver))
elif transfer.receiver in aliases:
line = f"[->(1){aliases[transfer.receiver]}:{label}"
left, right = 0, 2 * positions[transfer.receiver] + 1
else:
line = f"{aliases[transfer.sender]}->(1)]:{label}"
left, right = 2 * positions[transfer.sender] + 1, len(availability) - 1
emit(line, left, right, ARROW_HEIGHT)
receiver = transfer.receiver
if (
receiver in aliases
and receiver not in active
and transfer.receiver_ordinal == first_receives[receiver]
):
lines.append(f"activate {aliases[receiver]}")
active.add(receiver)
sender = transfer.sender
if (
sender in active
and programs[sender].ends_with_send
and transfer.sender_ordinal + transfer.count == len(programs[sender].events)
):
lines.append(f"deactivate {aliases[sender]}")
active.remove(sender)
emit_notes(transfer_index)
lines.extend(f"deactivateafter {aliases[core]}" for core in cores if core in active)
return "\n".join(lines) + "\n"
def main() -> int:
parser = argparse.ArgumentParser(
description="Generate SequenceDiagram.org text for PIM communication and intervening work."
)
parser.add_argument("pim_dir", type=Path, help="Directory containing core_<id>.json or core_<id>.pim files")
selection = parser.add_mutually_exclusive_group(required=True)
selection.add_argument("--cores", nargs="+", type=int, help="Core lifelines to display, in column order")
selection.add_argument("--all-cores", action="store_true", help="Display every used core, ordered by core ID")
parser.add_argument("--format", choices=("auto", "json", "pim"), default="auto")
parser.add_argument("-o", "--output", type=Path, help="Text output path (default: stdout)")
args = parser.parse_args()
try:
if not args.pim_dir.is_dir():
raise ValueError(f"{args.pim_dir}: not a directory")
if args.cores is not None and len(set(args.cores)) != len(args.cores):
raise ValueError("--cores contains duplicates")
programs = read_programs(args.pim_dir, args.format)
cores = (
[core for core in sorted(programs) if programs[core].operations]
if args.all_cores
else args.cores
)
missing = [core for core in cores if core not in programs]
if missing:
raise ValueError(f"missing artifact for selected core(s): {', '.join(map(str, missing))}")
transfers = visible_transfers(match_transfers(programs), set(cores), programs)
diagram = render_text(cores, transfers, programs)
if args.output:
args.output.write_text(diagram, encoding="utf-8")
else:
sys.stdout.write(diagram)
except (OSError, ValueError) as error:
parser.error(str(error))
return 0
if __name__ == "__main__":
raise SystemExit(main())