315 lines
13 KiB
Python
315 lines
13 KiB
Python
import argparse
|
|
import json
|
|
import numpy as np
|
|
import subprocess
|
|
import shutil
|
|
import sys
|
|
from pathlib import Path
|
|
from colorama import Style, Fore
|
|
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 gen_network_runner import gen_network_runner
|
|
from subprocess_utils import run_command_with_reporter
|
|
|
|
|
|
STAGE_COUNT = 6
|
|
GENERATED_DIR_NAMES = ("inputs", "outputs", "raptor", "runner", "simulation")
|
|
|
|
|
|
class ProgressReporter:
|
|
def __init__(self, total_models, stages_per_model=STAGE_COUNT):
|
|
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.current_label = ""
|
|
self.enabled = True
|
|
self.columns = shutil.get_terminal_size((100, 20)).columns
|
|
self.suspended = False
|
|
|
|
def _clear(self):
|
|
if self.enabled:
|
|
sys.stdout.write("\033[2K\r")
|
|
|
|
def _render(self):
|
|
if not self.enabled or self.suspended:
|
|
return
|
|
bar_width = 24
|
|
filled = int(bar_width * self.completed_steps / self.total_steps)
|
|
prefix_text = f"[{'#' * filled}{'-' * (bar_width - filled)}] {self.completed_steps}/{self.total_steps}"
|
|
if len(prefix_text) > self.columns:
|
|
prefix_text = f"{self.completed_steps}/{self.total_steps}"
|
|
|
|
label = f" {self.current_label}" if self.current_label else ""
|
|
available_label_width = max(0, self.columns - len(prefix_text))
|
|
label = label[:available_label_width]
|
|
|
|
if prefix_text.startswith("["):
|
|
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
|
|
else:
|
|
prefix = Fore.CYAN + prefix_text + Style.RESET_ALL
|
|
|
|
sys.stdout.write("\r" + prefix + label + Style.RESET_ALL)
|
|
sys.stdout.flush()
|
|
|
|
def log(self, message="", color=None):
|
|
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 suspend(self):
|
|
self.suspended = True
|
|
self._clear()
|
|
sys.stdout.flush()
|
|
|
|
def resume(self):
|
|
self.suspended = False
|
|
|
|
def finish(self):
|
|
if self.enabled:
|
|
self.suspended = True
|
|
self._clear()
|
|
sys.stdout.flush()
|
|
|
|
|
|
def run_command(cmd, cwd=None, reporter=None):
|
|
run_command_with_reporter(cmd, cwd=cwd, reporter=reporter)
|
|
|
|
|
|
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 = {
|
|
"Compile ONNX": Fore.BLUE,
|
|
"Build Runner": Fore.MAGENTA,
|
|
"Generate Inputs": Fore.YELLOW,
|
|
"Run Reference": Fore.GREEN,
|
|
"Compile PIM": Fore.CYAN,
|
|
"Run Simulator": Fore.MAGENTA,
|
|
"Compare Outputs": Fore.YELLOW,
|
|
}
|
|
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):
|
|
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"], reporter=reporter)
|
|
run_command([raptor_path, network_onnx_path, "-o", runner_base], reporter=reporter)
|
|
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):
|
|
run_command(["cmake", source_dir], cwd=build_dir, reporter=reporter)
|
|
run_command(["cmake", "--build", ".", "-j"], cwd=build_dir, reporter=reporter)
|
|
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):
|
|
run_command(
|
|
["cargo", "run", "--release", "--package", "pim-simulator", "--bin", "pim-simulator", "--",
|
|
"-f", str(pim_dir), "-o", str(output_bin_path), "-d", dump_ranges],
|
|
cwd=simulator_dir,
|
|
reporter=reporter,
|
|
)
|
|
|
|
|
|
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=1e-3):
|
|
all_passed = True
|
|
rows = []
|
|
for sim_array, (oi, name, _, shape) in zip(sim_arrays, outputs_descriptor):
|
|
csv_name = f"output{oi}_{name}.csv"
|
|
runner_array = np.loadtxt(runner_out_dir / csv_name, delimiter=',', dtype=np.float32).reshape(shape)
|
|
max_diff = float(np.max(np.abs(sim_array.astype(np.float64) - runner_array.astype(np.float64))))
|
|
passed = max_diff <= threshold
|
|
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}-+"
|
|
|
|
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=64, crossbar_count=8, 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()
|
|
onnx_include_dir = Path(onnx_include_dir).resolve()
|
|
simulator_dir = Path(simulator_dir).resolve()
|
|
owns_reporter = reporter is None
|
|
reporter = reporter or ProgressReporter(model_total)
|
|
|
|
workspace_dir = network_onnx_path.parent
|
|
clean_workspace_artifacts(workspace_dir, network_onnx_path.stem)
|
|
raptor_dir = workspace_dir / "raptor"
|
|
runner_dir = workspace_dir / "runner"
|
|
runner_build_dir = runner_dir / "build"
|
|
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
|
|
|
|
try:
|
|
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)
|
|
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, out=runner_dir / "runner.c", verbose=False)
|
|
runner_path = build_onnx_runner(runner_dir, runner_build_dir, reporter=reporter)
|
|
print_info(reporter, f"Runner built at {runner_path}")
|
|
reporter.advance()
|
|
|
|
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)
|
|
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)
|
|
print_info(reporter, f"Reference outputs saved to {out_dir}")
|
|
reporter.advance()
|
|
|
|
print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compile PIM")
|
|
compile_with_raptor(
|
|
network_mlir_path, raptor_path, raptor_dir / network_onnx_path.stem,
|
|
crossbar_size, crossbar_count,
|
|
cwd=raptor_dir, reporter=reporter)
|
|
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 Simulator")
|
|
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)
|
|
print_info(reporter, f"Simulator 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)
|
|
reporter.resume()
|
|
reporter.advance()
|
|
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 passed
|
|
except Exception:
|
|
failed_with_exception = True
|
|
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()
|
|
|
|
|
|
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 else 1)
|