| |
| """The TT device dataflow of FRNet (PORT_LOG.md "trace layout"), run on the CPU in fp32 and proven equal to the |
| literal graph (``reference/model.py``) by ``tests/test_reference_rewrites_host.py``. Each rewrite is exact in real |
| arithmetic; in float32 they differ only by summation order: |
| |
| - **points sorted by cell and padded** to ``N_cap`` (``host/preprocess.py`` ``FramePlan``): the network is |
| permutation-equivariant in the points; padded rows form one trailing sentinel segment whose results nothing reads; |
| - **host features**: the 8-d decoration and ``pre_norm`` in float64, one rounding (``decorate_points``); |
| - **segmented max** (6 sites: ``ScatterElements(max)``) as the exact **log-step** form of PLAN.md 1.3 (the K1 |
| kernel's stock-op fallback): round ``r`` takes ``x = max(x, x[idx_r])`` with the host shift table ``idx_r``, |
| then one gather of each segment's last row (``seg_last``) -> ``(M_cap, C)``; :func:`segmax_logstep`; |
| - **frustum2pixel** (6 ``ScatterND``) as a gather: ``image[pixel] = table[pix2vox_s[pixel]]`` from the voxel table |
| with one appended **zero row** (index ``M_cap``) for empty pixels -- ``ttaw.ops.gather.scatter_rows`` semantics |
| (C19); duplicates at s > 1 resolved on the host (last voxel wins); |
| - **pixel2point** as a gather of NHWC rows (``point2pix_s``); |
| - **concat-free convs**: ``conv(concat(a, b))`` = ``conv_a(a) + conv_b(b)`` for the fusion convs (256 = 128 + 128) |
| and the 5-way K-split ``fuse_layer1`` (640 = 5 x 128, probe P14: 8.3 vs 14.2 ms); |
| - **concat-free point layers**: ``linear(concat(a, b))`` = ``a @ W_a + b @ W_b + bias`` (point stem, point fusion, |
| ``point_fuse_layer1``), which also lets the two halves keep different dtypes; |
| - **align-corners bilinear** as interpolation matrices ``A_h X A_w^T`` (C18 ``interp_matrix``, float64 -> float32). |
| |
| :class:`Emulation` optionally rounds layer operands the way the device stores them (bf16, or the TF32-like |
| mantissa of the device's fp32 matmul, probe P12) -- a first-order CPU model to size the fp32 islands of the |
| port before the device runs (``code/scripts/precision_study.py``); it models no math-fidelity phases or tile |
| layouts, so its numbers are indicative only. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
| from typing import Dict, List, Optional, Sequence |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| from ..host.preprocess import FramePlan, segmax_rounds, segmax_shift_tables |
| from ..ttaw.golden import NULL_TAPS, TapRegistry |
| from ..ttaw.ops.upsample import interp_matrix |
| from .model import TorchLayers, _t |
| from .weights import Conv, FRNetWeights, Linear |
|
|
| __all__ = ["Emulation", "round_mode", "segmax_logstep", "FRNetPlanTorch", "plan_taps_to_literal"] |
|
|
| MODES = ("fp32", "bf16", "tf32") |
|
|
|
|
| def round_mode(x: torch.Tensor, mode: str) -> torch.Tensor: |
| """``fp32``: unchanged; ``bf16``: round to nearest even bfloat16; ``tf32``: truncate to 10 mantissa bits (the |
| probe P12 truncation model of the device's fp32 matmul operands; the device measured 1.4-1.75x its error).""" |
| if mode == "fp32": |
| return x |
| if mode == "bf16": |
| return x.to(torch.bfloat16).to(torch.float32) |
| if mode == "tf32": |
| return (x.contiguous().view(torch.int32) & ~0x1FFF).view(torch.float32) |
| raise ValueError(f"unknown rounding mode {mode!r}; one of {MODES}") |
|
|
|
|
| @dataclass(frozen=True) |
| class Emulation: |
| """Operand rounding per site (``fp32`` = exact float32). ``conv_act`` / ``conv_weight``: inputs and weights |
| (+ bias) of the 53 convs; ``point_act`` / ``point_weight``: inputs and weights of the 17 point / voxel linears; |
| ``gather``: rows produced by the embedding gathers (frustum2pixel, pixel2point: ``ttnn.embedding`` tables are |
| bf16-only, so ``bf16`` is the stock-op device); ``upsample``: the decoder's interpolation operands. ``overrides``: |
| per-layer modes by layer name (e.g. ``{"head.logits": ("tf32", "fp32")}``: (act, weight)).""" |
|
|
| conv_act: str = "fp32" |
| conv_weight: str = "fp32" |
| point_act: str = "fp32" |
| point_weight: str = "fp32" |
| gather: str = "fp32" |
| upsample: str = "fp32" |
| overrides: Dict[str, tuple] = field(default_factory=dict) |
|
|
| @classmethod |
| def parse(cls, spec: str) -> "Emulation": |
| """``"conv_act=bf16,point_act=tf32,gather=bf16"`` (unset sites stay fp32).""" |
| kw = {} |
| for part in filter(None, (p.strip() for p in spec.split(","))): |
| k, _, v = part.partition("=") |
| if k not in cls.__dataclass_fields__ or k == "overrides" or v not in MODES: |
| raise ValueError(f"bad emulation item {part!r}") |
| kw[k] = v |
| return cls(**kw) |
|
|
| @property |
| def exact(self) -> bool: |
| return all(getattr(self, k) == "fp32" for k in ("conv_act", "conv_weight", "point_act", "point_weight", |
| "gather", "upsample")) and not self.overrides |
|
|
| def label(self) -> str: |
| items = [f"{k}={getattr(self, k)}" for k in ("conv_act", "conv_weight", "point_act", "point_weight", |
| "gather", "upsample") if getattr(self, k) != "fp32"] |
| items += [f"{k}:{v}" for k, v in sorted(self.overrides.items())] |
| return ",".join(items) or "fp32" |
|
|
|
|
| def segmax_logstep(x: torch.Tensor, plan: FramePlan, rounds: Optional[int] = None) -> torch.Tensor: |
| """Per-voxel max of the sorted, padded rows ``x`` ``(N_cap, C)`` -> ``(M_cap, C)``: ``rounds`` (default: just |
| enough for the frame's longest segment) log-step rounds over the host shift tables, then a gather of every |
| segment's last row (PLAN.md 1.3). Exact (max is order-independent).""" |
| rounds = segmax_rounds(plan.max_segment_length) if rounds is None else int(rounds) |
| for idx in segmax_shift_tables(plan.seg_id, rounds): |
| x = torch.maximum(x, x[_t(idx.astype(np.int64))]) |
| return x[_t(plan.seg_last.astype(np.int64))] |
|
|
|
|
| def _with_zero_row(table: torch.Tensor) -> torch.Tensor: |
| return torch.cat([table, table.new_zeros(1, table.shape[1])], 0) |
|
|
|
|
| def _nhwc(x: torch.Tensor) -> torch.Tensor: |
| _, c, h, w = x.shape |
| return x.permute(0, 2, 3, 1).reshape(h * w, c) |
|
|
|
|
| def _nchw(rows: torch.Tensor, h: int, w: int) -> torch.Tensor: |
| return rows.reshape(1, h, w, -1).permute(0, 3, 1, 2).contiguous() |
|
|
|
|
| class FRNetPlanTorch: |
| """``FRNetPlanTorch(weights)(plan, taps)`` -> ``pred_probs`` of the sorted, padded rows ``(N_cap, 27)`` |
| (``plan.unsort`` gives the network-input order). Taps use the names of ``reference/model.py``; point taps are in |
| sorted padded order, voxel taps have ``M_cap`` rows (:func:`plan_taps_to_literal` aligns them).""" |
|
|
| def __init__(self, weights: FRNetWeights, *, emulation: Optional[Emulation] = None, |
| rounds: Optional[int] = None): |
| self.weights = weights |
| self.layers = TorchLayers(weights) |
| self.emu = emulation or Emulation() |
| self.rounds = rounds |
| H, W = weights.frustum_hw |
| self._interp = {} |
| for st in [1] + [s.stride for s in weights.stages]: |
| ah = interp_matrix(H // st, H, mode="linear", align_corners=True) |
| aw = interp_matrix(W // st, W, mode="linear", align_corners=True) |
| self._interp[st] = (_t(ah.astype(np.float32)), _t(aw.astype(np.float32))) |
| |
| self._split_lin = {} |
| for lin, parts in ((weights.point_stem, (128, 256)), (weights.point_fuse[0], (128,) * 5), |
| *((st.point_fusion, (128, 128)) for st in weights.stages)): |
| w_t, b_t, _ = self.layers.lin[lin.name] |
| if w_t.shape[0] != sum(parts): |
| raise ValueError(f"{lin.name}: {w_t.shape[0]} inputs, expected {sum(parts)}") |
| self._split_lin[lin.name] = (list(torch.split(w_t, list(parts), 0)), b_t) |
| self._split_conv = {} |
| for conv, parts in ((weights.fusion_stem, (128, 128)), (weights.fuse_layer1, (128,) * 5), |
| *((st.pixel_fusion, (128, 128)) for st in weights.stages)): |
| w_t, b_t, stride, pad, act = self.layers.conv[conv.name] |
| if w_t.shape[1] != sum(parts): |
| raise ValueError(f"{conv.name}: {w_t.shape[1]} inputs, expected {sum(parts)}") |
| self._split_conv[conv.name] = (list(torch.split(w_t, list(parts), 1)), b_t, stride, pad, act) |
|
|
| |
| def _modes(self, name: str, act: str, weight: str) -> tuple: |
| return self.emu.overrides.get(name, (act, weight)) |
|
|
| def linear(self, parts: Sequence[torch.Tensor], layer: Linear) -> torch.Tensor: |
| act_m, w_m = self._modes(layer.name, self.emu.point_act, self.emu.point_weight) |
| if layer.name in self._split_lin: |
| ws, b = self._split_lin[layer.name] |
| else: |
| w, b, _ = self.layers.lin[layer.name] |
| ws = [w] |
| if len(parts) != len(ws): |
| raise ValueError(f"{layer.name}: {len(parts)} input parts for {len(ws)} weight parts") |
| y = round_mode(b, w_m).expand(parts[0].shape[0], -1).clone() |
| for x, w in zip(parts, ws): |
| y = y + round_mode(x, act_m) @ round_mode(w, w_m) |
| return F.relu(y) if layer.relu else y |
|
|
| def conv(self, parts: Sequence[torch.Tensor], layer: Conv) -> torch.Tensor: |
| act_m, w_m = self._modes(layer.name, self.emu.conv_act, self.emu.conv_weight) |
| if layer.name in self._split_conv: |
| ws, b, stride, pad, act = self._split_conv[layer.name] |
| else: |
| w, b, stride, pad, act = self.layers.conv[layer.name] |
| ws = [w] |
| if len(parts) != len(ws): |
| raise ValueError(f"{layer.name}: {len(parts)} input parts for {len(ws)} weight parts") |
| y = None |
| for i, (x, w) in enumerate(zip(parts, ws)): |
| part = F.conv2d(round_mode(x, act_m), round_mode(w, w_m), round_mode(b, w_m) if i == 0 else None, |
| stride=stride, padding=pad) |
| y = part if y is None else y + part |
| if act == "hardswish": |
| y = F.hardswish(y) |
| elif act == "sigmoid": |
| y = torch.sigmoid(y) |
| return y |
|
|
| def gather(self, table: torch.Tensor, index: np.ndarray) -> torch.Tensor: |
| return round_mode(table, self.emu.gather)[_t(index.astype(np.int64))] |
|
|
| def upsample(self, x: torch.Tensor, stride: int) -> torch.Tensor: |
| if stride == 1: |
| return x |
| ah, aw = self._interp[stride] |
| m = self.emu.upsample |
| x = round_mode(x, m) |
| return torch.einsum("hi,bcij,wj->bchw", round_mode(ah, m), x, round_mode(aw, m)) |
|
|
| |
| def __call__(self, plan: FramePlan, taps: TapRegistry = NULL_TAPS) -> np.ndarray: |
| return self.forward(plan, taps) |
|
|
| @torch.no_grad() |
| def forward(self, plan: FramePlan, taps: TapRegistry = NULL_TAPS) -> np.ndarray: |
| Wt = self.weights |
| H, W = plan.frustum_hw |
| if (H, W) != Wt.frustum_hw: |
| raise ValueError(f"plan grid {plan.frustum_hw} != weights grid {Wt.frustum_hw}") |
|
|
| def tap(name: str, value: torch.Tensor) -> torch.Tensor: |
| if taps.wants(name): |
| taps.tap(name, value.numpy()) |
| return value |
|
|
| def to_image(table: torch.Tensor, s: int) -> torch.Tensor: |
| return _nchw(self.gather(_with_zero_row(table), plan.pix2vox[s]), H // s, W // s) |
|
|
| def to_points(img: torch.Tensor, s: int) -> torch.Tensor: |
| return self.gather(_nhwc(img), plan.point2pix[s]) |
|
|
| x = tap("encoder.input", _t(plan.features)) |
| pf: List[torch.Tensor] = [] |
| for i, lin in enumerate(Wt.ffe): |
| x = self.linear([x], lin) |
| pf.append(tap(f"encoder.pf{i}", x)) |
| vmax = tap("encoder.voxel_max", segmax_logstep(x, plan, self.rounds)) |
| vfeat = tap("encoder.voxel_feats", self.linear([vmax], Wt.compression)) |
| img = tap("backbone.frustum", to_image(vfeat, 1)) |
| for c in Wt.stem: |
| img = self.conv([img], c) |
| tap("backbone.stem", img) |
| p = tap("backbone.point_stem", self.linear([to_points(img, 1), pf[-1]], Wt.point_stem)) |
| seg = tap("backbone.segmax0", segmax_logstep(p, plan, self.rounds)) |
| x = tap("backbone.fusion_stem", self.conv([to_image(seg, 1), img], Wt.fusion_stem)) |
| maps, pts_out = [x], [p] |
| for st in Wt.stages: |
| li, s = st.index, st.stride |
| for bi, (c1, c2, ds) in enumerate(st.blocks): |
| y = self.conv([self.conv([x], c1)], c2) |
| idt = self.conv([x], ds) if ds is not None else x |
| x = tap(f"backbone.layer{li}.block{bi}", F.hardswish(y + idt)) |
| tap(f"backbone.layer{li}", x) |
| p = tap(f"backbone.point_fusion{li}", self.linear([to_points(x, s), p], st.point_fusion)) |
| seg = tap(f"backbone.segmax{li}", segmax_logstep(p, plan, self.rounds)) |
| fuse = tap(f"backbone.pixel_fusion{li}", self.conv([to_image(seg, s), x], st.pixel_fusion)) |
| att = tap(f"backbone.attention{li}", self.conv([self.conv([fuse], st.attention[0])], st.attention[1])) |
| x = tap(f"backbone.stage{li}", fuse * att + x) |
| maps.append(x) |
| pts_out.append(p) |
| strides = [1] + [st.stride for st in Wt.stages] |
| ups = [tap(f"decoder.up{k}", self.upsample(mp, s)) for k, (mp, s) in enumerate(zip(maps, strides))] |
| y = tap("decoder.fuse1", self.conv(ups, Wt.fuse_layer1)) |
| y = tap("decoder.fuse2", self.conv([y], Wt.fuse_layer2)) |
| q = tap("decoder.point_fuse1", self.linear(pts_out, Wt.point_fuse[0])) |
| q = tap("decoder.point_fuse2", self.linear([q], Wt.point_fuse[1])) |
| h = to_points(y, 1) |
| for i, (lin, skip) in enumerate(zip(Wt.head_mlps, (q, pf[2], pf[1], pf[0]))): |
| h = tap(f"head.mlp{i}", self.linear([h], lin) + skip) |
| logits = tap("head.logits", self.linear([h], Wt.conv_seg)) |
| return tap("pred_probs", torch.softmax(logits, -1)).numpy() |
|
|
|
|
| def plan_taps_to_literal(name: str, value: np.ndarray, plan: FramePlan) -> np.ndarray: |
| """A tap of :class:`FRNetPlanTorch` in the layout of ``reference/model.py``: point tensors (``N_cap`` rows) |
| unsorted to the network-input order, voxel tensors (``M_cap`` rows) cut to ``M``, maps unchanged.""" |
| v = np.asarray(value) |
| if v.ndim == 2 and v.shape[0] == plan.n_cap: |
| return plan.unsort(v) |
| if v.ndim == 2 and v.shape[0] == plan.m_cap: |
| return v[: plan.m] |
| return v |
|
|
