diff --git a/validation/raptor.py b/validation/raptor.py index 8b7d8cb..0746cdf 100644 --- a/validation/raptor.py +++ b/validation/raptor.py @@ -4,10 +4,13 @@ from colorama import Fore, Style from subprocess_utils import run_command_with_reporter -def compile_with_raptor(network_path, raptor_onnx_path: Path, crossbar_size, crossbar_count, reporter=None): +def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path, + crossbar_size, crossbar_count, cwd=None, reporter=None): # Define the arguments, with the possibility to set crossbar size and count args = [ network_path, + "-o", + output_base, "--maccel=PIM", "--EmitPimCodegen", # "--use-experimental-conv-impl=true", @@ -18,6 +21,7 @@ def compile_with_raptor(network_path, raptor_onnx_path: Path, crossbar_size, cro try: run_command_with_reporter( [str(raptor_onnx_path)] + [str(arg) for arg in args], + cwd=cwd, reporter=reporter, ) if reporter is None: diff --git a/validation/validate.py b/validation/validate.py index 697fca5..c959360 100644 --- a/validation/validate.py +++ b/validation/validate.py @@ -5,19 +5,21 @@ import subprocess import sys from pathlib import Path from colorama import Style, Fore -from validate_one import ProgressReporter, validate_network +from validate_one import ProgressReporter, clean_workspace_artifacts, validate_network def main(): ap = argparse.ArgumentParser(description="Validate all ONNX operations under the operations/ directory.") - ap.add_argument("--raptor-path", required=True, help="Path to the Raptor compiler binary.") - ap.add_argument("--onnx-include-dir", required=True, help="Path to OnnxMlirRuntime include 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 pim-simulator crate root (default: auto-detected relative to script).") ap.add_argument("--threshold", type=float, default=1e-3, help="Max allowed diff per output element.") ap.add_argument("--crossbar-size", type=int, default=64) ap.add_argument("--crossbar-count", type=int, default=8) + ap.add_argument("--clean", action="store_true", + help="Remove generated validation artifacts under each model workspace and exit.") a = ap.parse_args() script_dir = Path(__file__).parent.resolve() @@ -35,6 +37,21 @@ def main(): 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) diff --git a/validation/validate_one.py b/validation/validate_one.py index 8f6b6e7..9fab85b 100644 --- a/validation/validate_one.py +++ b/validation/validate_one.py @@ -13,6 +13,7 @@ from subprocess_utils import run_command_with_reporter STAGE_COUNT = 6 +GENERATED_DIR_NAMES = ("inputs", "outputs", "raptor", "runner", "simulation") class ProgressReporter: @@ -88,6 +89,27 @@ def run_command(cmd, cwd=None, reporter=None): run_command_with_reporter(cmd, cwd=cwd, reporter=reporter) +def clean_workspace_artifacts(workspace_dir, model_stem): + workspace_dir = Path(workspace_dir) + removed_paths = [] + + def remove_path(path): + if path.is_symlink() or path.is_file(): + path.unlink(missing_ok=True) + removed_paths.append(path) + elif path.is_dir(): + shutil.rmtree(path) + removed_paths.append(path) + + for name in GENERATED_DIR_NAMES: + remove_path(workspace_dir / name) + + for suffix in (".onnx.mlir", ".so", ".tmp"): + remove_path(workspace_dir / f"{model_stem}{suffix}") + + return removed_paths + + def print_stage(reporter, model_index, model_total, model_name, title): stage_colors = { "Compile ONNX": Fore.BLUE, @@ -108,19 +130,15 @@ def print_info(reporter, message): def compile_onnx_network(network_onnx_path, raptor_path, raptor_dir, runner_dir, reporter=None): - run_command([raptor_path, network_onnx_path, "--EmitONNXIR"], reporter=reporter) - run_command([raptor_path, network_onnx_path], reporter=reporter) - parent = network_onnx_path.parent stem = network_onnx_path.stem - so_path = parent / f"{stem}.so" - mlir_path = parent / f"{stem}.onnx.mlir" - tmp_path = parent / f"{stem}.tmp" - moved_so = runner_dir / so_path.name - moved_mlir = raptor_dir / mlir_path.name - so_path.rename(moved_so) - mlir_path.rename(moved_mlir) - tmp_path.unlink(missing_ok=True) - return moved_so, moved_mlir + onnx_ir_base = raptor_dir / stem + runner_base = runner_dir / stem + run_command([raptor_path, network_onnx_path, "-o", onnx_ir_base, "--EmitONNXIR"], reporter=reporter) + run_command([raptor_path, network_onnx_path, "-o", runner_base], reporter=reporter) + network_so_path = runner_base.with_suffix(".so") + network_mlir_path = onnx_ir_base.with_suffix(".onnx.mlir") + onnx_ir_base.with_suffix(".tmp").unlink(missing_ok=True) + return network_so_path, network_mlir_path def build_onnx_runner(source_dir, build_dir, reporter=None): @@ -200,6 +218,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir, reporter = reporter or ProgressReporter(model_total) workspace_dir = network_onnx_path.parent + clean_workspace_artifacts(workspace_dir, network_onnx_path.stem) raptor_dir = workspace_dir / "raptor" runner_dir = workspace_dir / "runner" runner_build_dir = runner_dir / "build" @@ -241,7 +260,9 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir, print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compile PIM") compile_with_raptor( - network_mlir_path, raptor_path, crossbar_size, crossbar_count, reporter=reporter) + network_mlir_path, raptor_path, raptor_dir / network_onnx_path.stem, + crossbar_size, crossbar_count, + cwd=raptor_dir, reporter=reporter) print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}") reporter.advance()