# SPDX-License-Identifier: Apache-2.0 """C21: top-k query selection on the device (PLAN.md section 1.2 row C21; probe P8, common/PROBES.md). The query-based LiDAR detectors (TransFusion, BEVFusion, PTv3) pick the ``K`` best (class, cell) proposals of a local-max heatmap ``[C, H*W]``. On the p150 that is ``ttnn.experimental.topk_large_indices`` (Blackhole only): one bf16 ROW_MAJOR row, ``k`` a multiple of 16 in [16, 2048], uint32 indices in descending value order, ~4.3 ns per element on one core (0.71 / 0.80 / 1.98 ms over 162k / 184k / 459k values, P8). What the helpers here guarantee: - **The row is class-major** (:func:`class_major_row`): index ``c * cells + p`` for class ``c`` and cell ``p`` -- the ONNX ``TopK`` order of a ``[1, C, H*W]`` reshape, so reference (teacher-forcing) indices need no conversion. The NHWC heatmap ``[1, 1, cells, C]`` (TILE) is transposed (one program), reshaped to one row while still tiled, and untilized: no ROW_MAJOR reshape of wide pages, whose L1 staging (2 x the destination page, API.md 13) would not fit PT's 0.9 MB row (TF's is 0.37 MB). - **Exact decode** (:func:`decode_class_major`): the uint32 indices are cast to fp32 (exact below 2**24) and decoded with comparisons only, ``cls = sum_c (idx >= c * cells)``, ``pos = idx - cls * cells`` -- no division, so no device rounding can move a boundary -- clamped into range (a 0xFFFFFFFF sentinel can never become an out-of-range gather index) and cast back to uint32 ROW_MAJOR ``[1, K]`` rows (``ttnn.embedding`` indices). - **Ties are unspecified** (P8: the device does not break ties by lowest index), so every gate after the selection is set-based (:func:`selection_metrics`) or teacher-forced; the host oracle :func:`topk_numpy` has ONNX semantics (stable descending: ties to the lower index). - **Never ``-inf`` masks**: surviving ``-inf`` values come back as index 0xFFFFFFFF (P8); mask with :data:`MASK_VALUE` (-1e30, finite in bf16) instead. - ``k`` is padded to a multiple of 16 (:func:`padded_k`; 500 -> 512); keep the first ``K`` of the descending output. The bf16 ``ttnn.topk`` composite also returns the values (P8: +8-98 us, exact values): :func:`topk_with_values`. An fp32 value re-gather costs ~1 us per index through ``ttnn.gather`` (P13); a two-stage top-k (row-split first stage) is optimization work (PLAN.md 5.3). numpy only at import; ttnn is imported inside the device functions. """ from __future__ import annotations from typing import Any, Dict, List, Optional, Sequence, Tuple import numpy as np __all__ = ["K_MULTIPLE", "K_MIN", "K_MAX", "MASK_VALUE", "SENTINEL", "padded_k", "class_major_row", "topk_indices", "decode_class_major", "topk_with_values", "TopKSelect", "topk_numpy", "class_major_numpy", "decode_class_major_numpy", "selection_metrics", "near_tie_report", "local_max_margins"] K_MULTIPLE = 16 K_MIN = 16 K_MAX = 2048 MASK_VALUE = -1.0e30 # masking value for scores (never -inf: P8 sentinel) SENTINEL = 0xFFFFFFFF # index returned for a surviving -inf FP32_EXACT = 1 << 24 # integers below are exact in fp32 (the decode's range) def padded_k(k: int) -> int: """``k`` rounded up to the multiple of 16 ``topk_large_indices`` takes (500 -> 512); 16 <= result <= 2048.""" k = int(k) if k <= 0: raise ValueError(f"k must be positive, got {k}") kp = max(K_MIN, -(-k // K_MULTIPLE) * K_MULTIPLE) if kp > K_MAX: raise ValueError(f"k={k} exceeds topk_large_indices' maximum {K_MAX}") return kp # ------------------------------------------------------------------------------------------------- device def class_major_row(heat_nhwc: Any, classes: int, cells: int, *, memory_config: Any = None): """NHWC heatmap ``[1, 1, cells, classes]`` (TILE, bf16) -> the class-major bf16 ROW_MAJOR row ``[1, classes * cells]`` of the top-k (transpose, TILE reshape, untilize: traceable).""" import ttnn shape = tuple(int(s) for s in heat_nhwc.shape) if shape != (1, 1, int(cells), int(classes)): raise ValueError(f"heat must be [1, 1, {cells}, {classes}], got {shape}") if heat_nhwc.dtype != ttnn.bfloat16: raise ValueError(f"topk_large_indices takes bf16 scores, got {heat_nhwc.dtype}") kw = {} if memory_config is None else {"memory_config": memory_config} n = int(classes) * int(cells) t = ttnn.transpose(heat_nhwc, 2, 3, **kw) # [1, 1, classes, cells] TILE flat = ttnn.reshape(t, (1, 1, 1, n)) # the row while still tiled (TILE reshape) ttnn.deallocate(t, False) rm = ttnn.to_layout(flat, ttnn.ROW_MAJOR_LAYOUT) # one row, untilized ttnn.deallocate(flat, False) return ttnn.reshape(rm, (1, n)) def topk_indices(row: Any, k: int): """``topk_large_indices`` of a bf16 ROW_MAJOR ``[rows, n]`` tensor -> uint32 ROW_MAJOR ``[rows, padded_k(k)]``, descending (ties unspecified). Refuses an ``n`` beyond fp32's exact integers (the decode's range).""" import ttnn n = int(row.shape[-1]) kp = padded_k(k) if n < kp: raise ValueError(f"the row has {n} values, fewer than k={kp}") if n > FP32_EXACT: raise ValueError(f"row of {n} values: indices above 2**24 are not exact in the fp32 decode") if row.dtype != ttnn.bfloat16 or row.layout != ttnn.ROW_MAJOR_LAYOUT: raise ValueError("topk_large_indices takes a bf16 ROW_MAJOR row") return ttnn.experimental.topk_large_indices(row, k=kp) def decode_class_major(indices: Any, cells: int, classes: int, *, clamp: bool = True) -> Tuple[Any, Any]: """Class-major flat indices (uint32, ROW_MAJOR ``[1, K]`` or TILE) -> ``(pos, cls)`` uint32 ROW_MAJOR ``[1, K]`` with ``idx = cls * cells + pos`` (module docstring: fp32 comparisons, exact; ``clamp`` keeps them in range).""" import ttnn cells, classes = int(cells), int(classes) if cells * classes > FP32_EXACT: raise ValueError("cells * classes exceeds fp32's exact integers") it = indices if indices.layout == ttnn.TILE_LAYOUT else ttnn.to_layout(indices, ttnn.TILE_LAYOUT) f = ttnn.typecast(it, ttnn.float32) if it is not indices: ttnn.deallocate(it, False) cls = None for c in range(1, classes): ge = ttnn.ge(f, float(c * cells)) if cls is None: cls = ge else: s = ttnn.add(cls, ge) ttnn.deallocate(cls, False) ttnn.deallocate(ge, False) cls = s if cls is None: # one class: cls = 0, pos = idx pos = f cls = ttnn.multiply(f, 0.0) else: off = ttnn.multiply(cls, float(cells)) pos = ttnn.subtract(f, off) ttnn.deallocate(off, False) ttnn.deallocate(f, False) if clamp: p2 = ttnn.clamp(pos, 0.0, float(cells - 1)) c2 = ttnn.clamp(cls, 0.0, float(classes - 1)) ttnn.deallocate(pos, False) ttnn.deallocate(cls, False) pos, cls = p2, c2 out = [] for t in (pos, cls): u = ttnn.typecast(t, ttnn.uint32) ttnn.deallocate(t, False) r = ttnn.to_layout(u, ttnn.ROW_MAJOR_LAYOUT) ttnn.deallocate(u, False) out.append(r) return out[0], out[1] def topk_with_values(scores_tile: Any, k: int): """The bf16 ``ttnn.topk`` composite on a TILE ``[.., n]`` row: ``(values bf16, indices)`` of the top ``k`` (P8: exact values, in-range indices even with ``-inf``; ties unspecified).""" import ttnn return ttnn.topk(scores_tile, k=int(k), dim=-1, largest=True, sorted=True) class TopKSelect: """Top-``k`` of a ``[1, 1, cells, classes]`` bf16 TILE heatmap -> uint32 class-major indices ``[1, padded_k]`` (descending) plus the decoded ``(pos, cls)`` rows: :func:`class_major_row`, :func:`topk_indices`, :func:`decode_class_major`. ``keep`` (default ``k``) proposals are meaningful; the rows hold ``padded_k(k)``.""" def __init__(self, cells: int, classes: int, k: int, *, name: str = "topk"): self.cells, self.classes, self.k = int(cells), int(classes), int(k) self.kp = padded_k(self.k) if self.cells * self.classes < self.kp: raise ValueError(f"{name}: {self.cells * self.classes} scores for k={self.kp}") self.name = name def indices(self, heat_nhwc: Any): row = class_major_row(heat_nhwc, self.classes, self.cells) idx = topk_indices(row, self.kp) import ttnn ttnn.deallocate(row, False) return idx def decode(self, indices: Any) -> Tuple[Any, Any]: return decode_class_major(indices, self.cells, self.classes) def __call__(self, heat_nhwc: Any) -> Tuple[Any, Any, Any]: """-> ``(indices, pos, cls)``, each uint32 ROW_MAJOR ``[1, padded_k]``.""" idx = self.indices(heat_nhwc) pos, cls = self.decode(idx) return idx, pos, cls def describe(self) -> Dict[str, Any]: return {"name": self.name, "cells": self.cells, "classes": self.classes, "k": self.k, "padded_k": self.kp, "order": "class-major (idx = cls * cells + pos)", "ties": "unspecified (P8)"} # --------------------------------------------------------------------------------------------------- host def topk_numpy(scores: Any, k: int) -> Tuple[np.ndarray, np.ndarray]: """ONNX ``TopK(largest=1, sorted=1)`` of a 1-D score vector: ``(values, indices)``, values descending, ties in ascending index order (a stable descending sort).""" s = np.asarray(scores).reshape(-1) order = np.argsort(-s.astype(np.float64), kind="stable")[: int(k)] return s[order], order.astype(np.int64) def class_major_numpy(heat_nhwc: Any) -> np.ndarray: """NHWC rows ``[cells, classes]`` -> the class-major flat vector of :func:`class_major_row`.""" h = np.asarray(heat_nhwc) return np.ascontiguousarray(h.reshape(-1, h.shape[-1]).T).reshape(-1) def decode_class_major_numpy(indices: Any, cells: int) -> Tuple[np.ndarray, np.ndarray]: """Host decode: ``(pos, cls)`` int64 of class-major flat indices.""" i = np.asarray(indices).astype(np.int64).reshape(-1) return i % int(cells), i // int(cells) def selection_metrics(device_indices: Any, reference_indices: Any, *, reference_scores: Optional[Any] = None, flat_scores: Optional[Any] = None, strong: float = 0.1, k: Optional[int] = None ) -> Dict[str, Any]: """Set-based agreement of a device top-k with the reference's (PLAN.md 2.5 / 0.3 item 1). ``device_indices`` / ``reference_indices``: flat indices (the first ``k`` of each are compared, default the reference length). ``reference_scores``: the reference values of ``reference_indices`` (same order) or ``flat_scores``: the reference score of every flat index. Returns ``shared`` (count) / ``shared_frac``, and with scores the strong proposals (reference score >= ``strong``): ``strong_total``, ``strong_kept``, ``strong_overlap`` (fraction; 1.0 when there is none), ``strong_missed`` (indices).""" d = np.asarray(device_indices).astype(np.int64).reshape(-1) r = np.asarray(reference_indices).astype(np.int64).reshape(-1) k = len(r) if k is None else int(k) d, r = d[:k], r[:k] ds = set(d.tolist()) out: Dict[str, Any] = {"k": k, "shared": len(ds & set(r.tolist())), "unique": len(ds) == len(d)} out["device_only"] = sorted(ds - set(r.tolist())) out["shared_frac"] = out["shared"] / max(k, 1) if reference_scores is not None or flat_scores is not None: if reference_scores is not None: vals = np.asarray(reference_scores, np.float64).reshape(-1)[:k] strong_set = set(r[vals >= strong].tolist()) else: fs = np.asarray(flat_scores, np.float64).reshape(-1) strong_set = set(np.flatnonzero(fs >= strong).tolist()) missed = sorted(strong_set - ds) out.update({"strong_threshold": float(strong), "strong_total": len(strong_set), "strong_kept": len(strong_set) - len(missed), "strong_overlap": 1.0 if not strong_set else (len(strong_set) - len(missed)) / len(strong_set), "strong_missed": missed}) return out def near_tie_report(missed: Sequence[int], device_indices: Any, reference_heat: Any, height: int, width: int, *, kernel: int = 3) -> List[Dict[str, Any]]: """Why a reference proposal is missing from a device selection (diagnostics of :func:`selection_metrics`). ``reference_heat``: the reference's class-major sigmoid heat BEFORE the local max (``[C * H * W]`` or ``[C, H, W]``). For each missed class-major index: its heat, the strongest same-class cell of its ``kernel x kernel`` window that the DEVICE selected, and the relative gap ``(heat_missed - heat_neighbour) / heat_missed`` (a few 0.1 % = a local-max near-tie: the neighbour won the device's 3x3 comparison; a bf16 heat holds 3 significant digits). ``neighbour`` is None when no same-class neighbour was selected.""" heat = np.asarray(reference_heat, np.float64).reshape(-1) cells = int(height) * int(width) sel = set(np.asarray(device_indices).astype(np.int64).reshape(-1).tolist()) r = int(kernel) // 2 out: List[Dict[str, Any]] = [] for idx in missed: c, p = divmod(int(idx), cells) y, x = divmod(p, int(width)) best = None for dy in range(-r, r + 1): for dx in range(-r, r + 1): yy, xx = y + dy, x + dx if (dy or dx) and 0 <= yy < height and 0 <= xx < width: q = c * cells + yy * int(width) + xx if q in sel and (best is None or heat[q] > heat[best]): best = q rec: Dict[str, Any] = {"index": int(idx), "class": c, "cell": [y, x], "heat": float(heat[idx]), "neighbour": None if best is None else int(best)} if best is not None: rec["neighbour_heat"] = float(heat[best]) rec["relative_gap"] = float((heat[idx] - heat[best]) / max(heat[idx], 1e-30)) out.append(rec) return out def local_max_margins(indices: Any, heat: Any, height: int, width: int, *, kernel: int = 3) -> np.ndarray: """How decisive the reference's local-max decision is for each class-major proposal (the selection metric of a "decisive" strong-proposal gate): ``(heat[p] - max_neighbour) / heat[p]`` over its same-class ``kernel x kernel`` neighbours in ``heat`` (class-major ``[C * H * W]`` or ``[C, H, W]``, BEFORE the local max). A margin of at most a few bf16 steps (0.39 % each) is a near-tie that a bf16 device heatmap cannot resolve; 1.0 for a cell without neighbours. TransFusion gates the strong proposals with a margin above 0.5 % (``strong_overlap_decisive``) and reports the rest with :func:`near_tie_report`.""" hv = np.asarray(heat, np.float64).reshape(-1) h, w = int(height), int(width) cells = h * w r = int(kernel) // 2 out = [] for idx in np.asarray(indices).astype(np.int64).reshape(-1): c, p = divmod(int(idx), cells) y, x = divmod(p, w) nb = [hv[c * cells + yy * w + xx] for yy in range(y - r, y + r + 1) for xx in range(x - r, x + r + 1) if (yy, xx) != (y, x) and 0 <= yy < h and 0 <= xx < w] out.append((hv[idx] - max(nb)) / max(hv[idx], 1e-30) if nb else 1.0) return np.asarray(out, np.float64)