# SPDX-License-Identifier: Apache-2.0 """C25: the CenterPoint dense head on the device (PLAN.md section 1.2 row C25; CenterPoint, PointPainting). mmdet3d ``CenterHead`` (single task, as Autoware's exports): ``shared = relu(conv3x3(concat[d_0..d_n]))``, then per output head ``out_h = conv1x1(relu(conv1x1(shared)))``. Two exact rewrites (each differs from the literal graph only by summation order; probe P14 and research/centerpoint/SPEC.md 8.1): - **K-split shared conv** (default; C17 ``KSplitConv``): ``conv(concat[d_i], W) == sum_i conv(d_i, W_i)``, the bias on the first partial conv and the ReLU fused into the last add. The concat (CenterPoint base: 177 MB) never exists; P14: 3.6 ms against 14.8 ms for the full-K conv, and more accurate. ``accumulate_dtype="float32"`` keeps the partials in fp32; - **merged sibling heads**: the six 1x1 -> ReLU -> 1x1 heads become two matmuls on NHWC rows (:class:`..second. RowLinear`): ``hidden = relu(x @ Wh + bh)`` with ``Wh = [Wh_1 | ... | Wh_6]`` (``64 -> 6*64``), then ``out = hidden @ Wo + bo`` with the block-diagonal ``Wo`` (``6*64 -> sum C_h``, zero-padded to a multiple of 32 columns so the output tensor is tile-exact). The output is ONE ``[1, 1, H*W, C_pad]`` tensor (fp32 by default: the logits feed host thresholds), i.e. one readback; :meth:`CenterHead.unpack` splits the host rows into ``{name: (C, H, W)}``. Module names for the precision policy: ``.shared``, ``.hidden``, ``.out``. """ from __future__ import annotations from typing import Any, Dict, List, Optional, Sequence, Tuple import numpy as np from ..ops.conv import Conv2d, FeatureMap, KSplitConv from ..precision import PrecisionPolicy from .second import ConvSpec, RowLinear, _precision, _to_dram, free_maps __all__ = ["CenterHead", "HeadSpec", "merge_heads", "centerhead_numpy"] HeadSpec = Tuple[str, ConvSpec, ConvSpec] # (output name, hidden 1x1 + ReLU, final 1x1) def merge_heads(heads: Sequence[HeadSpec], *, pad_to: int = 32 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, List[Tuple[str, int, int]]]: """``(Wh [Cin, sum hidden], bh, Wo [sum hidden, C_pad], bo [C_pad], slices)`` of the merged heads; ``slices`` lists ``(name, start, stop)`` of each head's output columns; columns past ``sum C_h`` are zero.""" if not heads: raise ValueError("no heads") cin = heads[0][1].in_channels wh, bh, rows, cols = [], [], 0, 0 for name, hidden, final in heads: if hidden.kind != "conv" or final.kind != "conv" or hidden.kernel != 1 or final.kernel != 1: raise ValueError(f"head {name!r}: merged heads need 1x1 convs") if hidden.in_channels != cin or final.in_channels != hidden.out_channels: raise ValueError(f"head {name!r}: channel mismatch") if not hidden.relu or final.relu: raise ValueError(f"head {name!r}: expected a ReLU after the hidden conv only") wh.append(np.asarray(hidden.weight, np.float32)[:, :, 0, 0].T) bh.append(np.zeros(hidden.out_channels, np.float32) if hidden.bias is None else np.asarray(hidden.bias, np.float32)) rows += hidden.out_channels cols += final.out_channels cpad = -(-cols // pad_to) * pad_to wo = np.zeros((rows, cpad), np.float32) bo = np.zeros(cpad, np.float32) slices, r, c = [], 0, 0 for name, hidden, final in heads: co = final.out_channels wo[r:r + hidden.out_channels, c:c + co] = np.asarray(final.weight, np.float32)[:, :, 0, 0].T if final.bias is not None: bo[c:c + co] = np.asarray(final.bias, np.float32) slices.append((name, c, c + co)) r += hidden.out_channels c += co return (np.ascontiguousarray(np.concatenate(wh, axis=1)), np.concatenate(bh).astype(np.float32), wo, bo, slices) class CenterHead: """Shared conv (K-split over ``split``) + merged heads; see the module docstring. ``shared``: the 3x3 :class:`ConvSpec` over the concatenated inputs (ReLU after); ``split``: channels of each input, in concat order; ``heads``: ``(name, hidden, final)`` in output order. ``k_split=False`` runs one conv on a device concat instead (A/B only). ``output_dtype``: of the packed head output (``"float32"``). The shared conv's output dtype is its precision's ``activations`` (bf16 by default); ``accumulate_dtype`` (default: the same) is the dtype of the K-split partial sums.""" def __init__(self, shared: ConvSpec, split: Sequence[int], heads: Sequence[HeadSpec], *, policy: Optional[PrecisionPolicy] = None, prefix: str = "head", output_dtype: str = "float32", k_split: bool = True, accumulate_dtype: Optional[str] = None): if shared.kind != "conv" or shared.stride != 1 or not shared.relu: raise ValueError("the shared conv must be a stride-1 conv followed by a ReLU") self.prefix = prefix self.split = [int(s) for s in split] if sum(self.split) != shared.in_channels: raise ValueError(f"split {self.split} does not cover the {shared.in_channels} shared-conv inputs") self.k_split = bool(k_split) and len(self.split) > 1 pad = shared.kernel // 2 if shared.padding is None else int(shared.padding) prec = _precision(policy, f"{prefix}.shared") if self.k_split: self.shared = KSplitConv(shared.weight, shared.bias, self.split, activation="relu", padding=pad, precision=prec, accumulate_dtype=accumulate_dtype or prec.activations, output_dtype=prec.activations, name=f"{prefix}.shared") else: self.shared = Conv2d(shared.weight, shared.bias, padding=pad, activation="relu", precision=prec, output_dtype=prec.activations, name=f"{prefix}.shared") wh, bh, wo, bo, self.slices = merge_heads(heads) self.out_channels = self.slices[-1][2] self.padded_channels = int(wo.shape[1]) hidden_prec = _precision(policy, f"{prefix}.hidden") self.hidden = RowLinear(wh, bh, activation="relu", precision=hidden_prec, output_dtype=hidden_prec.activations, name=f"{prefix}.hidden") self.out = RowLinear(wo, bo, activation=None, precision=_precision(policy, f"{prefix}.out"), output_dtype=output_dtype, name=f"{prefix}.out") def shared_forward(self, inputs: Sequence[FeatureMap]) -> FeatureMap: """``relu(conv3x3(concat(inputs)))`` (inputs are kept), DRAM interleaved.""" import ttnn if [x.channels for x in inputs] != self.split: raise ValueError(f"{self.prefix}: inputs have {[x.channels for x in inputs]} channels, " f"expected {self.split}") if self.k_split: return _to_dram(self.shared(inputs)) cat = inputs[0].with_tensor(ttnn.concat([x.tensor for x in inputs], dim=-1), sum(self.split)) out = _to_dram(self.shared(cat)) free_maps(cat) return out def heads_forward(self, shared: FeatureMap) -> FeatureMap: """Merged heads on the shared map -> ``[1, 1, H*W, C_pad]`` (``channels`` = the padded width).""" hidden = self.hidden(shared.tensor) out = self.out(hidden) free_maps(hidden) return shared.with_tensor(out, self.padded_channels) def __call__(self, inputs: Sequence[FeatureMap]) -> FeatureMap: shared = self.shared_forward(inputs) out = self.heads_forward(shared) free_maps(shared) return out def unpack(self, rows: Any, height: int, width: int) -> Dict[str, np.ndarray]: """Host rows ``[H*W, >= C]`` (a read-back of the head output) -> ``{name: (C_h, H, W)}`` float32.""" a = np.asarray(rows, dtype=np.float32) a = a.reshape(-1, a.shape[-1])[: height * width] return {name: np.ascontiguousarray(a[:, s:e].T.reshape(e - s, height, width)) for name, s, e in self.slices} def release(self) -> None: self.shared.release() self.hidden.release() self.out.release() def describe(self) -> Dict[str, Any]: return {"shared": self.shared.describe(), "k_split": self.k_split, "hidden": self.hidden.describe(), "out": self.out.describe(), "slices": [list(s) for s in self.slices]} def centerhead_numpy(shared: ConvSpec, heads: Sequence[HeadSpec], inputs_nchw: Sequence[Any]) -> Dict[str, np.ndarray]: """fp32 oracle (torch, the literal graph: concat, conv, six heads) -> ``{"shared": NCHW, name: NCHW}``.""" import torch from .second import _apply_torch with torch.no_grad(): x = torch.cat([torch.as_tensor(np.asarray(t, dtype=np.float32)) for t in inputs_nchw], dim=1) s = _apply_torch(shared, x) out = {"shared": s.numpy().copy()} for name, hidden, final in heads: out[name] = _apply_torch(final, _apply_torch(hidden, s)).numpy().copy() return out