rework actually broken dcp merge + compute re-batching (still to refine)

validation/raptor.py

@@ -1,4 +1,5 @@
 import re
+import shlex
 import subprocess
 from pathlib import Path
 from colorama import Fore, Style
@@ -37,8 +38,12 @@ def _parse_pim_pass_timings(output_text):
     return pass_timings
 
 
+def _format_command(cmd):
+    return shlex.join(str(arg) for arg in cmd)
+
+
 def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
-                        crossbar_size, crossbar_count, cwd=None, reporter=None):
+                        crossbar_size, crossbar_count, core_count=None, cwd=None, reporter=None):
     # Define the arguments, with the possibility to set crossbar size and count
     args = [
         network_path,
@@ -51,10 +56,18 @@ def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
         f"--crossbar-count={crossbar_count}",
         "--enable-timing",
     ]
+    if core_count is not None:
+        args.append(f"--core-count={core_count}")
+
+    cmd = [str(raptor_onnx_path)] + [str(arg) for arg in args]
+    if reporter is not None:
+        reporter.log(f"  Raptor command: {_format_command(cmd)}")
+    else:
+        print(f"Raptor command: {_format_command(cmd)}")
 
     try:
         output_text = run_command_with_reporter(
-            [str(raptor_onnx_path)] + [str(arg) for arg in args],
+            cmd,
             cwd=cwd,
             reporter=reporter,
             capture_output=True,

validation/validate.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
 import argparse
-import shlex
 import signal
 import subprocess
 import sys
@@ -11,12 +10,6 @@ from validate_one import ProgressReporter, clean_workspace_artifacts, validate_n
 from raptor import PIM_PASS_LABELS
 
 
-def format_command(cmd):
-    if isinstance(cmd, (list, tuple)):
-        return shlex.join(str(arg) for arg in cmd)
-    return str(cmd)
-
-
 def format_return_status(returncode):
     if returncode < 0:
         signal_num = -returncode
@@ -34,8 +27,6 @@ def print_validation_error(reporter, rel, exc):
           file=sys.stderr, flush=True)
     if isinstance(exc, subprocess.CalledProcessError):
         print(format_return_status(exc.returncode), file=sys.stderr, flush=True)
-        print("Retry command:", file=sys.stderr, flush=True)
-        print(format_command(exc.cmd), file=sys.stderr, flush=True)
     else:
         print(f"{type(exc).__name__}: {exc}", file=sys.stderr, flush=True)
     print("=" * 72, file=sys.stderr, flush=True)
@@ -65,6 +56,8 @@ def main():
     ap.add_argument("--threshold", type=float, default=1e-3, help="Max allowed diff per output element.")
     ap.add_argument("--crossbar-size", type=int, default=64)
     ap.add_argument("--crossbar-count", type=int, default=8)
+    ap.add_argument("--core-count", type=int, default=None,
+                    help="Core count to pass to Raptor. If omitted, Raptor uses its default.")
     ap.add_argument("--clean", action="store_true",
                     help="Remove generated validation artifacts under each model workspace and exit.")
     a = ap.parse_args()
@@ -114,7 +107,7 @@ def main():
         try:
             result = validate_network(
                 onnx_path, a.raptor_path, a.onnx_include_dir, simulator_dir,
-                crossbar_size=a.crossbar_size, crossbar_count=a.crossbar_count,
+                crossbar_size=a.crossbar_size, crossbar_count=a.crossbar_count, core_count=a.core_count,
                 threshold=a.threshold,
                 reporter=reporter,
                 model_index=index,

validation/validate_one.py

@@ -1,4 +1,3 @@
-import argparse
 import json
 import numpy as np
 import subprocess
@@ -258,7 +257,7 @@ def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold=1
 
 
 def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
-                     simulator_dir, crossbar_size=64, crossbar_count=8, threshold=1e-3,
+                     simulator_dir, crossbar_size=64, crossbar_count=8, core_count=None, threshold=1e-3,
                      reporter=None, model_index=1, model_total=1):
     network_onnx_path = Path(network_onnx_path).resolve()
     raptor_path = Path(raptor_path).resolve()
@@ -313,7 +312,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
         print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compile PIM")
         pim_pass_timings = compile_with_raptor(
             network_mlir_path, raptor_path, raptor_dir / network_onnx_path.stem,
-            crossbar_size, crossbar_count,
+            crossbar_size, crossbar_count, core_count=core_count,
             cwd=raptor_dir, reporter=reporter)
         print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
         reporter.advance()
@@ -350,18 +349,3 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
             reporter.log("=" * 72)
         if owns_reporter:
             reporter.finish()
-
-
-if __name__ == '__main__':
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--network-onnx", required=True)
-    ap.add_argument("--raptor-path", required=True)
-    ap.add_argument("--onnx-include-dir", required=True)
-    a = ap.parse_args()
-
-    simulator_dir = Path(__file__).parent.resolve() / ".." / "backend-simulators" / "pim" / "pim-simulator"
-
-    passed = validate_network(
-        a.network_onnx, a.raptor_path, a.onnx_include_dir, simulator_dir
-    )
-    raise SystemExit(0 if passed.passed else 1)