simplify pimcomp compare workflow

validation/tools/compare_raptor_pimcomp.py

@@ -21,6 +21,7 @@ from typing import Any
 
 import numpy as np
 import onnx
+from colorama import Fore, Style
 
 
 REPO = Path(__file__).resolve().parents[2]
@@ -37,7 +38,11 @@ from raptor_validation.onnx_utils import (  # noqa: E402
     write_inputs_to_memory_bin,
 )
 from raptor_validation.raptor import compile_with_raptor  # noqa: E402
-from raptor_validation.validate_one import build_dump_ranges, parse_pim_simulator_outputs  # noqa: E402
+from raptor_validation.validate_one import (  # noqa: E402
+    STAGE_COLORS,
+    build_dump_ranges,
+    parse_pim_simulator_outputs,
+)
 
 
 @dataclass
@@ -73,8 +78,14 @@ def shell_join(cmd: list[str]) -> str:
     return shlex.join(str(arg) for arg in cmd)
 
 
-def print_step(name: str, cmd: list[str] | None = None, cwd: Path | None = None):
-    print(f"\n[{name}]")
+def print_step(
+    name: str,
+    cmd: list[str] | None = None,
+    cwd: Path | None = None,
+    stage: str | None = None,
+):
+    color = STAGE_COLORS.get(stage or name, Fore.WHITE)
+    print("\n" + Style.BRIGHT + color + f"[{name}]" + Style.RESET_ALL)
     if cmd is not None:
         print(f"  cwd: {cwd or REPO}")
         print(f"  $ {shell_join(cmd)}")
@@ -108,9 +119,12 @@ def exception_message(exc: BaseException) -> str:
 
 def print_failure(name: str, exc: BaseException | str) -> None:
     message = exc if isinstance(exc, str) else exception_message(exc)
-    print(f"\n[{name} FAILED]")
+    print(
+        "\n" + Style.BRIGHT + Fore.RED + f"[{name} FAILED]" + Style.RESET_ALL,
+        file=sys.stderr,
+    )
     for line in message.splitlines()[:20]:
-        print(f"  {line}")
+        print(Fore.RED + f"  {line}" + Style.RESET_ALL, file=sys.stderr)
 
 
 def run_logged(
@@ -120,8 +134,9 @@ def run_logged(
     cwd: Path,
     timeout_sec: float,
     steps: list[StepRecord],
+    stage: str | None = None,
 ) -> str:
-    print_step(name, cmd, cwd)
+    print_step(name, cmd, cwd, stage)
     start = time.perf_counter()
     command = shell_join(cmd)
     try:
@@ -202,29 +217,29 @@ def load_saved_inputs(
 
 
 def prepare_pimcomp_model(model_path: Path, out_dir: Path) -> Path:
+    model = onnx.load(model_path)
+    if not any(node.op_type == "BatchNormalization" for node in model.graph.node):
+        return model_path
+
     out_dir.mkdir(parents=True, exist_ok=True)
     output_path = out_dir / f"{model_path.stem}_pimcomp.onnx"
-    model = onnx.load(model_path)
+    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):
-        from onnxsim import simplify
+        import onnxruntime as ort
 
-        model, equivalent = simplify(model, check_n=1)
-        if not equivalent:
-            raise RuntimeError("Conv+BatchNormalization folding changed the model output")
+        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):
-            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)
+            raise RuntimeError("PIMCOMP model preparation did not eliminate BatchNormalization")
     else:
-        shutil.copy2(model_path, output_path)
+        onnx.save(model, output_path)
     return output_path
 
 
@@ -259,8 +274,7 @@ def sanitize_output_name(name: str) -> str:
 
 
 def load_effective_hardware(args: argparse.Namespace) -> dict[str, int]:
-    config_path = args.pimcomp_dir / "config.json"
-    with open(config_path, "r", encoding="utf-8") as f:
+    with open(args.pimcomp_config, "r", encoding="utf-8") as f:
         config = json.load(f)
     rows, cols = config["chip_config"]["network_config"]["layout"]
     xbar_h, xbar_w = config["chip_config"]["core_config"]["matrix_config"]["xbar_size"]
@@ -304,7 +318,7 @@ def compile_reference(
     work_dir: Path,
     steps: list[StepRecord],
 ) -> Path:
-    raptor_dir = work_dir / "reference"
+    raptor_dir = work_dir / "raptor"
     runner_dir = work_dir / "runner"
     build_dir = runner_dir / "build"
     raptor_dir.mkdir(parents=True, exist_ok=True)
@@ -320,6 +334,7 @@ def compile_reference(
         cwd=REPO,
         timeout_sec=args.timeout_seconds,
         steps=steps,
+        stage="Compile ONNX",
     )
     run_logged(
         "Reference Native Compile",
@@ -327,10 +342,11 @@ def compile_reference(
         cwd=REPO,
         timeout_sec=args.timeout_seconds,
         steps=steps,
+        stage="Compile ONNX",
     )
     network_so = runner_base.with_suffix(".so")
 
-    print_step("Generate Runner Source")
+    print_step("Generate Runner Source", stage="Build Runner")
     gen_network_runner(
         model_path,
         network_so,
@@ -346,6 +362,7 @@ def compile_reference(
         cwd=build_dir,
         timeout_sec=args.timeout_seconds,
         steps=steps,
+        stage="Build Runner",
     )
     run_logged(
         "Build Runner",
@@ -367,7 +384,7 @@ def generate_reference_outputs(
     out_dir: Path,
 ) -> Path:
     inputs_dir = out_dir / "inputs"
-    reference_dir = out_dir / "reference_outputs"
+    reference_dir = out_dir / "outputs"
     inputs_dir.mkdir(parents=True, exist_ok=True)
     reference_dir.mkdir(parents=True, exist_ok=True)
     flags, _ = save_inputs_to_files(model_path, arrays_in_order, inputs_dir)
@@ -402,7 +419,7 @@ def compile_raptor_target(
         "--pim-emit-json",
         *args.raptor_extra_arg,
     ]
-    print_step("Compile Raptor PIM", cmd, REPO)
+    print_step("Compile Raptor PIM", cmd, REPO, "Compile PIM")
     start = time.perf_counter()
     command = shell_join(cmd)
     raptor_extra_args = ["--pim-emit-json", *args.raptor_extra_arg]
@@ -446,6 +463,7 @@ def run_rust_validation(
     label: str,
     pim_dir: Path,
     config_path: Path,
+    output_bin: Path,
     outputs_desc: list[tuple[int, str, int, list[int]]],
     reference_dir: Path,
     steps: list[StepRecord],
@@ -453,7 +471,6 @@ def run_rust_validation(
     *,
     channel_last: bool = False,
 ) -> CompareResult:
-    output_bin = pim_dir.parent / "semantic_validation" / "out.bin"
     dump_ranges = build_dump_ranges(config_path, outputs_desc)
     cmd = [
         "cargo",
@@ -472,14 +489,14 @@ def run_rust_validation(
         "-d",
         dump_ranges,
     ]
-    simulation_dir = pim_dir.parent / "semantic_validation"
-    simulation_dir.mkdir(parents=True, exist_ok=True)
+    output_bin.parent.mkdir(parents=True, exist_ok=True)
     run_logged(
         label,
         cmd,
         cwd=args.pim_simulator_dir,
         timeout_sec=args.timeout_seconds,
         steps=steps,
+        stage="Run Functional Simulation",
     )
     return compare_simulator_outputs(
         output_bin,
@@ -503,16 +520,20 @@ def compile_pimcomp(
     out_dir: Path,
     steps: list[StepRecord],
 ) -> tuple[Path, Path]:
-    out_dir.mkdir(parents=True, exist_ok=True)
-    pimcomp_output_dir = args.pimcomp_dir / "output"
+    frontend_json_dir = out_dir / "models/JSON"
+    frontend_json_dir.mkdir(parents=True, exist_ok=True)
+    runtime_config = out_dir / "config.json"
+    if args.pimcomp_config != runtime_config:
+        shutil.copy2(args.pimcomp_config, runtime_config)
+    pimcomp_output_dir = out_dir / "output"
     pimcomp_output_dir.mkdir(parents=True, exist_ok=True)
     for name in ("SimulationInfo.gz", "VerificationInfo.json", "MappingResult.txt"):
         (pimcomp_output_dir / name).unlink(missing_ok=True)
     model_name = args.pimcomp_model_name or f"compare_{model_path.stem}"
-    frontend_json = args.pimcomp_dir / "models/JSON" / f"{model_name}.json"
+    frontend_json = frontend_json_dir / f"{model_name}.json"
     frontend_cmd = [
         sys.executable,
-        "frontend.py",
+        str(args.pimcomp_dir / "frontend/frontend.py"),
         "--model_path",
         str(model_path),
         "--save_path",
@@ -524,6 +545,7 @@ def compile_pimcomp(
         cwd=args.pimcomp_dir / "frontend",
         timeout_sec=args.timeout_seconds,
         steps=steps,
+        stage="Compile PIM",
     )
     backend_cmd = [
         str(args.pimcomp_dir / "build" / "PIMCOMP-NN"),
@@ -536,12 +558,13 @@ def compile_pimcomp(
     run_logged(
         "PIMCOMP Backend",
         backend_cmd,
-        cwd=args.pimcomp_dir / "build",
+        cwd=frontend_json_dir.parent,
         timeout_sec=args.timeout_seconds,
         steps=steps,
+        stage="Compile PIM",
     )
-    copy_pimcomp_outputs(pimcomp_output_dir, out_dir)
-    return out_dir / "VerificationInfo.json", out_dir / "SimulationInfo.gz"
+    remove_tree(frontend_json_dir.parent)
+    return pimcomp_output_dir / "VerificationInfo.json", pimcomp_output_dir / "SimulationInfo.gz"
 
 
 def export_pimcomp_for_pimsim_nn(simulation_info: Path, output_dir: Path) -> Path:
@@ -836,6 +859,7 @@ def run_pimsim_nn(
         cwd=args.pimsim_nn_build_dir,
         timeout_sec=args.timeout_seconds * 10.0,
         steps=steps,
+        stage="Run Non-functional Simulation",
     )
     return parse_pimsim_nn_report(output)
 
@@ -1170,6 +1194,11 @@ def main():
     parser.add_argument("--raptor-path", default=REPO / "build_release/Release/bin/onnx-mlir", type=Path)
     parser.add_argument("--onnx-include-dir", default=REPO / "onnx-mlir/include", type=Path)
     parser.add_argument("--pimcomp-dir", default=REPO / "third_party/PIMCOMP-NN", type=Path)
+    parser.add_argument(
+        "--pimcomp-config",
+        type=Path,
+        help="PIMCOMP hardware config (default: <pimcomp-dir>/config.json).",
+    )
     parser.add_argument("--pim-simulator-dir", default=REPO / "backend-simulators/pim/pim-simulator", type=Path)
     parser.add_argument("--pimsim-nn-build-dir", default=REPO / "backend-simulators/pim/pimsim-nn/build", type=Path)
     parser.add_argument("--seed", type=int, default=0)
@@ -1213,6 +1242,12 @@ def main():
         args.pimcomp_pipeline = "element" if args.pimsim_mode == "latency" else "batch"
 
     model_path = args.model.resolve()
+    args.pimcomp_dir = args.pimcomp_dir.resolve()
+    args.pimcomp_config = (
+        args.pimcomp_config.resolve()
+        if args.pimcomp_config is not None
+        else args.pimcomp_dir / "config.json"
+    )
     out_dir = args.out_dir.resolve()
     out_dir.mkdir(parents=True, exist_ok=True)
 
@@ -1265,14 +1300,19 @@ def main():
             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")
+        arrays_in_order, runtime_inputs = load_saved_inputs(
+            model_path,
+            inputs_desc,
+            reference_dir.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}")
+        print_step("Reuse Raptor")
+        print(f"  Report: {reuse_report_path}")
 
-    expected_runner_path = out_dir / "runner" / "build" / "runner"
+    expected_runner_path = out_dir / "runner/build/runner"
 
     if not reuse_raptor:
         reference_compile = try_stage(
@@ -1289,7 +1329,14 @@ def main():
         else:
             if expected_runner_path.exists():
                 runner_path = expected_runner_path
-                print(f"\n[Continue] Reusing partial runner: {runner_path}")
+                print(
+                    "\n"
+                    + Style.BRIGHT
+                    + Fore.YELLOW
+                    + "[Continue]"
+                    + Style.RESET_ALL
+                    + f" Reusing partial runner: {runner_path}"
+                )
 
     if not reuse_raptor and runner_path is not None and runner_path.exists() and model_io is not None:
         generated_reference = try_stage(
@@ -1350,6 +1397,7 @@ def main():
                 "Rust Validation Raptor",
                 raptor_pim_dir,
                 raptor_pim_dir / "config.json",
+                out_dir / "simulation/out.bin",
                 outputs_desc,
                 reference_dir,
                 steps,
@@ -1370,7 +1418,7 @@ def main():
         "Prepare PIMCOMP model",
         prepare_pimcomp_model,
         model_path,
-        out_dir / "pimcomp_model",
+        out_dir / "pimcomp/model",
     )
 
     if args.reuse_pimcomp_dir is not None:
@@ -1385,7 +1433,8 @@ def main():
         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}")
+            print_step("Reuse PIMCOMP")
+            print(f"  Directory: {reused_pimcomp_dir}")
     else:
         compiled_pimcomp = try_stage(
             failures,
@@ -1408,7 +1457,7 @@ def main():
             verification_info,
             simulation_info,
             runtime_inputs,
-            out_dir / "pimcomp_exported",
+            out_dir / "pimcomp/exported",
         )
         if exported is not None:
             pimcomp_export_dir = exported
@@ -1433,6 +1482,7 @@ def main():
             "Rust Validation PIMCOMP",
             pimcomp_export_dir,
             pimcomp_export_dir / "config.json",
+            out_dir / "simulation/pimcomp.out.bin",
             outputs_desc,
             reference_dir,
             steps,
@@ -1494,7 +1544,7 @@ def main():
                 "Export PIMCOMP for pimsim-nn",
                 export_pimcomp_for_pimsim_nn,
                 simulation_info,
-                out_dir / "pimcomp_pimsim_nn",
+                out_dir / "pimcomp/pimsim_nn",
             )
             if pimcomp_pimsim_dir is not None:
                 perf = try_stage(
@@ -1525,7 +1575,7 @@ def main():
     else:
         pimcomp_instr = empty_instruction_summary("PIMCOMP SimulationInfo.gz is not available")
 
-    report_path = out_dir / "comparison_report.md"
+    report_path = out_dir / "pimcomp/comparison_report.md"
     write_report(
         report_path,
         model_path=model_path,
@@ -1569,22 +1619,30 @@ def main():
             "report_markdown": str(report_path),
         },
     }
-    json_path = out_dir / "comparison_report.json"
+    json_path = out_dir / "pimcomp/comparison_report.json"
     with open(json_path, "w", encoding="utf-8") as f:
         json.dump(json_report, f, indent=2)
         f.write("\n")
 
-    print(f"\n[Done]")
-    print(f"  Report: {report_path}")
-    print(f"  JSON:   {json_path}")
-    if failures or any(step.status != "passed" for step in steps):
-        print(f"  Completed with {len(failures)} recorded failure/skipped stage(s).")
-
+    failed_steps = 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):
+    failed = bool(failures or failed_steps or semantic_failure)
+    result = "FAIL" if args.fail_on_error and failed else "DONE" if failed else "PASS"
+    color = Fore.RED if result == "FAIL" else Fore.YELLOW if result == "DONE" else Fore.GREEN
+    print("\n" + Style.BRIGHT + f"Result: {color}{result}" + Style.RESET_ALL)
+    print(f"  Report: {report_path}")
+    print(f"  JSON:   {json_path}")
+    if failures or failed_steps:
+        print(
+            Fore.YELLOW
+            + f"  Completed with {len(failures)} recorded failure/skipped stage(s)."
+            + Style.RESET_ALL
+        )
+
+    if args.fail_on_error and failed:
         raise SystemExit(1)