diff --git a/backend-simulators/pim/pim-simulator/src/bin/pim-simulator/main.rs b/backend-simulators/pim/pim-simulator/src/bin/pim-simulator/main.rs index 04504e2..2bbe082 100644 --- a/backend-simulators/pim/pim-simulator/src/bin/pim-simulator/main.rs +++ b/backend-simulators/pim/pim-simulator/src/bin/pim-simulator/main.rs @@ -152,7 +152,12 @@ fn get_crossbars(config: &Value, args: &Args) -> anyhow::Result//comparison_report.{md,json}`. +Use `--models vgg8` to run one model, `--resume` after an interruption, or +`--dry-run` to inspect every command. The runner continues after a failed model +so all reports are produced, then returns a nonzero status if any comparison +failed. + +Arch-A low-latency example: + +```bash +RAPTOR_ROOT=$PWD +PIMCOMP="$RAPTOR_ROOT/third_party/PIMCOMP-NN" +CONFIG_BACKUP=$(mktemp) +cp "$PIMCOMP/config.json" "$CONFIG_BACKUP" +trap 'cp "$CONFIG_BACKUP" "$PIMCOMP/config.json"' EXIT +cp "$RAPTOR_ROOT/validation/networks/pimcomp_models/configs/arch-a.json" "$PIMCOMP/config.json" + +"$RAPTOR_ROOT/.venv/bin/python" "$RAPTOR_ROOT/validation/tools/compare_raptor_pimcomp.py" \ + --model "$RAPTOR_ROOT/validation/networks/pimcomp_models/vgg8/vgg8-mnist-reconstructed.onnx" \ + --out-dir /tmp/compare-vgg8-arch-a-ll \ + --core-count 168 \ + --crossbar-count 96 \ + --crossbar-size 128 \ + --mesh-rows 12 \ + --mesh-cols 14 \ + --pimsim-mode latency \ + --pimcomp-pipeline element \ + --fail-on-error +``` + +For Arch-A high throughput, use `--pimsim-mode throughput +--pimcomp-pipeline batch`. For Arch-B, use 138 cores, 128 crossbars, a +`6x23` mesh, and `configs/arch-b.json`. + +If only semantic and instruction comparison is required, add +`--skip-pimsim-nn`. This exact VGG-8 Arch-A LL smoke test passed both semantic +validations with maximum output differences below `5e-10`. + +Current Raptor status: + +- VGG-8, ResNet-18, fixed-batch ResNet-34, and GoogLeNet compile on Arch-A. +- Use `googlenet-12-no-softmax.onnx` for the paper-matched latency comparison. + The original model's final `vsoftmax` is supported by Raptor's functional + simulator but not by `pimsim-nn`; PIMCOMP does not schedule that operation. +- Raptor currently accepts one square `--crossbar-size`; Arch-C's rectangular + `512x1024` arrays can therefore be compiled by PIMCOMP but not compared + exactly with Raptor. + +Do not change the hardware profile to bypass either limitation; that would no +longer be a paper-matched comparison. + +## Monolith fallback + +The local `monolith` SSH alias points to the high-memory host. Copy only this +suite and the comparison driver; `-L` materializes the ResNet-18 symlink because +the canonical `resnetv2/depth_68` file may not exist remotely: + +```bash +REMOTE_REPO=/home/gmagnani/Project/Raptor +rsync -azL validation/networks/pimcomp_models/ \ + "monolith:$REMOTE_REPO/validation/networks/pimcomp_models/" +rsync -az validation/tools/compare_raptor_pimcomp.py \ + "monolith:$REMOTE_REPO/validation/tools/compare_raptor_pimcomp.py" +rsync -az validation/tools/run_pimcomp_paper_latency.py \ + "monolith:$REMOTE_REPO/validation/tools/run_pimcomp_paper_latency.py" +rsync -az --exclude=.git --exclude=build --exclude=output \ + third_party/PIMCOMP-NN/ \ + "monolith:$REMOTE_REPO/third_party/PIMCOMP-NN/" +``` + +Then use the same commands over SSH: + +```bash +ssh monolith +cd /home/gmagnani/Project/Raptor + +# One-time setup if the repository virtual environment is absent. +python3 -m venv .venv +.venv/bin/python -m pip install numpy onnx onnxruntime onnxsim colorama + +# Run every latency comparison serially. +.venv/bin/python validation/tools/run_pimcomp_paper_latency.py \ + --out-dir /tmp/raptor-pimcomp-paper-latency +``` + +Copy reports back without transferring large compiler artifacts: + +```bash +rsync -az --include='*/' --include='comparison_report.*' --exclude='*' \ + monolith:/tmp/raptor-pimcomp-paper-latency/ \ + /tmp/raptor-pimcomp-paper-latency/ +``` diff --git a/validation/networks/pimcomp_models/RESULTS.md b/validation/networks/pimcomp_models/RESULTS.md new file mode 100644 index 0000000..8d53254 --- /dev/null +++ b/validation/networks/pimcomp_models/RESULTS.md @@ -0,0 +1,33 @@ +# Raptor vs PIMCOMP latency results + +## GoogLeNet, Arch-A, low latency + +Measured with `googlenet-12-no-softmax.onnx`, batch 1, Raptor's best current +schedule, and PIMCOMP's GA/element artifacts. Both instruction streams were +simulated by the same `pimsim-nn` build using the complete Arch-A timing and +precision configuration. + +| Compiler | Latency (ms) | Instructions | Sends | Receives | MVMUL | +| --- | ---: | ---: | ---: | ---: | ---: | +| Raptor | 1132.458315 | 77,717,248 | 29,115 | 29,115 | 157,158 | +| PIMCOMP | 41.790450 | 3,398,070 | 110,334 | 110,334 | 113,639 | + +PIMCOMP is 27.10x faster in this latency simulation. + +Semantic validation did not pass the comparison driver's strict default +tolerance: the maximum logit differences from the native ONNX reference were +`0.01995039` for Raptor and `7.768404` for PIMCOMP. Treat these as performance +results, not as a correctness-equivalent comparison. + +No throughput experiment was run. + +## Reproduce + +```bash +.venv/bin/python validation/tools/run_pimcomp_paper_latency.py \ + --out-dir /tmp/raptor-pimcomp-paper-latency \ + --models googlenet +``` + +See [README.md](README.md) for model provenance, limitations, and monolith +instructions. diff --git a/validation/networks/pimcomp_models/configs/arch-a.json b/validation/networks/pimcomp_models/configs/arch-a.json new file mode 100644 index 0000000..6e2b409 --- /dev/null +++ b/validation/networks/pimcomp_models/configs/arch-a.json @@ -0,0 +1,54 @@ +{ + "chip_config": { + "core_config": { + "period": 1, + "matrix_config": { + "xbar_array_count": 96, + "period": 1, + "pipeline_mode": true, + "dac_resolution": 1, + "dac_count": 128, + "xbar_size": [128, 128], + "cell_precision": 2, + "xbar_latency": 7, + "xbar_read_power": 22.24, + "sample_hold_latency_cycle": 1, + "adc_resolution": 8, + "adc_latency_cycle": 20, + "adc_static_power": 0.322, + "adc_dynamic_power": 1.135, + "adc_count": 16, + "shift_adder_latency_cycle": 1, + "output_buffer_latency_cycle": 1 + }, + "vector_width": 32, + "vector_latency_cycle": 4, + "local_memory_config": { + "data_width": 64, + "period": 1, + "write_latency_cycle": 30, + "read_latency_cycle": 20 + }, + "global_memory_switch_id": -10 + }, + "global_memory_config": { + "data_width": 32, + "period": 1, + "write_latency_cycle": 50, + "read_latency_cycle": 40 + }, + "network_config": { + "bus_topology": "mesh", + "bus_width": 32, + "layout": [12, 14], + "net_config_file_path": "network_mesh_168.json" + }, + "core_cnt": 168, + "global_memory_switch_id": -10 + }, + "sim_config": { + "sim_mode": 0, + "sim_time": 200, + "report_verbose_level": 0 + } +} diff --git a/validation/networks/pimcomp_models/configs/arch-b.json b/validation/networks/pimcomp_models/configs/arch-b.json new file mode 100644 index 0000000..0c4092e --- /dev/null +++ b/validation/networks/pimcomp_models/configs/arch-b.json @@ -0,0 +1,54 @@ +{ + "chip_config": { + "core_config": { + "period": 1, + "matrix_config": { + "xbar_array_count": 128, + "period": 1, + "pipeline_mode": true, + "dac_resolution": 1, + "dac_count": 128, + "xbar_size": [128, 128], + "cell_precision": 2, + "xbar_latency": 7, + "xbar_read_power": 22.24, + "sample_hold_latency_cycle": 1, + "adc_resolution": 8, + "adc_latency_cycle": 20, + "adc_static_power": 0.322, + "adc_dynamic_power": 1.135, + "adc_count": 16, + "shift_adder_latency_cycle": 1, + "output_buffer_latency_cycle": 1 + }, + "vector_width": 32, + "vector_latency_cycle": 4, + "local_memory_config": { + "data_width": 64, + "period": 1, + "write_latency_cycle": 30, + "read_latency_cycle": 20 + }, + "global_memory_switch_id": -10 + }, + "global_memory_config": { + "data_width": 32, + "period": 1, + "write_latency_cycle": 50, + "read_latency_cycle": 40 + }, + "network_config": { + "bus_topology": "mesh", + "bus_width": 32, + "layout": [6, 23], + "net_config_file_path": "network_mesh_138.json" + }, + "core_cnt": 138, + "global_memory_switch_id": -10 + }, + "sim_config": { + "sim_mode": 0, + "sim_time": 200, + "report_verbose_level": 0 + } +} diff --git a/validation/networks/pimcomp_models/configs/arch-c.json b/validation/networks/pimcomp_models/configs/arch-c.json new file mode 100644 index 0000000..e6803c1 --- /dev/null +++ b/validation/networks/pimcomp_models/configs/arch-c.json @@ -0,0 +1,54 @@ +{ + "chip_config": { + "core_config": { + "period": 1, + "matrix_config": { + "xbar_array_count": 8, + "period": 1, + "pipeline_mode": true, + "dac_resolution": 1, + "dac_count": 128, + "xbar_size": [512, 1024], + "cell_precision": 2, + "xbar_latency": 7, + "xbar_read_power": 22.24, + "sample_hold_latency_cycle": 1, + "adc_resolution": 8, + "adc_latency_cycle": 20, + "adc_static_power": 0.322, + "adc_dynamic_power": 1.135, + "adc_count": 16, + "shift_adder_latency_cycle": 1, + "output_buffer_latency_cycle": 1 + }, + "vector_width": 32, + "vector_latency_cycle": 4, + "local_memory_config": { + "data_width": 64, + "period": 1, + "write_latency_cycle": 30, + "read_latency_cycle": 20 + }, + "global_memory_switch_id": -10 + }, + "global_memory_config": { + "data_width": 32, + "period": 1, + "write_latency_cycle": 50, + "read_latency_cycle": 40 + }, + "network_config": { + "bus_topology": "mesh", + "bus_width": 32, + "layout": [8, 8], + "net_config_file_path": "network_mesh_64.json" + }, + "core_cnt": 64, + "global_memory_switch_id": -10 + }, + "sim_config": { + "sim_mode": 0, + "sim_time": 200, + "report_verbose_level": 0 + } +} diff --git a/validation/networks/pimcomp_models/googlenet/googlenet-12-no-softmax.onnx b/validation/networks/pimcomp_models/googlenet/googlenet-12-no-softmax.onnx new file mode 100644 index 0000000..3f41b0d Binary files /dev/null and b/validation/networks/pimcomp_models/googlenet/googlenet-12-no-softmax.onnx differ diff --git a/validation/networks/pimcomp_models/resnet18/resnet18-v1-7.onnx b/validation/networks/pimcomp_models/resnet18/resnet18-v1-7.onnx new file mode 120000 index 0000000..408f533 --- /dev/null +++ b/validation/networks/pimcomp_models/resnet18/resnet18-v1-7.onnx @@ -0,0 +1 @@ +../../resnetv2/depth_68/resnetv2_depth_68.onnx \ No newline at end of file diff --git a/validation/networks/pimcomp_models/resnet34/resnet34-v1-7.onnx b/validation/networks/pimcomp_models/resnet34/resnet34-v1-7.onnx new file mode 100644 index 0000000..6c942db Binary files /dev/null and b/validation/networks/pimcomp_models/resnet34/resnet34-v1-7.onnx differ diff --git a/validation/networks/pimcomp_models/vgg8/vgg8-mnist-reconstructed.onnx b/validation/networks/pimcomp_models/vgg8/vgg8-mnist-reconstructed.onnx new file mode 100644 index 0000000..b2a1b9e Binary files /dev/null and b/validation/networks/pimcomp_models/vgg8/vgg8-mnist-reconstructed.onnx differ diff --git a/validation/networks/resnet/resnet18_torchvision.onnx b/validation/networks/resnet/resnet18_torchvision.onnx index 7d91dbd..6b8884e 100644 Binary files a/validation/networks/resnet/resnet18_torchvision.onnx and b/validation/networks/resnet/resnet18_torchvision.onnx differ diff --git a/validation/tools/compare_raptor_pimcomp.py b/validation/tools/compare_raptor_pimcomp.py index 7bf5b6f..d796e92 100644 --- a/validation/tools/compare_raptor_pimcomp.py +++ b/validation/tools/compare_raptor_pimcomp.py @@ -214,8 +214,17 @@ def prepare_pimcomp_model(model_path: Path, out_dir: Path) -> Path: if not equivalent: raise RuntimeError("Conv+BatchNormalization folding changed the model output") if any(node.op_type == "BatchNormalization" for node in model.graph.node): - raise RuntimeError("PIMCOMP model preparation did not eliminate BatchNormalization") - onnx.save(model, output_path) + import onnxruntime as ort + + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC + options.optimized_model_filepath = str(output_path) + ort.InferenceSession(str(model_path), options, providers=["CPUExecutionProvider"]) + model = onnx.load(output_path) + if any(node.op_type == "BatchNormalization" for node in model.graph.node): + raise RuntimeError("PIMCOMP model preparation did not eliminate BatchNormalization") + else: + onnx.save(model, output_path) else: shutil.copy2(model_path, output_path) return output_path @@ -271,9 +280,9 @@ def load_effective_hardware(args: argparse.Namespace) -> dict[str, int]: def write_pimsim_config(args: argparse.Namespace, out_dir: Path, hardware: dict[str, int]) -> Path: mesh_builder = load_mesh_builder() - example_config = REPO / "backend-simulators/pim/pimsim-nn/example/config/latency_config.json" - with open(example_config, "r", encoding="utf-8") as f: + with open(args.pimcomp_dir / "config.json", "r", encoding="utf-8") as f: config = json.load(f) + config["chip_config"]["core_config"].setdefault("rob_size", 1) config["chip_config"]["core_config"]["matrix_config"]["xbar_array_count"] = hardware["crossbar_count"] config["chip_config"]["core_config"]["matrix_config"]["xbar_size"] = [ hardware["crossbar_size"], @@ -491,10 +500,10 @@ def run_rust_validation( ) -def copy_pimcomp_outputs(args: argparse.Namespace, out_dir: Path): +def copy_pimcomp_outputs(source_dir: Path, out_dir: Path): out_dir.mkdir(parents=True, exist_ok=True) for name in ("SimulationInfo.gz", "VerificationInfo.json", "MappingResult.txt"): - shutil.copy2(args.pimcomp_dir / "output" / name, out_dir / name) + shutil.copy2(source_dir / name, out_dir / name) def compile_pimcomp( @@ -507,7 +516,7 @@ def compile_pimcomp( model_name = args.pimcomp_model_name or f"compare_{model_path.stem}" frontend_json = args.pimcomp_dir / "models/JSON" / f"{model_name}.json" frontend_cmd = [ - "python3", + sys.executable, "frontend.py", "--model_path", str(model_path), @@ -536,7 +545,7 @@ def compile_pimcomp( timeout_sec=args.timeout_seconds, steps=steps, ) - copy_pimcomp_outputs(args, out_dir) + copy_pimcomp_outputs(args.pimcomp_dir / "output", out_dir) return out_dir / "VerificationInfo.json", out_dir / "SimulationInfo.gz" @@ -798,7 +807,7 @@ def export_pimcomp_for_rust( def parse_pimsim_nn_report(output: str) -> dict[str, float | int | str]: patterns = { "output_count": r"output count:\s+([0-9]+)\s+samples", - "throughput": r"throughput:\s+([0-9.]+)\s+samples/s", + "throughput": r"throughput:\s+([0-9.eE+-]+)\s+samples/s", "average_latency_ms": r"average latency:\s+([0-9.eE+-]+)\s+ms", "latency_ms": r"latency:\s+([0-9.eE+-]+)\s+ms", "average_power_mw": r"average power:\s+([0-9.eE+-]+)\s+mW", @@ -1192,6 +1201,11 @@ def main(): type=Path, help="Reuse Raptor artifacts and results from an existing comparison_report.json.", ) + parser.add_argument( + "--reuse-pimcomp-dir", + type=Path, + help="Reuse a directory containing PIMCOMP SimulationInfo.gz, VerificationInfo.json, and MappingResult.txt.", + ) parser.add_argument("--skip-pimsim-nn", action="store_true") parser.add_argument("--verbose-raptor-compile", action="store_true") parser.add_argument("--raptor-extra-arg", action="append", default=[]) @@ -1365,17 +1379,31 @@ def main(): out_dir / "pimcomp_model", ) - compiled_pimcomp = try_stage( - failures, - "Compile PIMCOMP", - compile_pimcomp, - args, - pimcomp_model_path, - out_dir / "pimcomp", - steps, - ) if pimcomp_model_path is not None else None - if compiled_pimcomp is not None: - verification_info, simulation_info = compiled_pimcomp + if args.reuse_pimcomp_dir is not None: + reused_pimcomp_dir = args.reuse_pimcomp_dir.resolve() + copied_pimcomp = try_stage_success( + failures, + "Reuse PIMCOMP outputs", + copy_pimcomp_outputs, + reused_pimcomp_dir, + out_dir / "pimcomp", + ) + if copied_pimcomp: + verification_info = out_dir / "pimcomp/VerificationInfo.json" + simulation_info = out_dir / "pimcomp/SimulationInfo.gz" + print(f"\n[Reuse PIMCOMP]\n Directory: {reused_pimcomp_dir}") + else: + compiled_pimcomp = try_stage( + failures, + "Compile PIMCOMP", + compile_pimcomp, + args, + pimcomp_model_path, + out_dir / "pimcomp", + steps, + ) if pimcomp_model_path is not None else None + if compiled_pimcomp is not None: + verification_info, simulation_info = compiled_pimcomp if verification_info is not None and simulation_info is not None and model_io is not None: exported = try_stage( diff --git a/validation/tools/run_pimcomp_paper_latency.py b/validation/tools/run_pimcomp_paper_latency.py new file mode 100644 index 0000000..68fff5e --- /dev/null +++ b/validation/tools/run_pimcomp_paper_latency.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import re +import shlex +import shutil +import subprocess +import sys +from pathlib import Path + + +REPO = Path(__file__).resolve().parents[2] +SUITE = REPO / "validation/networks/pimcomp_models" +PIMCOMP_SOURCE = REPO / "third_party/PIMCOMP-NN" +COMPARE = REPO / "validation/tools/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-no-softmax.onnx", +} + + +def run(command: list[str], *, dry_run: bool, check: bool = True) -> int: + print(f"$ {shlex.join(command)}", flush=True) + if dry_run: + return 0 + return subprocess.run(command, cwd=REPO, check=check).returncode + + +def prepare_pimcomp(work_dir: Path) -> None: + shutil.copytree( + PIMCOMP_SOURCE, + work_dir, + dirs_exist_ok=True, + ignore=shutil.ignore_patterns(".git", "build", "output"), + ) + header = work_dir / "backend/GeneticAlgorithm.h" + source = header.read_text(encoding="utf-8") + if "int population_num = 200;" not in source: + raise RuntimeError("PIMCOMP GA population is not 200") + source, replacements = re.subn( + r"int max_iteration = \d+;", + "int max_iteration = 1000;", + source, + ) + if replacements != 1: + raise RuntimeError("Could not set PIMCOMP GA max_iteration") + header.write_text(source, encoding="utf-8") + shutil.copy2(SUITE / "configs/arch-a.json", work_dir / "config.json") + + +def comparison_command(model: Path, result_dir: Path, pimcomp_dir: Path, timeout: float) -> list[str]: + return [ + sys.executable, + str(COMPARE), + "--model", + str(model), + "--out-dir", + str(result_dir), + "--pimcomp-dir", + str(pimcomp_dir), + "--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", required=True, type=Path) + 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="Keep the existing work tree and 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() + work_dir = out_dir / "pimcomp-ga1000" + if not args.dry_run and out_dir.exists() and any(out_dir.iterdir()) and not args.resume: + parser.error(f"{out_dir} is not empty; choose a fresh directory or pass --resume") + + 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.dry_run: + print(f"# prepare isolated PIMCOMP GA build in {work_dir}") + else: + out_dir.mkdir(parents=True, exist_ok=True) + prepare_pimcomp(work_dir) + + run(["cmake", "--build", str(REPO / "build_release")], dry_run=args.dry_run) + run( + ["cmake", "-S", str(work_dir), "-B", str(work_dir / "build")], + dry_run=args.dry_run, + ) + run( + ["cmake", "--build", str(work_dir / "build"), "--target", "PIMCOMP-NN"], + dry_run=args.dry_run, + ) + + failed = [] + for name in args.models: + result_dir = out_dir / name + if args.resume and (result_dir / "comparison_report.json").exists(): + print(f"[{name}] completed report exists; skipping", flush=True) + continue + print(f"\n[{name}] Arch-A latency comparison", flush=True) + returncode = run( + comparison_command(MODELS[name], result_dir, work_dir, args.timeout_seconds), + dry_run=args.dry_run, + check=False, + ) + if returncode: + failed.append(name) + + if failed: + print(f"\nCompleted with failed comparisons: {', '.join(failed)}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/tools/split_onnx_prefixes.py b/validation/tools/split_onnx_prefixes.py new file mode 100644 index 0000000..22bf006 --- /dev/null +++ b/validation/tools/split_onnx_prefixes.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python3 + +import argparse +from pathlib import Path + +import onnx + + +def split_prefixes(model_path: Path, output_dir: Path, name: str) -> None: + model = onnx.shape_inference.infer_shapes(onnx.load(model_path)) + initializer_names = {initializer.name for initializer in model.graph.initializer} + input_names = [value.name for value in model.graph.input if value.name not in initializer_names] + extractor = onnx.utils.Extractor(model) + + output_dir.mkdir(parents=True, exist_ok=True) + for depth, node in enumerate(model.graph.node): + output_name = next(output for output in node.output if output) + prefix = extractor.extract_model(input_names, [output_name]) + prefix.ir_version = max(prefix.ir_version, 4) + onnx.checker.check_model(prefix) + + depth_name = f"depth_{depth:02d}" + depth_dir = output_dir / depth_name + depth_dir.mkdir(parents=True, exist_ok=True) + output_path = depth_dir / f"{name}_{depth_name}.onnx" + onnx.save(prefix, output_path) + print(f"{depth_name}: {node.op_type} -> {output_name} ({len(prefix.graph.node)} nodes)") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Split an ONNX graph into one ancestor prefix per node.") + parser.add_argument("model", type=Path) + parser.add_argument("output_dir", type=Path) + parser.add_argument("--name", required=True) + args = parser.parse_args() + split_prefixes(args.model, args.output_dir, args.name) + + +if __name__ == "__main__": + main()