| |
| """K1 ``segment_reduce``: per-segment reductions over rows sorted by segment (PLAN.md section 1.3 row K1; owned by the |
| FRNet port), with the exact stock-op fallback (the log-step segmented max) and their host twins. |
| |
| A *segmented* tensor holds rows sorted by segment id: segment ``s`` is the contiguous row range |
| ``[offsets[s], offsets[s + 1])`` (CSR offsets, non-decreasing; rows past ``offsets[-1]`` are padding that nothing |
| reads). FRNet's six ``ScatterElements(max)`` sites (points sorted by frustum cell, PLAN.md 2.8), PTv3 pooling and |
| BEVPool are of this form once the host sorts the rows. |
| |
| - **K1 on the device** (:class:`SegmentReduce`, kernel ``kernels/segment_reduce_dm.cpp``): ``out[s] = |
| max(x[offsets[s]:offsets[s + 1]])`` for ``s < num_segments``, empty segments -> ``empty_value`` (FRNet: the zero |
| row of the frustum2pixel gather table comes out of the kernel itself). Input ``[1, 1, R, C]`` TILE or ROW_MAJOR, |
| bf16 or fp32 (``R % 32 == 0``, ``C % 32 == 0``); offsets uint32 ``[1, K]`` ROW_MAJOR (``K >= num_segments + 1``, |
| an RT-dev tensor: rewritten per frame, no recapture); output ``[1, 1, num_segments, C]`` ROW_MAJOR, bf16 (round to |
| nearest even) or fp32 -- the layout ``ttnn.embedding`` takes as a table. **Exact**: the maximum is taken on |
| integer keys of the float bit patterns (signed order == float order, -0 < +0), so the result is bit-exact for any |
| input dtype, and fp32 -> bf16 output equals rounding first (a monotonic rounding commutes with max). One program, |
| common runtime args only (probe P3 rule), deterministic, no atomics, traceable. Work split: 2 data-movement |
| RISC-Vs per core, each a contiguous run of segments balanced on rows + 1 per segment, found on the device by binary |
| search over the offsets (no per-frame host split, nothing grid-dependent in the host tables). |
| - **max only on the device for now.** ``sum`` / ``mean`` need float adds, which the data-movement RISC-Vs do not have |
| (no F extension); they are host-oracle only here (:func:`segment_reduce_numpy`) and come with their consumers |
| (PTv3 / BEVPool, optimization phase). |
| - **Stock-op fallback** (:func:`logstep_segment_max`, PLAN.md 1.3): ``R`` Hillis-Steele rounds ``x = max(x, x[idx_r])`` |
| with host shift tables ``idx_r[i] = i - 2^r`` when row ``i - 2^r`` is in the same segment, else ``i`` (max is |
| idempotent), then one gather of each segment's last row: exact for segments of up to ``2^R`` rows, ``3R + 1`` |
| programs (untilize, ``ttnn.embedding``, ``ttnn.maximum`` per round). ``ttnn.embedding`` is bf16-only, so the |
| fallback is bf16 in and out. ``R`` is COMPILE (it sets the op count): bucket it per frame (:func:`segment_rounds`). |
| |
| Host side (numpy only): :func:`segment_offsets` (CSR offsets from sorted segment ids), :func:`segment_rounds`, |
| :func:`segment_shift_tables`, :func:`segment_last_rows`, oracles :func:`segment_reduce_numpy` and |
| :func:`logstep_segment_max_numpy`, and :func:`worker_segments` (the kernel's work split, for tests and diagnostics). |
| numpy only at import; ttnn is imported inside the device functions. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Any, Dict, List, Optional, Sequence, Tuple |
|
|
| import numpy as np |
|
|
| from . import kernel_path |
|
|
| __all__ = [ |
| "OPS", |
| "DEVICE_OPS", |
| "segment_offsets", |
| "check_offsets", |
| "segment_rounds", |
| "segment_shift_tables", |
| "segment_last_rows", |
| "segment_reduce_numpy", |
| "logstep_segment_max_numpy", |
| "worker_segments", |
| "SegmentReduceSpec", |
| "SegmentReduce", |
| "segment_reduce", |
| "logstep_segment_max", |
| "KERNEL", |
| ] |
|
|
| OPS = ("max", "min", "sum", "mean") |
| DEVICE_OPS = ("max",) |
| KERNEL = "segment_reduce_dm.cpp" |
| UNIT_ROWS = 32 |
| _OFF_CHUNK = 256 |
| _OUT_RING = 8 |
|
|
|
|
| |
|
|
| def segment_offsets(seg_id: Any, num_segments: int, *, num_rows: Optional[int] = None) -> np.ndarray: |
| """CSR offsets ``uint32 [num_segments + 1]`` of rows sorted by segment id: segment ``s`` = rows |
| ``[off[s], off[s + 1])``. ``seg_id`` (``int [R]``, non-decreasing) may contain ids ``>= num_segments`` (padding |
| rows, e.g. FRNet's trailing sentinel segment): they must come last and belong to no segment. ``num_rows`` |
| (default ``len(seg_id)``) bounds the offsets.""" |
| ids = np.asarray(seg_id).reshape(-1).astype(np.int64) |
| if ids.size and (np.any(np.diff(ids) < 0) or ids[0] < 0): |
| raise ValueError("segment ids must be non-negative and sorted (non-decreasing)") |
| counts = np.bincount(ids[ids < num_segments], minlength=num_segments).astype(np.int64) |
| off = np.zeros(num_segments + 1, np.int64) |
| np.cumsum(counts, out=off[1:]) |
| rows = len(ids) if num_rows is None else int(num_rows) |
| if off[-1] > rows: |
| raise ValueError(f"segments hold {off[-1]} rows, more than num_rows={rows}") |
| return off.astype(np.uint32) |
|
|
|
|
| def check_offsets(offsets: Any, *, num_segments: int, num_rows: int, length: Optional[int] = None) -> np.ndarray: |
| """Validate CSR offsets for the device and return them as a contiguous ``uint32 [1, length]`` row (``length`` |
| defaults to ``num_segments + 1``; extra entries repeat the last offset). Raises ``ValueError`` unless they start |
| at 0, never decrease and end at most at ``num_rows`` (the kernel clamps anyway; a bad table is a host bug).""" |
| a = np.asarray(offsets).reshape(-1) |
| if a.dtype.kind not in "iu": |
| raise ValueError(f"offsets must be integral, got dtype {a.dtype}") |
| a = a.astype(np.int64) |
| if a.size < num_segments + 1: |
| raise ValueError(f"{a.size} offsets for {num_segments} segments (need num_segments + 1)") |
| a = a[: num_segments + 1] |
| if a[0] != 0 or np.any(np.diff(a) < 0) or a[-1] > num_rows: |
| raise ValueError(f"offsets must start at 0, be non-decreasing and end <= num_rows={num_rows}") |
| n = num_segments + 1 if length is None else int(length) |
| if n < num_segments + 1: |
| raise ValueError(f"length {n} < num_segments + 1") |
| out = np.full((1, n), a[-1], np.uint32) |
| out[0, : a.size] = a |
| return out |
|
|
|
|
| def segment_rounds(max_segment_length: int) -> int: |
| """Log-step rounds that cover segments of up to ``max_segment_length`` rows: ``ceil(log2(L))`` (0 for L <= 1).""" |
| L = int(max_segment_length) |
| return int(np.ceil(np.log2(L))) if L > 1 else 0 |
|
|
|
|
| def segment_shift_tables(seg_id: Any, rounds: int) -> List[np.ndarray]: |
| """The gather tables of the log-step segmented max: round ``r`` uses ``idx[i] = i - 2^r`` if row ``i - 2^r`` has |
| the same segment id, else ``i`` -> ``rounds`` x ``uint32 [R]``. After ``rounds`` rounds the last row of every |
| segment of at most ``2^rounds`` rows holds the segment max (padding rows only ever see padding rows).""" |
| seg = np.asarray(seg_id).reshape(-1) |
| i = np.arange(len(seg), dtype=np.int64) |
| out = [] |
| for r in range(int(rounds)): |
| d = 1 << r |
| same = np.zeros(len(seg), bool) |
| same[d:] = seg[d:] == seg[:-d] |
| out.append(np.where(same, i - d, i).astype(np.uint32)) |
| return out |
|
|
|
|
| def segment_last_rows(offsets: Any, *, empty_row: int) -> np.ndarray: |
| """``uint32 [S]``: the last row of each segment of ``offsets`` (``S + 1`` entries), ``empty_row`` for an empty |
| segment. For :func:`logstep_segment_max` pass the number of input rows: the fallback appends one zero row there |
| (padding rows themselves are reduced too, so they do not stay zero).""" |
| off = np.asarray(offsets).reshape(-1).astype(np.int64) |
| last = off[1:] - 1 |
| return np.where(off[1:] > off[:-1], last, int(empty_row)).astype(np.uint32) |
|
|
|
|
| |
|
|
| def _float_keys(a: np.ndarray) -> np.ndarray: |
| """float32 -> int32 keys whose signed order is the float order with -0 < +0 (an involution on the bits).""" |
| b = np.ascontiguousarray(a, np.float32).view(np.int32) |
| return b ^ ((b >> 31) & np.int32(0x7FFFFFFF)) |
|
|
|
|
| def _from_keys(k: np.ndarray) -> np.ndarray: |
| return (k ^ ((k >> 31) & np.int32(0x7FFFFFFF))).astype(np.int32).view(np.float32) |
|
|
|
|
| def segment_reduce_numpy(x: Any, offsets: Any, num_segments: Optional[int] = None, *, op: str = "max", |
| empty_value: float = 0.0) -> np.ndarray: |
| """Host oracle: ``out[s] = op(x[off[s]:off[s + 1]], axis=0)`` (``x`` ``[R, C]`` or ``[..., R, C]``), empty |
| segments -> ``empty_value``. ``max`` / ``min`` keep the input dtype (exact) and, for float32 / bf16-valued |
| inputs, order the floats as the K1 kernel does (signed zeros: ``-0 < +0``; numpy's ``maximum`` would keep |
| whichever zero comes first); ``sum`` / ``mean`` accumulate in float64 and return float32.""" |
| if op not in OPS: |
| raise ValueError(f"op {op!r}: one of {OPS}") |
| a = np.asarray(x) |
| a = a.reshape(-1, a.shape[-1]) |
| off = np.asarray(offsets).reshape(-1).astype(np.int64) |
| S = len(off) - 1 if num_segments is None else int(num_segments) |
| off = off[: S + 1] |
| out_dtype = a.dtype if op in ("max", "min") else np.float32 |
| out = np.full((S, a.shape[1]), empty_value, dtype=out_dtype) |
| nonempty = np.flatnonzero(off[1:] > off[:-1]) |
| if nonempty.size: |
| starts = off[:-1][nonempty] |
| if op in ("max", "min") and a.dtype == np.float32: |
| keys = _float_keys(a[: off[-1]]) |
| red = _from_keys((np.maximum if op == "max" else np.minimum).reduceat(keys, starts, axis=0)) |
| elif op in ("max", "min"): |
| red = (np.maximum if op == "max" else np.minimum).reduceat(a[: off[-1]], starts, axis=0) |
| else: |
| red = np.add.reduceat(a[: off[-1]].astype(np.float64), starts, axis=0) |
| if op == "mean": |
| red = red / (off[1:][nonempty] - starts)[:, None] |
| out[nonempty] = red.astype(out_dtype) |
| return out |
|
|
|
|
| def logstep_segment_max_numpy(x: Any, shift_tables: Sequence[Any], last_rows: Any) -> np.ndarray: |
| """Host twin of :func:`logstep_segment_max` (any dtype; exact): the rounds, then a gather of ``last_rows`` from |
| the result with one zero row appended (index ``R``: the empty segments).""" |
| a = np.asarray(x) |
| a = a.reshape(-1, a.shape[-1]) |
| for idx in shift_tables: |
| a = np.maximum(a, a[np.asarray(idx).reshape(-1).astype(np.int64)]) |
| a = np.concatenate([a, np.zeros((1, a.shape[1]), a.dtype)], axis=0) |
| return a[np.asarray(last_rows).reshape(-1).astype(np.int64)] |
|
|
|
|
| def worker_segments(offsets: Any, num_segments: int, num_workers: int, *, num_rows: Optional[int] = None |
| ) -> List[Tuple[int, int]]: |
| """The kernel's work split (host twin of ``segment_reduce_dm.cpp``): worker ``w`` owns the segments whose cost |
| ``off[s] + s`` lies in ``[w T / W, (w + 1) T / W)`` with ``T = off[S] + S`` -> ``[(s_lo, s_hi)]`` per worker. |
| Offsets are clamped to ``num_rows`` as the kernel clamps them.""" |
| off = np.asarray(offsets).reshape(-1).astype(np.int64)[: num_segments + 1] |
| if num_rows is not None: |
| off = np.minimum(off, int(num_rows)) |
| cost = off + np.arange(num_segments + 1) |
| total = int(cost[-1]) |
| out = [] |
| for w in range(num_workers): |
| t0 = (w * total) // num_workers |
| t1 = ((w + 1) * total) // num_workers |
| lo = 0 if w == 0 else int(np.searchsorted(cost, t0, side="left")) |
| hi = num_segments if w + 1 == num_workers else int(np.searchsorted(cost, t1, side="left")) |
| out.append((lo, max(lo, hi))) |
| return out |
|
|
|
|
| |
|
|
| def _f32_bits(value: float) -> int: |
| return int(np.array([value], np.float32).view(np.uint32)[0]) |
|
|
|
|
| @dataclass(frozen=True) |
| class SegmentReduceSpec: |
| """The static shape of one K1 call site (COMPILE): ``num_rows`` x ``channels`` input, ``num_segments`` outputs, |
| dtypes / layout, the empty value, and how many workers (data-movement RISC-Vs) per core run it.""" |
|
|
| num_rows: int |
| channels: int |
| num_segments: int |
| input_dtype: str = "float32" |
| input_layout: str = "tile" |
| output_dtype: str = "float32" |
| op: str = "max" |
| empty_value: float = 0.0 |
| workers_per_core: int = 2 |
|
|
| def __post_init__(self) -> None: |
| if self.op not in DEVICE_OPS: |
| raise NotImplementedError(f"K1 runs {DEVICE_OPS} on the device; {self.op!r} is host-oracle only " |
| "(segment_reduce_numpy): sum / mean need float adds the data-movement RISC-Vs " |
| "do not have") |
| if self.num_rows <= 0 or self.num_rows % UNIT_ROWS: |
| raise ValueError(f"num_rows must be a positive multiple of {UNIT_ROWS}, got {self.num_rows}") |
| if self.channels <= 0 or self.channels % 32: |
| raise ValueError(f"channels must be a positive multiple of 32, got {self.channels}") |
| if self.num_segments <= 0: |
| raise ValueError("num_segments must be positive") |
| for name, value in (("input_dtype", self.input_dtype), ("output_dtype", self.output_dtype)): |
| if value not in ("float32", "bfloat16"): |
| raise ValueError(f"{name} {value!r}: 'float32' or 'bfloat16'") |
| if self.input_layout not in ("tile", "row_major"): |
| raise ValueError(f"input_layout {self.input_layout!r}: 'tile' or 'row_major'") |
| if self.workers_per_core not in (1, 2): |
| raise ValueError("workers_per_core must be 1 or 2") |
| if not np.isfinite(self.empty_value): |
| raise ValueError("empty_value must be finite") |
|
|
| @property |
| def in_bytes(self) -> int: |
| return 4 if self.input_dtype == "float32" else 2 |
|
|
| @property |
| def out_bytes(self) -> int: |
| return 4 if self.output_dtype == "float32" else 2 |
|
|
| @property |
| def in_page_bytes(self) -> int: |
| return (1024 if self.input_layout == "tile" else self.channels) * self.in_bytes |
|
|
| @property |
| def unit_bytes(self) -> int: |
| pages = self.channels // 32 if self.input_layout == "tile" else UNIT_ROWS |
| return pages * self.in_page_bytes |
|
|
| @property |
| def out_row_bytes(self) -> int: |
| return self.channels * self.out_bytes |
|
|
| def l1_bytes_per_worker(self) -> Dict[str, int]: |
| """The L1 scratch (CB) sizes of one worker: 2 data units, offsets window, probe block, accumulator, output |
| ring (each + 64 B so the kernel can align to the 64-byte DRAM read alignment).""" |
| return {"data": 2 * self.unit_bytes + 64, "offsets": _OFF_CHUNK * 4 + 64, "probe": 128, |
| "acc": self.channels * 4, "out": _OUT_RING * self.out_row_bytes + 64} |
|
|
| def check_l1(self, budget: int = 600 * 1024) -> None: |
| """Refuse configurations whose per-core scratch exceeds ``budget`` (default 600 KiB of ~1.4 MiB).""" |
| per_core = self.workers_per_core * sum(self.l1_bytes_per_worker().values()) |
| if per_core > budget: |
| raise ValueError(f"K1 scratch {per_core} B per core exceeds {budget} B") |
|
|
|
|
| class SegmentReduce: |
| """K1 for one call site (static :class:`SegmentReduceSpec`). ``__call__(x, offsets, *, output=None)`` runs one |
| ``ttnn.generic_op`` (eagerly or inside a trace capture) and returns the ``[1, 1, num_segments, C]`` ROW_MAJOR |
| DRAM output (allocated per call unless ``output`` is given: a persistent buffer of that spec). |
| |
| ``x``: ``[1, 1, num_rows, C]`` interleaved DRAM tensor of the spec's dtype / layout, rows sorted by segment; |
| ``offsets``: uint32 ROW_MAJOR ``[1, K]`` DRAM (``K >= num_segments + 1``; :func:`check_offsets`), e.g. a |
| ``TraceRunner`` input rewritten per frame. The core grid is read from the device (never hard-coded); the split |
| across cores happens on the device, so ETH (12x10) and WORKER (11x10) runs give identical outputs.""" |
|
|
| def __init__(self, spec: SegmentReduceSpec, *, name: str = "segment_reduce", grid: Optional[Tuple[int, int]] = None): |
| self.spec = spec |
| self.name = name |
| self.grid = grid |
| spec.check_l1() |
|
|
| |
| def _grid(self, device: Any) -> Tuple[int, int]: |
| if self.grid is not None: |
| return int(self.grid[0]), int(self.grid[1]) |
| g = device.compute_with_storage_grid_size() |
| return int(g.x), int(g.y) |
|
|
| def program(self, x: Any, offsets: Any, out: Any): |
| """The ``ttnn.ProgramDescriptor`` of one call (buffer addresses in common runtime args).""" |
| import ttnn |
|
|
| spec = self.spec |
| gx, gy = self._grid(x.device()) |
| rect = ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1)) |
| crs = ttnn.CoreRangeSet([rect]) |
| slots = spec.workers_per_core |
| workers = gx * gy * slots |
| off_len = int(offsets.shape[-1]) |
| sizes = spec.l1_bytes_per_worker() |
| cbs, kernels = [], [] |
| configs = [ttnn.WriterConfigDescriptor(), ttnn.ReaderConfigDescriptor()][:slots] |
| accessors = (list(ttnn.TensorAccessorArgs(x).get_compile_time_args()) |
| + list(ttnn.TensorAccessorArgs(offsets).get_compile_time_args()) |
| + list(ttnn.TensorAccessorArgs(out).get_compile_time_args())) |
| for slot, config in enumerate(configs): |
| ids = [5 * slot + k for k in range(5)] |
| for cb_id, size in zip(ids, (sizes["data"], sizes["offsets"], sizes["probe"], sizes["acc"], |
| sizes["out"])): |
| cbs.append(ttnn.CBDescriptor(total_size=size, core_ranges=crs, format_descriptors=[ |
| ttnn.CBFormatDescriptor(buffer_index=cb_id, data_format=ttnn.uint32, page_size=size)])) |
| ct = [1 if spec.input_layout == "tile" else 0, spec.in_bytes, spec.out_bytes, spec.channels, |
| spec.num_rows, spec.num_segments, workers, slots, slot, 0, 0, gx, *ids, _OUT_RING, _OFF_CHUNK, |
| _f32_bits(spec.empty_value), off_len * 4] + accessors |
| kernels.append(ttnn.KernelDescriptor( |
| kernel_source=kernel_path(KERNEL), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH, |
| core_ranges=crs, compile_time_args=ct, defines=[], runtime_args=ttnn.RuntimeArgs(), |
| common_runtime_args=[x.buffer_address(), offsets.buffer_address(), out.buffer_address()], |
| config=config)) |
| return ttnn.ProgramDescriptor(kernels=kernels, semaphores=[], cbs=cbs) |
|
|
| |
| def _check_inputs(self, x: Any, offsets: Any) -> None: |
| import ttnn |
|
|
| spec = self.spec |
| shape = tuple(int(s) for s in x.shape) |
| if shape != (1, 1, spec.num_rows, spec.channels): |
| raise ValueError(f"{self.name}: x is {shape}, expected (1, 1, {spec.num_rows}, {spec.channels})") |
| want_dtype = ttnn.float32 if spec.input_dtype == "float32" else ttnn.bfloat16 |
| if x.dtype != want_dtype: |
| raise ValueError(f"{self.name}: x dtype {x.dtype}, expected {spec.input_dtype}") |
| want_layout = ttnn.TILE_LAYOUT if spec.input_layout == "tile" else ttnn.ROW_MAJOR_LAYOUT |
| if x.layout != want_layout: |
| raise ValueError(f"{self.name}: x layout {x.layout}, expected {spec.input_layout}") |
| oshape = tuple(int(s) for s in offsets.shape) |
| if len(oshape) != 2 or oshape[0] != 1 or oshape[1] < spec.num_segments + 1: |
| raise ValueError(f"{self.name}: offsets must be [1, K >= {spec.num_segments + 1}], got {oshape}") |
| if offsets.dtype != ttnn.uint32 or offsets.layout != ttnn.ROW_MAJOR_LAYOUT: |
| raise ValueError(f"{self.name}: offsets must be uint32 ROW_MAJOR") |
| for t, what in ((x, "x"), (offsets, "offsets")): |
| mc = t.memory_config() |
| if getattr(mc, "buffer_type", ttnn.BufferType.DRAM) != ttnn.BufferType.DRAM or \ |
| getattr(mc, "memory_layout", ttnn.TensorMemoryLayout.INTERLEAVED) != \ |
| ttnn.TensorMemoryLayout.INTERLEAVED: |
| raise ValueError(f"{self.name}: {what} must be DRAM interleaved") |
|
|
| def output_spec(self): |
| """``(shape, ttnn dtype, layout)`` of the output.""" |
| import ttnn |
|
|
| dt = ttnn.float32 if self.spec.output_dtype == "float32" else ttnn.bfloat16 |
| return (1, 1, self.spec.num_segments, self.spec.channels), dt, ttnn.ROW_MAJOR_LAYOUT |
|
|
| def allocate_output(self, device: Any): |
| import ttnn |
|
|
| shape, dt, layout = self.output_spec() |
| return ttnn.allocate_tensor_on_device(ttnn.Shape(list(shape)), dt, layout, device, ttnn.DRAM_MEMORY_CONFIG) |
|
|
| def __call__(self, x: Any, offsets: Any, *, output: Any = None): |
| import ttnn |
|
|
| self._check_inputs(x, offsets) |
| out = self.allocate_output(x.device()) if output is None else output |
| shape, dt, layout = self.output_spec() |
| if tuple(int(s) for s in out.shape) != shape or out.dtype != dt or out.layout != layout: |
| raise ValueError(f"{self.name}: output must be {shape} {dt} ROW_MAJOR") |
| return ttnn.generic_op([x, offsets, out], self.program(x, offsets, out)) |
|
|
| def describe(self) -> Dict[str, Any]: |
| s = self.spec |
| return {"name": self.name, "kernel": KERNEL, "op": s.op, "num_rows": s.num_rows, "channels": s.channels, |
| "num_segments": s.num_segments, "input": f"{s.input_dtype}/{s.input_layout}", |
| "output": f"{s.output_dtype}/row_major", "empty_value": s.empty_value, |
| "workers_per_core": s.workers_per_core, "grid": list(self.grid) if self.grid else "device", |
| "l1_bytes_per_worker": s.l1_bytes_per_worker()} |
|
|
|
|
| def segment_reduce(x: Any, offsets: Any, *, num_segments: int, op: str = "max", output_dtype: Optional[str] = None, |
| empty_value: float = 0.0, output: Any = None): |
| """One-shot K1 (builds a :class:`SegmentReduce` from the tensors; ports build it once per call site).""" |
| import ttnn |
|
|
| in_dtype = "float32" if x.dtype == ttnn.float32 else "bfloat16" |
| spec = SegmentReduceSpec(num_rows=int(x.shape[-2]), channels=int(x.shape[-1]), num_segments=int(num_segments), |
| input_dtype=in_dtype, input_layout="tile" if x.layout == ttnn.TILE_LAYOUT else |
| "row_major", output_dtype=output_dtype or in_dtype, op=op, empty_value=empty_value) |
| return SegmentReduce(spec)(x, offsets, output=output) |
|
|
|
|
| |
|
|
| def logstep_segment_max(x: Any, shift_tables: Sequence[Any], last_rows: Any, *, layout: str = "tile"): |
| """The exact log-step segmented max with stock ops (PLAN.md 1.3; the K1 fallback): ``x`` bf16 ``[1, 1, R, C]`` |
| (TILE or ROW_MAJOR, ``C % 32 == 0``), ``shift_tables`` = ``rounds`` uint32 ``[1, R]`` device tensors |
| (:func:`segment_shift_tables`), ``last_rows`` uint32 ``[1, S]`` (:func:`segment_last_rows` with ``empty_row=R``) |
| -> bf16 ``[1, 1, S, C]`` (``layout``). Per round: untilize, ``ttnn.embedding`` gather, ``ttnn.maximum`` (3 |
| programs); then one PADDED gather from the result with a zero row appended at ``R`` (empty segments -> 0). |
| bf16 only (``ttnn.embedding`` tables are bf16); exact for segments of up to ``2^rounds`` rows.""" |
| import ttnn |
|
|
| from .gather import gather_rows, with_zero_rows |
|
|
| if x.dtype != ttnn.bfloat16: |
| raise ValueError("logstep_segment_max is bf16-only (ttnn.embedding); use K1 (SegmentReduce) for fp32") |
| cur = x if x.layout == ttnn.TILE_LAYOUT else ttnn.to_layout(x, ttnn.TILE_LAYOUT) |
| for idx in shift_tables: |
| table = ttnn.to_layout(cur, ttnn.ROW_MAJOR_LAYOUT) |
| shifted = gather_rows(idx, table, layout="tile") |
| nxt = ttnn.maximum(cur, shifted) |
| for t in (table, shifted): |
| ttnn.deallocate(t, False) |
| if cur is not x: |
| ttnn.deallocate(cur, False) |
| cur = nxt |
| rows = int(cur.shape[-2]) |
| table = with_zero_rows(cur) |
| out = gather_rows(last_rows, table, sentinel=rows, layout=layout) |
| ttnn.deallocate(table, False) |
| if cur is not x: |
| ttnn.deallocate(cur, False) |
| return out |
|
|
