#!/usr/bin/env python3 """Save exact attention-tap arrays and quantify the MatMul error sources.""" import argparse import hashlib import json import sys from pathlib import Path import numpy as np import onnx from onnx import numpy_helper REPO_ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO_ROOT / "validation")) from raptor_validation.onnx_utils import onnx_io # noqa: E402 from raptor_validation.validate_one import ( # noqa: E402 parse_pim_simulator_outputs, sanitize_output_name, ) TAP_NAMES = { "v": "/model.10/m/m.0/attn/Split_output_2", "raw": "/model.10/m/m.0/attn/MatMul_output_0", "scaled": "/model.10/m/m.0/attn/Mul_output_0", "rhs": "/model.10/m/m.0/attn/Transpose_1_output_0", "c": "/model.10/m/m.0/attn/MatMul_1_output_0", } ABSOLUTE_TOLERANCE = 1e-3 RELATIVE_TOLERANCE = 1e-5 def sha256(path): digest = hashlib.sha256() with Path(path).open("rb") as stream: for block in iter(lambda: stream.read(1 << 20), b""): digest.update(block) return digest.hexdigest() def metric(actual, expected): difference = np.abs(actual.astype(np.float64) - expected.astype(np.float64)) allowed = ABSOLUTE_TOLERANCE + RELATIVE_TOLERANCE * np.abs(expected.astype(np.float64)) return { "max_abs": float(np.max(difference)), "mean_abs": float(np.mean(difference)), "rms": float(np.sqrt(np.mean(np.square(difference)))), "elements_over_validator_limit": int(np.count_nonzero(difference > allowed)), } def f32_matmul(lhs, rhs): return np.matmul(lhs.astype(np.float32), rhs.astype(np.float32)).astype(np.float32) def f64_matmul(lhs, rhs): return np.matmul(lhs.astype(np.float64), rhs.astype(np.float64)).astype(np.float64) def load_constant(model, output_name): for initializer in model.graph.initializer: if initializer.name == output_name: return float(numpy_helper.to_array(initializer).reshape(-1)[0]) for node in model.graph.node: if output_name not in node.output: continue for attribute in node.attribute: if attribute.name == "value" and attribute.HasField("t"): return float(numpy_helper.to_array(attribute.t).reshape(-1)[0]) raise ValueError(f"could not find ONNX Constant producing {output_name}") def load_arrays(workspace, model_path): model = onnx.load(model_path) descriptors = onnx_io(model_path) output_descriptors = {name: (index, dtype, shape) for index, name, dtype, shape in descriptors[1]} missing = sorted(set(TAP_NAMES.values()) - set(output_descriptors)) if missing: raise ValueError("tap model is missing outputs: " + ", ".join(missing)) sim_arrays = parse_pim_simulator_outputs( workspace / "simulation" / "out.bin", descriptors[1] ) reference = {} simulated = {} input_files = {} for key, name in TAP_NAMES.items(): index, _dtype, shape = output_descriptors[name] csv_path = workspace / "outputs" / f"output{index}_{sanitize_output_name(name)}.csv" reference[key] = np.loadtxt(csv_path, delimiter=",", dtype=np.float32).reshape(shape) simulated[key] = np.asarray(sim_arrays[index], dtype=np.float32).reshape(shape) input_files[key] = csv_path return reference, simulated, input_files def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--workspace", type=Path, required=True, help="validator workspace containing inputs, outputs, and simulation") parser.add_argument("--model", type=Path, required=True, help="five-output ONNX tap model") parser.add_argument("--output-dir", type=Path, required=True, help="directory for arrays.npz, metadata.json, and decomposition.json") args = parser.parse_args() args.output_dir.mkdir(parents=True, exist_ok=True) model = onnx.load(args.model) reference, simulated, source_files = load_arrays(args.workspace, args.model) scale = np.float32(load_constant(model, "/model.10/m/m.0/attn/Constant_1_output_0")) ref_v, ref_raw, ref_scaled, ref_rhs, ref_c = (reference[key] for key in ("v", "raw", "scaled", "rhs", "c")) sim_v, sim_raw, sim_scaled, sim_rhs, sim_c = (simulated[key] for key in ("v", "raw", "scaled", "rhs", "c")) ref_score_transpose = np.swapaxes(ref_scaled, -1, -2) sim_score_transpose = np.swapaxes(sim_scaled, -1, -2) ref_ss_f32 = f32_matmul(ref_v, ref_rhs) sim_ss_f32 = f32_matmul(sim_v, sim_rhs) ref_ss_f64 = f64_matmul(ref_v, ref_rhs) sim_ss_f64 = f64_matmul(sim_v, sim_rhs) ref_split_f32 = (f32_matmul(ref_v, np.swapaxes(ref_raw, -1, -2)) * scale).astype(np.float32) sim_split_f32 = (f32_matmul(sim_v, np.swapaxes(sim_raw, -1, -2)) * scale).astype(np.float32) ref_split_f64 = f64_matmul(ref_v, np.swapaxes(ref_raw, -1, -2)) * np.float64(scale) sim_split_f64 = f64_matmul(sim_v, np.swapaxes(sim_raw, -1, -2)) * np.float64(scale) arrays = { **{f"ref_{key}": value for key, value in reference.items()}, **{f"sim_{key}": value for key, value in simulated.items()}, "ref_ss_f32": ref_ss_f32, "sim_ss_f32": sim_ss_f32, "ref_ss_f64": ref_ss_f64, "sim_ss_f64": sim_ss_f64, "ref_split_f32": ref_split_f32, "sim_split_f32": sim_split_f32, "ref_split_f64": ref_split_f64, "sim_split_f64": sim_split_f64, } arrays_path = args.output_dir / "arrays.npz" np.savez_compressed(arrays_path, **arrays) metrics = { "validator_policy": { "absolute_tolerance": ABSOLUTE_TOLERANCE, "relative_tolerance": RELATIVE_TOLERANCE, }, "scale": float(scale), "shape": list(ref_c.shape), "tap_differences": {key: metric(simulated[key], reference[key]) for key in TAP_NAMES}, "rhs_transpose_consistency": metric(ref_rhs, ref_score_transpose), "sim_rhs_transpose_consistency": metric(sim_rhs, sim_score_transpose), "c_sim_vs_ss_f32": metric(sim_c, ref_ss_f32), "c_ref_vs_ss_f32": metric(ref_c, ref_ss_f32), "c_sim_vs_simulated_inputs_ss_f32": metric(sim_c, sim_ss_f32), "v_drift_only": metric(f32_matmul(sim_v, ref_rhs), ref_ss_f32), "rhs_drift_only": metric(f32_matmul(ref_v, sim_rhs), ref_ss_f32), "joint_input_drift": metric(sim_ss_f32, ref_ss_f32), "scale_reassociation_reference": metric(ref_split_f32, ref_ss_f32), "scale_reassociation_simulated": metric(sim_split_f32, sim_ss_f32), "reference_accumulation_f32_vs_f64": metric(ref_ss_f32, ref_ss_f64), "simulated_accumulation_f32_vs_f64": metric(sim_ss_f32, sim_ss_f64), "split_accumulation_reference_f32_vs_f64": metric(ref_split_f32, ref_split_f64), "split_accumulation_simulated_f32_vs_f64": metric(sim_split_f32, sim_split_f64), } decomposition_path = args.output_dir / "decomposition.json" decomposition_path.write_text(json.dumps(metrics, indent=2) + "\n", encoding="utf-8") metadata = { "model": str(args.model), "model_sha256": sha256(args.model), "workspace": str(args.workspace), "arrays_sha256": sha256(arrays_path), "source_sha256": {key: sha256(path) for key, path in source_files.items()}, "simulator_output_sha256": sha256(args.workspace / "simulation" / "out.bin"), "input_sha256": sha256(args.workspace / "inputs" / "in0.csv"), "outputs": TAP_NAMES, "arrays": {key: {"dtype": str(value.dtype), "shape": list(value.shape)} for key, value in arrays.items()}, } (args.output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8") print(json.dumps(metrics, indent=2)) if __name__ == "__main__": main()