File size: 6,733 Bytes
4d9b003 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | # SPDX-License-Identifier: Apache-2.0
"""C19: row gathers with ``ttnn.embedding`` and the gather form of scatter (PLAN.md section 1.2 row C19; probe P4).
A scatter of ``P`` source rows into an ``N``-row destination (a BEV canvas, a to-dense map) is written as a
gather: the host builds ``index[n]`` = the source row of destination row ``n``, or ``sentinel`` (an explicit
all-zero row appended to the table) where nothing lands. Every shape is static, so it traces, and nothing has to
clear the destination first. Probe P4 (common/PROBES.md) measured it bit-exact at the real shapes, with:
- the table ROW_MAJOR bf16 ``[1, 1, V, D]`` (``D % 32 == 0`` for the fused tilized output), UINT32 indices
``[1, N]`` (ROW_MAJOR);
- **the default call** ``ttnn.embedding(index, table, layout=TILE, padding_idx=sentinel,
embeddings_type=PADDED)``: 0.19-0.54 ms per canvas (CenterPoint 480 x 480: 0.21 ms); GENERIC is 5.4-14.6x slower
on canvases (most indices hit the sentinel);
- **PADDED returns the table's own pad row, not zeros** (S:ptv3:304): the zero row must really be in the table
(:func:`with_zero_rows`).
``ttnn.embedding`` is bf16-only; fp32 row gathers are ``ttnn.gather`` (bit-exact but ~1 us per index, probe P13),
not covered here. Indices must lie in ``[0, V)``: an out-of-range index reads outside the table (and can hang the
chip), so :func:`check_index` validates host index arrays before they are uploaded.
numpy only at import; ttnn is imported inside the device functions.
"""
from __future__ import annotations
from typing import Any, Optional, Tuple
import numpy as np
__all__ = [
"check_index",
"gather_rows",
"gather_rows_numpy",
"scatter_rows",
"scatter_rows_numpy",
"with_zero_rows",
]
def check_index(index: Any, *, num_rows: int, sentinel: Optional[int] = None) -> np.ndarray:
"""Validate a host gather index and return it as a contiguous ``uint32`` array of shape ``[1, N]``.
``num_rows`` is the number of real source rows; valid entries are ``0 <= i < num_rows`` or ``i == sentinel``
(the zero row, conventionally ``num_rows``). Raises ``ValueError`` for anything else (negative, too large,
non-integral, non-finite), because an out-of-range index makes the device read outside the table."""
a = np.asarray(index)
if a.dtype.kind == "f":
if not np.all(np.isfinite(a)) or np.any(a != np.floor(a)):
raise ValueError("gather index holds non-integral or non-finite values")
elif a.dtype.kind not in "iu":
raise ValueError(f"gather index must be integral, got dtype {a.dtype}")
flat = a.reshape(-1).astype(np.int64)
valid = (flat >= 0) & (flat < int(num_rows))
if sentinel is not None:
valid |= flat == int(sentinel)
bad = ~valid
if np.any(bad):
i = int(np.flatnonzero(bad)[0])
raise ValueError(f"gather index [{i}] = {int(flat[i])} is outside [0, {num_rows})"
+ ("" if sentinel is None else f" and is not the sentinel {sentinel}"))
return np.ascontiguousarray(flat.astype(np.uint32).reshape(1, -1))
def gather_rows_numpy(index: Any, table: Any) -> np.ndarray:
"""Host oracle: ``out[n] = table[index[n]]`` (``table`` ``[V, D]``)."""
t = np.asarray(table)
return t.reshape(-1, t.shape[-1])[np.asarray(index).reshape(-1).astype(np.int64)]
def scatter_rows_numpy(rows: Any, index: Any, *, sentinel: int) -> np.ndarray:
"""Host oracle of :func:`scatter_rows`: ``rows`` ``[P, D]`` with a zero row appended at ``sentinel``
(``sentinel >= P``; rows between ``P`` and ``sentinel`` are zero too), gathered by ``index``."""
r = np.asarray(rows, dtype=np.float32)
r = r.reshape(-1, r.shape[-1])
table = np.zeros((max(int(sentinel) + 1, r.shape[0]), r.shape[1]), np.float32)
table[: r.shape[0]] = r
return gather_rows_numpy(index, table)
def _rm(t: Any):
import ttnn
return ttnn.to_layout(t, ttnn.ROW_MAJOR_LAYOUT) if t.layout != ttnn.ROW_MAJOR_LAYOUT else t
def with_zero_rows(rows: Any, count: int = 1):
"""A ROW_MAJOR gather table ``[1, 1, P + count, D]``: the device rows ``rows`` (``[1, 1, P, D]``, TILE or
ROW_MAJOR, bf16) followed by ``count`` zero rows (the sentinel rows of :func:`scatter_rows`). One ``ttnn.pad``
on the device (plus an untilize for a TILE input); traceable."""
import ttnn
if count < 1:
raise ValueError("count must be >= 1")
shape = tuple(int(s) for s in rows.shape)
if len(shape) != 4 or shape[0] != 1 or shape[1] != 1:
raise ValueError(f"rows must be [1, 1, P, D], got {shape}")
rm = _rm(rows)
return ttnn.pad(rm, [(0, 0), (0, 0), (0, int(count)), (0, 0)], 0.0)
def gather_rows(index: Any, table: Any, *, sentinel: Optional[int] = None, layout: str = "tile",
memory_config: Any = None):
"""``out[0, 0, n] = table[0, 0, index[0, n]]`` on the device -> ``[1, 1, N, D]``.
``index``: UINT32 ROW_MAJOR ``[1, N]`` device tensor; ``table``: bf16 ROW_MAJOR ``[1, 1, V, D]`` (a TILE table
is untilized by ttnn first: one extra program). ``sentinel``: the index of an all-zero table row; given, the
call is P4's fast PADDED form (the table must really hold zeros there). ``layout="tile"`` (default) writes
the tilized output directly when ``N % 32 == 0`` and ``D % 32 == 0``."""
import ttnn
lay = ttnn.TILE_LAYOUT if layout == "tile" else ttnn.ROW_MAJOR_LAYOUT
kw = {"layout": lay}
if sentinel is not None:
kw.update(padding_idx=int(sentinel), embeddings_type=ttnn.EmbeddingsType.PADDED)
if memory_config is not None:
kw["memory_config"] = memory_config
out = ttnn.embedding(index, table, **kw) # [1, N, D]
n, d = int(out.shape[-2]), int(out.shape[-1])
return ttnn.reshape(out, (1, 1, n, d))
def scatter_rows(rows: Any, index: Any, *, sentinel: Optional[int] = None, layout: str = "tile",
memory_config: Any = None) -> Tuple[Any, Any]:
"""Gather-form scatter: ``canvas[n] = rows[index[n]]``, or zeros where ``index[n] == sentinel``.
``rows``: ``[1, 1, P, D]`` device rows (TILE or ROW_MAJOR bf16); ``sentinel`` defaults to ``P`` (the zero
row appended by :func:`with_zero_rows`). Returns ``(canvas, table)``: the ``[1, 1, N, D]`` canvas (TILE by
default, e.g. the NHWC input of a conv) and the table, which the caller may deallocate."""
p = int(rows.shape[-2])
sentinel = p if sentinel is None else int(sentinel)
if sentinel < p:
raise ValueError(f"sentinel {sentinel} must be >= the number of source rows {p}")
table = with_zero_rows(rows, sentinel - p + 1)
return gather_rows(index, table, sentinel=sentinel, layout=layout, memory_config=memory_config), table
|