even faster on pimcomp models

validation/raptor_validation/pimsim_nn.py

@@ -1,8 +1,11 @@
 import json
 import re
 import shutil
+import struct
 from pathlib import Path
 
+import numpy as np
+
 
 _METRIC_PATTERNS = {
     "output_count": r"output count:\s+([0-9]+)\s+samples",
@@ -14,6 +17,14 @@ _METRIC_PATTERNS = {
 }
 
 
+def read_raptor_instruction_count(path):
+    with open(path, "rb") as f:
+        magic, version, count = struct.unpack("<4sII", f.read(12))
+    if magic != b"PIMB" or version != 1:
+        raise ValueError(f"Invalid Raptor binary instruction file: {path}")
+    return count
+
+
 def parse_pimsim_nn_metrics(output):
     metrics = {"raw_output": output}
     for name, pattern in _METRIC_PATTERNS.items():
@@ -45,6 +56,46 @@ def export_raptor_latency_artifact(pim_dir, output_dir):
         json.dump(config, f, separators=(",", ":"))
         f.write("\n")
 
+    binary_sources = sorted(pim_dir.glob("core_*.pim"), key=lambda path: int(path.stem.split("_")[1]))
+    if binary_sources:
+        record_dtype = np.dtype([
+            ("opcode", "u1"), ("rd", "u1"), ("r1", "u1"), ("flags", "u1"),
+            ("r2_or_imm", "<i4"), ("generic1", "<i4"),
+            ("generic2", "<i4"), ("generic3", "<i4"),
+        ])
+        for source in binary_sources:
+            destination = output_dir / source.name
+            shutil.copyfile(source, destination)
+            count = read_raptor_instruction_count(destination)
+            if destination.stat().st_size != 12 + count * record_dtype.itemsize:
+                raise ValueError(f"Invalid Raptor binary instruction file: {source}")
+            records = np.memmap(destination, dtype=record_dtype, mode="r+", offset=12, shape=(count,))
+            scalar_sldi = np.zeros(count, dtype=bool)
+            vmv = records["opcode"] == 22
+            sldi = records["opcode"] == 1
+            for register in np.unique(records["r2_or_imm"][vmv]):
+                definitions = np.flatnonzero(sldi & (records["rd"] == register))
+                consumers = np.flatnonzero(vmv & (records["r2_or_imm"] == register))
+                reaching = np.searchsorted(definitions, consumers) - 1
+                scalar_sldi[definitions[reaching[reaching >= 0]]] = True
+            for start in range(0, count, 1_000_000):
+                chunk = records[start:start + 1_000_000]
+                setbw = chunk["opcode"] == 8
+                chunk["generic1"][setbw] = 8
+                chunk["generic2"][setbw] = 8
+                address_sldi = (chunk["opcode"] == 1) & ~scalar_sldi[start:start + len(chunk)]
+                if np.any(chunk["r2_or_imm"][address_sldi] % 4):
+                    raise ValueError(f"Raptor address is not aligned to its fp32 element width: {source}")
+                chunk["r2_or_imm"][address_sldi] //= 4
+                transfer = (chunk["opcode"] >= 25) & (chunk["opcode"] <= 30)
+                if np.any(chunk["generic2"][transfer] % 4) or np.any(chunk["generic3"][transfer] % 4):
+                    raise ValueError(f"Raptor transfer field is not aligned to its fp32 element width: {source}")
+                chunk["generic2"][transfer] //= 4
+                chunk["generic3"][transfer] //= 4
+            records.flush()
+            del records
+        return output_dir
+
     byte_size_fields = {
         "ld": "size",
         "st": "size",
@@ -56,12 +107,19 @@ def export_raptor_latency_artifact(pim_dir, output_dir):
     for source in sorted(pim_dir.glob("core_*.json"), key=lambda path: int(path.stem.split("_")[1])):
         with open(source, encoding="utf-8") as f:
             instructions = json.load(f)
-        for instruction in instructions:
+        last_sldi = {}
+        scalar_sldi = set()
+        for index, instruction in enumerate(instructions):
+            if instruction["op"] == "sldi":
+                last_sldi[instruction["rd"]] = index
+            elif instruction["op"] == "vmv" and instruction["rs2"] in last_sldi:
+                scalar_sldi.add(last_sldi[instruction["rs2"]])
+        for index, instruction in enumerate(instructions):
             op = instruction["op"]
             if op == "setbw":
                 instruction["ibiw"] = 8
                 instruction["obiw"] = 8
-            elif op == "sldi":
+            elif op == "sldi" and index not in scalar_sldi:
                 instruction["imm"] = int8_bytes(instruction["imm"], "address")
             if field := byte_size_fields.get(op):
                 instruction[field] = int8_bytes(instruction[field], f"{op} {field}")

validation/raptor_validation/validate_one.py

@@ -1,8 +1,10 @@
 import json
 import os
+import re
 import shutil
 import subprocess
 import sys
+import time
 import numpy as np
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -10,7 +12,7 @@ from colorama import Style, Fore
 from .gen_network_runner import gen_network_runner
 from .onnx_utils import gen_random_inputs, save_inputs_to_files, onnx_io, write_inputs_to_memory_bin, _ONNX_TO_NP
 from .raptor import compile_with_raptor
-from .pimsim_nn import export_raptor_latency_artifact, parse_pimsim_nn_metrics
+from .pimsim_nn import export_raptor_latency_artifact, parse_pimsim_nn_metrics, read_raptor_instruction_count
 from .subprocess_utils import run_command_with_reporter
 
 STAGE_TITLES = (
@@ -80,6 +82,35 @@ class ValidationResult:
     pimsim_power_mw: float | None = None
     pimsim_energy_pj: float | None = None
     pimsim_status: str = PIMSIM_SKIPPED
+    compile_time_s: float | None = None
+    host_memory_bytes: int | None = None
+    cores_memory_bytes: int | None = None
+    used_core_count: int | None = None
+    used_crossbar_count: int | None = None
+
+
+_MEMORY_UNITS = {"B": 1, "KB": 1 << 10, "MB": 1 << 20, "GB": 1 << 30}
+
+
+def collect_pim_resource_metrics(pim_dir):
+    pim_dir = Path(pim_dir)
+    report_path = pim_dir.parent / "reports" / "memory_report.txt"
+    report = report_path.read_text(encoding="utf-8") if report_path.exists() else ""
+
+    def memory_bytes(label):
+        match = re.search(rf"^\s*{re.escape(label)}:\s+([0-9.]+)\s+(B|KB|MB|GB)$", report, re.MULTILINE)
+        return round(float(match.group(1)) * _MEMORY_UNITS[match.group(2)]) if match else None
+
+    with open(pim_dir / "config.json", encoding="utf-8") as f:
+        config = json.load(f)
+    used_cores = sum(read_raptor_instruction_count(path) > 0 for path in pim_dir.glob("core_*.pim"))
+    used_crossbars = sum(sum(groups) for groups in config.get("array_group_map", {}).values())
+    return {
+        "host_memory_bytes": memory_bytes("Host memory"),
+        "cores_memory_bytes": memory_bytes("Local memory after reuse"),
+        "used_core_count": used_cores,
+        "used_crossbar_count": used_crossbars,
+    }
 
 
 class ProgressReporter:
@@ -416,8 +447,6 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
         pimsim_nn_build_dir = Path(pimsim_nn_build_dir).resolve()
         pimsim_config_path = Path(pimsim_config_path).resolve()
     compile_extra_args = list(raptor_extra_args or [])
-    if pimsim_enabled and "--pim-emit-json" not in compile_extra_args:
-        compile_extra_args.append("--pim-emit-json")
     owns_reporter = reporter is None
     reporter = reporter or ProgressReporter(model_total, stages_per_model=len(MODE_STAGE_TITLES[mode]), verbose=verbose)
 
@@ -435,6 +464,8 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
                  f" {Style.BRIGHT}Validating {network_onnx_path.name}{Style.RESET_ALL}")
     failed_with_exception = False
     pim_pass_timings = {}
+    compile_time_s = None
+    resource_metrics = {}
 
     try:
         stem = network_onnx_path.stem
@@ -443,6 +474,19 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
         runner_path = runner_build_dir / "runner"
         pim_output_base = raptor_dir / stem
 
+        def compile_pim():
+            nonlocal compile_time_s, resource_metrics
+            started = time.perf_counter()
+            timings = compile_with_raptor(
+                network_onnx_path, raptor_path, pim_output_base, crossbar_size,
+                crossbar_count, core_count=core_count,
+                raptor_extra_args=compile_extra_args, cwd=raptor_dir,
+                verbose=verbose, reporter=reporter,
+                timeout_sec=command_timeout_seconds)
+            compile_time_s = time.perf_counter() - started
+            resource_metrics = collect_pim_resource_metrics(raptor_dir / "pim")
+            return timings
+
         if mode != MODE_RUN_ONLY:
             print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compile ONNX")
             network_so_path, network_mlir_path = compile_onnx_network(
@@ -468,16 +512,14 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
 
             if mode == MODE_COMPILE_ONLY:
                 print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compile PIM")
-                pim_pass_timings = compile_with_raptor(
-                    network_onnx_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
-                    core_count=core_count,
-                    raptor_extra_args=compile_extra_args,
-                    cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds)
+                pim_pass_timings = compile_pim()
                 print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
                 reporter.advance()
                 reporter.record_result(True)
                 reporter.log(Style.BRIGHT + f"Result: {Fore.GREEN}PASS{Style.RESET_ALL}" + Style.RESET_ALL)
-                return ValidationResult(passed=True, pim_pass_timings=pim_pass_timings)
+                return ValidationResult(
+                    passed=True, pim_pass_timings=pim_pass_timings,
+                    compile_time_s=compile_time_s, **resource_metrics)
 
         if mode == MODE_RUN_ONLY:
             required_paths = [
@@ -489,6 +531,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
             missing = [f"{description} at {path}" for path, description in required_paths if not path.exists()]
             if missing:
                 raise FileNotFoundError("run-only mode requires existing artifacts:\n  " + "\n  ".join(missing))
+            resource_metrics = collect_pim_resource_metrics(raptor_dir / "pim")
 
         print_stage(reporter, model_index, model_total, network_onnx_path.name, "Generate Inputs")
         inputs_descriptor, outputs_descriptor = onnx_io(network_onnx_path)
@@ -508,11 +551,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
 
         if mode != MODE_RUN_ONLY:
             print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compile PIM")
-            pim_pass_timings = compile_with_raptor(
-                network_onnx_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
-                core_count=core_count,
-                raptor_extra_args=compile_extra_args,
-                cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds)
+            pim_pass_timings = compile_pim()
             print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
             reporter.advance()
 
@@ -587,6 +626,8 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
             pimsim_power_mw=pimsim_power_mw,
             pimsim_energy_pj=pimsim_energy_pj,
             pimsim_status=pimsim_status,
+            compile_time_s=compile_time_s,
+            **resource_metrics,
         )
     except Exception:
         failed_with_exception = True