diff --git a/validation/tools/compare_raptor_pimcomp.py b/validation/tools/compare_raptor_pimcomp.py index 4127410..9782931 100644 --- a/validation/tools/compare_raptor_pimcomp.py +++ b/validation/tools/compare_raptor_pimcomp.py @@ -186,26 +186,39 @@ def remove_tree(path: Path) -> None: def load_model_inputs(model_path: Path, seed: int): - model = onnx.load(model_path) - initializer_names = {init.name for init in model.graph.initializer} - initializer_values = { - init.name: onnx.numpy_helper.to_array(init) for init in model.graph.initializer - } inputs_desc, outputs_desc = onnx_io(model_path) - runtime_desc = [desc for desc in inputs_desc if desc[1] not in initializer_names] - runtime_arrays, _ = gen_random_inputs(runtime_desc, seed=seed) + arrays_in_order, _ = gen_random_inputs(inputs_desc, seed=seed) + return inputs_desc, outputs_desc, arrays_in_order, arrays_in_order - runtime_by_name = { - desc[1]: arr for desc, arr in zip(runtime_desc, runtime_arrays) - } - arrays_in_order = [] - for _, name, elem_type, _ in inputs_desc: - if name in initializer_values: - arrays_in_order.append(initializer_values[name].astype(_ONNX_TO_NP[elem_type], copy=False)) - else: - arrays_in_order.append(runtime_by_name[name]) - runtime_only = [arr for desc, arr in zip(inputs_desc, arrays_in_order) if desc[1] not in initializer_names] - return inputs_desc, outputs_desc, arrays_in_order, runtime_only + +def load_saved_inputs( + model_path: Path, + inputs_desc: list[tuple[int, str, int, list[int]]], + inputs_dir: Path, +) -> tuple[list[np.ndarray], list[np.ndarray]]: + arrays = [] + for idx, name, elem_type, shape in inputs_desc: + array = np.loadtxt(inputs_dir / f"in{idx}.csv", delimiter=",", dtype=_ONNX_TO_NP[elem_type]).reshape(shape) + arrays.append(array) + return arrays, arrays + + +def prepare_pimcomp_model(model_path: Path, out_dir: Path) -> Path: + out_dir.mkdir(parents=True, exist_ok=True) + output_path = out_dir / f"{model_path.stem}_pimcomp.onnx" + model = onnx.load(model_path) + if any(node.op_type == "BatchNormalization" for node in model.graph.node): + from onnxsim import simplify + + model, equivalent = simplify(model, check_n=1) + 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) + else: + shutil.copy2(model_path, output_path) + return output_path def compare_simulator_outputs( @@ -215,11 +228,15 @@ def compare_simulator_outputs( *, threshold: float, rtol: float, + channel_last: bool = False, ) -> CompareResult: sim_arrays = parse_pim_simulator_outputs(output_bin, outputs_desc) max_diffs: dict[str, float] = {} passed = True for sim_array, (idx, name, _, shape) in zip(sim_arrays, outputs_desc): + if channel_last and len(shape) == 4: + n, c, h, w = shape + sim_array = sim_array.reshape(n, h, w, c).transpose(0, 3, 1, 2) csv_name = reference_dir / f"output{idx}_{sanitize_output_name(name)}.csv" ref = np.loadtxt(csv_name, delimiter=",", dtype=np.float32).reshape(shape) diff = np.abs(sim_array.astype(np.float64) - ref.astype(np.float64)) @@ -306,7 +323,14 @@ def compile_reference( run_logged( "Reference Emit ONNX IR", - [str(args.raptor_path), str(model_path), "-o", str(onnx_ir_base), "--EmitONNXIR"], + [ + str(args.raptor_path), + str(model_path), + "-o", + str(onnx_ir_base), + "--EmitONNXIR", + "--enable-conv-opt-pass=false", + ], cwd=REPO, timeout_sec=args.timeout_seconds, steps=steps, @@ -433,6 +457,8 @@ def run_rust_validation( reference_dir: Path, steps: list[StepRecord], args: argparse.Namespace, + *, + channel_last: bool = False, ) -> CompareResult: output_bin = pim_dir.parent / "semantic_validation" / "out.bin" dump_ranges = build_dump_ranges(config_path, outputs_desc) @@ -468,6 +494,7 @@ def run_rust_validation( reference_dir, threshold=args.threshold, rtol=args.rtol, + channel_last=channel_last, ) @@ -484,7 +511,7 @@ def compile_pimcomp( steps: list[StepRecord], ) -> tuple[Path, Path]: out_dir.mkdir(parents=True, exist_ok=True) - model_name = f"compare_{model_path.stem}" + 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", @@ -504,7 +531,8 @@ def compile_pimcomp( backend_cmd = [ str(args.pimcomp_dir / "build" / "PIMCOMP-NN"), f"-m={model_name}", - "-p=batch", + f"-r={args.pimcomp_replication}", + f"-p={args.pimcomp_pipeline}", "-v=YES", "-s=YES", ] @@ -527,24 +555,20 @@ def export_pimcomp_for_pimsim_nn(simulation_info: Path, output_dir: Path) -> Pat output_dir.mkdir(parents=True, exist_ok=True) sim_config = sim_info["config"] - present_core_indices = sorted( - int(key[4:]) for key, value in sim_info.items() if key.startswith("core") and isinstance(value, list) and value - ) - if not present_core_indices: - raise ValueError("PIMCOMP SimulationInfo.gz does not contain any non-empty core instruction streams") - expected_core_indices = list(range(present_core_indices[-1] + 1)) - if present_core_indices != expected_core_indices: - raise ValueError(f"PIMCOMP core numbering is not contiguous: {present_core_indices}") + core_count = int(sim_config["core_cnt"]) + if core_count <= 0: + raise ValueError("PIMCOMP SimulationInfo.gz must configure at least one core") + core_indices = range(core_count) config = { - "core_cnt": len(present_core_indices), + "core_cnt": core_count, "xbar_size": sim_config["xbar_size"], "xbar_array_count": sim_config["xbar_array_count"], "cell_precision": sim_config["cell_precision"], "adc_count": sim_config["adc_count"], "array_group_map": {}, } - for core_idx in present_core_indices: + for core_idx in core_indices: core_name = f"core{core_idx}" config["array_group_map"][core_name] = sim_config["array_group_map"].get(core_name, []) @@ -552,9 +576,9 @@ def export_pimcomp_for_pimsim_nn(simulation_info: Path, output_dir: Path) -> Pat json.dump(config, f, separators=(",", ":")) f.write("\n") - for core_idx in present_core_indices: + for core_idx in core_indices: core_key = f"core{core_idx}" - instructions = sim_info[core_key] + instructions = sim_info.get(core_key, []) or [{"op": "lldi", "rd": 0, "imm": 0, "len": 0}] with open(output_dir / f"core_{core_idx}.json", "w", encoding="utf-8") as f: json.dump(instructions, f, separators=(",", ":")) f.write("\n") @@ -652,11 +676,12 @@ def export_pimcomp_for_rust( core_dir.mkdir(parents=True, exist_ok=True) local_to_global = local_group_map.get(core_idx, {}) ag_counts = sim_info["config"]["array_group_map"].get(core_name, []) - group_prefix = [] + local_group_to_physical = {} total_crossbars = 0 - for count in ag_counts: - group_prefix.append(total_crossbars) - total_crossbars += count + for local_group in sorted(local_to_global): + width = ag_counts[local_group] + local_group_to_physical[local_group] = total_crossbars + total_crossbars += width config["array_group_map"][core_name] = list(range(total_crossbars)) for local_group, global_ag in sorted(local_to_global.items()): @@ -664,7 +689,7 @@ def export_pimcomp_for_rust( weight_name = output_to_weight[info["node_name"]] matrix = gemm_weights[weight_name] row_slice = slice(info["height_start"], info["height_end"] + 1) - first_physical = group_prefix[local_group] + first_physical = local_group_to_physical[local_group] for crossbar_idx, crossbar in enumerate(info["crossbar"]): col_slice = slice(crossbar["width_start"], crossbar["width_end"] + 1) tile = np.zeros((xbar_size, col_slice.stop - col_slice.start), dtype=np.float32) @@ -700,7 +725,13 @@ def export_pimcomp_for_rust( ) if op == "ld": if ver_inst["stage"] == "INPUT": - src = input_addr + exporter.byte_offset(ver_inst["source_offset"]) + if ver_inst["node_index"] == 1: + src = input_addr + exporter.byte_offset(ver_inst["source_offset"]) + else: + provider_index = -ver_inst["source_address"] + src = output_base + exporter.byte_offset( + provider_index * max_output + ver_inst["source_offset"] + ) elif ver_inst["stage"] == "BIAS": src = bias_addrs[node_list[ver_inst["node_index"]]["name"]] + exporter.byte_offset(ver_inst["source_offset"]) else: @@ -742,7 +773,7 @@ def export_pimcomp_for_rust( elif op == "mvmul": local_group = sim_inst["group"] global_ag = local_to_global[local_group] - first_physical = group_prefix[local_group] + first_physical = local_group_to_physical[local_group] widths = [ crossbar["width_end"] - crossbar["width_start"] + 1 for crossbar in ag_info[global_ag]["crossbar"] @@ -987,12 +1018,16 @@ def write_report( pimcomp_instr: dict[str, Any], raptor_pass_timings: dict[str, float], pimsim_mode: str, + pimcomp_pipeline: str, + pimcomp_replication: str, ): lines = [ "# Raptor vs PIMCOMP Comparison Report", "", f"- Model: `{model_path}`", f"- Hardware: `{hardware.get('core_count', 'n/a')} cores`, `{hardware.get('crossbar_count', 'n/a')} xbars/core`, `{hardware.get('crossbar_size', 'n/a')}x{hardware.get('crossbar_size', 'n/a')}` crossbars, mesh `{hardware.get('mesh_rows', 'n/a')}x{hardware.get('mesh_cols', 'n/a')}`", + f"- PIMCOMP pipeline: `{pimcomp_pipeline}`", + f"- PIMCOMP replication: `{pimcomp_replication}`", "", ] @@ -1152,15 +1187,29 @@ def main(): parser.add_argument("--mesh-cols", type=int) parser.add_argument("--pimsim-time-ms", type=int, default=1000) parser.add_argument("--pimsim-mode", choices=["latency", "throughput"], default="latency") + parser.add_argument("--pimcomp-pipeline", choices=["element", "batch"]) + parser.add_argument("--pimcomp-model-name", help="Use a PIMCOMP built-in model name such as vgg16.") + parser.add_argument( + "--pimcomp-replication", + choices=["balance", "W0H0", "uniform", "GA"], + default="balance", + ) + parser.add_argument( + "--reuse-raptor-report", + type=Path, + help="Reuse Raptor artifacts and results from an existing comparison_report.json.", + ) 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=[]) parser.add_argument( "--fail-on-error", action="store_true", - help="Return a non-zero process status after writing the reports if any compilation/run stage failed.", + help="Return a non-zero status if a stage or semantic validation fails.", ) args = parser.parse_args() + if args.pimcomp_pipeline is None: + args.pimcomp_pipeline = "element" if args.pimsim_mode == "latency" else "batch" model_path = args.model.resolve() out_dir = args.out_dir.resolve() @@ -1188,7 +1237,9 @@ def main(): verification_info: Path | None = None simulation_info: Path | None = None pimcomp_export_dir: Path | None = None + pimcomp_model_path: Path | None = None pimsim_config: Path | None = None + reuse_raptor = args.reuse_raptor_report is not None raptor_validation = skipped_validation("Raptor validation did not run") pimcomp_validation = skipped_validation("PIMCOMP validation did not run") @@ -1205,29 +1256,46 @@ def main(): if model_io is not None: inputs_desc, outputs_desc, arrays_in_order, runtime_inputs = model_io + if reuse_raptor and model_io is not None: + reuse_report_path = args.reuse_raptor_report.resolve() + with open(reuse_report_path, "r", encoding="utf-8") as f: + reused = json.load(f) + reused_hardware = reused["hardware"] + if reused_hardware != hardware: + raise ValueError(f"Reused Raptor hardware differs: {reused_hardware} != {hardware}") + reference_dir = Path(reused["paths"]["reference_outputs"]) + raptor_pim_dir = Path(reused["paths"]["raptor_pim"]) + arrays_in_order, runtime_inputs = load_saved_inputs(model_path, inputs_desc, reuse_report_path.parent / "inputs") + raptor_validation = CompareResult(**reused["raptor_validation"]) + raptor_perf = reused["raptor_performance"] + raptor_instr = reused["raptor_instruction_summary"] + raptor_pass_timings = reused["raptor_pass_timings"] + print(f"\n[Reuse Raptor]\n Report: {reuse_report_path}") + expected_network_mlir = out_dir / "reference" / f"{model_path.stem}.onnx.mlir" expected_runner_path = out_dir / "runner" / "build" / "runner" - reference_compile = try_stage( - failures, - "Compile reference", - compile_reference, - args, - model_path, - out_dir, - steps, - ) - if reference_compile is not None: - network_mlir, _, runner_path = reference_compile - else: - if expected_network_mlir.exists(): - network_mlir = expected_network_mlir - print(f"\n[Continue] Reusing partial ONNX MLIR: {network_mlir}") - if expected_runner_path.exists(): - runner_path = expected_runner_path - print(f"\n[Continue] Reusing partial runner: {runner_path}") + if not reuse_raptor: + reference_compile = try_stage( + failures, + "Compile reference", + compile_reference, + args, + model_path, + out_dir, + steps, + ) + if reference_compile is not None: + network_mlir, _, runner_path = reference_compile + else: + if expected_network_mlir.exists(): + network_mlir = expected_network_mlir + print(f"\n[Continue] Reusing partial ONNX MLIR: {network_mlir}") + if expected_runner_path.exists(): + runner_path = expected_runner_path + print(f"\n[Continue] Reusing partial runner: {runner_path}") - if runner_path is not None and runner_path.exists() and model_io is not None: + if not reuse_raptor and runner_path is not None and runner_path.exists() and model_io is not None: generated_reference = try_stage( failures, "Run reference", @@ -1242,14 +1310,14 @@ def main(): ) if generated_reference is not None: reference_dir = generated_reference - else: + elif not reuse_raptor: record_failure( failures, "Skip reference outputs", "Reference outputs were skipped because the native runner or model inputs are not available.", ) - if network_mlir is not None and network_mlir.exists() and hardware["core_count"] > 0: + if not reuse_raptor and network_mlir is not None and network_mlir.exists() and hardware["core_count"] > 0: compiled_raptor = try_stage( failures, "Compile Raptor PIM", @@ -1262,14 +1330,14 @@ def main(): ) if compiled_raptor is not None: raptor_pim_dir, raptor_pass_timings = compiled_raptor - else: + elif not reuse_raptor: record_failure( failures, "Skip Raptor PIM compile", "Raptor PIM compile was skipped because the ONNX MLIR or hardware configuration is not available.", ) - if raptor_pim_dir is not None: + if not reuse_raptor and raptor_pim_dir is not None: wrote_inputs = try_stage_success( failures, "Write Raptor inputs", @@ -1298,18 +1366,26 @@ def main(): raptor_validation = skipped_validation("Output descriptors are not available") else: raptor_validation = skipped_validation("Raptor input materialization failed") - else: + elif not reuse_raptor: raptor_validation = skipped_validation("Raptor PIM compilation did not produce a PIM directory") + pimcomp_model_path = try_stage( + failures, + "Prepare PIMCOMP model", + prepare_pimcomp_model, + model_path, + out_dir / "pimcomp_model", + ) + compiled_pimcomp = try_stage( failures, "Compile PIMCOMP", compile_pimcomp, args, - model_path, + 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 @@ -1318,7 +1394,7 @@ def main(): failures, "Export PIMCOMP for Rust", export_pimcomp_for_rust, - model_path, + pimcomp_model_path, verification_info, simulation_info, runtime_inputs, @@ -1351,6 +1427,7 @@ def main(): reference_dir, steps, args, + channel_last=True, ) pimcomp_validation = validation if validation is not None else failed_validation("PIMCOMP validation failed") elif pimcomp_export_dir is None: @@ -1379,13 +1456,15 @@ def main(): ) if args.skip_pimsim_nn: - raptor_perf = skipped_perf("Skipped by --skip-pimsim-nn") + if not reuse_raptor: + raptor_perf = skipped_perf("Skipped by --skip-pimsim-nn") pimcomp_perf = skipped_perf("Skipped by --skip-pimsim-nn") elif pimsim_config is None: - raptor_perf = skipped_perf("pimsim-nn config is not available") + if not reuse_raptor: + raptor_perf = skipped_perf("pimsim-nn config is not available") pimcomp_perf = skipped_perf("pimsim-nn config is not available") else: - if raptor_pim_dir is not None: + if not reuse_raptor and raptor_pim_dir is not None: perf = try_stage( failures, "pimsim-nn Raptor", @@ -1398,7 +1477,7 @@ def main(): args, ) raptor_perf = perf if perf is not None else failed_perf("pimsim-nn Raptor failed") - else: + elif not reuse_raptor: raptor_perf = skipped_perf("Raptor PIM directory is not available") if simulation_info is not None: @@ -1427,10 +1506,10 @@ def main(): else: pimcomp_perf = skipped_perf("PIMCOMP SimulationInfo.gz is not available") - if raptor_pim_dir is not None and raptor_pim_dir.exists(): + if not reuse_raptor and raptor_pim_dir is not None and raptor_pim_dir.exists(): parsed = try_stage(failures, "Parse Raptor instructions", parse_raptor_instructions, raptor_pim_dir) raptor_instr = parsed if parsed is not None else empty_instruction_summary(error="Failed to parse Raptor instructions") - else: + elif not reuse_raptor: raptor_instr = empty_instruction_summary("Raptor PIM directory is not available") if simulation_info is not None and simulation_info.exists(): @@ -1454,12 +1533,17 @@ def main(): pimcomp_instr=pimcomp_instr, raptor_pass_timings=raptor_pass_timings, pimsim_mode=args.pimsim_mode, + pimcomp_pipeline=args.pimcomp_pipeline, + pimcomp_replication=args.pimcomp_replication, ) json_report = { "model": str(model_path), "hardware": hardware, "pimsim_mode": args.pimsim_mode, + "pimcomp_pipeline": args.pimcomp_pipeline, + "pimcomp_replication": args.pimcomp_replication, + "reused_raptor_report": optional_path(args.reuse_raptor_report.resolve()) if reuse_raptor else None, "failures": failures, "steps": [asdict(step) for step in steps], "raptor_validation": asdict(raptor_validation), @@ -1489,7 +1573,11 @@ def main(): if failures or any(step.status != "passed" for step in steps): print(f" Completed with {len(failures)} recorded failure/skipped stage(s).") - if args.fail_on_error and (failures or any(step.status != "passed" for step in steps)): + semantic_failure = any( + result.status == "done" and not result.passed + for result in (raptor_validation, pimcomp_validation) + ) + if args.fail_on_error and (failures or any(step.status != "passed" for step in steps) or semantic_failure): raise SystemExit(1)