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