#!/usr/bin/env python3 import argparse import os import signal import subprocess import sys from concurrent.futures import ProcessPoolExecutor, as_completed from contextlib import redirect_stderr, redirect_stdout from itertools import groupby from pathlib import Path from tempfile import TemporaryDirectory from colorama import Style, Fore from raptor_validation.validate_one import ( MODE_COMPILE_ONLY, MODE_FULL, MODE_RUN_ONLY, MODE_STAGE_TITLES, PIMSIM_DONE, PIMSIM_FAILED, PIMSIM_NOT_RUN, PIMSIM_SKIPPED, PIMSIM_UNSUPPORTED, ProgressReporter, ValidationResult, clean_workspace_artifacts, load_pimcomp_hardware, pimcomp_compatibility_errors, validate_network, ) from raptor_validation.raptor import PIM_PASS_LABELS DEFAULT_PIMCOMP_CONFIG = "arch-a" PIMCOMP_CONFIG_CHOICES = ("arch-a", "arch-b", "arch-c") def run_validation_job(job): models, options = job completed = [] for index, onnx_path, rel, log_path in models: reporter = ProgressReporter( options["model_total"], stages_per_model=len(MODE_STAGE_TITLES[options["mode"]]), enabled=False, verbose=options["verbose"], ) sys.stdout.flush() sys.stderr.flush() 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): try: result = validate_network( onnx_path, reporter=reporter, model_index=index, model_total=options["model_total"], verbose=options["verbose"], mode=options["mode"], **options["validation_kwargs"], ) except Exception as exc: result = ValidationResult(False, pimsim_status=PIMSIM_NOT_RUN) print_validation_error(reporter, rel, exc) finally: reporter.finish() finally: os.dup2(saved_stdout, 1) os.dup2(saved_stderr, 2) os.close(saved_stdout) os.close(saved_stderr) completed.append((str(rel), result, str(log_path))) return completed def format_return_status(returncode): if returncode < 0: signal_num = -returncode try: signal_name = signal.Signals(signal_num).name except ValueError: signal_name = "UNKNOWN" return f"Program terminated by signal {signal_name} ({signal_num})." return f"Program exited with code {returncode}." def print_validation_error(reporter, rel, exc): reporter.suspend() print(Style.BRIGHT + Fore.RED + f"Exception while validating {rel}" + Style.RESET_ALL, file=sys.stderr, flush=True) if isinstance(exc, subprocess.CalledProcessError): print(format_return_status(exc.returncode), file=sys.stderr, flush=True) if getattr(exc, "output_already_streamed", False): print("Failure log already printed above.", file=sys.stderr, flush=True) elif exc.output: output_text = exc.output.decode("utf-8", errors="replace") if isinstance(exc.output, bytes) else str(exc.output) if output_text: print(output_text, file=sys.stderr, end="" if output_text.endswith("\n") else "\n", flush=True) else: print(f"{type(exc).__name__}: {exc}", file=sys.stderr, flush=True) print("=" * 72, file=sys.stderr, flush=True) reporter.resume() def print_average_pim_pass_timings(pass_timing_sums, pass_timing_counts, total_timing_sum, timed_benchmark_count): if timed_benchmark_count == 0: return print("\n" + Style.BRIGHT + Fore.CYAN + "Average PIM Pass Timings" + Style.RESET_ALL) for _, label in PIM_PASS_LABELS: count = pass_timing_counts[label] if count == 0: continue print(f" {label.ljust(28)} {pass_timing_sums[label] / count:.4f}s") print(f" {'Total'.ljust(28)} {total_timing_sum / timed_benchmark_count:.4f}s") def format_pimsim_metric(result, value, unit): if result.pimsim_status == PIMSIM_DONE: return f"{value:.6f} {unit}" return result.pimsim_status def main(): script_dir = Path(__file__).parent.resolve() pimcomp_configs_dir = script_dir / "pimsim_configs" / "pimcomp" default_pimsim_config = pimcomp_configs_dir / DEFAULT_PIMCOMP_CONFIG / "latency_config.json" default_hardware = load_pimcomp_hardware(default_pimsim_config) ap = argparse.ArgumentParser(description="Validate all ONNX operations under the operations/ directory.") ap.add_argument("--raptor-path", help="Path to the Raptor compiler binary.") ap.add_argument("--onnx-include-dir", help="Path to OnnxMlirRuntime include directory.") ap.add_argument("--operations-dir", default=None, help="Root of the operations tree (default: operations).") ap.add_argument("--simulator-dir", default=None, help="Path to the functional pim-simulator crate root " "(default: auto-detected relative to script).") ap.add_argument("--non-functional-simulator-build-dir", metavar="PATH", default=None, help="Path to the non-functional simulator build directory " "(default: auto-detected relative to script).") ap.add_argument("--pimcomp-config", choices=PIMCOMP_CONFIG_CHOICES, default=DEFAULT_PIMCOMP_CONFIG, help="Hardware/timing profile for non-functional simulation " "(default: arch-a).") ap.add_argument("--skip-non-functional-simulation", action="store_true", help="Skip non-functional simulation.") ap.add_argument("--threshold", type=float, default=1e-3, help="Absolute tolerance for per-element output comparison.") ap.add_argument("--relative-threshold", type=float, default=1e-5, help="Relative tolerance for per-element output comparison.") ap.add_argument("--seed", type=int, default=0, help="RNG seed for generated validation inputs.") ap.add_argument("--crossbar-size", type=int, default=default_hardware["crossbar_rows"], help=f"Crossbar size (default: {default_hardware['crossbar_rows']}).") ap.add_argument("--crossbar-count", type=int, default=default_hardware["crossbar_count"], help=f"Crossbars per core (default: {default_hardware['crossbar_count']}).") ap.add_argument("--core-count", type=int, default=default_hardware["core_count"], help=f"Core count (default: {default_hardware['core_count']}).") ap.add_argument("--raptor-extra-arg", action="append", default=[], help="Additional argument to pass through to the Raptor compiler. Repeat as needed.") ap.add_argument("--command-timeout-seconds", type=float, default=1000000.0, help="Per-subprocess timeout in seconds for compiler, runner, and simulation commands.") ap.add_argument("-j", "--jobs", type=int, default=os.cpu_count() or 1, help="Number of model validations to run in parallel (default: all available CPUs).") ap.add_argument("--clean", action="store_true", help="Remove generated validation artifacts under each model workspace and exit.") mode_group = ap.add_mutually_exclusive_group() mode_group.add_argument("--compile-only", action="store_true", help="Compile reference and PIM artifacts only; do not run reference execution, " "simulations, or comparison.") mode_group.add_argument("--run-only", action="store_true", help="Reuse existing compiled artifacts and only run inputs, reference execution, " "simulations, and comparison.") ap.add_argument("--verbose", action="store_true", help="Print per-stage progress and subprocess logs for passing validations too.") a = ap.parse_args() if a.jobs < 1: ap.error("--jobs must be at least 1") operations_dir = Path(a.operations_dir).resolve() if a.operations_dir else script_dir / "operations" simulator_dir = Path(a.simulator_dir).resolve() if a.simulator_dir else ( script_dir / ".." / "backend-simulators" / "pim" / "pim-simulator" ) pimsim_nn_build_dir = ( Path(a.non_functional_simulator_build_dir).resolve() if a.non_functional_simulator_build_dir else script_dir / ".." / "backend-simulators" / "pim" / "pimsim-nn" / "build" ) pimsim_config_path = pimcomp_configs_dir / a.pimcomp_config / "latency_config.json" if not operations_dir.is_dir(): print(Fore.RED + f"Operations directory not found: {operations_dir}" + Style.RESET_ALL) sys.exit(1) onnx_files = sorted(operations_dir.rglob("*.onnx")) if not onnx_files: print(Fore.YELLOW + f"No .onnx files found under {operations_dir}" + Style.RESET_ALL) sys.exit(1) if a.clean: removed_count = 0 for onnx_path in onnx_files: removed_count += len(clean_workspace_artifacts(onnx_path.parent, onnx_path.stem)) print(Style.BRIGHT + f"Removed {removed_count} generated artifact path(s)." + Style.RESET_ALL) sys.exit(0) missing_args = [] if not a.raptor_path: missing_args.append("--raptor-path") if not a.onnx_include_dir: missing_args.append("--onnx-include-dir") if missing_args: ap.error("the following arguments are required unless --clean is used: " + ", ".join(missing_args)) print(Style.BRIGHT + f"Found {len(onnx_files)} ONNX file(s) to validate." + Style.RESET_ALL) print(f"Operations root: {operations_dir}") print(f"Parallel jobs: {a.jobs}") print("=" * 72) mode = MODE_FULL if a.compile_only: mode = MODE_COMPILE_ONLY elif a.run_only: mode = MODE_RUN_ONLY selected_pimsim_config = None if not a.skip_non_functional_simulation: compatibility_errors = pimcomp_compatibility_errors( pimsim_config_path, core_count=a.core_count, crossbar_count=a.crossbar_count, crossbar_size=a.crossbar_size, ) if compatibility_errors: print( Fore.RED + "Non-functional simulation disabled: " + "; ".join(compatibility_errors) + Style.RESET_ALL, file=sys.stderr, ) else: selected_pimsim_config = pimsim_config_path results = {} # relative_path -> ValidationResult pass_timing_sums = {label: 0.0 for _, label in PIM_PASS_LABELS} pass_timing_counts = {label: 0 for _, label in PIM_PASS_LABELS} total_timing_sum = 0.0 timed_benchmark_count = 0 reporter = ProgressReporter(len(onnx_files), stages_per_model=1, verbose=a.verbose) validation_kwargs = { "raptor_path": a.raptor_path, "onnx_include_dir": a.onnx_include_dir, "simulator_dir": simulator_dir, "crossbar_size": a.crossbar_size, "crossbar_count": a.crossbar_count, "core_count": a.core_count, "raptor_extra_args": a.raptor_extra_arg, "pimsim_nn_build_dir": pimsim_nn_build_dir, "pimsim_config_path": selected_pimsim_config, "command_timeout_seconds": a.command_timeout_seconds, "threshold": a.threshold, "rtol": a.relative_threshold, "seed": a.seed, } indexed_files = list(enumerate(onnx_files, start=1)) workspace_groups = [ list(group) for _, group in groupby(indexed_files, key=lambda indexed_path: indexed_path[1].parent) ] with TemporaryDirectory(prefix="raptor-validation-") as log_dir: jobs = [ ( [ ( index, onnx_path, onnx_path.relative_to(operations_dir), Path(log_dir) / f"{index}.log", ) for index, onnx_path in workspace_group ], { "model_total": len(onnx_files), "mode": mode, "verbose": a.verbose, "validation_kwargs": validation_kwargs, }, ) for workspace_group in workspace_groups ] with ProcessPoolExecutor(max_workers=a.jobs) as executor: futures = [executor.submit(run_validation_job, job) for job in jobs] for future in as_completed(futures): completed_group = future.result() for rel, result, log_path in completed_group: reporter.suspend() output = Path(log_path).read_text(encoding="utf-8", errors="replace") if output: print(output, end="" if output.endswith("\n") else "\n") reporter.resume() reporter.advance() reporter.record_result(result.passed) results[rel] = result if result.pim_pass_timings: benchmark_total = 0.0 for label, duration in result.pim_pass_timings.items(): pass_timing_sums[label] += duration pass_timing_counts[label] += 1 benchmark_total += duration total_timing_sum += benchmark_total timed_benchmark_count += 1 reporter.finish() results = dict(sorted(results.items())) # Summary n_passed = sum(1 for result in results.values() if result.passed) n_total = len(results) status_width = len("Result") path_width = max(len("Operation"), *(len(rel) for rel in results)) formatted_metrics = { rel: ( format_pimsim_metric(result, result.pimsim_latency_ms, "ms"), format_pimsim_metric(result, result.pimsim_power_mw, "mW"), ) for rel, result in results.items() } latency_width = max(len("Latency"), *(len(metrics[0]) for metrics in formatted_metrics.values())) power_width = max(len("Power"), *(len(metrics[1]) for metrics in formatted_metrics.values())) separator = ( f"+-{'-' * path_width}-+-{'-' * status_width}-+-{'-' * latency_width}" f"-+-{'-' * power_width}-+") print(separator) print( f"| {'Operation'.ljust(path_width)} | {'Result'.ljust(status_width)} | " f"{'Latency'.rjust(latency_width)} | {'Power'.rjust(power_width)} |" ) print(separator) for rel, result in results.items(): plain_status = "PASS" if result.passed else "FAIL" status = Fore.GREEN + plain_status.ljust(status_width) + Style.RESET_ALL if result.passed else \ Fore.RED + plain_status.ljust(status_width) + Style.RESET_ALL latency, power = formatted_metrics[rel] print( f"| {rel.ljust(path_width)} | {status} | {latency.rjust(latency_width)} | " f"{power.rjust(power_width)} |") print(separator) print("\n" + Style.BRIGHT + Fore.CYAN + "Summary" + Style.RESET_ALL) print(Style.BRIGHT + f"Passed: {n_passed}" + Style.RESET_ALL) print(Style.BRIGHT + f"Failed: {n_total - n_passed}" + Style.RESET_ALL) measured_latencies = [ result.pimsim_latency_ms for result in results.values() if result.pimsim_status == PIMSIM_DONE ] pimsim_failed = sum( result.pimsim_status == PIMSIM_FAILED for result in results.values() ) pimsim_skipped = sum( result.pimsim_status in (PIMSIM_SKIPPED, PIMSIM_NOT_RUN) for result in results.values() ) pimsim_unsupported = sum( result.pimsim_status == PIMSIM_UNSUPPORTED for result in results.values() ) print( Style.BRIGHT + f"pimsim-nn: {len(measured_latencies)} measured, " f"{pimsim_failed} failed, {pimsim_unsupported} unsupported, " f"{pimsim_skipped} skipped" + Style.RESET_ALL ) if measured_latencies: print( Style.BRIGHT + f"Total latency: {sum(measured_latencies):.6f} ms" + Style.RESET_ALL ) if a.verbose: print_average_pim_pass_timings( pass_timing_sums, pass_timing_counts, total_timing_sum, timed_benchmark_count, ) sys.exit(0 if n_passed == n_total else 1) if __name__ == "__main__": main()