code/scripts/precision_study.py

#!/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())