| |
| """C24 companion: the PillarFeatureNet (PFN) of the pillar detectors on the device, and the host staging of its |
| inputs (PLAN.md section 1.2; CenterPoint, PointPainting, TransFusion). Promoted by the PointPainting port from the |
| CenterPoint bundle's ``tt/pfn.py`` + ``tt/model.py FrameStaging`` (device-verified there: PFN PCC 0.99997) with the |
| TransFusion bundle's generalisation (any hidden / output width, pillars of fewer than 32 points); those bundles keep |
| their copies until they re-vendor. |
| |
| The encoder (mmdet3d ``PFNLayer`` x 2 as Autoware exports it; BN folded into both MatMuls by the port's loader): |
| |
| h = relu(x @ W0 + b0) x: (P, K, F) decorated points, padded slots all zero h: (P, K, H) |
| m = max over the K slots of h padded slots included: their rows are relu(b0) |
| y = max over the slots of relu([h, m] @ W1 + b1) = relu(max_slots(h @ W1a) + m @ W1b + b1) y: (P, O) |
| |
| Device form. One pillar is one 32-row tile (its point slots), so a slot max is a tile-column reduction (``ttnn.max`` |
| over dim -2). Because ``relu(. + c)`` is monotonic, the per-pillar term ``c = m @ W1b + b1`` is added AFTER the max |
| (exact, also in floating point: rounding is monotonic), so no slot broadcast is needed; two identity blocks fold the |
| rest into the matmuls (exact: a value times 1.0 with HiFi4 and fp32 accumulation): |
| |
| x [P*32, 16] --linear W0 (16 x H_pad) + b0, ReLU--> h [P*32, H_pad] |
| h --linear Wz = [I_H | W1a] (H_pad x Z)--> z [P*32, Z] = [h | h @ W1a | 0] |
| z as [P, 32, Z] --max over dim -2--> mz [P, 1, Z] = [m | max(h @ W1a) | 0] |
| mz --linear Wy = [W1b ; I_O] (Z x O) + b1, ReLU--> y [P, 1, O] |
| y --untilize, view--> [1, 1, P, O] ROW_MAJOR bf16 rows |
| |
| (``H_pad`` = H rounded up to 32, ``Z`` = H + O rounded up to 32, ``O % 32 == 0``: the rows are a C19 gather table.) |
| The rows feed the gather-form scatter (C19 ``ops.gather.scatter_rows``, which appends the zero sentinel row). |
| |
| **Pillars of K < 32 points** (TransFusion: 20): slots K..31 repeat slot 0 (:func:`pack_features`; every kept |
| pillar has a point in slot 0), so both slot maxima are unchanged (max is idempotent), whereas zero rows there would |
| add ``relu(b0)`` to the maxima of full pillars. With K = 32 (CenterPoint, PointPainting) the padded slots are the |
| model's own zero slots. |
| |
| **Input**: the host's decorated features as bf16 ROW_MAJOR ``[1, 1, P, 32 * 16]``: one 1 KB row per pillar (32 |
| slots x F <= 16 features zero-padded to 16). One row per pillar keeps the host-to-device copy at P pages (CenterPoint: |
| 5.1 ms for 40,000 pages of 1 KB vs 32 ms as 1.28 M pages of 32 B); the device reshapes it to ``[1, 1, P*32, 16]`` and |
| tilizes it. Every pillar row is computed, also past the frame's pillar count (stale or zero input): the scatter index |
| never points at those rows, so no pillar count reaches the device. :class:`FrameStaging` writes these inputs in place |
| into persistent host tensors (``tensors.HostStaging``): no per-frame ``from_torch`` of 20 M elements (CenterPoint: 188 |
| ms) and no allocation; only the kept pillars are converted (torch fp32 -> bf16, RNE). |
| |
| **Precision** (``ttaw.precision.Precision``): HiFi4 is required (the identity blocks are exact only with it); fp32 |
| accumulation; weights per ``Precision.weights`` (fp32: a device "fp32" matmul behaves like TF32, probe P12); |
| ``Precision.activations`` is the dtype of the hidden tensors ``h`` and ``z`` (bf16 by default; fp32 keeps the PFN's |
| internal roundings off: PointPainting measured on the CPU that the bf16 rounding of ``h`` / ``z`` is the largest PFN |
| error term, rel. L2 0.0045 vs 0.0017 for the bf16 output rounding alone). The output rows are always bf16 (the |
| ``ttnn.embedding`` table of C19 is bf16-only). |
| |
| **Two-term PFN** (``terms=2``; TransFusion's default since its fix rounds 1 / 2, ttaw 0.21.0): every linear is a |
| :class:`TwoTermLinear` (weights and biases as two bf16 terms ``hi + lo``, fp32 outputs); ``h``, ``z`` and the slot |
| max stay fp32 and are split into two bf16 terms (:func:`split_terms`) before the next linear; only the output rows |
| are rounded to bf16. With ``bias_add=True`` (the default) the products are bias-free ``ttnn.matmul`` calls |
| (:class:`RowMatmul`) and each fp32 bias is added by one exact fp32 ``ttnn.add``: ``ttnn.linear``'s fused bias (also |
| an all-zero one) rounds the fp32 output to about 11 significant bits on this tree (TransFusion probes T4 / T5). On |
| TransFusion's golden frames the device PFN output error (relative to fp64 on the same bf16 features) fell from 5.5e-3 |
| (fused bias) to 2.0e-3 and a 0.48 % heatmap local-max near-tie stayed a local maximum |
| (``bundles/transfusion-p150/PORT_LOG.md`` "Fix round 1" / "Fix round 2"). Cost: 3 matmuls + 3 adds per linear and |
| two term splits (P = 65,536: see the TransFusion OPT_BASELINE). ``terms=1`` is the 0.19.0 module unchanged. |
| |
| numpy only at import; ttnn / torch are imported inside the device and staging functions. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Any, Dict, Optional |
|
|
| import numpy as np |
|
|
| from ..precision import Precision |
| from .second import RowLinear, free_maps |
|
|
| __all__ = ["PFN_IN_FEATURES", "PILLAR_SLOTS", "PfnPlan", "TtPillarFeatureNet", "FrameStaging", "features_shape", |
| "pack_features", "pfn_reference_numpy", "pack_inputs", "DEFAULT_PRECISION", "TWO_TERM_PRECISION", |
| "split_terms", "RowMatmul", "TwoTermLinear"] |
|
|
| PFN_IN_FEATURES = 16 |
| PILLAR_SLOTS = 32 |
| TILE = 32 |
| DEFAULT_PRECISION = "HiFi4+fp32:w=fp32" |
| TWO_TERM_PRECISION = "HiFi4+fp32:w=bf16" |
|
|
|
|
| def _round_up(n: int, m: int = TILE) -> int: |
| return -(-int(n) // m) * m |
|
|
|
|
| @dataclass(frozen=True) |
| class PfnPlan: |
| """The device matrices of the PFN (float32 host arrays; see the module docstring).""" |
|
|
| w0: np.ndarray |
| b0: np.ndarray |
| wz: np.ndarray |
| wy: np.ndarray |
| b1: np.ndarray |
| in_features: int |
| hidden: int |
| out_features: int |
|
|
| @classmethod |
| def from_weights(cls, w0: Any, b0: Any, w1: Any, b1: Any, *, in_padded: int = PFN_IN_FEATURES) -> "PfnPlan": |
| """From the BN-folded encoder weights: ``w0`` (F, H) and ``w1`` (2H, O), ``[in, out]`` layout; layer 1 reads |
| ``concat[h, max_slots h]`` (rows 0..H-1 of ``w1`` act on ``h``, rows H..2H-1 on the max).""" |
| w0 = np.asarray(w0, np.float32) |
| w1 = np.asarray(w1, np.float32) |
| if w0.ndim != 2 or w1.ndim != 2: |
| raise ValueError(f"w0 {w0.shape} / w1 {w1.shape} must be [in, out] matrices") |
| f, hdim = w0.shape |
| out = int(w1.shape[1]) |
| if w1.shape[0] != 2 * hdim: |
| raise ValueError(f"w1 {w1.shape} does not take concat[h, max h] of width {2 * hdim}") |
| if f > in_padded: |
| raise ValueError(f"{f} features do not fit the {in_padded}-wide slot rows") |
| if out % TILE: |
| raise ValueError(f"PFN output width {out}: the C19 gather table needs a multiple of {TILE}") |
| b0 = np.asarray(b0, np.float32).reshape(-1) |
| b1 = np.asarray(b1, np.float32).reshape(-1) |
| if b0.shape != (hdim,) or b1.shape != (out,): |
| raise ValueError(f"biases {b0.shape} / {b1.shape} do not match widths {hdim} / {out}") |
| hp, zp = _round_up(hdim), _round_up(hdim + out) |
| W0 = np.zeros((in_padded, hp), np.float32) |
| W0[:f, :hdim] = w0 |
| B0 = np.zeros(hp, np.float32) |
| B0[:hdim] = b0 |
| Wz = np.zeros((hp, zp), np.float32) |
| Wz[:hdim, :hdim] = np.eye(hdim, dtype=np.float32) |
| Wz[:hdim, hdim:hdim + out] = w1[:hdim] |
| Wy = np.zeros((zp, out), np.float32) |
| Wy[:hdim] = w1[hdim:] |
| Wy[hdim:hdim + out] = np.eye(out, dtype=np.float32) |
| return cls(W0, B0, Wz, Wy, b1.copy(), int(f), int(hdim), out) |
|
|
| def forward_numpy(self, features: Any, *, hidden_round: Any = None) -> np.ndarray: |
| """float32 host emulation of the device graph (``(P, S, F)`` slots -> ``(P, out_features)``); give it the |
| 32-slot layout of :func:`pack_features` to emulate the device's slots exactly. ``hidden_round``: a function |
| applied to ``h`` and ``z`` (e.g. a bf16 rounding) to model the hidden dtype.""" |
| r = hidden_round or (lambda a: a) |
| x = np.asarray(features, np.float32) |
| p, slots, f = x.shape |
| xp = np.zeros((p, slots, self.w0.shape[0]), np.float32) |
| xp[:, :, :f] = x |
| h = r(np.maximum(xp @ self.w0 + self.b0, 0.0).astype(np.float32)) |
| z = r((h @ self.wz).astype(np.float32)) |
| mz = z.max(axis=1) |
| return np.maximum(mz @ self.wy + self.b1, 0.0).astype(np.float32) |
|
|
|
|
| def pfn_reference_numpy(features: Any, w0: Any, b0: Any, w1: Any, b1: Any) -> np.ndarray: |
| """The literal encoder in float64 (concat form, max over the given slots, padded slots included) -> ``(P, O)`` |
| float32: the oracle of :class:`PfnPlan` and of the device module.""" |
| x = np.asarray(features, np.float64) |
| w0, w1 = np.asarray(w0, np.float64), np.asarray(w1, np.float64) |
| h = np.maximum(x @ w0 + np.asarray(b0, np.float64), 0.0) |
| m = h.max(axis=1, keepdims=True) |
| y = np.maximum(np.concatenate([h, np.broadcast_to(m, h.shape)], axis=2) @ w1 + np.asarray(b1, np.float64), 0.0) |
| return y.max(axis=1).astype(np.float32) |
|
|
|
|
| def features_shape(num_pillars: int): |
| """The device input layout of the PFN: ``(1, 1, num_pillars, 32 * 16)`` (one bf16 row per pillar).""" |
| return (1, 1, int(num_pillars), PILLAR_SLOTS * PFN_IN_FEATURES) |
|
|
|
|
| def pack_features(features: Any, num_pillars: int, *, out: Optional[np.ndarray] = None) -> np.ndarray: |
| """Decorated features ``(P_kept, K, F)`` (K <= 32 point slots, F <= 16; the model's padded slots zero) -> the |
| device input :func:`features_shape` as float32: features zero-padded to 16, slots K..31 copies of slot 0 when |
| K < 32 (module docstring), rows past ``P_kept`` zero. ``out`` (an array of that shape) is reused when given.""" |
| x = np.asarray(features, np.float32) |
| if x.ndim != 3: |
| raise ValueError(f"features must be (pillars, slots, features), got {x.shape}") |
| p, slots, f = x.shape |
| if slots > PILLAR_SLOTS or f > PFN_IN_FEATURES or p > num_pillars: |
| raise ValueError(f"features {x.shape} do not fit ({num_pillars}, {PILLAR_SLOTS}, <= {PFN_IN_FEATURES})") |
| shape = features_shape(num_pillars) |
| if out is None: |
| out = np.zeros(shape, np.float32) |
| elif out.shape != shape or out.dtype != np.float32: |
| raise ValueError(f"out must be float32 {shape}") |
| else: |
| out.fill(0.0) |
| v = out.reshape(int(num_pillars), PILLAR_SLOTS, PFN_IN_FEATURES) |
| v[:p, :slots, :f] = x |
| if 0 < slots < PILLAR_SLOTS: |
| v[:p, slots:, :f] = x[:, :1, :] |
| return out |
|
|
|
|
| def pack_inputs(features: Any, canvas_index: Any, num_pillars: int, *, capacity: int, |
| features_out: Optional[np.ndarray] = None) -> Dict[str, np.ndarray]: |
| """The two PFN-graph inputs as numpy (tests and tools; the API uses :class:`FrameStaging`): ``{"features": |
| (1, 1, capacity, 512) float32, "canvas_index": (1, N) uint32}`` from the kept pillars' features |
| ``features[:num_pillars]`` and a cell -> pillar-row index (``capacity`` = the zero sentinel row). The index is |
| range-checked: a bad index would read outside the gather table.""" |
| from ..ops.gather import check_index |
|
|
| n = int(num_pillars) |
| feats = pack_features(np.asarray(features)[:n], capacity, out=features_out) |
| return {"features": feats, "canvas_index": check_index(canvas_index, num_rows=n, sentinel=capacity)} |
|
|
|
|
| def split_terms(t: Any): |
| """An fp32 device tensor -> two bf16 tensors ``hi = bf16(t)``, ``lo = bf16(t - hi)`` (``hi + lo`` carries ~16 |
| significant bits; the subtraction is exact in fp32). ``t`` is kept.""" |
| import ttnn |
|
|
| hi = ttnn.typecast(t, ttnn.bfloat16) |
| hi32 = ttnn.typecast(hi, ttnn.float32) |
| diff = ttnn.subtract(t, hi32) |
| ttnn.deallocate(hi32, False) |
| lo = ttnn.typecast(diff, ttnn.bfloat16) |
| ttnn.deallocate(diff, False) |
| return hi, lo |
|
|
|
|
| class RowMatmul(RowLinear): |
| """C24 :class:`RowLinear` without a bias: ``ttnn.matmul`` with an fp32 output, so no bias epilogue touches the fp32 |
| accumulator (``ttnn.linear``'s fused bias, also an all-zero one, degrades an fp32 output to about 11 significant |
| bits: TransFusion probes T4 / T5); add the bias afterwards with an exact fp32 ``ttnn.add`` |
| (:class:`TwoTermLinear` ``bias_add``).""" |
|
|
| def __init__(self, weight: Any, *, precision: Any, name: str): |
| super().__init__(weight, None, activation=None, precision=precision, output_dtype="float32", name=name) |
|
|
| def __call__(self, x: Any) -> Any: |
| import ttnn |
|
|
| if int(x.shape[-1]) != self.in_channels: |
| raise ValueError(f"{self.name}: input has {int(x.shape[-1])} channels, expected {self.in_channels}") |
| weight, _, grid, cfg = self._device or self._upload(x.device()) |
| return ttnn.matmul(x, weight, compute_kernel_config=cfg, core_grid=grid, memory_config=ttnn.DRAM_MEMORY_CONFIG, |
| dtype=ttnn.float32) |
|
|
|
|
| class TwoTermLinear: |
| """``y = act(x @ W + b)`` on a two-term input ``x = hi + lo`` (bf16 TILE tensors ``[..., R, K]``; ``lo`` may be |
| None for a bf16 input) with ``W`` and ``b`` as two bf16 terms: ``hi @ W_hi``, ``hi @ W_lo`` and ``lo @ W_hi`` |
| (bf16 x bf16 products, exact under HiFi4, fp32 accumulation and fp32 outputs) summed in fp32, the activation in |
| fp32. Returns the fp32 TILE tensor (DRAM); the inputs are kept. The dropped ``lo @ W_lo`` term is ~2^-16 |
| relative. Identity blocks in ``W`` stay exact (1.0 is a bf16 hi term). |
| |
| ``bias_add=True`` (default): the three products are bias-free :class:`RowMatmul` calls and the fp32 bias is added |
| by one exact fp32 ``ttnn.add`` (with the ReLU); ``bias_add=False``: three C24 :class:`RowLinear` with the bias |
| terms fused into the matmuls, whose epilogue rounds the fp32 output to about 11 significant bits.""" |
|
|
| def __init__(self, weight: Any, bias: Any = None, *, activation: Optional[str] = None, |
| precision: Any = TWO_TERM_PRECISION, name: str = "linear", bias_add: bool = True): |
| from ..tensors import round_to_bf16 |
|
|
| prec = Precision.parse(precision) |
| if prec.weights != "bfloat16" or prec.fidelity != "HiFi4" or not prec.fp32_acc: |
| raise ValueError(f"{name}: two-term linears need HiFi4, fp32 accumulation and bf16 weight terms, got " |
| f"{prec.label}") |
| w = np.asarray(weight, np.float32) |
| wh = round_to_bf16(w) |
| wl = round_to_bf16(w - wh) |
| b = np.zeros(w.shape[1], np.float32) if bias is None else np.asarray(bias, np.float32).reshape(-1) |
| if activation not in (None, "relu"): |
| raise ValueError(f"{name}: activation {activation!r}: None or 'relu'") |
| self.name = name |
| self.activation = activation |
| self.in_channels, self.out_channels = (int(n) for n in w.shape) |
| self.precision = prec |
| self.bias_add = bool(bias_add) |
| self._bias = None |
| if self.bias_add: |
| self.main = RowMatmul(wh, precision=prec, name=f"{name}.hi") |
| self.corr = RowMatmul(wl, precision=prec, name=f"{name}.wlo") |
| self.low = self.main |
| self.bias_host = b.reshape(1, -1).copy() |
| else: |
| bh = round_to_bf16(b) |
| bl = round_to_bf16(b - bh) |
| self.main = RowLinear(wh, bh, precision=prec, output_dtype="float32", name=f"{name}.hi") |
| self.corr = RowLinear(wl, bl, precision=prec, output_dtype="float32", name=f"{name}.wlo") |
| self.low = RowLinear(wh, None, precision=prec, output_dtype="float32", name=f"{name}.lo") |
|
|
| def _bias_tensor(self, device: Any) -> Any: |
| import torch |
| import ttnn |
|
|
| if self._bias is None: |
| self._bias = ttnn.from_torch(torch.from_numpy(self.bias_host), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, |
| device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| return self._bias |
|
|
| def __call__(self, hi: Any, lo: Any = None) -> Any: |
| import ttnn |
|
|
| from ..ops.conv import unary_with_param |
|
|
| relu = [unary_with_param("relu")] if self.activation else [] |
| if self.bias_add: |
| bias = self._bias_tensor(hi.device()) |
| y = self.corr(hi) |
| if lo is not None: |
| c = self.low(lo) |
| y2 = ttnn.add(y, c, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| free_maps(y, c) |
| y = y2 |
| a = self.main(hi) |
| y2 = ttnn.add(a, y, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| free_maps(a, y) |
| y = ttnn.add(y2, bias, activations=relu, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| free_maps(y2) |
| return y |
| a = self.main(hi) |
| b = self.corr(hi) |
| last = lo is None |
| y = ttnn.add(a, b, activations=relu if last else [], memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| free_maps(a, b) |
| if lo is not None: |
| c = self.low(lo) |
| y2 = ttnn.add(y, c, activations=relu, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| free_maps(y, c) |
| y = y2 |
| return y |
|
|
| def release(self) -> None: |
| for layer in {id(m): m for m in (self.main, self.corr, self.low)}.values(): |
| layer.release() |
| if self._bias is not None: |
| free_maps(self._bias) |
| self._bias = None |
|
|
| def describe(self) -> Dict[str, Any]: |
| return {"name": self.name, "terms": 2, "cin": self.in_channels, "cout": self.out_channels, |
| "activation": self.activation, "precision": self.precision.label, "output_dtype": "float32", |
| "bias": "fp32 add" if self.bias_add else "fused"} |
|
|
|
|
| class TtPillarFeatureNet: |
| """The PFN on ``num_pillars`` pillars: ``[1, 1, P, 32*16]`` bf16 ROW_MAJOR -> ``[1, 1, P, O]`` bf16 ROW_MAJOR |
| rows (module docstring). ``precision``: a :class:`..precision.Precision` or its spec; its ``activations`` dtype is |
| that of the hidden tensors ``h`` / ``z`` (``terms=1``). ``terms=2``: the two-term PFN (module docstring; needs |
| ``w=bf16``, e.g. :data:`TWO_TERM_PRECISION`; hidden tensors fp32), ``bias_add`` as in :class:`TwoTermLinear`.""" |
|
|
| def __init__(self, plan: PfnPlan, *, num_pillars: int, precision: Any = DEFAULT_PRECISION, name: str = "pfn", |
| terms: int = 1, bias_add: bool = True): |
| self.plan = plan |
| self.name = name |
| self.num_pillars = int(num_pillars) |
| self.terms = int(terms) |
| if self.terms not in (1, 2): |
| raise ValueError(f"{name}: terms must be 1 or 2, got {terms}") |
| self.precision = Precision.parse(precision) |
| if self.precision.fidelity != "HiFi4": |
| raise ValueError(f"{name}: the PFN identity blocks are exact only with HiFi4 (x * 1.0)") |
| self.bias_add = bool(bias_add) |
| if self.terms == 2: |
| kw = dict(precision=self.precision, bias_add=self.bias_add) |
| self.l0: Any = TwoTermLinear(plan.w0, plan.b0, activation="relu", name=f"{name}.l0", **kw) |
| self.lz: Any = TwoTermLinear(plan.wz, None, name=f"{name}.lz", **kw) |
| self.ly: Any = TwoTermLinear(plan.wy, plan.b1, activation="relu", name=f"{name}.ly", **kw) |
| else: |
| hidden = self.precision.activations |
| self.l0 = RowLinear(plan.w0, plan.b0, activation="relu", precision=self.precision, output_dtype=hidden, |
| name=f"{name}.l0") |
| self.lz = RowLinear(plan.wz, None, precision=self.precision, output_dtype=hidden, name=f"{name}.lz") |
| self.ly = RowLinear(plan.wy, plan.b1, activation="relu", precision=self.precision, |
| output_dtype="bfloat16", name=f"{name}.ly") |
| self.out_features = plan.out_features |
|
|
| @property |
| def features_shape(self): |
| return features_shape(self.num_pillars) |
|
|
| def __call__(self, features_rm: Any) -> Any: |
| import ttnn |
|
|
| p = self.num_pillars |
| rows = ttnn.reshape(features_rm, (1, 1, p * PILLAR_SLOTS, PFN_IN_FEATURES)) |
| x = ttnn.to_layout(rows, ttnn.TILE_LAYOUT) |
| free_maps(rows) |
| if self.terms == 2: |
| y = self._two_term(x) |
| else: |
| h = self.l0(x) |
| free_maps(x) |
| z = self.lz(h) |
| free_maps(h) |
| z4 = ttnn.reshape(z, (1, p, PILLAR_SLOTS, int(z.shape[-1]))) |
| mz = ttnn.max(z4, dim=2, keepdim=True) |
| free_maps(z) |
| y = self.ly(mz) |
| free_maps(mz) |
| y_rm = ttnn.to_layout(y, ttnn.ROW_MAJOR_LAYOUT) |
| free_maps(y) |
| return ttnn.reshape(y_rm, (1, 1, p, self.out_features)) |
|
|
| def _two_term(self, x: Any) -> Any: |
| """``terms=2``: fp32 ``h`` / ``z`` / slot max, each split into two bf16 terms for the next linear -> bf16 |
| ``[1, P, 1, O]`` (TILE).""" |
| import ttnn |
|
|
| p = self.num_pillars |
| h = self.l0(x) |
| free_maps(x) |
| hh, hl = split_terms(h) |
| free_maps(h) |
| z = self.lz(hh, hl) |
| free_maps(hh, hl) |
| z4 = ttnn.reshape(z, (1, p, PILLAR_SLOTS, int(z.shape[-1]))) |
| mz = ttnn.max(z4, dim=2, keepdim=True) |
| free_maps(z) |
| mh, ml = split_terms(mz) |
| free_maps(mz) |
| y32 = self.ly(mh, ml) |
| free_maps(mh, ml) |
| y = ttnn.typecast(y32, ttnn.bfloat16) |
| free_maps(y32) |
| return y |
|
|
| def release(self) -> None: |
| for layer in (self.l0, self.lz, self.ly): |
| layer.release() |
|
|
| def describe(self) -> Dict[str, Any]: |
| return {"name": self.name, "num_pillars": self.num_pillars, "in_features": self.plan.in_features, |
| "hidden": self.plan.hidden, "out": self.plan.out_features, "z_width": int(self.plan.wz.shape[1]), |
| "slots": PILLAR_SLOTS, "precision": self.precision.label, "terms": self.terms, |
| "hidden_dtype": "float32" if self.terms == 2 else self.precision.activations, |
| **({"bias": "fp32 add" if self.bias_add else "fused"} if self.terms == 2 else {})} |
|
|
|
|
| class FrameStaging: |
| """Persistent host buffers of the two PFN-graph inputs: bf16 ``features`` (:func:`features_shape`) and uint32 |
| ``canvas_index`` ``[1, num_cells]`` (``tensors.HostStaging``: the ttnn host tensor's own memory, written in |
| place). :meth:`write` converts only the kept pillars (torch fp32 -> bf16, RNE) into their rows (slots K..31 = |
| slot 0 when K < 32), zeroes the rows the previous frame used beyond them, copies the range-checked index, and |
| returns the ``TraceRunner`` inputs (host tensors, uploaded as they are). Without a zero-copy alias (another |
| tt-metal build) it falls back to :func:`pack_features` + ``HostStaging.write`` (same values).""" |
|
|
| def __init__(self, num_pillars: int, num_cells: int): |
| from ..tensors import HostStaging |
|
|
| self.num_pillars = int(num_pillars) |
| self.num_cells = int(num_cells) |
| self.features = HostStaging(features_shape(self.num_pillars), "bfloat16") |
| self.index = HostStaging((1, self.num_cells), "uint32") |
| self._used = self.num_pillars |
|
|
| @property |
| def zero_copy(self) -> bool: |
| return bool(self.features.zero_copy and self.index.zero_copy) |
|
|
| def write(self, features: Any, canvas_index: Any, num_pillars: Optional[int] = None) -> Dict[str, Any]: |
| """``features``: ``(P, K, F)`` float32 decorated pillars (only the first ``num_pillars`` rows are used; |
| default all); ``canvas_index``: ``num_cells`` entries in ``[0, num_pillars)`` or the sentinel |
| ``self.num_pillars``.""" |
| from ..ops.gather import check_index |
|
|
| f_all = np.asarray(features) |
| p = int(f_all.shape[0] if num_pillars is None else num_pillars) |
| if p > self.num_pillars: |
| raise ValueError(f"{p} pillars exceed the staging capacity {self.num_pillars}") |
| feats = np.asarray(f_all[:p], np.float32) |
| index = check_index(canvas_index, num_rows=p, sentinel=self.num_pillars) |
| if index.size != self.num_cells: |
| raise ValueError(f"canvas index of {index.size} cells, expected {self.num_cells}") |
| if not self.zero_copy: |
| return {"features": self.features.write(pack_features(feats, self.num_pillars)), |
| "canvas_index": self.index.write(index)} |
| import torch |
|
|
| if feats.ndim != 3 or feats.shape[1] > PILLAR_SLOTS or feats.shape[2] > PFN_IN_FEATURES: |
| raise ValueError(f"features {feats.shape} do not fit (P, <= {PILLAR_SLOTS}, <= {PFN_IN_FEATURES})") |
| view = self.features.view.reshape(self.num_pillars, PILLAR_SLOTS, PFN_IN_FEATURES).view(np.int16) |
| dst = torch.from_numpy(view) |
| slots, f = feats.shape[1], feats.shape[2] |
| src = torch.from_numpy(np.ascontiguousarray(feats)).to(torch.bfloat16).view(torch.int16) |
| dst[:p, :slots, :f] = src |
| if f < PFN_IN_FEATURES: |
| dst[:p, :, f:] = 0 |
| if 0 < slots < PILLAR_SLOTS: |
| dst[:p, slots:, :f] = src[:, :1, :] |
| if p < self._used: |
| dst[p:self._used] = 0 |
| self._used = p |
| self.index.view[...] = index |
| return {"features": self.features.tensor, "canvas_index": self.index.tensor} |
|
|
