Blame for code/tt_diffusion_planner/ttaw/models/pillars.py

# SPDX-License-Identifier: Apache-2.0
"""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        # decorated features zero-padded to 16 bf16 = 32-byte slot rows
PILLAR_SLOTS = 32           # one 32-row tile per pillar
TILE = 32
DEFAULT_PRECISION = "HiFi4+fp32:w=fp32"
TWO_TERM_PRECISION = "HiFi4+fp32:w=bf16"     # terms=2: the weight terms are bf16-exact


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            # (16, H_pad)
    b0: np.ndarray            # (H_pad,)
    wz: np.ndarray            # (H_pad, Z) = [I_H | W1a | 0]
    wy: np.ndarray            # (Z, O)     = [W1b ; I_O ; 0]
    b1: np.ndarray            # (O,)
    in_features: int          # F
    hidden: int               # H (layer-0 width)
    out_features: int         # O

    @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                                     # lo @ W_hi: the same device weights
            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)                                        # the small terms first, then the main product
            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))   # 1 KB rows -> 32 B slot rows
        x = ttnn.to_layout(rows, ttnn.TILE_LAYOUT)                      # [1, 1, P*32, 16]
        free_maps(rows)
        if self.terms == 2:
            y = self._two_term(x)
        else:
            h = self.l0(x)                                              # [1, 1, P*32, H_pad]
            free_maps(x)
            z = self.lz(h)                                              # [1, 1, P*32, Z]
            free_maps(h)
            z4 = ttnn.reshape(z, (1, p, PILLAR_SLOTS, int(z.shape[-1])))  # one pillar = one tile row block (a view)
            mz = ttnn.max(z4, dim=2, keepdim=True)                      # [1, P, 1, Z]
            free_maps(z)
            y = self.ly(mz)                                             # [1, P, 1, O] bf16
            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)                                                  # fp32 [1, 1, P*32, H_pad]
        free_maps(x)
        hh, hl = split_terms(h)
        free_maps(h)
        z = self.lz(hh, hl)                                             # fp32 [1, 1, P*32, Z]
        free_maps(hh, hl)
        z4 = ttnn.reshape(z, (1, p, PILLAR_SLOTS, int(z.shape[-1])))
        mz = ttnn.max(z4, dim=2, keepdim=True)                          # fp32 [1, P, 1, Z]
        free_maps(z)
        mh, ml = split_terms(mz)
        free_maps(mz)
        y32 = self.ly(mh, ml)                                           # fp32 [1, P, 1, O]
        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            # rows that may hold data (all, until the first write)

    @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:                                           # no host alias: plain writes
            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}