Python script for compare

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(