# 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 ```` 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)