add non-functional validation
Validate Operations / validate-operations (push) Has been cancelled

refactors
This commit is contained in:
NiccoloN
2026-07-27 10:54:59 +02:00
parent 4964e889da
commit 620e381cfb
30 changed files with 677 additions and 164 deletions
+1
View File
@@ -0,0 +1 @@
"""Reusable Raptor validation support."""
@@ -0,0 +1,243 @@
#!/usr/bin/env python3
import argparse, os, pathlib, textwrap
from .onnx_utils import onnx_io
from onnx import TensorProto
# ONNX dtype -> (ctype, printf, ONNX_TYPE_*)
DTYPES = {
TensorProto.FLOAT: ("float", "%.9g", "ONNX_TYPE_FLOAT"),
TensorProto.DOUBLE: ("double", "%.17g", "ONNX_TYPE_DOUBLE"),
TensorProto.INT64: ("int64_t", "%lld", "ONNX_TYPE_INT64"),
TensorProto.INT32: ("int32_t", "%d", "ONNX_TYPE_INT32"),
TensorProto.UINT8: ("uint8_t", "%u", "ONNX_TYPE_UINT8"),
TensorProto.INT8: ("int8_t", "%d", "ONNX_TYPE_INT8"),
TensorProto.BOOL: ("uint8_t", "%u", "ONNX_TYPE_BOOL"),
TensorProto.FLOAT16: ("uint16_t", "%u", "ONNX_TYPE_FLOAT16"),
TensorProto.BFLOAT16:("uint16_t", "%u", "ONNX_TYPE_BFLOAT16"),
}
def esc(s): return s.replace("\\","\\\\").replace('"','\\"')
def gen_c(inputs, outputs, entry, so_name):
in_blocks=[]
for i,name,et,shape in inputs:
if et not in DTYPES:
raise ValueError(f"Unsupported dtype for input '{name}': {et}")
cty, pfmt, onnx_ty = DTYPES[et]
shp_list = ", ".join(str(d) for d in shape) if shape else ""
rank = len(shape)
in_blocks.append(textwrap.dedent(f"""
// ---- Input {i}: "{esc(name)}" ({cty}) ----
const char *in{i}_csv=NULL, *in{i}_csv_file=NULL, *in{i}_shape_str=NULL;
char *in{i}_csv_buf=NULL; // holds file contents if --in{i}-csv-file used
int has_in{i}=0; double in{i}_fill=0.0; int in{i}_fill_set=0;
for (int ai=1; ai<argc; ++ai) {{
if (strncmp(argv[ai],"--in{i}-csv",11)==0 && ai+1<argc) {{ in{i}_csv=argv[ai+1]; has_in{i}=1; }}
else if (strncmp(argv[ai],"--in{i}-csv-file",16)==0 && ai+1<argc) {{ in{i}_csv_file=argv[ai+1]; has_in{i}=1; }}
else if (strncmp(argv[ai],"--in{i}-fill",12)==0 && ai+1<argc) {{ in{i}_fill=atof(argv[ai+1]); in{i}_fill_set=1; has_in{i}=1; }}
else if (strncmp(argv[ai],"--in{i}-shape",13)==0 && ai+1<argc) {{ in{i}_shape_str=argv[ai+1]; }}
}}
if (!has_in{i}) {{
fprintf(stderr,"ERROR: provide one of --in{i}-csv/--in{i}-csv-file/--in{i}-fill for input {i}.\\n");
return 2;
}}
// If a CSV file was provided, read it fully and use its content as the CSV string.
if (in{i}_csv_file && !in{i}_csv) {{
FILE *f=fopen(in{i}_csv_file,"rb");
if(!f){{perror("fopen --in{i}-csv-file"); return 2;}}
fseek(f, 0, SEEK_END);
long sz = ftell(f); if (sz < 0) {{ perror("ftell"); fclose(f); return 2; }}
fseek(f, 0, SEEK_SET);
in{i}_csv_buf = (char*)malloc((size_t)sz + 1);
if(!in{i}_csv_buf){{fprintf(stderr,"OOM reading --in{i}-csv-file.\\n"); fclose(f); return 2;}}
size_t got = fread(in{i}_csv_buf, 1, (size_t)sz, f);
fclose(f);
if (got != (size_t)sz) {{ fprintf(stderr,"ERROR: short read for --in{i}-csv-file.\\n"); free(in{i}_csv_buf); return 2; }}
in{i}_csv_buf[sz] = '\\0';
in{i}_csv = in{i}_csv_buf;
}}
int64_t *in{i}_shape=NULL; int in{i}_rank=0;
if (in{i}_shape_str) {{
char *tmp=strdup(in{i}_shape_str);
for(char*p=tmp; *p; ++p) if(*p=='x'||*p=='X') in{i}_rank++;
in{i}_rank++;
in{i}_shape=(int64_t*)malloc(sizeof(int64_t)*in{i}_rank);
int di=0; char *tok=strtok(tmp,"xX");
while(tok && di<in{i}_rank) {{ in{i}_shape[di++]=atoll(tok); tok=strtok(NULL,"xX"); }}
free(tmp);
}} else {{
in{i}_rank={rank};
in{i}_shape=(int64_t*)malloc(sizeof(int64_t)*in{i}_rank);
int64_t def_shape[]={{{shp_list}}};
for(int k=0;k<in{i}_rank;k++) in{i}_shape[k]=def_shape[k];
}}
long long in{i}_nelem=1; for(int k=0;k<in{i}_rank;k++) in{i}_nelem*=in{i}_shape[k];
size_t in{i}_bytes = sizeof({cty}) * (size_t)in{i}_nelem;
void *in{i}_buf = malloc(in{i}_bytes);
if(!in{i}_buf){{fprintf(stderr,"OOM for input {i}.\\n"); if(in{i}_csv_buf) free(in{i}_csv_buf); return 2;}}
if (in{i}_csv) {{
char *buf=strdup(in{i}_csv); long long idx=0; char *tok=strtok(buf,",\\n\\r\\t ");
while(tok) {{
if(idx>=in{i}_nelem) break;
double v=atof(tok);
(({cty}*)in{i}_buf)[idx++] = ({cty})v;
tok=strtok(NULL,",\\n\\r\\t ");
}}
free(buf);
if(idx!=in{i}_nelem){{fprintf(stderr,"ERROR: CSV provided %lld values, expected %lld.\\n",(long long)idx,in{i}_nelem); if(in{i}_csv_buf) free(in{i}_csv_buf); return 2;}}
}} else if (in{i}_fill_set) {{
{cty} vv=({cty})in{i}_fill; for(long long t=0;t<in{i}_nelem;t++) (({cty}*)in{i}_buf)[t]=vv;
}} else {{
fprintf(stderr,"ERROR: no data source for input {i}.\\n"); if(in{i}_csv_buf) free(in{i}_csv_buf); return 2;
}}
OMTensor *in{i}_tensor = omTensorCreateWithOwnership(in{i}_buf, in{i}_shape, in{i}_rank, {onnx_ty}, /*owning=*/1);
if(in{i}_csv_buf) free(in{i}_csv_buf);
if(!in{i}_tensor){{fprintf(stderr,"ERROR: omTensorCreateWithOwnership failed for input {i}.\\n");return 2;}}
"""))
# Optional per-output CSV dump
csv_write_blocks=[]
for oi,name,et,shape in outputs:
if et not in DTYPES:
raise ValueError(f"Unsupported dtype for output '{name}': {et}")
cty, pfmt, _ = DTYPES[et]
safe = esc(name)
csv_write_blocks.append(textwrap.dedent(f"""
if (save_csv_dir) {{
// Build "DIR/output{oi}_<sanitized name>.csv"
char fname[512];
// simple sanitizer: copy name => replace non [A-Za-z0-9_.-] with '_'
char clean[256]; int ci=0; const char *src="{safe}";
for (; src[ci] && ci < 255; ++ci) {{
char ch = src[ci];
int ok = (ch>='A'&&ch<='Z')||(ch>='a'&&ch<='z')||(ch>='0'&&ch<='9')||ch=='_'||ch=='-'||ch=='.';
clean[ci] = ok ? ch : '_';
}}
clean[ci] = '\\0';
snprintf(fname, sizeof(fname), "%s/output{oi}_%s.csv", save_csv_dir, clean);
FILE *csv = fopen(fname, "w");
if (!csv) {{ perror("fopen --save-csv-dir"); }}
else {{
OMTensor *t = omTensorListGetOmtByIndex(out_list, {oi});
int64_t rank = omTensorGetRank(t);
int64_t const *shape = omTensorGetShape(t);
long long numel = 1; for (int64_t k=0;k<rank;k++) numel *= shape[k];
{cty} *p = ({cty}*)omTensorGetDataPtr(t);
if (rank == 2) {{
int64_t R = shape[0], C = shape[1];
for (int64_t r=0; r<R; ++r) {{
for (int64_t c=0; c<C; ++c) {{
long long idx = r*C + c;
fprintf(csv, "{pfmt}%s", p[idx], (c+1<C)?",":"");
}}
fprintf(csv, "\\n");
}}
}} else {{
for (long long i=0;i<numel;i++) {{
fprintf(csv, "{pfmt}%s", p[i], (i+1<numel)?",":"");
}}
fprintf(csv, "\\n");
}}
fclose(csv);
}}
}}
"""))
n_in=len(inputs)
build_inputs="\n".join([f" arr[{i}] = in{i}_tensor;" for i,_,_,_ in inputs])
return f"""\
// Auto-generated onnx network runner
#include "OnnxMlirRuntime.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <inttypes.h>
OMTensorList *{entry}(OMTensorList *inputs);
int main(int argc, char **argv) {{
// optional: --save-csv-dir <DIR> (directory must exist)
const char *save_csv_dir = NULL;
for (int ai=1; ai<argc; ++ai) {{
if (strncmp(argv[ai], "--save-csv-dir", 14)==0 && ai+1 < argc) {{
save_csv_dir = argv[ai+1];
}}
}}
if (argc == 1) {{
fprintf(stderr,
"Usage: %s "
"[--inK-csv \\\"v1,v2,...\\\" | --inK-csv-file path | --inK-fill c] "
"[--inK-shape 1x...xD] "
"[--save-csv-dir /path/to/dir]\\n"
"Repeat for K=0..%d.\\n",
argv[0], {max(0, n_in-1)});
return 1;
}}
{"".join(in_blocks)}
OMTensor *arr[{n_in}];
{build_inputs}
OMTensorList *in_list = omTensorListCreate(arr, {n_in});
if(!in_list){{fprintf(stderr,"ERROR: omTensorListCreate failed.\\n");return 2;}}
OMTensorList *out_list = {entry}(in_list);
if(!out_list){{fprintf(stderr,"ERROR: model returned NULL.\\n");omTensorListDestroy(in_list);return 3;}}
// ---- Optional per-output CSV dump ----
{"".join(csv_write_blocks)}
// ---- Cleanup ----
omTensorListDestroy(in_list);
// Some debug-heavy models return aliased outputs. This runner is a short-
// lived process, so destroy only the list wrapper and let process exit
// reclaim the output tensors safely.
omTensorListDestroyShallow(out_list);
return 0;
}}
"""
def gen_network_runner(network_onnx, network_so, onnx_include_dir, entry, out, verbose):
ins, outs = onnx_io(network_onnx)
out_c = out or "runner.c"
so_abs = os.path.abspath(network_so)
onnx_include_dir = str(onnx_include_dir)
csrc = gen_c(ins, outs, entry, pathlib.Path(so_abs).name)
pathlib.Path(out_c).write_text(csrc)
cmake=f"""\
cmake_minimum_required(VERSION 3.15)
project(onnx_mlir_runner C)
add_executable({pathlib.Path(out_c).stem} {pathlib.Path(out_c).name})
target_include_directories({pathlib.Path(out_c).stem} PUBLIC {esc(onnx_include_dir)})
add_library(model_so SHARED IMPORTED)
set_target_properties(model_so PROPERTIES IMPORTED_LOCATION {esc(so_abs)})
target_link_libraries({pathlib.Path(out_c).stem} PUBLIC model_so)
"""
pathlib.Path(out_c).with_name("CMakeLists.txt").write_text(cmake)
if verbose:
print(f"[OK] Wrote {out_c}")
print("[OK] Wrote CMakeLists.txt")
if __name__=="__main__":
ap=argparse.ArgumentParser()
ap.add_argument("--network-onnx", required=True)
ap.add_argument("--network-so", required=True)
ap.add_argument("--onnx-include-dir", required=True)
ap.add_argument("--entry", default="run_main_graph")
ap.add_argument("--out", default=None)
a=ap.parse_args()
gen_network_runner(a.network_onnx, a.network_so, a.onnx_include_dir, a.entry, a.out, True)
+203
View File
@@ -0,0 +1,203 @@
import csv
import onnx
import json
import pathlib
import numpy as np
from onnx import TensorProto
_ONNX_TO_NP = {
TensorProto.FLOAT: np.float32,
TensorProto.DOUBLE: np.float64,
TensorProto.INT64: np.int64,
TensorProto.INT32: np.int32,
TensorProto.UINT8: np.uint8,
TensorProto.INT8: np.int8,
TensorProto.BOOL: np.uint8, # store as 0/1 bytes
TensorProto.FLOAT16: np.float16, # generate in f32 then cast
TensorProto.BFLOAT16: getattr(np, "bfloat16", np.float32), # cast if available
}
def onnx_io(path):
m = onnx.load(path)
g = m.graph
def shp(tt):
s = []
if tt.HasField("shape"):
for d in tt.shape.dim:
s.append(int(d.dim_value) if d.HasField("dim_value") else 1)
return s
ins, outs = [], []
initializer_names = {initializer.name for initializer in g.initializer}
for v in g.input:
if v.name in initializer_names:
continue
t = v.type.tensor_type
ins.append((len(ins), v.name, t.elem_type, shp(t)))
for i, v in enumerate(g.output):
t = v.type.tensor_type
outs.append((i, v.name, t.elem_type, shp(t)))
return ins, outs
def onnx_io_bitsize(io):
idx, name, elem_type, shape = io
num_elements = shape[0]
for dim in shape[1:]:
num_elements *= dim
return num_elements * _ONNX_TO_NP[elem_type]().itemsize * 8
def _dtype_bounds(np_dtype):
"""Return (min, max) inclusive bounds for integer dtypes; None for floats."""
if np_dtype in (np.int8, np.int16, np.int32, np.int64):
info = np.iinfo(np_dtype)
return int(info.min), int(info.max)
if np_dtype in (np.uint8, np.uint16, np.uint32, np.uint64):
info = np.iinfo(np_dtype)
return int(info.min), int(info.max)
return None
def gen_random_inputs(
onnx_inputs,
*,
shape_overrides: dict | None = None,
float_range: tuple[float, float] = (-1.0, 1.0),
int_range: tuple[int, int] = (-3, 3),
dyn_dim_default: int = 1,
seed: int | None = None,
):
"""
Generate random NumPy arrays for each ONNX input.
Params
------
shape_overrides:
Dict mapping input index OR input name -> tuple/list of dims.
Overrides the shape inferred from the model (useful for dynamic dims).
float_range:
Range for floats (uniform).
int_range:
Range for integers (uniform integers, inclusive of low/high with np.integers semantics).
dyn_dim_default:
If a dim is dynamic/unknown, use this value (unless shape_overrides provides one).
seed:
RNG seed for reproducibility.
Returns
-------
inputs_list: list[np.ndarray]
Arrays in graph input order (index-sorted).
inputs_dict: dict[str, np.ndarray]
Mapping input_name -> array in the ONNX-declared dtype.
"""
rng = np.random.default_rng(seed)
ins = onnx_inputs
# Normalize overrides to support both index and name keys.
shape_overrides = shape_overrides or {}
name_overrides = {k: tuple(v) for k, v in shape_overrides.items() if isinstance(k, str)}
idx_overrides = {int(k): tuple(v) for k, v in shape_overrides.items() if isinstance(k, int)}
arrays_by_name = {}
arrays_in_order = []
for idx, name, elem_type, shape in ins:
# Resolve dtype
if elem_type not in _ONNX_TO_NP:
raise ValueError(f"Unsupported ONNX dtype for input '{name}': {elem_type}")
np_dtype = _ONNX_TO_NP[elem_type]
# Resolve shape: model -> replace unknowns with dyn_dim_default -> apply overrides
resolved_shape = list(shape or [])
if not resolved_shape:
resolved_shape = [dyn_dim_default] # scalar-like: treat as 1-dim with size dyn_dim_default
# If your onnx_io already sets unknown dims to 1, we still allow overriding:
if idx in idx_overrides:
resolved_shape = list(idx_overrides[idx])
elif name in name_overrides:
resolved_shape = list(name_overrides[name])
# Make sure no zeros
resolved_shape = [int(d if d and d > 0 else dyn_dim_default) for d in resolved_shape]
size = int(np.prod(resolved_shape))
# Generate data
if np.issubdtype(np_dtype, np.floating):
lo, hi = float_range
# generate in float32/64 and cast as needed
base_dtype = np.float32 if np_dtype in (np.float16, getattr(np, "bfloat16", np.float32)) else np_dtype
arr = rng.uniform(lo, hi, size=size).astype(base_dtype).reshape(resolved_shape)
# cast to f16/bf16 if required
if np_dtype is np.float16:
arr = arr.astype(np.float16)
elif getattr(np, "bfloat16", None) is not None and np_dtype is np.bfloat16:
arr = arr.astype(np.bfloat16)
elif np_dtype == np.uint8 and elem_type == TensorProto.BOOL:
# Bool as 0/1 bytes
arr = (rng.random(size=size) < 0.5).astype(np.uint8).reshape(resolved_shape)
elif np.issubdtype(np_dtype, np.integer):
lo, hi = int_range
bounds = _dtype_bounds(np_dtype)
if bounds is not None:
lo = max(lo, bounds[0])
hi = min(hi, bounds[1])
# np.random.integers is exclusive of high; add 1 for int range
arr = rng.integers(lo, hi + 1, size=size, dtype=np_dtype).reshape(resolved_shape)
else:
raise ValueError(f"Unhandled dtype mapping for input '{name}' (elem_type={elem_type}).")
arrays_by_name[name] = arr
arrays_in_order.append(arr)
return arrays_in_order, arrays_by_name
def save_inputs_to_files(onnx_path, arrays_in_order, out_dir):
"""
Save arrays to CSV files. Returns (flags, files) where flags is a list
like ["--in0-csv-file", "...", "--in0-shape", "Dx...xD", ...]
and files is the list of created paths.
"""
out_dir = pathlib.Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
ins, _ = onnx_io(onnx_path)
flags = []
files = []
for idx, _name, _et, shape in ins:
arr = arrays_in_order[idx]
csv_path = out_dir / f"in{idx}.csv"
# Write row-major flattened values, comma-separated, with newlines allowed
with open(csv_path, "w", newline="") as f:
writer = csv.writer(f)
# For 2D, write each row; otherwise write flattened single row for clarity
if arr.ndim == 2:
for r in range(arr.shape[0]):
writer.writerow(arr[r].reshape(-1))
else:
writer.writerow(arr.flatten())
shape_str = "x".join(str(d) for d in arr.shape)
flags += [f"--in{idx}-csv-file", str(csv_path), f"--in{idx}-shape", shape_str]
files.append(str(csv_path))
return flags, files
def write_inputs_to_memory_bin(memory_bin_path, config_json_path, arrays_in_order):
"""Overwrite input regions in memory.bin at addresses from config.json."""
with open(config_json_path) as f:
config = json.load(f)
input_addresses = config["inputs_addresses"]
assert len(input_addresses) == len(arrays_in_order), \
f"Address/input count mismatch: {len(input_addresses)} vs {len(arrays_in_order)}"
with open(memory_bin_path, "r+b") as f:
for addr, arr in zip(input_addresses, arrays_in_order):
native = arr.astype(arr.dtype.newbyteorder("="), copy=False)
f.seek(addr)
f.write(native.tobytes(order="C"))
+83
View File
@@ -0,0 +1,83 @@
import re
import shlex
import subprocess
from pathlib import Path
from colorama import Fore, Style
from .subprocess_utils import run_command_with_reporter
PIM_PASS_LABELS = (
("ONNXToSpatialPass", "ONNX to Spatial"),
("MergeComputeNodesPass", "Merge Compute Nodes"),
("SpatialToPimPass", "Spatial to PIM"),
("PimBufferizationPass", "Bufferize PIM"),
("HostConstantFoldingPass", "Fold Host Constants"),
("PimLocalMemoryPlanningPass", "Plan Local Memory"),
("VerificationPass", "Verify PIM"),
("EmitPimCodePass", "Emit PIM Code"),
)
PIM_PASS_LABEL_BY_SUFFIX = dict(PIM_PASS_LABELS)
TIMING_LINE_RE = re.compile(r"^\s*([0-9]+\.[0-9]+)\s+\(\s*[0-9.]+%\)\s+(.+?)\s*$")
def _parse_pim_pass_timings(output_text):
pass_timings = {}
for line in output_text.splitlines():
match = TIMING_LINE_RE.match(line)
if not match:
continue
duration = float(match.group(1))
pass_name = match.group(2)
for suffix, label in PIM_PASS_LABEL_BY_SUFFIX.items():
if pass_name.endswith(suffix):
pass_timings[label] = pass_timings.get(label, 0.0) + duration
break
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, core_count,
raptor_extra_args, cwd, verbose, reporter, timeout_sec):
# Define the arguments, with the possibility to set crossbar size and count
args = [
network_path,
"-o",
output_base,
"--maccel=PIM",
"--EmitPimCodegen",
f"--crossbar-size={crossbar_size}",
f"--crossbar-count={crossbar_count}",
]
if core_count is not None:
args.append(f"--core-count={core_count}")
if raptor_extra_args:
args.extend(str(arg) for arg in raptor_extra_args)
if verbose:
args.append("--enable-timing")
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(
cmd,
cwd=cwd,
reporter=reporter,
capture_output=True,
timeout_sec=timeout_sec,
)
if reporter is None:
print(Fore.GREEN + "Raptor execution successful" + Style.RESET_ALL)
return _parse_pim_pass_timings(output_text)
except subprocess.CalledProcessError:
if reporter is None:
print(Fore.RED + "Raptor execution failed" + Style.RESET_ALL)
raise
@@ -0,0 +1,119 @@
import errno
import os
import pty
import selectors
import subprocess
import time
MAX_ERROR_OUTPUT_BYTES = 8192
def _read_chunk(fd, treat_eio_as_eof=False):
try:
return os.read(fd, 4096)
except OSError as exc:
if treat_eio_as_eof and exc.errno == errno.EIO:
return b""
raise
def _stream_output(fd, process, reporter, treat_eio_as_eof=False, stream_output=True, timeout_sec=None):
selector = selectors.DefaultSelector()
recent_output = bytearray()
captured_output = bytearray()
deadline = None if timeout_sec is None else time.monotonic() + timeout_sec
try:
selector.register(fd, selectors.EVENT_READ)
while selector.get_map():
select_timeout = None
if deadline is not None:
remaining = deadline - time.monotonic()
if remaining <= 0:
process.kill()
process.wait()
raise subprocess.TimeoutExpired(process.args, timeout_sec, output=bytes(captured_output))
select_timeout = min(1.0, remaining)
for key, _ in selector.select(select_timeout):
data = _read_chunk(key.fileobj, treat_eio_as_eof=treat_eio_as_eof)
if not data:
selector.unregister(key.fileobj)
os.close(key.fileobj)
continue
if stream_output:
reporter._clear()
os.write(1, data)
reporter._render()
captured_output.extend(data)
if stream_output:
recent_output.extend(data)
if len(recent_output) > MAX_ERROR_OUTPUT_BYTES:
del recent_output[:-MAX_ERROR_OUTPUT_BYTES]
finally:
selector.close()
return_code = process.wait()
if return_code != 0:
error_output = captured_output if not stream_output else recent_output
exc = subprocess.CalledProcessError(return_code, process.args, output=bytes(error_output))
exc.output_already_streamed = stream_output and bool(captured_output)
raise exc
return bytes(captured_output)
def run_command_with_reporter(cmd, cwd=None, reporter=None, capture_output=False, timeout_sec=None):
if reporter is None:
if capture_output:
completed = subprocess.run(
cmd,
cwd=cwd,
check=True,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
timeout=timeout_sec,
)
return completed.stdout.decode("utf-8", errors="replace")
subprocess.run(cmd, cwd=cwd, check=True, timeout=timeout_sec)
return None
stream_output = bool(getattr(reporter, "verbose", False))
if not stream_output:
completed = subprocess.run(
cmd,
cwd=cwd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
timeout=timeout_sec,
)
if completed.returncode != 0:
raise subprocess.CalledProcessError(completed.returncode, completed.args, output=completed.stdout)
return completed.stdout.decode("utf-8", errors="replace") if capture_output else None
try:
master_fd, slave_fd = pty.openpty()
except OSError:
process = subprocess.Popen(
cmd,
cwd=cwd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
)
assert process.stdout is not None
output = _stream_output(process.stdout.fileno(), process, reporter, timeout_sec=timeout_sec)
return output.decode("utf-8", errors="replace") if capture_output else None
try:
process = subprocess.Popen(
cmd,
cwd=cwd,
stdout=slave_fd,
stderr=slave_fd,
)
finally:
os.close(slave_fd)
output = _stream_output(master_fd, process, reporter, treat_eio_as_eof=True, timeout_sec=timeout_sec)
return output.decode("utf-8", errors="replace") if capture_output else None
@@ -0,0 +1,583 @@
import json
import os
import re
import shutil
import sys
import numpy as np
from dataclasses import dataclass, field
from pathlib import Path
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 .subprocess_utils import run_command_with_reporter
STAGE_TITLES = (
"Compile ONNX",
"Build Runner",
"Generate Inputs",
"Run Reference",
"Compile PIM",
"Run Functional Simulation",
"Compare Outputs",
"Run Non-functional Simulation",
)
STAGE_COUNT = len(STAGE_TITLES)
GENERATED_DIR_NAMES = ("inputs", "outputs", "raptor", "runner", "simulation")
MODE_FULL = "full"
MODE_COMPILE_ONLY = "compile_only"
MODE_RUN_ONLY = "run_only"
MODE_STAGE_TITLES = {
MODE_FULL: STAGE_TITLES,
MODE_COMPILE_ONLY: (
"Compile ONNX",
"Build Runner",
"Compile PIM",
),
MODE_RUN_ONLY: (
"Generate Inputs",
"Run Reference",
"Run Functional Simulation",
"Compare Outputs",
"Run Non-functional Simulation",
),
}
PIMSIM_DONE = "DONE"
PIMSIM_FAILED = "ERROR"
PIMSIM_SKIPPED = "SKIP"
PIMSIM_NOT_RUN = "-"
PIMSIM_LATENCY_RE = re.compile(r"\blatency:\s+([0-9.eE+-]+)\s+ms")
PIMSIM_POWER_RE = re.compile(r"\baverage power:\s+([0-9.eE+-]+)\s+mW")
def sanitize_output_name(name):
return "".join(ch if ch.isalnum() or ch in "_.-" else "_" for ch in name[:255])
@dataclass
class ValidationResult:
passed: bool
pim_pass_timings: dict[str, float] = field(default_factory=dict)
pimsim_latency_ms: float | None = None
pimsim_power_mw: float | None = None
pimsim_status: str = PIMSIM_SKIPPED
class ProgressReporter:
def __init__(self, total_models, stages_per_model=STAGE_COUNT, enabled=None, verbose=False):
self.total_models = total_models
self.stages_per_model = stages_per_model
self.total_steps = max(1, total_models * stages_per_model)
self.completed_steps = 0
self.passed_models = 0
self.failed_models = 0
self.current_label = ""
self.enabled = (
sys.stdout.isatty() and "CODEX_CI" not in os.environ
if enabled is None else enabled
)
self.verbose = verbose
self.columns = max(1, shutil.get_terminal_size((100, 20)).columns)
self.suspended = False
self.rendered_width = 0
self.rendered_rows = 0
def _clear(self):
if self.enabled and self.rendered_rows:
columns = max(1, shutil.get_terminal_size((100, 20)).columns)
rows = max(self.rendered_rows, (self.rendered_width + columns - 1) // columns)
sys.stdout.write("\r\033[2K")
for _ in range(rows - 1):
sys.stdout.write("\033[1A\r\033[2K")
sys.stdout.flush()
self.rendered_width = 0
self.rendered_rows = 0
def _render(self):
if not self.enabled or self.suspended:
return
self.columns = max(1, shutil.get_terminal_size((100, 20)).columns)
bar_width = min(24, max(4, self.columns - 24))
filled = int(bar_width * self.completed_steps / self.total_steps)
counts_text = f"P:{self.passed_models} F:{self.failed_models}"
prefix_text = f"[{'#' * filled}{'-' * (bar_width - filled)}] {self.completed_steps}/{self.total_steps}"
bar = Fore.GREEN + ("#" * filled) + Fore.CYAN + ("-" * (bar_width - filled))
prefix = Fore.CYAN + f"[{bar}{Fore.CYAN}] {self.completed_steps}/{self.total_steps}" + Style.RESET_ALL
counts = (
" "
+ Style.BRIGHT
+ Fore.GREEN
+ f"P:{self.passed_models}"
+ Style.RESET_ALL
+ " "
+ Style.BRIGHT
+ Fore.RED
+ f"F:{self.failed_models}"
+ Style.RESET_ALL
)
model_counter = ""
label = ""
if self.current_label.startswith("[") and "] " in self.current_label:
model_counter, label = self.current_label.split("] ", 1)
model_counter = f" {model_counter}]"
label = f" {label}"
elif self.current_label:
label = f" {self.current_label}"
fixed_width = len(prefix_text) + len(model_counter) + len(counts_text) + 2
if fixed_width > self.columns:
model_counter = ""
fixed_width = len(prefix_text) + len(counts_text) + 2
if fixed_width > self.columns:
prefix_text = f"{self.completed_steps}/{self.total_steps}"
prefix = Fore.CYAN + prefix_text + Style.RESET_ALL
fixed_width = len(prefix_text) + len(counts_text) + 2
if fixed_width > self.columns:
counts = ""
counts_text = ""
fixed_width = len(prefix_text) + 1
available_label_width = max(0, self.columns - fixed_width)
label = label[:available_label_width]
plain_counts = f" {counts_text}" if counts_text else ""
plain_line = prefix_text + model_counter + plain_counts + label
rendered_line = prefix + model_counter + counts + label + Style.RESET_ALL
self._clear()
sys.stdout.write(rendered_line)
sys.stdout.flush()
self.rendered_width = len(plain_line)
self.rendered_rows = max(1, (self.rendered_width + self.columns - 1) // self.columns)
def log(self, message="", color=None):
if not self.verbose:
self._render()
return
if self.enabled:
self._clear()
if color:
print(color + message + Style.RESET_ALL)
else:
print(message)
self._render()
def set_stage(self, model_index, model_total, model_name, stage_name):
self.current_label = f"[{model_index}/{model_total}] {model_name} · {stage_name}"
self._render()
def advance(self):
self.completed_steps = min(self.total_steps, self.completed_steps + 1)
self._render()
def record_result(self, passed):
if passed:
self.passed_models += 1
else:
self.failed_models += 1
self._render()
def suspend(self):
if self.enabled:
self._clear()
self.suspended = True
def resume(self):
self.suspended = False
self._render()
def finish(self):
if self.enabled:
self.suspended = True
self._clear()
def run_command(cmd, cwd=None, reporter=None, timeout_sec=None, capture_output=False):
return run_command_with_reporter(
cmd,
cwd=cwd,
reporter=reporter,
timeout_sec=timeout_sec,
capture_output=capture_output,
)
def load_pimcomp_hardware(config_path):
with open(config_path, encoding="utf-8") as f:
config = json.load(f)
matrix = config["chip_config"]["core_config"]["matrix_config"]
rows, cols = config["chip_config"]["network_config"]["layout"]
xbar_rows, xbar_cols = matrix["xbar_size"]
return {
"core_count": config["chip_config"]["core_cnt"],
"crossbar_count": matrix["xbar_array_count"],
"crossbar_rows": xbar_rows,
"crossbar_cols": xbar_cols,
"mesh_rows": rows,
"mesh_cols": cols,
}
def pimcomp_compatibility_errors(config_path, *, core_count, crossbar_count, crossbar_size):
hardware = load_pimcomp_hardware(config_path)
errors = []
if hardware["mesh_rows"] * hardware["mesh_cols"] != hardware["core_count"]:
errors.append(
f"config layout {hardware['mesh_rows']}x{hardware['mesh_cols']} does not match "
f"{hardware['core_count']} cores"
)
if core_count != hardware["core_count"]:
errors.append(f"--core-count={core_count}, config requires {hardware['core_count']}")
if crossbar_count != hardware["crossbar_count"]:
errors.append(
f"--crossbar-count={crossbar_count}, config requires {hardware['crossbar_count']}"
)
if (
hardware["crossbar_rows"] != hardware["crossbar_cols"]
or crossbar_size != hardware["crossbar_rows"]
):
errors.append(
f"--crossbar-size={crossbar_size}, config requires "
f"{hardware['crossbar_rows']}x{hardware['crossbar_cols']}"
)
return errors
def run_pimsim_nn(pimsim_nn_build_dir, pim_dir, config_path, reporter=None, timeout_sec=None):
output = run_command(
[pimsim_nn_build_dir / "ChipTest", pim_dir, config_path, "false"],
cwd=pimsim_nn_build_dir,
reporter=reporter,
timeout_sec=timeout_sec,
capture_output=True,
)
latency_match = PIMSIM_LATENCY_RE.search(output)
power_match = PIMSIM_POWER_RE.search(output)
if not latency_match or not power_match:
raise RuntimeError("pimsim-nn output did not contain latency and average power")
return float(latency_match.group(1)), float(power_match.group(1))
def clean_workspace_artifacts(workspace_dir, model_stem):
workspace_dir = Path(workspace_dir)
removed_paths = []
def remove_path(path):
if path.is_symlink() or path.is_file():
path.unlink(missing_ok=True)
removed_paths.append(path)
elif path.is_dir():
shutil.rmtree(path)
removed_paths.append(path)
for name in GENERATED_DIR_NAMES:
remove_path(workspace_dir / name)
for suffix in (".onnx.mlir", ".so", ".tmp"):
remove_path(workspace_dir / f"{model_stem}{suffix}")
return removed_paths
def print_stage(reporter, model_index, model_total, model_name, title):
stage_colors = {
STAGE_TITLES[0]: Fore.BLUE,
STAGE_TITLES[1]: Fore.MAGENTA,
STAGE_TITLES[2]: Fore.YELLOW,
STAGE_TITLES[3]: Fore.GREEN,
STAGE_TITLES[4]: Fore.CYAN,
STAGE_TITLES[5]: Fore.MAGENTA,
STAGE_TITLES[6]: Fore.YELLOW,
STAGE_TITLES[7]: Fore.BLUE,
}
color = stage_colors.get(title, Fore.WHITE)
reporter.log(Style.BRIGHT + color + f"[{title}]" + Style.RESET_ALL)
reporter.set_stage(model_index, model_total, model_name, title)
def print_info(reporter, message):
reporter.log(f" {message}")
def compile_onnx_network(network_onnx_path, raptor_path, raptor_dir, runner_dir, reporter=None, timeout_sec=None):
stem = network_onnx_path.stem
onnx_ir_base = raptor_dir / stem
runner_base = runner_dir / stem
run_command([raptor_path, network_onnx_path, "-o", onnx_ir_base, "--EmitONNXIR",
"--mlir-elide-elementsattrs-if-larger=16"],
reporter=reporter, timeout_sec=timeout_sec)
run_command([raptor_path, network_onnx_path, "-o", runner_base], reporter=reporter, timeout_sec=timeout_sec)
network_so_path = runner_base.with_suffix(".so")
network_mlir_path = onnx_ir_base.with_suffix(".onnx.mlir")
onnx_ir_base.with_suffix(".tmp").unlink(missing_ok=True)
return network_so_path, network_mlir_path
def build_onnx_runner(source_dir, build_dir, reporter=None, timeout_sec=None):
run_command(["cmake", source_dir], cwd=build_dir, reporter=reporter, timeout_sec=timeout_sec)
run_command(["cmake", "--build", ".", "-j"], cwd=build_dir, reporter=reporter, timeout_sec=timeout_sec)
return build_dir / "runner"
def build_dump_ranges(config_path, outputs_descriptor):
with open(config_path) as f:
output_addresses = json.load(f)["outputs_addresses"]
ranges = []
for addr, (_, _, dtype_code, shape) in zip(output_addresses, outputs_descriptor):
byte_size = int(np.prod(shape)) * np.dtype(_ONNX_TO_NP[dtype_code]).itemsize
ranges.append(f"{addr},{byte_size}")
return ",".join(ranges)
def run_pim_simulator(simulator_dir, pim_dir, output_bin_path, dump_ranges, reporter=None, timeout_sec=None):
run_command(
["cargo", "run", "--no-default-features", "--release", "--package", "pim-simulator", "--bin", "pim-simulator",
"--",
"-f", str(pim_dir), "-o", str(output_bin_path), "-d", dump_ranges],
cwd=simulator_dir,
reporter=reporter,
timeout_sec=timeout_sec,
)
def parse_pim_simulator_outputs(output_bin_path, outputs_descriptor):
raw = output_bin_path.read_bytes()
arrays = []
offset = 0
for _, _, dtype_code, shape in outputs_descriptor:
dtype = np.dtype(_ONNX_TO_NP[dtype_code])
count = int(np.prod(shape))
array = np.frombuffer(raw, dtype=dtype, count=count, offset=offset).reshape(shape)
offset += count * dtype.itemsize
arrays.append(array)
return arrays
def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold, rtol, verbose):
all_passed = True
rows = []
for sim_array, (oi, name, _, shape) in zip(sim_arrays, outputs_descriptor):
csv_name = f"output{oi}_{sanitize_output_name(name)}.csv"
runner_array = np.loadtxt(runner_out_dir / csv_name, delimiter=',', dtype=np.float32).reshape(shape)
sim_array64 = sim_array.astype(np.float64)
runner_array64 = runner_array.astype(np.float64)
abs_diff = np.abs(sim_array64 - runner_array64)
allowed_diff = threshold + rtol * np.abs(runner_array64)
max_diff = float(np.max(abs_diff))
passed = bool(np.all(abs_diff <= allowed_diff))
rows.append((name, f"{max_diff:.6e}", passed))
if not passed:
all_passed = False
name_width = max(len("Output"), *(len(name) for name, _, _ in rows))
diff_width = max(len("Max diff"), *(len(diff) for _, diff, _ in rows))
result_width = len("Result")
separator = f" +-{'-' * name_width}-+-{'-' * diff_width}-+-{'-' * result_width}-+"
if verbose or not all_passed:
print(separator)
print(f" | {'Output'.ljust(name_width)} | {'Max diff'.ljust(diff_width)} | {'Result'} |")
print(separator)
for name, diff_text, passed in rows:
status_text = ("PASS" if passed else "FAIL").ljust(result_width)
status = Fore.GREEN + status_text + Style.RESET_ALL if passed else Fore.RED + status_text + Style.RESET_ALL
print(f" | {name.ljust(name_width)} | {diff_text.ljust(diff_width)} | {status} |")
print(separator)
return all_passed
def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
simulator_dir, crossbar_size, crossbar_count, core_count,
raptor_extra_args,
pimsim_nn_build_dir, pimsim_config_path,
threshold, rtol,
seed, reporter, model_index, model_total, verbose,
command_timeout_seconds, mode):
network_onnx_path = Path(network_onnx_path).resolve()
raptor_path = Path(raptor_path).resolve()
onnx_include_dir = Path(onnx_include_dir).resolve()
simulator_dir = Path(simulator_dir).resolve()
pimsim_enabled = pimsim_nn_build_dir is not None and pimsim_config_path is not None
if pimsim_enabled:
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)
workspace_dir = network_onnx_path.parent
raptor_dir = workspace_dir / "raptor"
runner_dir = workspace_dir / "runner"
runner_build_dir = runner_dir / "build"
if mode != MODE_RUN_ONLY:
clean_workspace_artifacts(workspace_dir, network_onnx_path.stem)
Path.mkdir(raptor_dir, exist_ok=True)
Path.mkdir(runner_build_dir, parents=True, exist_ok=True)
reporter.log(Fore.CYAN + f"[{model_index}/{model_total}]" + Style.RESET_ALL +
f" {Style.BRIGHT}Validating {network_onnx_path.name}{Style.RESET_ALL}")
failed_with_exception = False
pim_pass_timings = {}
try:
stem = network_onnx_path.stem
network_so_path = runner_dir / f"{stem}.so"
network_mlir_path = raptor_dir / f"{stem}.onnx.mlir"
runner_path = runner_build_dir / "runner"
pim_output_base = raptor_dir / stem
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(
network_onnx_path, raptor_path, raptor_dir, runner_dir, reporter=reporter,
timeout_sec=command_timeout_seconds)
print_info(reporter, f"MLIR saved to {network_mlir_path}")
print_info(reporter, f"Shared library saved to {network_so_path}")
reporter.advance()
print_stage(reporter, model_index, model_total, network_onnx_path.name, "Build Runner")
gen_network_runner(
network_onnx_path,
network_so_path,
onnx_include_dir,
entry="run_main_graph",
out=runner_dir / "runner.c",
verbose=False,
)
runner_path = build_onnx_runner(runner_dir, runner_build_dir, reporter=reporter,
timeout_sec=command_timeout_seconds)
print_info(reporter, f"Runner built at {runner_path}")
reporter.advance()
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)
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)
if mode == MODE_RUN_ONLY:
required_paths = [
(network_so_path, "compiled reference shared library"),
(network_mlir_path, "exported ONNX MLIR"),
(runner_path, "built reference runner"),
(raptor_dir / "pim" / "config.json", "compiled PIM artifacts"),
]
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))
print_stage(reporter, model_index, model_total, network_onnx_path.name, "Generate Inputs")
inputs_descriptor, outputs_descriptor = onnx_io(network_onnx_path)
inputs_list, _inputs_dict = gen_random_inputs(inputs_descriptor, seed=seed)
flags, _files = save_inputs_to_files(network_onnx_path, inputs_list, out_dir=workspace_dir / "inputs")
print_info(reporter, f"Saved {len(inputs_list)} input file(s) to {workspace_dir / 'inputs'}")
reporter.advance()
print_stage(reporter, model_index, model_total, network_onnx_path.name, "Run Reference")
out_dir = workspace_dir / "outputs"
Path.mkdir(out_dir, exist_ok=True)
run_cmd = [runner_path, *flags]
run_cmd += ["--save-csv-dir", f"{out_dir}"]
run_command(run_cmd, cwd=runner_build_dir, reporter=reporter, timeout_sec=command_timeout_seconds)
print_info(reporter, f"Reference outputs saved to {out_dir}")
reporter.advance()
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)
print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
reporter.advance()
print_stage(
reporter, model_index, model_total, network_onnx_path.name,
"Run Functional Simulation")
pim_dir = raptor_dir / "pim"
write_inputs_to_memory_bin(pim_dir / "memory.bin", pim_dir / "config.json", inputs_list)
simulation_dir = workspace_dir / "simulation"
Path.mkdir(simulation_dir, exist_ok=True)
dump_ranges = build_dump_ranges(pim_dir / "config.json", outputs_descriptor)
output_bin_path = simulation_dir / "out.bin"
run_pim_simulator(simulator_dir, pim_dir, output_bin_path, dump_ranges, reporter=reporter,
timeout_sec=command_timeout_seconds)
print_info(reporter, f"Functional simulation output saved to {output_bin_path}")
reporter.advance()
print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compare Outputs")
sim_arrays = parse_pim_simulator_outputs(output_bin_path, outputs_descriptor)
reporter.suspend()
passed = validate_outputs(sim_arrays, out_dir, outputs_descriptor, threshold, rtol=rtol, verbose=verbose)
reporter.resume()
reporter.advance()
print_stage(
reporter, model_index, model_total, network_onnx_path.name,
"Run Non-functional Simulation")
pimsim_latency_ms = None
pimsim_power_mw = None
pimsim_status = PIMSIM_SKIPPED
if pimsim_enabled:
try:
pimsim_latency_ms, pimsim_power_mw = run_pimsim_nn(
pimsim_nn_build_dir,
pim_dir,
pimsim_config_path,
reporter=reporter,
timeout_sec=command_timeout_seconds,
)
pimsim_status = PIMSIM_DONE
print_info(
reporter,
f"Latency: {pimsim_latency_ms:.6f} ms, "
f"Power: {pimsim_power_mw:.6f} mW")
except Exception as exc:
pimsim_status = PIMSIM_FAILED
reporter.suspend()
print(
Fore.RED
+ f"pimsim-nn non-functional simulation failed: {type(exc).__name__}: {exc}"
+ Style.RESET_ALL,
file=sys.stderr,
flush=True,
)
reporter.resume()
else:
print_info(reporter, "pimsim-nn non-functional simulation skipped")
reporter.advance()
passed = passed and pimsim_status != PIMSIM_FAILED
reporter.record_result(passed)
status = Fore.GREEN + "PASS" + Style.RESET_ALL if passed else Fore.RED + "FAIL" + Style.RESET_ALL
reporter.log(Style.BRIGHT + f"Result: {status}" + Style.RESET_ALL)
return ValidationResult(
passed=passed,
pim_pass_timings=pim_pass_timings,
pimsim_latency_ms=pimsim_latency_ms,
pimsim_power_mw=pimsim_power_mw,
pimsim_status=pimsim_status,
)
except Exception:
failed_with_exception = True
reporter.record_result(False)
reporter.log(Style.BRIGHT + Fore.RED + "Result: FAIL" + Style.RESET_ALL)
reporter.suspend()
raise
finally:
if not failed_with_exception:
reporter.log("=" * 72)
if owns_reporter:
reporter.finish()