#!/usr/bin/env python3 import argparse import signal import subprocess import sys from pathlib import Path 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, 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 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("--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() 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("=" * 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=len(MODE_STAGE_TITLES[mode]), verbose=a.verbose) for index, onnx_path in enumerate(onnx_files, start=1): rel = onnx_path.relative_to(operations_dir) try: result = validate_network( onnx_path, a.raptor_path, a.onnx_include_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, reporter=reporter, model_index=index, model_total=len(onnx_files), verbose=a.verbose, mode=mode, ) results[str(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 except subprocess.CalledProcessError as exc: results[str(rel)] = ValidationResult(False, pimsim_status=PIMSIM_NOT_RUN) print_validation_error(reporter, rel, exc) except Exception as exc: results[str(rel)] = ValidationResult(False, pimsim_status=PIMSIM_NOT_RUN) print_validation_error(reporter, rel, exc) reporter.finish() # 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'.ljust(latency_width)} | {'Power'.ljust(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.ljust(latency_width)} | " f"{power.ljust(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() ) print( Style.BRIGHT + f"pimsim-nn: {len(measured_latencies)} measured, " f"{pimsim_failed} failed, {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()