#!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 """First-order CPU emulation of the TT numerics of FRNet, to size the fp32 islands of the port before the device runs (research venv: torch + onnx; CPU only, no device). cd bundles/frnet-p150 && OMP_NUM_THREADS=4 PYTHONPATH=code \\ /home/ubuntu/experiments/tt-models/tools/research-venv/bin/python code/scripts/precision_study.py \\ --out /home/ubuntu/experiments/tt-models/logs/frnet/m0/precision_study.json The TT dataflow (``reference/rewrites.py`` ``FRNetPlanTorch``) runs with operand rounding per site (``Emulation``): ``bf16`` (RNE), ``tf32`` (truncation to 10 mantissa bits: probe P12's model of the device's fp32 matmul, whose measured error is 1.4-1.75x that model's, so ``tf32`` rows are optimistic). Configurations: - ``bf16_all``: bf16 everywhere (weights and activations): the research simulation's worst case; - ``plan``: the planned functional numerics -- conv activations bf16 with fp32 weights (PLAN.md 9.3 precision policy), the point / voxel layers as an fp32 island (device fp32 matmul ~ tf32), bf16 rows out of the ``ttnn.embedding`` gathers (bf16-only tables), fp32 interpolation matmuls (~ tf32); - ``plan_bf16_points``: ``plan`` with bf16 activations and fp32 weights in the point layers (the fast path); - ``plan_fp32_gather``: ``plan`` with exact fp32 gathers (what a fp32 row-gather kernel, K3, would buy); - ``plan_fp32_convs`` / ``plan_exact_points``: ``plan`` with fp32 conv activations, resp. exact fp32 point layers (e.g. two-term bf16 operands): which island to deepen first when a frame misses a gate. Metrics per frame against the stored ONNX Runtime ``pred_probs`` (the reference of record): PCC over all points and argmax agreement on the raw points (the PLAN.md 2.8 gates: >= 0.998 and >= 98.5 %), and the module PCC of the device modules' outputs against the exact fp32 dataflow (gate >= 0.999). Indicative only: no math-fidelity phases, no tile layouts, no on-device accumulation order. """ from __future__ import annotations import argparse import json import sys import time from pathlib import Path import numpy as np HERE = Path(__file__).resolve() sys.path.insert(0, str(HERE.parents[1])) WS = Path("/home/ubuntu/experiments/tt-models") CONFIGS = { "bf16_all": "conv_act=bf16,conv_weight=bf16,point_act=bf16,point_weight=bf16,gather=bf16,upsample=bf16", "plan": "conv_act=bf16,point_act=tf32,point_weight=tf32,gather=bf16,upsample=tf32", "plan_bf16_points": "conv_act=bf16,point_act=bf16,gather=bf16,upsample=tf32", "plan_fp32_gather": "conv_act=bf16,point_act=tf32,point_weight=tf32,upsample=tf32", "plan_fp32_convs": "point_act=tf32,point_weight=tf32,gather=bf16,upsample=tf32", "plan_exact_points": "conv_act=bf16,gather=bf16,upsample=tf32", # M1b device follow-up (PORT_LOG "M1 device results"): the device's float32 conv-weight path is not exact # (C17 / YOLOX: TF32-like, close to bf16 weights); the point island exact (two-term operands) "exact_points_tf32_conv_w": "conv_act=bf16,conv_weight=tf32,gather=bf16,upsample=tf32", "exact_points_bf16_conv_w": "conv_act=bf16,conv_weight=bf16,gather=bf16,upsample=tf32", } FRAMES = ["pandaset_019_40.ot128", "pandaset_029_40.ot128", "nuscenes_scene-0103_20.ot128", "nuscenes_scene-0916_20.ot128", "vls128_frame000_fronthalf.ot128", "pandaset_019_40.qt128", "vls128_frame000_fronthalf.qt128"] MODULES = ("encoder.pf3", "encoder.voxel_feats", "backbone.stem", "backbone.point_stem", "backbone.fusion_stem", "backbone.stage1", "backbone.stage2", "backbone.stage3", "backbone.stage4", "decoder.fuse2", "decoder.point_fuse2", "head.logits") def main() -> int: ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) ap.add_argument("--out", type=Path, required=True) ap.add_argument("--frames", nargs="*", default=FRAMES) ap.add_argument("--configs", nargs="*", default=list(CONFIGS)[:4]) ap.add_argument("--threads", type=int, default=4) ap.add_argument("--weights-dir", default=str(WS / "assets" / "frnet" / "hf_lidar_frnet")) a = ap.parse_args() import torch from tt_frnet.host.preprocess import build_plan from tt_frnet.reference.config import FRNetConfig from tt_frnet.reference.rewrites import Emulation, FRNetPlanTorch from tt_frnet.reference.weights import load_weights from tt_frnet.tests import _data from tt_frnet.ttaw.golden import TapRegistry from tt_frnet.ttaw.metrics import pcc torch.set_num_threads(a.threads) weights, configs = {}, {} result = {"configs": {k: CONFIGS[k] for k in a.configs}, "frames": {}, "note": __doc__.split("\n\n")[2]} for frame in a.frames: variant = frame.rpartition(".")[2] if variant not in weights: weights[variant] = load_weights(a.weights_dir, variant) configs[variant] = FRNetConfig.from_model_dir(a.weights_dir, variant) w = weights[variant] inputs, ort_probs, _ = _data.load_frame(frame) plan = build_plan(inputs, configs[variant], pre_norm=(w.pre_norm_scale, w.pre_norm_shift)) n_raw = inputs.num_points_raw ref_label = ort_probs.argmax(1) t0 = time.time() exact = TapRegistry(include=list(MODULES)) FRNetPlanTorch(w)(plan, taps=exact) rows = {} for name in a.configs: taps = TapRegistry(include=list(MODULES)) probs = plan.unsort(FRNetPlanTorch(w, emulation=Emulation.parse(CONFIGS[name]))(plan, taps=taps)) label = probs.argmax(1) rows[name] = {"pred_probs_pcc": float(pcc(probs, ort_probs)), "argmax_agreement_raw": float((label[:n_raw] == ref_label[:n_raw]).mean()), "argmax_agreement_all": float((label == ref_label).mean()), "modules": {m: float(pcc(taps[m][: plan.n] if taps[m].shape[0] == plan.n_cap else taps[m], exact[m][: plan.n] if exact[m].shape[0] == plan.n_cap else exact[m])) for m in MODULES}} r = rows[name] worst = min(r["modules"].items(), key=lambda kv: kv[1]) print(f"{frame} {name}: pcc {r['pred_probs_pcc']:.6f} argmax(raw) {r['argmax_agreement_raw']:.4%} " f"worst module {worst[0]} {worst[1]:.6f}", flush=True) result["frames"][frame] = {"N": inputs.num_points, "N_raw": n_raw, "M": inputs.num_voxels, "seconds": round(time.time() - t0, 1), "results": rows} a.out.parent.mkdir(parents=True, exist_ok=True) a.out.write_text(json.dumps(result, indent=1) + "\n") return 0 if __name__ == "__main__": sys.exit(main())