code/tt_diffusion_planner/ttaw/ops/upsample.py

# SPDX-License-Identifier: Apache-2.0
"""C18: up-sampling and interpolation helpers for NHWC feature maps (PLAN.md section 1.2 row C18).

- :func:`upsample_nearest`: integer-factor nearest up-sampling through ``ttnn.upsample`` (floor indexing, exactly
  ``F.interpolate(mode="nearest")`` and ONNX Resize nearest / asymmetric / floor: S:yolox section 4.3). ttnn's
  nearest path takes ROW_MAJOR (interleaved or height / block sharded) or a TILE input whose W and C are multiples
  of 32 (``upsample_common.cpp:25-65``) and always writes ROW_MAJOR, so the map goes TILE -> ROW_MAJOR ->
  ``[N, H, W, C]`` -> upsample -> ``[1, 1, N*sH*H*sW*W, C]`` -> TILE (``output_layout``). Bit-exact (a data move).
- :func:`interp_matrix`: the 1-D interpolation matrix ``A [out, in]`` of ``F.interpolate`` (``"nearest"`` floor
  rule, ``"linear"`` with ``align_corners`` True / False, the half-pixel rule clamped at 0 like torch), host numpy
  float64. A 2-D resize is ``Y_c = A_h X_c A_w^T`` per channel (S:frnet:449, S:bevdet:488, S:meteor:497).
- :class:`Resize2d`: that separable resize on the device for any sizes (integer or not), e.g. the
  ``align_corners=True`` bilinear of FRNet / BEVDet's FPN_LSS, which ``ttnn.upsample`` cannot do (its bilinear path
  is half-pixel only, integer scales, sharded bf16 input). The W pass is a matmul over the transposed rows
  (``[N*H, C, W] @ A_w^T``), the H pass ``A_h @ [H, W'*C]`` per image (``kron(I_N, A_h)`` for N > 1). Default
  operands are float32 with HiFi4 + fp32 accumulation, because bf16 cannot hold weights such as 511/1023 (S:frnet:449)
  and a device fp32 matmul is TF32-like (~1.5e-3 relative, probe P12); ``dtype="bfloat16"`` halves the traffic.
  Static non-integer resizes may also use a precomputed ``grid_sample`` grid (probe P9: int16 corner indices exact,
  bf16 weights; 1.7x faster than fp32 grids) through C23.

  **Tail (0.23.1).** The H pass leaves ``[1, 1, N*H', W'*C]``; the map wants ``[1, 1, N*H'*W', C]``. The original
  ``tail="rm"`` goes TILE -> ROW_MAJOR -> ``reshape`` -> TILE. That ROW_MAJOR reshape turns W'*C-wide pages into
  C-wide pages (e.g. 96,000 B -> 192 B): ttnn's dual-kernel RM reshape then writes small DRAM pages from both NOCs,
  which **deadlocks a Blackhole p150 under ETH dispatch** within a few hundred traced replays (Blackhole SYS-1419;
  ``logs/meteor/hangfix/ETH_DISPATCH_HANG_ISSUE.md``). ``tail="tile"`` reshapes the TILE tensor directly (no
  ROW_MAJOR round trip; bit-identical, a data move; ~4-5 ms slower at METEOR's 800 x 48000). ``tail="auto"``
  (default) uses ``tile`` when the RM destination page (C x element size) is < :data:`RESIZE_TAIL_SAFE_PAGE_BYTES`
  and the tt-metal tree lacks ``patches/tt-metal-reshape-rm-sys1419.patch`` (``ttaw.device.reshape_patch_present``;
  0.23.2), else ``rm`` (padded tiles included: bit-identical on the device for C 48 / 144 and odd row counts). With
  the patch the RM reshape takes ttnn's single-kernel path, which is safe and faster (METEOR ETH trace 556.8 ms with
  ``rm`` vs 594.7 ms with ``tile``). Environment override for A/B:
  ``TTAW_RESIZE2D_TAIL=rm|tile|auto`` (read at construction).
"""
from __future__ import annotations

from typing import Any, Dict, Optional, Sequence, Tuple, Union

import numpy as np

from ..precision import Precision
from ..tensors import dtype_name, ttnn_dtype
from .conv import FeatureMap, _layout, _memory_config

__all__ = ["upsample_nearest", "interp_matrix", "resize_matrices", "Resize2d", "RESIZE_TAIL_SAFE_PAGE_BYTES"]

# Smallest ROW_MAJOR destination page (bytes) the Resize2d tail may reshape into on DRAM (module docstring, 0.23.1).
RESIZE_TAIL_SAFE_PAGE_BYTES = 4096
_TAIL_MODES = ("auto", "rm", "tile")


def _pair(value: Union[int, Sequence[int]]) -> Tuple[int, int]:
    if isinstance(value, (int, np.integer)):
        return int(value), int(value)
    a, b = (int(x) for x in value)
    return a, b


def upsample_nearest(fm: FeatureMap, scale: Union[int, Sequence[int]] = 2, *, output_layout: Any = "tile",
                     memory_config: Any = "dram") -> FeatureMap:
    """Nearest up-sampling by integer factors (bit-exact). Returns ``[1, 1, N*sH*H*sW*W, C]`` in
    ``output_layout`` (TILE by default, what a following conv / concat wants)."""
    import ttnn

    sh, sw = _pair(scale)
    if sh < 1 or sw < 1:
        raise ValueError(f"scale must be positive integers, got {scale!r}")
    t = fm.tensor
    if t.layout != ttnn.ROW_MAJOR_LAYOUT:
        t = ttnn.to_layout(t, ttnn.ROW_MAJOR_LAYOUT)
    t = ttnn.reshape(t, (fm.batch, fm.height, fm.width, fm.channels))
    up = ttnn.upsample(t, (sh, sw), mode="nearest", memory_config=_memory_config(memory_config))
    out = ttnn.reshape(up, (1, 1, fm.batch * fm.height * sh * fm.width * sw, fm.channels))
    if _layout(output_layout) == ttnn.TILE_LAYOUT:
        out = ttnn.to_layout(out, ttnn.TILE_LAYOUT)
    return FeatureMap(out, fm.batch, fm.height * sh, fm.width * sw, fm.channels)


def interp_matrix(in_size: int, out_size: int, *, mode: str = "linear", align_corners: bool = False,
                  scale: Optional[float] = None) -> np.ndarray:
    """``A`` (float64, ``[out_size, in_size]``) with ``y = A @ x`` equal to 1-D ``F.interpolate`` of ``x``.

    ``mode``: ``"nearest"`` (``src = min(floor(dst * s), in - 1)``) or ``"linear"`` (``align_corners=True``:
    ``src = dst * (in - 1) / (out - 1)``; False: ``src = max((dst + 0.5) * s - 0.5, 0)``). ``s`` is ``in / out``,
    or ``1 / scale`` when the caller passed ``scale_factor=scale`` to torch (its default
    ``recompute_scale_factor=None`` keeps the given factor)."""
    if in_size < 1 or out_size < 1:
        raise ValueError(f"sizes must be positive, got {in_size} -> {out_size}")
    s = (1.0 / float(scale)) if scale is not None else in_size / out_size
    dst = np.arange(out_size, dtype=np.float64)
    a = np.zeros((out_size, in_size), np.float64)
    rows = np.arange(out_size)
    if mode == "nearest":
        if align_corners:
            raise ValueError("align_corners applies to linear modes only")
        # torch's nearest_idx: floorf(dst * (float)scale) in float32 (the exact 2x case equals dst >> 1)
        prod = (dst.astype(np.float32) * np.float32(s)).astype(np.float32)
        src = np.minimum(np.floor(prod).astype(np.int64), in_size - 1)
        a[rows, src] = 1.0
        return a
    if mode != "linear":
        raise ValueError(f"mode {mode!r}: expected 'nearest' or 'linear'")
    if align_corners:
        src = dst * ((in_size - 1) / (out_size - 1)) if out_size > 1 else np.zeros(out_size)
    else:
        src = np.maximum((dst + 0.5) * s - 0.5, 0.0)
    i0 = np.minimum(np.floor(src).astype(np.int64), in_size - 1)
    i1 = np.minimum(i0 + 1, in_size - 1)
    lam = src - i0
    np.add.at(a, (rows, i0), 1.0 - lam)
    np.add.at(a, (rows, i1), lam)
    return a


def resize_matrices(in_hw: Tuple[int, int], out_hw: Tuple[int, int], *, mode: str = "linear",
                    align_corners: bool = False, scale: Optional[Union[float, Sequence[float]]] = None
                    ) -> Tuple[np.ndarray, np.ndarray]:
    """``(A_h [H', H], A_w [W', W])`` of a 2-D resize: ``Y_c = A_h @ X_c @ A_w.T``."""
    sh, sw = (None, None) if scale is None else (_pair_f(scale))
    return (interp_matrix(in_hw[0], out_hw[0], mode=mode, align_corners=align_corners, scale=sh),
            interp_matrix(in_hw[1], out_hw[1], mode=mode, align_corners=align_corners, scale=sw))


def _pair_f(value: Union[float, Sequence[float]]) -> Tuple[float, float]:
    if isinstance(value, (int, float, np.floating, np.integer)):
        return float(value), float(value)
    a, b = (float(x) for x in value)
    return a, b


class Resize2d:
    """Separable 2-D resize of an NHWC feature map on the device (see the module docstring).

    Args:
        in_hw, out_hw: input / output spatial size. ``channels``, ``batch``: map size (static: one instance per
            call site).
        mode, align_corners, scale: as :func:`interp_matrix`.
        dtype: operand dtype of the two matmuls (``"float32"`` default; the input is typecast, the output cast
            back to ``output_dtype``).
        precision: compute config of the matmuls (``"accurate"``: HiFi4 + fp32 accumulation; never LoFi for fp32
            operands, probe P12).
    Matrices are uploaded on the first call (warm-up), never inside a capture."""

    def __init__(self, in_hw: Tuple[int, int], out_hw: Tuple[int, int], *, channels: int, batch: int = 1,
                 mode: str = "linear", align_corners: bool = False, scale: Optional[Any] = None,
                 dtype: Any = "float32", output_dtype: Any = "bfloat16", precision: Any = "accurate",
                 output_layout: Any = "tile", tail: str = "auto", name: str = "resize"):
        import os

        self.name = name
        tail = (os.environ.get("TTAW_RESIZE2D_TAIL") or tail).strip().lower()
        if tail not in _TAIL_MODES:
            raise ValueError(f"{name}: tail={tail!r}: expected one of {_TAIL_MODES}")
        self.tail_request = tail
        self.in_hw = (int(in_hw[0]), int(in_hw[1]))
        self.out_hw = (int(out_hw[0]), int(out_hw[1]))
        self.channels, self.batch = int(channels), int(batch)
        self.mode, self.align_corners = mode, bool(align_corners)
        self.a_h, self.a_w = resize_matrices(self.in_hw, self.out_hw, mode=mode, align_corners=align_corners,
                                             scale=scale)
        self.dtype = dtype_name(dtype)
        self.output_dtype = dtype_name(output_dtype)
        self.precision = Precision.parse(precision)
        self.output_layout = output_layout
        self._dev: Optional[Tuple[Any, Any]] = None
        self.tail = self._resolve_tail(tail)

    def _resolve_tail(self, tail: str) -> str:
        """``auto`` -> ``tile`` when the ROW_MAJOR tail would write small DRAM pages (module docstring), else ``rm``."""
        if tail != "auto":
            return tail
        from ..device import reshape_patch_present

        elem = 4 if self.output_dtype in ("float32", "uint32", "int32") else 2
        if self.channels * elem >= RESIZE_TAIL_SAFE_PAGE_BYTES:
            return "rm"
        # small pages: safe as ROW_MAJOR only on a tt-metal with the reshape patch (single kernel; ~5x faster than the
        # TILE reshape at METEOR's lift shape), else the TILE tail
        return "rm" if reshape_patch_present(min_page_bytes=RESIZE_TAIL_SAFE_PAGE_BYTES) else "tile"

    def host_matrices(self) -> Tuple[np.ndarray, np.ndarray]:
        """The device operands: ``A_w^T [W, W']`` and ``kron(I_N, A_h) [N*H', N*H]`` (float32)."""
        aw_t = np.ascontiguousarray(self.a_w.T, np.float32)
        ah = self.a_h if self.batch == 1 else np.kron(np.eye(self.batch), self.a_h)
        return aw_t, np.ascontiguousarray(ah, np.float32)

    def reference(self, nchw: np.ndarray) -> np.ndarray:
        """Host float64 result of the same matrices (tests)."""
        x = np.asarray(nchw, np.float64)
        return np.einsum("ph,nchw,qw->ncpq", self.a_h, x, self.a_w)

    def __call__(self, fm: FeatureMap) -> FeatureMap:
        import torch
        import ttnn

        n, (h, w), c = self.batch, self.in_hw, self.channels
        (h2, w2) = self.out_hw
        if (fm.batch, fm.height, fm.width, fm.channels) != (n, h, w, c):
            got = (fm.batch, fm.height, fm.width, fm.channels)
            raise ValueError(f"{self.name}: built for {(n, h, w, c)}, got {got}")
        dt = ttnn_dtype(self.dtype)
        cfg = self.precision.compute_kernel_config()
        if self._dev is None:
            aw_t, ah = self.host_matrices()
            dev = fm.tensor.device()
            self._dev = (ttnn.from_torch(torch.from_numpy(aw_t), dtype=dt, layout=ttnn.TILE_LAYOUT, device=dev),
                         ttnn.from_torch(torch.from_numpy(ah), dtype=dt, layout=ttnn.TILE_LAYOUT, device=dev))
        aw_t_dev, ah_dev = self._dev
        x = fm.tensor
        if x.layout != ttnn.ROW_MAJOR_LAYOUT:
            x = ttnn.to_layout(x, ttnn.ROW_MAJOR_LAYOUT)
        x = ttnn.reshape(x, (1, n * h, w, c))
        x = ttnn.to_layout(x, ttnn.TILE_LAYOUT)
        if x.dtype != dt:
            x = ttnn.typecast(x, dt)
        xt = ttnn.transpose(x, -2, -1)                                   # [1, N*H, C, W]
        y = ttnn.matmul(xt, aw_t_dev, compute_kernel_config=cfg)         # [1, N*H, C, W']
        y = ttnn.transpose(y, -2, -1)                                    # [1, N*H, W', C]
        y = ttnn.to_layout(y, ttnn.ROW_MAJOR_LAYOUT)
        y = ttnn.reshape(y, (1, 1, n * h, w2 * c))
        y = ttnn.to_layout(y, ttnn.TILE_LAYOUT)
        z = ttnn.matmul(ah_dev, y, compute_kernel_config=cfg)            # [1, 1, N*H', W'*C]
        if z.dtype != ttnn_dtype(self.output_dtype):
            z = ttnn.typecast(z, ttnn_dtype(self.output_dtype))          # TILE typecast (every dtype pair)
        if self.tail == "tile":                                          # TILE reshape: no small RM DRAM pages
            z = ttnn.reshape(z, (1, 1, n * h2 * w2, c))
            if _layout(self.output_layout) != ttnn.TILE_LAYOUT:
                z = ttnn.to_layout(z, ttnn.ROW_MAJOR_LAYOUT)
            return FeatureMap(z, n, h2, w2, c)
        z = ttnn.to_layout(z, ttnn.ROW_MAJOR_LAYOUT)
        z = ttnn.reshape(z, (1, 1, n * h2 * w2, c))
        if _layout(self.output_layout) == ttnn.TILE_LAYOUT:
            z = ttnn.to_layout(z, ttnn.TILE_LAYOUT)
        return FeatureMap(z, n, h2, w2, c)

    def release(self) -> None:
        import ttnn

        for t in self._dev or ():
            if t.is_allocated():
                ttnn.deallocate(t)
        self._dev = None

    def describe(self) -> Dict[str, Any]:
        return {"name": self.name, "in_hw": list(self.in_hw), "out_hw": list(self.out_hw), "channels": self.channels,
                "batch": self.batch, "mode": self.mode, "align_corners": self.align_corners, "dtype": self.dtype,
                "tail": self.tail}