code/tt_diffusion_planner/ttaw/models/sparse_encoder.py

# SPDX-License-Identifier: Apache-2.0
"""C28: the gather-GEMM sparse-conv encoder on the device (PLAN.md 1.2 row C28; probe P4; host rulebooks: C13
``ttaw.sparse``). Owned by the BEVFusion port; PTv3's sub-manifold stages are the second consumer.

Every spconv layer of the Autoware graphs (sub-manifold and strided alike) is a gather (C13): ``y[o] = sum_k
W[:, k, :] @ x[nmap[o, k]] + b``. On the device one layer is:

1. ONE ``ttnn.embedding`` (TILE + PADDED, P4) of the layer's tile-ordered im2col index (C13 ``im2col_index``,
   uint32 ``[rows / 32, K * 32]``) into the bf16 ROW_MAJOR table ``[1, 1, V + 1, Cin_pad]`` of the layer's input rows
   (the last row = the explicit zero sentinel ``V``, ``Cin_pad % 32 == 0``) -> ``[rows / 32, K * 32, Cin_pad]``
   whose tiles, in memory order, are the tiles of the im2col matrix ``[rows, K * Cin_pad]``;
2. ``ttnn.experimental.view`` to ``[1, 1, rows, K * Cin_pad]`` (metadata only);
3. a bias-free ``ttnn.matmul`` with the BN-folded im2col weight ``[K * Cin_pad, Cout_pad]`` into fp32 (C24
   ``RowMatmul``: no fused-bias epilogue rounds the fp32 accumulator), optionally a second one with the weight's
   bf16 low term (``weight_terms=2``);
4. one exact fp32 ``ttnn.add`` of the bias (with the ReLU fused when there is no residual), or bias + the fp32
   identity of a basic block followed by the ReLU fused into the second add;
5. the next table: the fp32 output cast to bf16, untilized and given its zero row (C19 ``with_zero_rows``).

**Activation terms.** ``ttnn.embedding`` is bf16-only, so a table rounds every layer's input to bf16 (8 significant
bits). A layer whose table is *two-term* (``in_terms=2``) carries ``hi = bf16(x)`` in channels ``[0, C)`` and ``lo =
bf16(x - hi)`` in ``[C, 2C)`` of the same table (~16 significant bits): the im2col weight repeats each tap's rows for
the ``lo`` block, so the same single gather and single matmul compute ``hi @ W + lo @ W``. Where ``2C <= 32`` (Cin 4
of a VFE input, 16 at BEVFusion's level 1) the ``lo`` block rides in channels that the tile padding carries anyway:
free. A producer with ``out_terms=2`` writes its output columns twice (``[W | W]``) and splits them on the device into
``[hi | lo]`` (``y - bf16(y) * mask``, then one bf16 cast: exact). Consecutive layers must agree (the producer's
``out_terms`` is the consumer's ``in_terms``; :class:`SparseEncoder` checks it).

Rows past a level's active count are computed (capacity padding: nothing data-dependent reaches the device) but
never gathered: every index points at an active row or at the sentinel, so no count is needed on the device.

numpy only at import (ttnn / torch inside the device calls). Host oracle: :func:`sparse_encoder_numpy`, with the
device's bf16 table rounding modelled (``round_tables=True``).
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple

import numpy as np

from ..precision import Precision, PrecisionPolicy
from .second import free_maps

__all__ = ["SparseConvSpec", "TtSparseConv", "SparseEncoder", "DEFAULT_PRECISION", "TILE_ROWS", "pad_channels",
           "pack_table", "table_rows_numpy", "sparse_layer_numpy", "sparse_encoder_numpy", "dense_gather",
           "im2col_weight"]

TILE_ROWS = 32
# fp32 weights (TF32-like products on this tree, P12), HiFi4, fp32 accumulation
DEFAULT_PRECISION = "HiFi4+fp32:w=fp32"


def pad_channels(c: int, multiple: int = TILE_ROWS) -> int:
    """``c`` rounded up to a multiple of 32 (the table width a TILE embedding needs, P4)."""
    return -(-int(c) // multiple) * multiple


def _bf16(a: np.ndarray) -> np.ndarray:
    from ..tensors import round_to_bf16

    return round_to_bf16(np.asarray(a, np.float32))


@dataclass(frozen=True)
class SparseConvSpec:
    """One BN-folded sparse conv: ``y = relu?(gather_conv(x, weight) + bias (+ identity))``.

    ``weight``: KRSC ``[Cout, k0, k1, k2, Cin]`` float32 (tap ``k = (a * k1 + b) * k2 + c``, C13 order) or already
    ``[Cout, K, Cin]``; ``index``: the name of the gather index it reads (several layers may share one: the
    sub-manifold convs of a level); ``residual``: the name of an earlier layer whose (fp32, pre-split) output is
    added before the ReLU; ``in_terms`` / ``out_terms``: 1 (bf16 table) or 2 (``[hi | lo]`` table, module
    docstring)."""

    name: str
    weight: np.ndarray
    bias: np.ndarray
    index: str
    residual: Optional[str] = None
    relu: bool = True
    in_terms: int = 1
    out_terms: int = 1

    def __post_init__(self) -> None:
        w = np.asarray(self.weight, np.float32)
        if w.ndim not in (3, 5):
            raise ValueError(f"{self.name}: weight must be KRSC [Cout, k0, k1, k2, Cin] or [Cout, K, Cin], got {w.shape}")
        if np.shape(self.bias) != (w.shape[0],):
            raise ValueError(f"{self.name}: bias must be ({w.shape[0]},), got {np.shape(self.bias)}")
        if self.in_terms not in (1, 2) or self.out_terms not in (1, 2):
            raise ValueError(f"{self.name}: terms must be 1 or 2")
        if not (np.all(np.isfinite(w)) and np.all(np.isfinite(self.bias))):
            raise ValueError(f"{self.name}: non-finite weights")

    @property
    def weight3(self) -> np.ndarray:
        w = np.asarray(self.weight, np.float32)
        return w.reshape(w.shape[0], -1, w.shape[-1])

    @property
    def kernel_volume(self) -> int:
        return int(self.weight3.shape[1])

    @property
    def in_channels(self) -> int:
        return int(self.weight3.shape[2])

    @property
    def out_channels(self) -> int:
        return int(self.weight3.shape[0])

    @property
    def cin_pad(self) -> int:
        """Table width of the input: ``in_terms * Cin`` rounded up to 32."""
        return pad_channels(self.in_terms * self.in_channels)

    @property
    def cout_pad(self) -> int:
        """Output / next-table width: ``out_terms * Cout`` rounded up to 32."""
        return pad_channels(self.out_terms * self.out_channels)

    def matrix(self) -> np.ndarray:
        """The im2col weight ``[K * cin_pad, cout_pad]`` float32 (:func:`im2col_weight`)."""
        return im2col_weight(self.weight3, cin_pad=self.cin_pad, cout_pad=self.cout_pad, in_terms=self.in_terms,
                             out_terms=self.out_terms)

    def bias_row(self) -> np.ndarray:
        """``(1, cout_pad)`` float32: the bias, repeated for the ``lo`` block when ``out_terms == 2``."""
        out = np.zeros((1, self.cout_pad), np.float32)
        c = self.out_channels
        for t in range(self.out_terms):
            out[0, t * c:(t + 1) * c] = np.asarray(self.bias, np.float32)
        return out

    def lo_mask(self) -> np.ndarray:
        """``(1, cout_pad)`` float32: 1 on the ``lo`` output columns (``out_terms == 2``), else 0."""
        out = np.zeros((1, self.cout_pad), np.float32)
        if self.out_terms == 2:
            out[0, self.out_channels:2 * self.out_channels] = 1.0
        return out


def im2col_weight(weight3: np.ndarray, *, cin_pad: int, cout_pad: int, in_terms: int = 1,
                  out_terms: int = 1) -> np.ndarray:
    """``[Cout, K, Cin]`` -> ``[K * cin_pad, cout_pad]``: row ``k * cin_pad + j`` = ``W[:, k, j mod Cin]`` for the
    ``in_terms`` channel blocks (the ``lo`` block repeats the ``hi`` weights), zero for padding; columns repeated for
    the ``out_terms`` blocks, zero padding beyond."""
    w = np.asarray(weight3, np.float32)
    cout, K, cin = w.shape
    if in_terms * cin > cin_pad or out_terms * cout > cout_pad:
        raise ValueError(f"cin_pad {cin_pad} / cout_pad {cout_pad} too small for {in_terms} x {cin} / "
                         f"{out_terms} x {cout}")
    m = np.zeros((K, cin_pad, cout_pad), np.float32)
    wt = np.transpose(w, (1, 2, 0))                                      # (K, Cin, Cout)
    for ti in range(in_terms):
        for to in range(out_terms):
            m[:, ti * cin:(ti + 1) * cin, to * cout:(to + 1) * cout] = wt
    return m.reshape(K * cin_pad, cout_pad)


# --------------------------------------------------------------------------------------------- host tables

def table_rows_numpy(x: np.ndarray, *, terms: int, width: int) -> np.ndarray:
    """``(N, C)`` float rows -> ``(N, width)`` float32 table rows with the device's bf16 values: ``terms=1``:
    ``bf16(x)`` in ``[0, C)``; ``terms=2``: ``hi = bf16(x)`` in ``[0, C)``, ``lo = bf16(x - hi)`` in ``[C, 2C)``;
    zero elsewhere."""
    x = np.asarray(x, np.float32)
    n, c = x.shape
    if terms * c > width:
        raise ValueError(f"{terms} x {c} channels do not fit a width of {width}")
    out = np.zeros((n, width), np.float32)
    hi = _bf16(x)
    out[:, :c] = hi
    if terms == 2:
        out[:, c:2 * c] = _bf16(x - hi)
    return out


def pack_table(x: np.ndarray, capacity: int, *, terms: int = 1, width: Optional[int] = None,
               out: Optional[np.ndarray] = None) -> np.ndarray:
    """The first layer's input table as a host array ``(1, 1, capacity + 1, width)`` float32 (bf16-exact values,
    upload as bf16 ROW_MAJOR): the ``N <= capacity`` rows of ``x`` (:func:`table_rows_numpy`), zero rows after them,
    and the zero sentinel row at ``capacity``. ``width`` defaults to ``terms * C`` rounded up to 32."""
    x = np.asarray(x, np.float32)
    n, c = x.shape
    if n > capacity:
        raise ValueError(f"{n} rows exceed the capacity {capacity}")
    width = pad_channels(terms * c) if width is None else int(width)
    shape = (1, 1, int(capacity) + 1, width)
    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[...] = 0.0
    out[0, 0, :n] = table_rows_numpy(x, terms=terms, width=width)
    return out


def sparse_layer_numpy(table: np.ndarray, nmap: np.ndarray, spec: SparseConvSpec, *,
                       identity: Optional[np.ndarray] = None) -> np.ndarray:
    """Host oracle of one device layer on a host table ``(V + 1, cin_pad)`` (the last row the zero sentinel):
    gathers by ``nmap`` (``-1`` -> sentinel), the im2col matmul in float64, bias, identity, ReLU -> ``(N, Cout)``
    float32 (the logical columns)."""
    t = np.asarray(table, np.float64).reshape(-1, spec.cin_pad)
    sentinel = t.shape[0] - 1
    idx = np.where(np.asarray(nmap) >= 0, nmap, sentinel).astype(np.int64)
    cols = t[idx].reshape(idx.shape[0], -1)                              # (N, K * cin_pad)
    m = spec.matrix().astype(np.float64)[:, :spec.out_channels]
    y = cols @ m + np.asarray(spec.bias, np.float64)
    if identity is not None:
        y = y + np.asarray(identity, np.float64)
    if spec.relu:
        y = np.maximum(y, 0.0)
    return y.astype(np.float32)


def sparse_encoder_numpy(specs: Sequence[SparseConvSpec], features: np.ndarray, maps: Mapping[str, np.ndarray], *,
                         round_tables: bool = True) -> Dict[str, np.ndarray]:
    """Host oracle of :class:`SparseEncoder`: ``features`` ``(N0, C0)`` (the first layer's input rows) and the
    neighbour maps by index name -> ``{layer name: (N, Cout) float32}``. ``round_tables=True`` models the device's
    tables (bf16 or two-term ``hi + lo``, per the specs' terms); False is the fp32 network (the reference)."""
    outs: Dict[str, np.ndarray] = {}
    x = np.asarray(features, np.float32)
    for spec in specs:
        if round_tables:
            table = table_rows_numpy(x, terms=spec.in_terms, width=spec.cin_pad)
        else:
            table = np.zeros((x.shape[0], spec.cin_pad), np.float32)    # an exact table: the lo block is zero
            table[:, :x.shape[1]] = x
        table = np.concatenate([table, np.zeros((1, spec.cin_pad), np.float32)], 0)
        ident = outs[spec.residual] if spec.residual else None
        y = sparse_layer_numpy(table, maps[spec.index], spec, identity=ident)
        outs[spec.name] = y
        x = y
    return outs


# --------------------------------------------------------------------------------------------- device layer

class TtSparseConv:
    """One device layer of :class:`SparseConvSpec` (module docstring, steps 1-4). ``weight_terms=2`` runs the
    im2col weight as two bf16 terms (two matmuls, ~16 significant bits) instead of fp32 weights (TF32-like
    products, P12). Device weights are uploaded by the first call (outside any capture)."""

    def __init__(self, spec: SparseConvSpec, *, precision: Any = DEFAULT_PRECISION, weight_terms: int = 1):
        from .pillars import RowMatmul

        self.spec = spec
        self.name = spec.name
        self.precision = Precision.parse(precision) if not isinstance(precision, Precision) else precision
        self.weight_terms = int(weight_terms)
        if self.weight_terms not in (1, 2):
            raise ValueError(f"{self.name}: weight_terms must be 1 or 2")
        m = spec.matrix()
        if self.weight_terms == 2:
            if self.precision.weights != "bfloat16":
                self.precision = self.precision.with_(weights="bfloat16")
            hi = _bf16(m)
            self.mm = [RowMatmul(hi, precision=self.precision, name=f"{self.name}.whi"),
                       RowMatmul(_bf16(m - hi), precision=self.precision, name=f"{self.name}.wlo")]
        else:
            self.mm = [RowMatmul(m, precision=self.precision, name=self.name)]
        self.k_in = int(m.shape[0])
        self._consts: Optional[Tuple[Any, Any]] = None

    @property
    def rows_per_block(self) -> int:
        return TILE_ROWS

    def _const(self, device: Any) -> Tuple[Any, Any]:
        import torch
        import ttnn

        if self._consts is None:
            def up(a):
                return ttnn.from_torch(torch.from_numpy(np.ascontiguousarray(a, np.float32)), dtype=ttnn.float32,
                                       layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG)

            mask = up(self.spec.lo_mask()) if self.spec.out_terms == 2 else None
            self._consts = (up(self.spec.bias_row()), mask)
        return self._consts

    def gather(self, table: Any, index: Any):
        """The im2col operand ``[1, 1, rows, K * cin_pad]`` bf16 TILE (one PADDED embedding + a view)."""
        import ttnn

        width = int(table.shape[-1])
        if width != self.spec.cin_pad:
            raise ValueError(f"{self.name}: table width {width} != cin_pad {self.spec.cin_pad}")
        k = self.spec.kernel_volume
        if int(index.shape[-1]) != k * TILE_ROWS:
            raise ValueError(f"{self.name}: index has {int(index.shape[-1])} columns, expected {k * TILE_ROWS}")
        rows = int(index.shape[0]) * TILE_ROWS
        sentinel = int(table.shape[-2]) - 1
        g = ttnn.embedding(index, table, layout=ttnn.TILE_LAYOUT, padding_idx=sentinel,
                           embeddings_type=ttnn.EmbeddingsType.PADDED)       # [rows / 32, K * 32, cin_pad]
        return ttnn.experimental.view(g, [1, 1, rows, k * width])

    def __call__(self, table: Any, index: Any, identity: Any = None):
        """``table`` bf16 ROW_MAJOR ``[1, 1, V + 1, cin_pad]`` (row V zero), ``index`` uint32 ``[rows / 32, K * 32]``,
        ``identity`` an fp32 TILE ``[1, 1, rows, cout_pad]`` (the residual source's output) or None -> the fp32 TILE
        output ``[1, 1, rows, cout_pad]`` (DRAM). The inputs are kept."""
        import ttnn

        from ..ops.conv import unary_with_param

        x = self.gather(table, index)
        ys = [mm(x) for mm in self.mm]
        free_maps(x)
        y = ys[-1]
        for extra in reversed(ys[:-1]):                                  # the small term first, then the main one
            s = ttnn.add(extra, y, memory_config=ttnn.DRAM_MEMORY_CONFIG)
            free_maps(extra, y)
            y = s
        bias, _ = self._const(table.device())
        relu = [unary_with_param("relu")] if self.spec.relu else []
        if identity is None:
            out = ttnn.add(y, bias, activations=relu, memory_config=ttnn.DRAM_MEMORY_CONFIG)
            free_maps(y)
            return out
        s = ttnn.add(y, bias, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        free_maps(y)
        out = ttnn.add(s, identity, activations=relu, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        free_maps(s)
        return out

    def to_table(self, y: Any):
        """The fp32 output -> the next layer's bf16 ROW_MAJOR table ``[1, 1, rows + 1, cout_pad]`` (zero last row;
        ``[hi | lo]`` split for ``out_terms == 2``). ``y`` is kept."""
        import ttnn

        from ..ops.gather import with_zero_rows

        if self.spec.out_terms == 2:
            _, mask = self._const(y.device())
            hi = ttnn.typecast(y, ttnn.bfloat16)
            h32 = ttnn.typecast(hi, ttnn.float32)
            free_maps(hi)
            hm = ttnn.multiply(h32, mask)
            free_maps(h32)
            r = ttnn.subtract(y, hm)
            free_maps(hm)
            t = ttnn.typecast(r, ttnn.bfloat16)
            free_maps(r)
        else:
            t = ttnn.typecast(y, ttnn.bfloat16)
        table = with_zero_rows(t, 1)
        free_maps(t)
        return table

    def release(self) -> None:
        for mm in self.mm:
            mm.release()
        if self._consts is not None:
            free_maps(*[c for c in self._consts if c is not None])
            self._consts = None

    def describe(self) -> Dict[str, Any]:
        s = self.spec
        return {"name": s.name, "index": s.index, "cin": s.in_channels, "cout": s.out_channels, "K": s.kernel_volume,
                "cin_pad": s.cin_pad, "cout_pad": s.cout_pad, "in_terms": s.in_terms, "out_terms": s.out_terms,
                "residual": s.residual, "relu": s.relu, "weight_terms": self.weight_terms,
                "precision": self.precision.label}


class SparseEncoder:
    """A chain of :class:`TtSparseConv` layers (in order; each reads the previous layer's table, the first one the
    given input table). ``policy`` resolves ``<prefix-free layer name>`` to a precision (default
    :data:`DEFAULT_PRECISION`); ``weight_terms`` 1 or 2 for every layer (or a ``{name: terms}`` mapping)."""

    def __init__(self, specs: Sequence[SparseConvSpec], *, policy: Optional[PrecisionPolicy] = None,
                 weight_terms: Any = 1, name: str = "sparse"):
        if not specs:
            raise ValueError("no layers")
        self.name = name
        names = [s.name for s in specs]
        if len(set(names)) != len(names):
            raise ValueError("layer names must be unique")
        for prev, cur in zip(specs[:-1], specs[1:]):
            if prev.out_terms != cur.in_terms or prev.cout_pad != cur.cin_pad:
                raise ValueError(f"{cur.name}: input terms / width {cur.in_terms} / {cur.cin_pad} != the output of "
                                 f"{prev.name} ({prev.out_terms} / {prev.cout_pad})")
            if prev.out_channels != cur.in_channels:
                raise ValueError(f"{cur.name}: {cur.in_channels} input channels after {prev.out_channels}")
        seen = set()
        for s in specs:
            if s.residual is not None:
                if s.residual not in seen:
                    raise ValueError(f"{s.name}: residual {s.residual!r} is not an earlier layer")
                src = next(p for p in specs if p.name == s.residual)
                if src.cout_pad != s.cout_pad or src.out_channels != s.out_channels or src.out_terms != s.out_terms:
                    raise ValueError(f"{s.name}: residual {s.residual!r} has another output layout")
            seen.add(s.name)
        self.specs = list(specs)
        self.layers: List[TtSparseConv] = []
        for s in specs:
            prec = policy.resolve(s.name) if policy is not None else Precision.parse(DEFAULT_PRECISION)
            wt = int(weight_terms.get(s.name, 1)) if isinstance(weight_terms, Mapping) else int(weight_terms)
            self.layers.append(TtSparseConv(s, precision=prec, weight_terms=wt))
        # the last use of every layer output (as a residual identity), to free it as early as possible
        self._last_use: Dict[str, int] = {}
        for i, s in enumerate(specs):
            if s.residual is not None:
                self._last_use[s.residual] = i

    @property
    def index_names(self) -> List[str]:
        """The gather indices the chain reads, in first-use order."""
        out: List[str] = []
        for s in self.specs:
            if s.index not in out:
                out.append(s.index)
        return out

    @property
    def input_width(self) -> int:
        return self.specs[0].cin_pad

    @property
    def input_terms(self) -> int:
        return self.specs[0].in_terms

    def __call__(self, table: Any, indices: Mapping[str, Any], taps: Optional[Dict[str, Any]] = None,
                 keep_table: bool = True):
        """``table``: the first layer's bf16 ROW_MAJOR input table (kept); ``indices``: ``{index name: uint32 device
        tensor}``; ``taps`` (a dict) collects every layer's fp32 output ``[1, 1, rows, cout_pad]`` by layer name.
        Returns ``(y, out_table)``: the last layer's fp32 output and (``keep_table``) its bf16 ROW_MAJOR table
        (zero sentinel row appended), else ``(y, None)``."""
        missing = [n for n in self.index_names if n not in indices]
        if missing:
            raise KeyError(f"{self.name}: missing gather indices {missing}")
        cur_table = table
        outs: Dict[str, Any] = {}
        y = None
        n = len(self.layers)
        for i, (layer, spec) in enumerate(zip(self.layers, self.specs)):
            ident = outs[spec.residual] if spec.residual else None
            y = layer(cur_table, indices[spec.index], ident)
            if cur_table is not table:
                free_maps(cur_table)
            # free outputs whose last residual use was this layer (unless tapped)
            for name, last in list(self._last_use.items()):
                if last == i and name in outs and (taps is None):
                    free_maps(outs.pop(name))
            outs[spec.name] = y
            if taps is not None:
                taps[spec.name] = y
            last_layer = i == n - 1
            if not last_layer or keep_table:
                cur_table = layer.to_table(y)
            if not last_layer and taps is None and spec.name not in self._last_use:
                free_maps(outs.pop(spec.name))
        return y, (cur_table if keep_table else None)

    def release(self) -> None:
        for layer in self.layers:
            layer.release()

    def describe(self) -> Dict[str, Any]:
        return {"name": self.name, "layers": [layer.describe() for layer in self.layers],
                "indices": self.index_names}


def dense_gather(table: Any, index: Any):
    """To-dense as a gather (C13 ``dense_gather_index``, one slice): ``table`` bf16 ROW_MAJOR ``[1, 1, V + 1, C]``
    (row V zero), ``index`` uint32 ``[1, cells]`` -> ``[1, 1, cells, C]`` bf16 TILE (empty cells exactly 0)."""
    from ..ops.gather import gather_rows

    return gather_rows(index, table, sentinel=int(table.shape[-2]) - 1)