| |
| """C23: ``ttnn.grid_sample`` helpers (PLAN.md section 1.2 row C23; probe P9, common/PROBES.md; first built by the |
| BEVDet port for AlignBEV and FPN_LSS). |
| |
| Rules from probe P9 (``common/probes/g4-attention-gridsample-msda.md``), enforced or encoded here: |
| |
| - **fp32 grids** for every data-dependent grid (bf16 grids cost PCC 0.9988 at 53 px and 0.992 at 128 px); |
| - **input C % 32 == 0** (C = 80 is rejected; zero pad channels stay exactly 0): :func:`padded_channels`, |
| :func:`pad_channels`; |
| - **per-batch grids**: grid N must equal input N (no broadcast over the batch); |
| - **ROW_MAJOR, interleaved (or height-sharded)** input and grid, ``padding_mode="zeros"`` only; |
| - **never a FLOAT32 input together with ``fp32_dest_acc_en``**: wrong, non-deterministic output (strict xfail in P9); |
| :func:`grid_sample` refuses it; |
| - **exact integer coordinates**: the kernel recovers pixel coordinates as ``x = g * W/2 + (W - 1)/2`` in fp32 |
| (``grid_sample_reader_common.hpp``, align_corners=False), so the align_corners=False normalisation |
| ``g = (2x + 1) / W - 1`` returns integer x exactly whenever ``(2x + 1) / W`` is exact in fp32 (``W`` a power of |
| two): an identity warp is then a bit-exact copy (P9: the literal align_corners=True normalisation leaves 2.8-7.6 % |
| of values off by up to 1/32). Zero padding acts in pixel space, so the normalisation choice does not change what |
| is sampled: callers state pixel coordinates (:func:`affine_grid`, :func:`resize_grid`) and always call the op with |
| ``align_corners=False``; |
| - the four bilinear weights are computed in fp32 and **truncated** (not rounded) to bf16 by the kernel: a 0.33-0.6 % |
| rrmse floor for fractional coordinates (plan risk R3; K5 with fp32 weights removes it). :func:`emulate_grid_sample` |
| reproduces that arithmetic on the host (the device test's oracle). |
| |
| Host functions are numpy only (grids are built on the host, per frame or once, and uploaded as fp32 ROW_MAJOR |
| tensors); :func:`grid_sample` is the device call with the checks above. DCNv2 as one grid_sample with K = 9 |
| (``batch_output_channels=True``, BEVFormer / METEOR) and the fused MSDA wrapper are added by their ports. |
| """ |
| from __future__ import annotations |
|
|
| from typing import Any, Optional, Sequence, Tuple |
|
|
| import numpy as np |
|
|
| __all__ = ["OUTSIDE", "padded_channels", "pad_channels", "pixel_to_grid", "kernel_pixels", "affine_pixel_coords", |
| "affine_grid", "resize_pixel_coords", "resize_grid", "emulate_grid_sample", "grid_sample"] |
|
|
| F32 = np.float32 |
| |
| OUTSIDE = F32(-4.0) |
|
|
|
|
| def padded_channels(c: int, multiple: int = 32) -> int: |
| """The channel count ``grid_sample`` accepts for ``c`` real channels (the next multiple of 32).""" |
| return -(-int(c) // int(multiple)) * int(multiple) |
|
|
|
|
| def pad_channels(x: np.ndarray, multiple: int = 32) -> np.ndarray: |
| """Zero-pad the last (channel) dimension of an NHWC array to a multiple of ``multiple`` (no copy when aligned).""" |
| x = np.asarray(x) |
| c = x.shape[-1] |
| cp = padded_channels(c, multiple) |
| if cp == c: |
| return x |
| out = np.zeros(x.shape[:-1] + (cp,), x.dtype) |
| out[..., :c] = x |
| return out |
|
|
|
|
| def pixel_to_grid(px: Any, size: int) -> np.ndarray: |
| """Pixel coordinate (0 = centre of the first pixel) -> normalised coordinate for ``align_corners=False``: |
| ``(2 px + 1) / size - 1`` in float32 (exact for integer ``px`` when ``size`` is a power of two).""" |
| p = np.asarray(px, F32) |
| return (F32(2.0) * p + F32(1.0)) / F32(size) - F32(1.0) |
|
|
|
|
| def kernel_pixels(g: Any, size: int) -> np.ndarray: |
| """The pixel coordinate the device kernel recovers from a normalised coordinate (align_corners=False): |
| ``g * size/2 + (size - 1)/2`` in float32.""" |
| g = np.asarray(g, F32) |
| return g * F32(size / 2.0) + F32((size - 1) / 2.0) |
|
|
|
|
| def affine_pixel_coords(transforms: np.ndarray, out_hw: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]: |
| """Per batch n and output pixel (h, w): ``ix = a w + b h + c``, ``iy = d w + e h + f`` in float32 with one |
| rounding per operation (``transforms`` [N, 6] = (a, b, c, d, e, f), the BEVDet AlignBEV / BEVFormer-style pixel |
| affine) -> (ix, iy) [N, H, W].""" |
| t = np.asarray(transforms, F32).reshape(-1, 6) |
| h, w = out_hw |
| hh, ww = np.meshgrid(np.arange(h, dtype=F32), np.arange(w, dtype=F32), indexing="ij") |
| a, b, c, d, e, f = (t[:, k, None, None] for k in range(6)) |
| ix = (a * ww + b * hh) + c |
| iy = (d * ww + e * hh) + f |
| return ix.astype(F32), iy.astype(F32) |
|
|
|
|
| def affine_grid(transforms: np.ndarray, in_hw: Tuple[int, int], out_hw: Optional[Tuple[int, int]] = None, *, |
| valid: Optional[Sequence[bool]] = None) -> np.ndarray: |
| """fp32 grid [N, H, W, 2] (x, y) sampling an ``in_hw`` input at :func:`affine_pixel_coords` (align_corners=False |
| normalisation: identity transforms are bit-exact copies for power-of-two input sizes). Batches with |
| ``valid[n] == False`` sample :data:`OUTSIDE` (exact zeros, e.g. a history slot that holds no frame yet).""" |
| out_hw = tuple(in_hw) if out_hw is None else tuple(out_hw) |
| ix, iy = affine_pixel_coords(transforms, out_hw) |
| grid = np.stack([pixel_to_grid(ix, in_hw[1]), pixel_to_grid(iy, in_hw[0])], axis=-1).astype(F32) |
| if valid is not None: |
| v = np.asarray(valid, bool).reshape(-1) |
| if v.size != grid.shape[0]: |
| raise ValueError(f"valid has {v.size} entries for {grid.shape[0]} batches") |
| grid[~v] = OUTSIDE |
| return np.ascontiguousarray(grid) |
|
|
|
|
| def resize_pixel_coords(in_size: int, out_size: int, align_corners: bool = True) -> np.ndarray: |
| """Source pixel of every output pixel of a 1-D bilinear resize (float64), torch's ``F.interpolate`` rules: |
| align_corners=True ``i (in - 1) / (out - 1)``; False ``(i + 0.5) in / out - 0.5`` (not clamped: the zero padding |
| of grid_sample differs from torch's edge clamp there, so prefer C18 ``Resize2d`` for align_corners=False).""" |
| i = np.arange(out_size, dtype=np.float64) |
| if align_corners: |
| return i * (in_size - 1) / max(out_size - 1, 1) |
| return (i + 0.5) * in_size / out_size - 0.5 |
|
|
|
|
| def resize_grid(in_hw: Tuple[int, int], out_hw: Tuple[int, int], *, align_corners: bool = True, |
| batch: int = 1) -> np.ndarray: |
| """Static fp32 grid [batch, Ho, Wo, 2] of a bilinear resize (source pixels in float64, normalised with the |
| align_corners=False formula; call :func:`grid_sample`). For align_corners=True every source pixel lies inside |
| the input, so the result equals ``F.interpolate(..., align_corners=True)`` up to the bf16 weight truncation.""" |
| sy = resize_pixel_coords(in_hw[0], out_hw[0], align_corners) |
| sx = resize_pixel_coords(in_hw[1], out_hw[1], align_corners) |
| gx = ((2.0 * sx + 1.0) / in_hw[1] - 1.0).astype(F32) |
| gy = ((2.0 * sy + 1.0) / in_hw[0] - 1.0).astype(F32) |
| grid = np.empty((out_hw[0], out_hw[1], 2), F32) |
| grid[..., 0] = gx[None, :] |
| grid[..., 1] = gy[:, None] |
| return np.ascontiguousarray(np.broadcast_to(grid, (int(batch),) + grid.shape)) |
|
|
|
|
| def _truncate_bf16(a: np.ndarray) -> np.ndarray: |
| bits = np.ascontiguousarray(a, F32).view(np.uint32) & np.uint32(0xFFFF0000) |
| return bits.view(F32) |
|
|
|
|
| def emulate_grid_sample(x: np.ndarray, grid: np.ndarray) -> np.ndarray: |
| """Host emulation of the device's bilinear ``grid_sample`` (align_corners=False, zero padding): pixel recovery |
| as :func:`kernel_pixels`, ``floor``, fp32 fractions and weights truncated to bf16, taps outside the input skipped. |
| ``x`` [N, H, W, C], ``grid`` [N, Ho, Wo, 2] -> [N, Ho, Wo, C] float64 (the device sums the four bf16 products |
| in its own order, so compare with a tolerance; an exact identity / outside sample is exact on both).""" |
| x = np.asarray(x, np.float64) |
| g = np.asarray(grid, F32) |
| n, h, w, _ = x.shape |
| px = kernel_pixels(g[..., 0], w) |
| py = kernel_pixels(g[..., 1], h) |
| x0 = np.floor(px).astype(np.int64) |
| y0 = np.floor(py).astype(np.int64) |
| fx = (px - x0.astype(F32)).astype(F32) |
| fy = (py - y0.astype(F32)).astype(F32) |
| one = F32(1.0) |
| out = np.zeros(g.shape[:-1] + (x.shape[-1],), np.float64) |
| bidx = np.arange(n)[:, None, None] |
| for dy, wy in ((0, one - fy), (1, fy)): |
| for dx, wx in ((0, one - fx), (1, fx)): |
| yy, xx = y0 + dy, x0 + dx |
| ok = (yy >= 0) & (yy < h) & (xx >= 0) & (xx < w) |
| wgt = _truncate_bf16((wy * wx).astype(F32)).astype(np.float64) * ok |
| vals = x[bidx, np.clip(yy, 0, h - 1), np.clip(xx, 0, w - 1)] |
| out += vals * wgt[..., None] |
| return out |
|
|
|
|
| def grid_sample(x: Any, grid: Any, *, compute_kernel_config: Any = None, memory_config: Any = None, |
| batch_output_channels: bool = False): |
| """``ttnn.grid_sample(x, grid)``: bilinear, zero padding, ``align_corners=False`` (pair it with |
| :func:`pixel_to_grid` / :func:`affine_grid` / :func:`resize_grid`), with the P9 rules checked before the call. |
| The output goes to DRAM by default.""" |
| import ttnn |
|
|
| if x.layout != ttnn.ROW_MAJOR_LAYOUT or grid.layout != ttnn.ROW_MAJOR_LAYOUT: |
| raise ValueError("grid_sample needs ROW_MAJOR input and grid") |
| if int(x.padded_shape[-1]) % 32: |
| raise ValueError(f"grid_sample input channels must be a multiple of 32, got {x.padded_shape[-1]} " |
| "(pad them: deform.pad_channels)") |
| if int(x.shape[0]) != int(grid.shape[0]): |
| raise ValueError(f"grid batch {grid.shape[0]} != input batch {x.shape[0]} (expand the grid per batch)") |
| if grid.dtype != ttnn.float32: |
| raise ValueError("use an fp32 grid for data-dependent sampling (probe P9: bf16 grids lose PCC)") |
| if x.dtype == ttnn.float32 and compute_kernel_config is not None and \ |
| getattr(compute_kernel_config, "fp32_dest_acc_en", False): |
| raise ValueError("a FLOAT32 input with fp32_dest_acc_en returns garbage (probe P9 defect)") |
| return ttnn.grid_sample(x, grid, mode="bilinear", padding_mode="zeros", align_corners=False, |
| use_precomputed_grid=False, batch_output_channels=batch_output_channels, |
| memory_config=ttnn.DRAM_MEMORY_CONFIG if memory_config is None else memory_config, |
| compute_kernel_config=compute_kernel_config) |
|
|
