code/tt_diffusion_planner/ttaw/ops/segment.py

# 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