validation

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]
+    arrays_in_order, _ = gen_random_inputs(inputs_desc, seed=seed)
-    runtime_arrays, _ = gen_random_inputs(runtime_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)
+def load_saved_inputs(
-    }
+    model_path: Path,
-    arrays_in_order = []
+    inputs_desc: list[tuple[int, str, int, list[int]]],
-    for _, name, elem_type, _ in inputs_desc:
+    inputs_dir: Path,
-        if name in initializer_values:
+) -> tuple[list[np.ndarray], list[np.ndarray]]:
-            arrays_in_order.append(initializer_values[name].astype(_ONNX_TO_NP[elem_type], copy=False))
+    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:
-            arrays_in_order.append(runtime_by_name[name])
+        shutil.copy2(model_path, output_path)
-    runtime_only = [arr for desc, arr in zip(inputs_desc, arrays_in_order) if desc[1] not in initializer_names]
+    return output_path
-    return inputs_desc, outputs_desc, arrays_in_order, runtime_only
 
 
 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(
+    core_count = int(sim_config["core_cnt"])
-        int(key[4:]) for key, value in sim_info.items() if key.startswith("core") and isinstance(value, list) and value
+    if core_count <= 0:
-    )
+        raise ValueError("PIMCOMP SimulationInfo.gz must configure at least one core")
-    if not present_core_indices:
+    core_indices = range(core_count)
-        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}")
 
     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:
+        for local_group in sorted(local_to_global):
-            group_prefix.append(total_crossbars)
+            width = ag_counts[local_group]
-            total_crossbars += count
+            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":
+                    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,9 +1256,26 @@ 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"
 
+    if not reuse_raptor:
         reference_compile = try_stage(
             failures,
             "Compile reference",
@@ -1227,7 +1295,7 @@ def main():
                 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:
+        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:
+        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)