| |
| """C22: heatmap peaks of the query-based LiDAR detectors -- sigmoid and the 3x3 local maximum (PLAN.md section 1.2 |
| row C22; probe P6, common/PROBES.md). |
| |
| TransFusion, BEVFusion and PTv3 keep a heatmap cell only where it is the maximum of its 3x3 neighbourhood, with |
| model-specific border and class rules: |
| |
| - ``"transfusion"`` (TF, S:transfusion:169): ``MaxPool 3x3 s1 p0`` padded back with a ZERO border, every class: the |
| border cells are always suppressed; |
| - ``"bevfusion"`` (BF, S:bevfusion:258): classes ``pooled_classes`` (default 0-3) as TF; the other classes are not |
| pooled (``local_max = heat``): every cell passes; |
| - ``"ptv3"`` (PT, S:ptv3:217): ``local_max`` starts as the heat itself and only the interior of the pooled classes is |
| replaced by the pool: border cells and the other classes pass. |
| |
| All three are one exact masked chain (P6, bit-exact on the device, TF 95 / BF 85 / PT 145 us): |
| |
| keep = m_pass + (heat == max_pool2d(heat, 3x3, stride 1, padding 1)) * m_eq out = heat * keep |
| |
| with two constant masks ``m_eq`` / ``m_pass`` (disjoint, values 0 / 1, laid out as the NHWC rows ``[H*W, C]`` of the |
| heatmap). The pool's padding never matters: a padded window only covers border cells, where ``m_eq`` is 0. Not the |
| literal ``p0`` pool + pad form (222 us for TF and two more programs). ``ttnn.max_pool2d`` is bf16 only (FLOAT32 is |
| rejected, ``pool_op.cpp:30``), so the heat is bf16: :func:`sigmoid_heat` takes fp32 logits through an fp32 sigmoid and |
| rounds once. |
| |
| Device layout: the heatmap is a TILE ``[1, 1, H*W, C]`` tensor (what a conv head writes); the result has the same |
| layout (the input of C21's class-major top-k row and of the query-heat gather of C26). |
| |
| numpy only at import; ttnn is imported inside the device functions. |
| """ |
| from __future__ import annotations |
|
|
| from typing import Any, Dict, Optional, Sequence, Tuple |
|
|
| import numpy as np |
|
|
| __all__ = ["RULES", "DEFAULT_POOLED_CLASSES", "local_max_masks", "local_max_numpy", "local_max_literal_numpy", |
| "sigmoid_heat", "sigmoid_heat_numpy", "LocalMax"] |
|
|
| RULES = ("transfusion", "bevfusion", "ptv3") |
| DEFAULT_POOLED_CLASSES = 4 |
|
|
|
|
| def _pooled(classes: int, rule: str, pooled_classes: Optional[Sequence[int]]) -> np.ndarray: |
| if rule not in RULES: |
| raise ValueError(f"rule {rule!r}: expected one of {RULES}") |
| if pooled_classes is None: |
| pooled_classes = range(classes) if rule == "transfusion" else range(min(DEFAULT_POOLED_CLASSES, classes)) |
| p = np.zeros(classes, np.float32) |
| for c in pooled_classes: |
| if not 0 <= int(c) < classes: |
| raise ValueError(f"pooled class {c} outside [0, {classes})") |
| p[int(c)] = 1.0 |
| if rule == "transfusion" and not p.all(): |
| raise ValueError("the transfusion rule pools every class") |
| return p |
|
|
|
|
| def local_max_masks(height: int, width: int, classes: int, *, rule: str = "transfusion", |
| pooled_classes: Optional[Sequence[int]] = None, kernel: int = 3) -> Tuple[np.ndarray, np.ndarray]: |
| """The constant masks ``(m_pass, m_eq)`` of the chain for ``rule`` (module docstring), float32 NHWC rows |
| ``[height * width, classes]``. ``kernel`` (odd): the pool size; the border is ``kernel // 2`` cells wide.""" |
| if kernel < 1 or kernel % 2 == 0: |
| raise ValueError(f"kernel must be odd and positive, got {kernel}") |
| r = kernel // 2 |
| if height <= 2 * r or width <= 2 * r: |
| raise ValueError(f"a {height}x{width} map has no interior for a {kernel}x{kernel} pool") |
| pooled = _pooled(classes, rule, pooled_classes) |
| interior = np.zeros((height, width, 1), np.float32) |
| interior[r:height - r, r:width - r] = 1.0 |
| m_eq = interior * pooled[None, None, :] |
| if rule == "transfusion": |
| m_pass = np.zeros_like(m_eq) |
| elif rule == "bevfusion": |
| m_pass = np.broadcast_to(1.0 - pooled[None, None, :], m_eq.shape).astype(np.float32) |
| else: |
| m_pass = 1.0 - m_eq |
| return (np.ascontiguousarray(m_pass.reshape(height * width, classes)), |
| np.ascontiguousarray(m_eq.reshape(height * width, classes))) |
|
|
|
|
| def _pool_numpy(heat_nhwc: np.ndarray, height: int, width: int, kernel: int) -> np.ndarray: |
| """3x3 (k x k) stride-1 max pool with ``-inf`` padding of NHWC rows ``[H*W, C]`` (exact on any float data).""" |
| c = heat_nhwc.shape[-1] |
| x = np.asarray(heat_nhwc, np.float64).reshape(height, width, c) |
| r = kernel // 2 |
| p = np.full((height + 2 * r, width + 2 * r, c), -np.inf) |
| p[r:r + height, r:r + width] = x |
| out = np.full_like(x, -np.inf) |
| for dy in range(kernel): |
| for dx in range(kernel): |
| np.maximum(out, p[dy:dy + height, dx:dx + width], out=out) |
| return out.reshape(height * width, c) |
|
|
|
|
| def local_max_numpy(heat_nhwc: Any, height: int, width: int, m_pass: Any, m_eq: Any, kernel: int = 3) -> np.ndarray: |
| """Host oracle of the masked chain on NHWC rows ``[H*W, C]`` (float32 out; exact on bf16-valued input).""" |
| h = np.asarray(heat_nhwc, np.float32).reshape(height * width, -1) |
| pooled = _pool_numpy(h, height, width, kernel) |
| keep = np.asarray(m_pass, np.float32) + (h == pooled).astype(np.float32) * np.asarray(m_eq, np.float32) |
| return (h * keep).astype(np.float32) |
|
|
|
|
| def local_max_literal_numpy(heat_nchw: Any, rule: str = "transfusion", *, |
| pooled_classes: Optional[Sequence[int]] = None, kernel: int = 3) -> np.ndarray: |
| """The models' literal forms (``heat * (heat == local_max)``, NCHW ``(C, H, W)`` or ``(1, C, H, W)``), the oracle |
| the masked chain is proven against: TF ``max_pool p0`` + zero border (all classes); BF pooled classes as TF and |
| the others ``local_max = heat``; PT ``local_max = heat`` with the interior of the pooled classes pooled.""" |
| x = np.asarray(heat_nchw, np.float32) |
| squeeze = x.ndim == 4 |
| if squeeze: |
| x = x[0] |
| c, h, w = x.shape |
| pooled = _pooled(c, rule, pooled_classes).astype(bool) |
| r = kernel // 2 |
| full = _pool_numpy(x.reshape(c, h * w).T, h, w, kernel).T.reshape(c, h, w) |
| interior = np.zeros((h, w), bool) |
| interior[r:h - r, r:w - r] = True |
| if rule == "transfusion": |
| local = np.where(interior[None], full, 0.0) |
| elif rule == "bevfusion": |
| local = np.where(pooled[:, None, None], np.where(interior[None], full, 0.0), x) |
| else: |
| local = np.where(pooled[:, None, None] & interior[None], full, x) |
| out = (x * (x == local)).astype(np.float32) |
| return out[None] if squeeze else out |
|
|
|
|
| def sigmoid_heat_numpy(logits: Any, dtype: str = "bfloat16") -> np.ndarray: |
| """Host model of :func:`sigmoid_heat`: float64 sigmoid of the float32 logits, rounded once to ``dtype`` (bf16 RNE).""" |
| s = (1.0 / (1.0 + np.exp(-np.asarray(logits, np.float64)))).astype(np.float32) |
| if dtype in ("bfloat16", "bf16"): |
| from ..tensors import round_to_bf16 |
|
|
| return round_to_bf16(s) |
| return s |
|
|
|
|
| def sigmoid_heat(logits: Any, *, dtype: str = "bfloat16", memory_config: Any = None): |
| """Device ``sigmoid(logits)`` computed in the logits' dtype (fp32 logits: an fp32 sigmoid) and cast once to |
| ``dtype`` (bf16: what ``max_pool2d`` and C21's ``topk_large_indices`` take). Returns ``(heat, sig)``: the cast |
| heat and the sigmoid before the cast (the same tensor when no cast is needed).""" |
| import ttnn |
|
|
| from ..tensors import ttnn_dtype |
|
|
| kw = {} if memory_config is None else {"memory_config": memory_config} |
| sig = ttnn.sigmoid(logits, **kw) |
| want = ttnn_dtype(dtype) |
| heat = sig if sig.dtype == want else ttnn.typecast(sig, want) |
| return heat, sig |
|
|
|
|
| def _mc(spec: Any): |
| import ttnn |
|
|
| if spec is None or spec == "dram": |
| return ttnn.DRAM_MEMORY_CONFIG |
| if spec == "l1": |
| return ttnn.L1_MEMORY_CONFIG |
| return spec |
|
|
|
|
| class LocalMax: |
| """The device chain for one heatmap geometry: ``heat`` (bf16 TILE ``[1, 1, H*W, C]``) -> ``heat * keep``. |
| |
| The masks are uploaded (bf16 TILE, exact 0 / 1) by the first call, outside any capture; with the TF rule |
| (``m_pass`` = 0) the chain is ``max_pool2d``, ``eq``, ``multiply`` (keep = eq * m_eq), ``multiply``; otherwise |
| ``addcmul(m_pass, eq, m_eq)`` replaces the first multiply. ``memory_config``: where the pool and the masks live |
| (``"dram"`` default; ``"l1"`` is P6's 95 us form for maps that fit).""" |
|
|
| def __init__(self, height: int, width: int, classes: int, *, rule: str = "transfusion", |
| pooled_classes: Optional[Sequence[int]] = None, kernel: int = 3, memory_config: Any = "dram", |
| name: str = "local_max"): |
| self.height, self.width, self.classes = int(height), int(width), int(classes) |
| self.rule = rule |
| self.kernel = int(kernel) |
| self.name = name |
| self.m_pass, self.m_eq = local_max_masks(self.height, self.width, self.classes, rule=rule, |
| pooled_classes=pooled_classes, kernel=self.kernel) |
| self.passes = bool(self.m_pass.any()) |
| self.memory_config = memory_config |
| self._dev: Optional[Dict[str, Any]] = None |
|
|
| @property |
| def cells(self) -> int: |
| return self.height * self.width |
|
|
| def _upload(self, device: Any) -> Dict[str, Any]: |
| import torch |
| import ttnn |
|
|
| def up(a: np.ndarray): |
| t = torch.from_numpy(np.ascontiguousarray(a.reshape(1, 1, self.cells, self.classes))) |
| return ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, |
| memory_config=_mc(self.memory_config)) |
|
|
| self._dev = {"m_eq": up(self.m_eq), "m_pass": up(self.m_pass) if self.passes else None} |
| return self._dev |
|
|
| def __call__(self, heat: Any): |
| import ttnn |
|
|
| shape = tuple(int(s) for s in heat.shape) |
| if shape != (1, 1, self.cells, self.classes): |
| raise ValueError(f"{self.name}: heat must be [1, 1, {self.cells}, {self.classes}], got {shape}") |
| if heat.dtype != ttnn.bfloat16: |
| raise ValueError(f"{self.name}: max_pool2d takes bf16 only (got {heat.dtype}): use sigmoid_heat") |
| dev = self._dev or self._upload(heat.device()) |
| r = self.kernel // 2 |
| pooled = ttnn.max_pool2d(input_tensor=heat, batch_size=1, input_h=self.height, input_w=self.width, |
| channels=self.classes, kernel_size=[self.kernel, self.kernel], stride=[1, 1], |
| padding=[r, r], dilation=[1, 1], ceil_mode=False, |
| memory_config=_mc(self.memory_config), output_layout=ttnn.TILE_LAYOUT) |
| if tuple(int(s) for s in pooled.shape) != shape: |
| full = pooled |
| pooled = ttnn.slice(full, [0, 0, 0, 0], list(shape)) |
| ttnn.deallocate(full, False) |
| eq = ttnn.eq(heat, pooled) |
| ttnn.deallocate(pooled, False) |
| if self.passes: |
| keep = ttnn.addcmul(dev["m_pass"], eq, dev["m_eq"]) |
| else: |
| keep = ttnn.multiply(eq, dev["m_eq"]) |
| ttnn.deallocate(eq, False) |
| out = ttnn.multiply(heat, keep) |
| ttnn.deallocate(keep, False) |
| return out |
|
|
| def numpy(self, heat_nhwc: Any) -> np.ndarray: |
| """Host oracle of this chain on NHWC rows (bf16-valued input gives the device result bit for bit).""" |
| return local_max_numpy(heat_nhwc, self.height, self.width, self.m_pass, self.m_eq, self.kernel) |
|
|
| def release(self) -> None: |
| import ttnn |
|
|
| if self._dev is not None: |
| for t in self._dev.values(): |
| if t is not None and t.is_allocated(): |
| ttnn.deallocate(t) |
| self._dev = None |
|
|
| def describe(self) -> Dict[str, Any]: |
| return {"name": self.name, "grid": [self.height, self.width], "classes": self.classes, "rule": self.rule, |
| "kernel": self.kernel, "pooled_classes": [int(c) for c in np.flatnonzero(self.m_eq.max(axis=0))], |
| "passes": self.passes, "memory_config": str(self.memory_config)} |
|
|
