code/tt_diffusion_planner/ttaw/ops/topk.py

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