142 lines
6.1 KiB
Python
142 lines
6.1 KiB
Python
import argparse
|
|
import json
|
|
import numpy as np
|
|
import subprocess
|
|
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
|
|
|
|
|
|
def compile_onnx_network(network_onnx_path, raptor_path, raptor_dir, runner_dir):
|
|
subprocess.run([raptor_path, network_onnx_path, "--EmitONNXIR"], check=True)
|
|
subprocess.run([raptor_path, network_onnx_path], check=True)
|
|
parent = network_onnx_path.parent
|
|
stem = network_onnx_path.stem
|
|
so_path = parent / f"{stem}.so"
|
|
mlir_path = parent / f"{stem}.onnx.mlir"
|
|
tmp_path = parent / f"{stem}.tmp"
|
|
moved_so = runner_dir / so_path.name
|
|
moved_mlir = raptor_dir / mlir_path.name
|
|
so_path.rename(moved_so)
|
|
mlir_path.rename(moved_mlir)
|
|
tmp_path.unlink(missing_ok=True)
|
|
return moved_so, moved_mlir
|
|
|
|
|
|
def build_onnx_runner(source_dir, build_dir):
|
|
subprocess.run(["cmake", source_dir], cwd=build_dir, check=True)
|
|
subprocess.run(["cmake", "--build", ".", "-j"], cwd=build_dir, check=True)
|
|
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):
|
|
subprocess.run(
|
|
["cargo", "run", "--release", "--package", "pim-simulator", "--bin", "pim-simulator", "--",
|
|
"-f", str(pim_dir), "-o", str(output_bin_path), "-d", dump_ranges],
|
|
cwd=simulator_dir, check=True
|
|
)
|
|
|
|
|
|
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
|
|
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
|
|
status = Fore.GREEN + "[PASS]" if passed else Fore.RED + "[FAIL]"
|
|
print(f" {name}: max diff = {max_diff:.6e} {status}" + Style.RESET_ALL)
|
|
if not passed:
|
|
all_passed = False
|
|
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):
|
|
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()
|
|
|
|
workspace_dir = network_onnx_path.parent
|
|
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)
|
|
|
|
print(Style.BRIGHT + "\nCompiling the onnx network:" + Style.RESET_ALL)
|
|
network_so_path, network_mlir_path = compile_onnx_network(network_onnx_path, raptor_path, raptor_dir, runner_dir)
|
|
|
|
print(Style.BRIGHT + "\nGenerating and building the runner:" + Style.RESET_ALL)
|
|
gen_network_runner(network_onnx_path, network_so_path, onnx_include_dir, out=runner_dir / "runner.c")
|
|
runner_path = build_onnx_runner(runner_dir, runner_build_dir)
|
|
|
|
print(Style.BRIGHT + "\nGenerating random inputs:" + Style.RESET_ALL)
|
|
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(Style.BRIGHT + "\nRunning inference with the runner:" + Style.RESET_ALL)
|
|
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}"]
|
|
subprocess.run(run_cmd, cwd=runner_build_dir, check=True)
|
|
|
|
print(Style.BRIGHT + "\nCompiling for PIM with Raptor:" + Style.RESET_ALL)
|
|
compile_with_raptor(network_mlir_path, raptor_path, crossbar_size, crossbar_count)
|
|
|
|
print(Style.BRIGHT + "\nRunning PIM simulation:" + Style.RESET_ALL)
|
|
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)
|
|
|
|
print(Style.BRIGHT + "\nValidating the results:" + Style.RESET_ALL)
|
|
sim_arrays = parse_pim_simulator_outputs(output_bin_path, outputs_descriptor)
|
|
return validate_outputs(sim_arrays, out_dir, outputs_descriptor, threshold)
|
|
|
|
|
|
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)
|