# SPDX-License-Identifier: Apache-2.0 """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 # rows per streamed unit (one tile-row): the input rows must be a multiple of it _OFF_CHUNK = 256 # offsets per L1 window (multiple of 16) _OUT_RING = 8 # output staging rows in flight per worker # ------------------------------------------------------------------------------------------------ host tables 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) # ------------------------------------------------------------------------------------------------- oracles 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 # ------------------------------------------------------------------------------------------------- device K1 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" # "float32" | "bfloat16" input_layout: str = "tile" # "tile" | "row_major" output_dtype: str = "float32" # "float32" | "bfloat16" (RNE) 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 # (x, y) cores; None = the device's compute grid spec.check_l1() # ---- descriptor 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)] # data, offsets, probe, acc, out 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) # ---- checks 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) # -------------------------------------------------------------------------------------- stock-op fallback 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) # [1, 1, R + 1, C] ROW_MAJOR, row R = 0 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