code/models/tt/nms_kernels.py

# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
"""Device NMS with two small custom data-movement kernels (``ttnn.generic_op``, JIT-compiled from
``code/kernels/sp_nms``; no tt-metal rebuild) around two standard ``max_pool2d`` calls.

    s_sm [B*h*w, 65] TILE  --nms_fold-->  P [H, SW+2r, 32] L1 height-sharded (strip layout)
    P --max_pool2d [2r+1,1] pad (r,0)--> --max_pool2d [1,2r+1] pad 0-->  M [H, SW, 32] (9x9 max)
    (M, P) --nms_unfold-->  NMS map [H, W] ROW_MAJOR DRAM interleaved

Lane ``l`` of the strip layout holds image columns ``l*SW - r .. l*SW + SW + r - 1`` (zero outside
the image; scores are >= 0 so zero padding never changes a window max), so the pools see every
pixel's full 9x9 window. Each core owns ``ROWS`` consecutive image rows in all three tensors, so
both custom kernels only touch their own L1 shards (plus the S tile reads / output row writes).
Everything is bit-exact: no arithmetic, only copies, max and an equality test.
"""
from __future__ import annotations

import os

import ttnn

_KDIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "kernels", "sp_nms")
LANES = 32
#: Keypoint candidates stored per slot (2 slots per core: BRISC rows, NCRISC rows). With a radius-4
#: NMS a 2-row band holds at most ~W/5 strict maxima; more (plateau ties) -> overflow -> host fallback.
CAND_CAP = 63


def threshold_bits(threshold: float) -> int:
    """bf16 bit pattern p with  float(x) > threshold  <=>  int16(bits(x)) > int16(p)  for bf16 x
    (the fp32 bits of the threshold truncated to 16; see postprocess._bf16_threshold_bits)."""
    import struct

    b = struct.unpack("<I", struct.pack("<f", float(threshold)))[0] >> 16
    return b


def _src(name: str) -> str:
    with open(os.path.join(_KDIR, name)) as f:
        return f.read()


class DeviceNms:
    """Static plan for one (H, W, radius) on one device; ``__call__`` is trace-safe (no host writes)."""

    CNT_PAGE = int(os.environ.get("SP_CNT_PAGE", "256"))  # bytes per page of the dense slot-count tensor (SP_SF_DIRECT)

    def __init__(self, device, H: int, W: int, radius: int, *, border: int | None = None,
                 threshold: float | None = None, cap: int = CAND_CAP):
        """``threshold``/``border`` set: the unfold step also emits the keypoint candidate list
        (``KeypointCandidates`` layout, see nms_unfold_kp.cpp) for that threshold and border."""
        if W % LANES or H % 8 or W % 8:
            raise ValueError(f"DeviceNms needs W % 32 == 0 and H % 8 == 0, got {H}x{W}")
        self.device, self.H, self.W, self.r = device, H, W, int(radius)
        self.SW = W // LANES
        self.PW = self.SW + 2 * self.r
        g = device.compute_with_storage_grid_size()
        ncores_max = g.x * g.y
        # smallest rows-per-core (most cores) that divides H
        self.rows = next(r for r in range(1, H + 1) if H % r == 0 and H // r <= ncores_max)
        self.ncores = H // self.rows
        self.grid = ttnn.num_cores_to_corerangeset(self.ncores, g, row_wise=True)
        self.cores = ttnn.corerange_to_cores(self.grid, row_wise=True)

        def mc(width_rows):
            spec = ttnn.ShardSpec(self.grid, [self.rows * width_rows, LANES], ttnn.ShardOrientation.ROW_MAJOR)
            return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, spec)

        self.p_mc = mc(self.PW)
        self.m_rows = self.rows * self.SW
        self._fold_src = _src("nms_fold.cpp")
        self._pool_dm_src = _src("nms_pool_dm.cpp")
        self._pool_c_src = _src("nms_pool_compute.cpp")
        self._unfold_src = _src("nms_unfold.cpp")
        self.kp = threshold is not None
        if self.kp:
            self._unfold_fn = _src("nms_unfold_fn.inc")
            self._unfold_src = self._unfold_fn + _src("nms_unfold_kp.cpp")
            self.border = int(border)
            self.threshold = float(threshold)
            self.thr_bits = threshold_bits(self.threshold)
            self.cap = int(cap)
            # per-request threshold / border (class C): a 64-byte parameter tensor read by the
            # unfold kernel; rewritten by set_params() only when the values change
            self._prm_vals = (self.thr_bits, self.border, 0)
            self.params = ttnn.from_torch(self._prm_host(*self._prm_vals), dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT,
                                          device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        self._split = (self.rows + 1) // 2  # BRISC rows [0, split), NCRISC rows [split, rows)
        #: SP_KP_REC=1: the unfold also writes final keypoint-header entries per slot (records tensor in
        #: the NMS cores' L1) and the keypoint list is a DMA gather (kp_compact3.cpp) instead of the
        #: one-core scalar expansion (kp_compact2.cpp)
        self.rec_mode = self.kp and os.environ.get("SP_KP_REC", "1") == "1" and self._rec_formula_ok()
        self.rec = None  # records tensor of the last unfold call (consumed by DeviceSampler.compact)
        #: SP_NMS_PMASK=1: the one-op window max also applies the NMS compare (SFPU) and outputs the strip-layout
        #: NMS map N; the unfold (word-store loop) then reads N only (needs the fused pool and an even SW)
        #: even radius only: an SFPU load at DST row offset 2 i covers one column parity per (offset / 2) % 2, so the
        #: centre shift by r positions (nms_center) is a pure shift only for even r
        self.pmask = (self.kp and self.pool_fused_ok() and self.SW % 2 == 0 and self.r % 2 == 0
                      and os.environ.get("SP_NMS_PMASK", "1") == "1" and os.environ.get("SP_NMS_UNF_W", "1") == "1")
        #: SP_NMS_POOL_UNF=1 (PMASK): the unfold + candidates run inside the window-max op, on its data-movement RISCs,
        #: reading the NMS rows from the op's output CB (one op less, M never written)
        self.pool_unf = self.pmask and os.environ.get("SP_NMS_POOL_UNF", "1") == "1"
        #: SP_SF_DIRECT=1: the unfold also writes the slot counts densely (16 B per slot, tensor `counts` on core (0, 0))
        #: and the sampler builds its keypoint list itself (no compaction op)
        self.cnt_mode = self.kp and self.rec_mode and os.environ.get("SP_SF_DIRECT", "1") == "1"
        self.counts = None

    def _rec_formula_ok(self) -> bool:
        """The unfold kernel's integer tap-cell formula equals postprocess.SampleTables for this H, W."""
        import torch

        from .postprocess import SampleTables

        t = SampleTables(self.H, self.W)
        for n, c, a0, a1 in ((self.W, t.wc, t.x0, t.x1), (self.H, t.hc, t.y0, t.y1)):
            v = torch.arange(n)
            q = torch.div((2 * v - 7) * (c - 1), 16 * c - 9, rounding_mode="floor")
            if not (torch.equal(q.clamp(0, c - 1), a0) and torch.equal((q + 1).clamp(0, c - 1), a1)):
                return False
        return True

    @staticmethod
    def _prm_host(thr_bits: int, border: int, spec: int = 0):
        import torch

        v = torch.zeros(1, 16, dtype=torch.int32)
        v[0, 0], v[0, 1], v[0, 2] = int(thr_bits), int(border), int(spec)
        return v

    def set_params(self, threshold: float, border: int, spec: int | None = None) -> None:
        """Make the next replay use this keypoint threshold / border (and, single-D2H mode, the
        speculative descriptor bucket ``spec``): H2D of 64 B, only on change."""
        if int(border) < 0 or 2 * int(border) >= min(self.H, self.W):
            raise ValueError(f"border {border} out of range")
        vals = (threshold_bits(threshold), int(border), self._prm_vals[2] if spec is None else int(spec))
        if vals != self._prm_vals:
            host = ttnn.from_torch(self._prm_host(*vals), dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT)
            ttnn.copy_host_to_device_tensor(host, self.params, cq_id=0)
            self._prm_vals = vals
        self.threshold, self.border = float(threshold), int(border)

    def _kernels(self, src, ct_args, rt_fn, extra_defs=()):
        ks = []
        defs = [(d, "1") for d in os.environ.get("SP_NMS_DEFS", "").split(",") if d] + list(extra_defs)
        for proc, (a, b) in ((0, (0, self._split)), (1, (self._split, self.rows))):
            rt = ttnn.RuntimeArgs()
            for idx, core in enumerate(self.cores):
                rt[core.x][core.y] = rt_fn(idx, a, b)
            cfg = ttnn.ReaderConfigDescriptor() if proc == 0 else ttnn.WriterConfigDescriptor()
            ks.append(
                ttnn.KernelDescriptor(
                    kernel_source=src,
                    source_type=ttnn.KernelDescriptor.SourceType.SOURCE_CODE,
                    core_ranges=self.grid,
                    compile_time_args=list(ct_args[proc]),
                    runtime_args=rt,
                    config=cfg,
                    defines=defs,
                )
            )
        return ks

    def _scratch(self, size):
        cbs = []
        for idx in (0, 1):
            fmt = ttnn.CBFormatDescriptor(buffer_index=idx, data_format=ttnn.bfloat16, page_size=size)
            cbs.append(ttnn.CBDescriptor(total_size=size, core_ranges=self.grid, format_descriptors=[fmt]))
        return cbs

    def fold5_ok(self) -> bool:
        """SP_NMS_FOLD5: all rows of a core in one cell row (one tile column), at most 5 rows per core."""
        return os.environ.get("SP_NMS_FOLD5", "1") == "1" and self.rows <= 5 and 8 % self.rows == 0

    def fold5(self, s_sm: ttnn.Tensor) -> ttnn.Tensor:
        """The fold on 5 RISCs per core (nms_fold5_*.cpp): the core's S tiles are read once, image row r of the core is
        folded by RISC r % 5 (2 data-movement RISCs + 3 TRISCs running plain scalar code). Same steps as fold()."""
        wc = self.W // 8
        tcols = s_sm.padded_shape[-1] // 32
        p = ttnn.allocate_tensor_on_device(
            ttnn.Shape([1, 1, self.H * self.PW, LANES]), ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT, self.device, self.p_mc
        )
        nt = (wc + 31) // 32 + 1
        CB_S = 0
        acc = ttnn.TensorAccessorArgs(s_sm).get_compile_time_args()
        SC = ttnn.KernelDescriptor.SourceType.SOURCE_CODE
        if not hasattr(self, "_fold5_src"):
            row = _src("nms_fold5_row.inc")
            self._fold5_src = (row + _src("nms_fold5_dm.cpp"), row + _src("nms_fold5_trisc.cpp"))
        ks = []
        for proc, cfg in ((0, ttnn.ReaderConfigDescriptor()), (1, ttnn.WriterConfigDescriptor())):
            rt = ttnn.RuntimeArgs()
            for idx, core in enumerate(self.cores):
                rt[core.x][core.y] = [s_sm.buffer_address(), p.buffer_address(), idx * self.rows]
            ks.append(ttnn.KernelDescriptor(kernel_source=self._fold5_src[0], source_type=SC, core_ranges=self.grid,
                                            compile_time_args=[proc, CB_S, wc, tcols, self.SW, self.r, self.rows, nt] + acc,
                                            runtime_args=rt, config=cfg))
        rt = ttnn.RuntimeArgs()
        for idx, core in enumerate(self.cores):
            rt[core.x][core.y] = [p.buffer_address(), idx * self.rows]
        ks.append(ttnn.KernelDescriptor(kernel_source=self._fold5_src[1], source_type=SC, core_ranges=self.grid,
                                        compile_time_args=[CB_S, wc, self.SW, self.r, self.rows, nt], runtime_args=rt,
                                        config=ttnn.ComputeConfigDescriptor()))
        f = ttnn.CBFormatDescriptor(buffer_index=CB_S, data_format=ttnn.bfloat16, page_size=2048)
        cbs = [ttnn.CBDescriptor(total_size=2048 * nt, core_ranges=self.grid, format_descriptors=[f])]
        return ttnn.generic_op([s_sm, p], ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))

    def fold(self, s_sm: ttnn.Tensor) -> ttnn.Tensor:
        """s_sm: [1, 1, H/8*W/8, 65] TILE bf16 interleaved -> P (L1 height-sharded strip layout)."""
        if self.fold5_ok():
            return self.fold5(s_sm)
        wc = self.W // 8
        tcols = s_sm.padded_shape[-1] // 32
        p = ttnn.allocate_tensor_on_device(
            ttnn.Shape([1, 1, self.H * self.PW, LANES]), ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT, self.device, self.p_mc
        )
        acc = ttnn.TensorAccessorArgs(s_sm).get_compile_time_args()
        ct = [[cb, wc, tcols, self.SW, self.r] + acc for cb in (0, 1)]
        s_addr, p_addr = s_sm.buffer_address(), p.buffer_address()
        rows = self.rows
        # SP_NMS_FOLD_LOCAL=1: the fold's dense row buffer in RISC-local memory (nms_fold.cpp FOLD_LOCAL)
        fdefs = [("FOLD_LOCAL", "1")] if os.environ.get("SP_NMS_FOLD_LOCAL", "1") == "1" else []
        if os.environ.get("SP_NMS_FOLD_SEG", "1") == "1":
            fdefs.append(("FOLD_SEG", "1"))  # gather in runs of cells with pointer steps
        ks = self._kernels(self._fold_src, ct, lambda idx, a, b: [s_addr, p_addr, idx * rows, a, b], fdefs)
        prog = ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=self._scratch(((wc + 31) // 32 + 1) * 2048 + ((self.W + 2 * self.r) * 2 + 31) // 32 * 32 + 32))
        return ttnn.generic_op([s_sm, p], prog)

    def pool_fused_ok(self) -> bool:
        """The one-op window max (nms_pool_*.cpp) needs the halo rows on the neighbouring cores only
        (r <= rows), and all input + output rows in the 16 bf16 DST tiles (rows + 2 r + rows <= 16)."""
        return (os.environ.get("SP_NMS_POOL1", "1") == "1" and self.r <= self.rows
                and 2 * self.rows + 2 * self.r <= 16 and self.PW <= 32)

    def pool_fused(self, p: ttnn.Tensor, unf: bool = False):
        """P -> M (9x9 window max, [H * SW, 32] L1 height-sharded like the two-pool path) in ONE op.
        ``unf`` (SP_NMS_POOL_UNF): the kp unfold runs in the same op on M's rows in CB_OUT -> (NMS map, candidates)."""
        if unf:
            m = None
            out = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, self.H, self.W]), ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT,
                                                 self.device, ttnn.DRAM_MEMORY_CONFIG)
            cand, uio, uct, urt, udefs, uscr = self._unfold_kp_parts(None, p, out)
        else:
            m = ttnn.allocate_tensor_on_device(
                ttnn.Shape([1, 1, self.H * self.SW, LANES]), ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT, self.device,
                ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1,
                                  ttnn.ShardSpec(self.grid, [self.m_rows, LANES], ttnn.ShardOrientation.ROW_MAJOR)))
        rows, r = self.rows, self.r
        nrow = rows + 2 * r
        CB_IN, CB_OUT, CB_VW, CB_VR = 0, 1, 2, 3
        # V scratch: ROWS 2 KB rows + one slack page (the +64 B re-reads run past the last row)
        vmc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1,
                                ttnn.ShardSpec(self.grid, [(rows + 1) * 32, LANES], ttnn.ShardOrientation.ROW_MAJOR))
        v = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, self.ncores * (rows + 1) * 32, LANES]), ttnn.bfloat16,
                                           ttnn.ROW_MAJOR_LAYOUT, self.device, vmc)
        n = self.ncores
        SC = ttnn.KernelDescriptor.SourceType.SOURCE_CODE
        ks = []
        CB_U0 = 5  # unfold scratch CBs 5 (PROC 0) and 6 (PROC 1)
        for proc, cfg in ((0, ttnn.ReaderConfigDescriptor()), (1, ttnn.WriterConfigDescriptor())):
            rt = ttnn.RuntimeArgs()
            a, b = (0, self._split) if proc == 0 else (self._split, self.rows)
            for i, core in enumerate(self.cores):
                pc = self.device.worker_core_from_logical_core(self.cores[max(i - 1, 0)])
                nc = self.device.worker_core_from_logical_core(self.cores[min(i + 1, n - 1)])
                rt[core.x][core.y] = [p.buffer_address(), m.buffer_address() if m is not None else 0, pc.x, pc.y, nc.x, nc.y,
                                      int(i > 0), int(i < n - 1)] + (urt(i, proc, a, b) if unf else [])
            ct = [proc, CB_IN, CB_OUT, rows, r, self.PW, self.SW] + ([CB_U0 + proc] + uct if unf else [])
            ks.append(ttnn.KernelDescriptor(kernel_source=(self._unfold_fn + self._pool_dm_src) if unf else self._pool_dm_src,
                                            source_type=SC, core_ranges=self.grid, compile_time_args=ct, runtime_args=rt,
                                            config=cfg, defines=((list(udefs) + [("POOL_UNF", "1"), ("UNF_SPLIT", str(self._split))]
                                                      + ([("UNF_ROW_WAIT(yl)", "cb_wait_front(get_compile_time_arg_val(2), (yl) + 1)")]
                                                         if os.environ.get("SP_NMS_ROWWAIT", "1") == "1" else [])) if unf else [])
                                            + ([("POOL_ZDMA", "1")] if os.environ.get("SP_NMS_ZDMA", "1") == "1" else [])))
        ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=False, dst_full_sync_en=True)
        defs = [(d, "1") for d in os.environ.get("SP_NMS_POOL_DEFS", "").split(",") if d]
        CB_C = 4
        if self.pmask:
            defs.append(("PMASK", "1"))
        if unf:
            defs.append(("POOL_UNF", "1"))
        ks.append(ttnn.KernelDescriptor(kernel_source=self._pool_c_src, source_type=SC, core_ranges=self.grid,
                                        compile_time_args=[CB_IN, CB_OUT, rows, r, self.PW, self.SW, CB_VW, CB_VR, CB_C], runtime_args=[],
                                        config=ccfg, defines=defs))

        def cb(i, pages):
            f = ttnn.CBFormatDescriptor(buffer_index=i, data_format=ttnn.bfloat16, page_size=2048)
            return ttnn.CBDescriptor(total_size=2048 * pages, core_ranges=self.grid, format_descriptors=[f])

        cb_vw = ttnn.cb_descriptor_from_sharded_tensor(CB_VW, v)
        cb_vw.format_descriptors = [ttnn.CBFormatDescriptor(buffer_index=CB_VW, data_format=ttnn.bfloat16, page_size=2048)]
        cb_vr = ttnn.cb_descriptor_from_sharded_tensor(CB_VR, v)
        cb_vr.format_descriptors = [ttnn.CBFormatDescriptor(buffer_index=CB_VR, data_format=ttnn.bfloat16, page_size=64)]
        cbs = [cb(CB_IN, nrow), cb(CB_OUT, rows), cb_vw, cb_vr] + ([cb(CB_C, rows)] if self.pmask else [])
        if unf:
            for i in (0, 1):
                f = ttnn.CBFormatDescriptor(buffer_index=CB_U0 + i, data_format=ttnn.bfloat16, page_size=uscr)
                cbs.append(ttnn.CBDescriptor(total_size=uscr, core_ranges=self.grid, format_descriptors=[f]))
            ttnn.generic_op([v] + uio, ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))
            ttnn.deallocate(v)
            return out, cand
        ttnn.generic_op([p, v, m], ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))
        ttnn.deallocate(v)
        return m

    def pool(self, p: ttnn.Tensor) -> ttnn.Tensor:
        if self.pool_fused_ok():
            return self.pool_fused(p)
        r, k = self.r, 2 * self.r + 1
        my = ttnn.max_pool2d(
            input_tensor=p, batch_size=1, input_h=self.H, input_w=self.PW, channels=LANES,
            kernel_size=[k, 1], stride=[1, 1], padding=[r, 0], dilation=[1, 1],
        )
        m = ttnn.max_pool2d(
            input_tensor=my, batch_size=1, input_h=self.H, input_w=self.PW, channels=LANES,
            kernel_size=[1, k], stride=[1, 1], padding=[0, 0], dilation=[1, 1], deallocate_input=True,
        )
        sp = m.memory_config().shard_spec
        if (not m.is_sharded()) or list(sp.shape) != [self.m_rows, LANES] or sp.grid != self.grid:
            raise RuntimeError(f"unexpected pool output sharding {m.memory_config()}")
        return m

    def unfold(self, m: ttnn.Tensor, p: ttnn.Tensor):
        out = ttnn.allocate_tensor_on_device(
            ttnn.Shape([1, 1, self.H, self.W]), ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT, self.device,
            ttnn.DRAM_MEMORY_CONFIG,
        )
        if self.kp:
            return self._unfold_kp(m, p, out)
        acc = ttnn.TensorAccessorArgs(out).get_compile_time_args()
        ct = [[cb, self.SW, self.r] + acc for cb in (0, 1)]
        m_addr, p_addr, o_addr = m.buffer_address(), p.buffer_address(), out.buffer_address()
        rows = self.rows
        ks = self._kernels(self._unfold_src, ct, lambda idx, a, b: [m_addr, p_addr, o_addr, idx * rows, a, b])
        prog = ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=self._scratch(2 * self.W * 2))
        return ttnn.generic_op([m, p, out], prog)

    def cand_memory_config(self):
        """Candidate slots: L1 of core (0, 0) (SP_CAND_L1=1, default; the one-core kp_compact kernel
        then reads them in place instead of NSLOT DRAM page reads) or DRAM interleaved."""
        if os.environ.get("SP_CAND_L1", "1") != "1":
            return ttnn.DRAM_MEMORY_CONFIG
        c0 = ttnn.CoreCoord(0, 0)
        spec = ttnn.ShardSpec(ttnn.CoreRangeSet([ttnn.CoreRange(c0, c0)]), [self.nslots, self.cap + 1], ttnn.ShardOrientation.ROW_MAJOR)
        return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, spec)

    @property
    def nslots(self) -> int:
        return 2 * self.ncores

    def slot_first_rows(self):
        """Image row of the first row of every candidate slot (slot = 2 * core + risc)."""
        rows = []
        for c in range(self.ncores):
            rows += [c * self.rows, c * self.rows + self._split]
        return rows

    def _unfold_kp(self, m, p, out):
        cand, io, ct, rt_fn, defs, scratch = self._unfold_kp_parts(m, p, out)
        ks = []
        for proc, (a, b) in ((0, (0, self._split)), (1, (self._split, self.rows))):
            rt = ttnn.RuntimeArgs()
            for idx, core in enumerate(self.cores):
                rt[core.x][core.y] = rt_fn(idx, proc, a, b)
            cfg = ttnn.ReaderConfigDescriptor() if proc == 0 else ttnn.WriterConfigDescriptor()
            ks.append(
                ttnn.KernelDescriptor(
                    kernel_source=self._unfold_src,
                    source_type=ttnn.KernelDescriptor.SourceType.SOURCE_CODE,
                    core_ranges=self.grid,
                    compile_time_args=[proc] + ct,
                    runtime_args=rt,
                    config=cfg,
                    defines=defs,
                )
            )
        prog = ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=self._scratch(scratch))
        ttnn.generic_op(io, prog)
        return out, cand

    def _unfold_kp_parts(self, m, p, out):
        """cand tensor, io tensors (without m when m is None), CT args after the scratch CB id, RT-arg fn
        (core idx, proc, first row, end row), defines, scratch CB size of the kp unfold."""
        cand = ttnn.allocate_tensor_on_device(
            ttnn.Shape([1, 1, self.nslots, self.cap + 1]), ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT, self.device,
            self.cand_memory_config(),
        )
        acc = (ttnn.TensorAccessorArgs(out).get_compile_time_args() + ttnn.TensorAccessorArgs(cand).get_compile_time_args()
               + ttnn.TensorAccessorArgs(self.params).get_compile_time_args())
        ct = [self.SW, self.r, self.H, 0, 0, self.cap] + acc
        m_addr = m.buffer_address() if m is not None else 0
        p_addr, o_addr, c_addr = p.buffer_address(), out.buffer_address(), cand.buffer_address()
        q_addr = self.params.buffer_address()
        rows = self.rows
        defs = [(d, "1") for d in os.environ.get("SP_NMS_DEFS", "").split(",") if d]
        if os.environ.get("SP_NMS_UNF_W", "1") == "1" and self.SW % 2 == 0:
            defs.append(("UNF_W", "1"))  # word-store compare loop (nms_unfold_kp.cpp)
        if self.pmask:
            defs.append(("PMASK", "1"))  # m holds the NMS output already
            if os.environ.get("SP_NMS_UNF_SPARSE", "1") == "1":
                defs.append(("UNF_SPARSE", "1"))  # DMA-zeroed row buffers, only non-zero words stored
                if os.environ.get("SP_NMS_UNF_BLK", "1") == "1":
                    defs.append(("UNF_BLK", "1"))  # 8 word loads per OR test
        io = ([m] if m is not None else []) + [p, out, cand, self.params]
        extra = []
        if self.rec_mode:
            # records: per NMS core 2 slots x (CAP + 1) x 16 B in its own L1
            spec = ttnn.ShardSpec(self.grid, [2, 4 * (self.cap + 1)], ttnn.ShardOrientation.ROW_MAJOR)
            rmc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, spec)
            self.rec = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, self.nslots, 4 * (self.cap + 1)]), ttnn.uint32,
                                                      ttnn.ROW_MAJOR_LAYOUT, self.device, rmc)
            io.append(self.rec)
            extra = [self.rec.buffer_address()]
            defs += [("KP_REC", "1"), ("KP_WC", str(self.W // 8)), ("KP_HC", str(self.H // 8))]
            if self.cnt_mode:
                # dense slot counts, 16 B per slot, L1-interleaved pages of CNT_PAGE bytes (spread over many cores:
                # every sampler RISC reads all of them at the same time)
                spp = self.CNT_PAGE // 16
                npg = (self.nslots + spp - 1) // spp
                self.counts = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, npg, self.CNT_PAGE // 4]), ttnn.uint32,
                                                             ttnn.ROW_MAJOR_LAYOUT, self.device, ttnn.L1_MEMORY_CONFIG)
                io.append(self.counts)
                extra.append(self.counts.buffer_address())
                ct += ttnn.TensorAccessorArgs(self.counts).get_compile_time_args()
                defs += [("KP_CNT", str(self.nslots)), ("KP_CNT_PAGE", str(self.CNT_PAGE))]
        def rt_fn(idx, proc, a, b):
            return [m_addr, p_addr, o_addr, idx * rows, a, b, c_addr, 2 * idx + proc, q_addr] + extra

        return cand, io, ct, rt_fn, defs, 2 * self.W * 2 + ((self.cap + 1) * 4 + 63) // 64 * 64 + 64

    def __call__(self, s_sm: ttnn.Tensor):
        """-> NMS map, or (NMS map, candidates) when built with a threshold."""
        p = self.fold(s_sm)
        if self.pool_unf and self.kp:
            res = self.pool_fused(p, unf=True)
            ttnn.deallocate(p)
            return res
        m = self.pool(p)
        out = self.unfold(m, p)
        ttnn.deallocate(m)
        ttnn.deallocate(p)
        return out


class DeviceSampler:
    """Keypoint list + bilinear descriptor sampling on device, bit-exact with the host path
    (``postprocess.extract_keypoints_bf16`` + ``sample_descriptors_nhwc`` before its final
    L2-normalisation, which stays on the host on the (n, 256) result).

    * ``compact`` (main trace, ``kp_compact.cpp``, one core): concatenates the candidate slots of the
      NMS unfold kernel into the keypoint header ``hdr`` (count, overflow, (y, x, score) list in raster
      order, at most KMAX).
    * ``sample(hdr, chunks)``: gather kernel (``sample_gather.cpp``, all cores: cell rows looked up
      through the per-axis table ``tab`` + tap weights from the per-pixel table ``wtab``, both
      uploaded once, written as TILE pages) -> ((G0*W0 + G1*W1) + G2*W2) + G3*W3 in fp32 with a
      column-broadcast weight (Tensix fp32 mul/add are IEEE-exact, verified bitwise; same operation
      order as the host) -> ``untilize_chunks.cpp`` writes the [n, 256] fp32 rows into ``CHUNK``-row
      tensors so the host reads back only ceil(n / CHUNK) of them. Runs in the main trace on
      ``hdr`` and in a small second trace on the host-uploaded ``hdr_b`` (top-k fallback when there
      are more than KMAX candidates)."""

    def __init__(self, device, height: int, width: int, nms: "DeviceNms", channels: int = 256, kmax: int = 1024, chunk: int = 64):
        from .postprocess import SampleTables
        import torch

        self.device, self.c, self.kmax, self.chunk = device, channels, kmax, chunk
        self.H, self.W, self.nms = height, width, nms
        DR, RM = ttnn.DRAM_MEMORY_CONFIG, ttnn.ROW_MAJOR_LAYOUT

        def alloc(shape, dt):
            return ttnn.allocate_tensor_on_device(ttnn.Shape(shape), dt, RM, device, DR)

        # persistent buffers, allocated before any trace capture
        self.hdr = alloc([1, 16 + 4 * kmax], ttnn.uint32)
        self.hdr_b = alloc([1, 16 + 4 * kmax], ttnn.uint32)
        # zero both headers once: capture_b's compile pass (and the first replay before any upload) reads hdr_b; left
        # uninitialised, its DRAM garbage gave keypoint cells far outside the descriptor map and an out-of-range
        # NoC read (watcher: NOC target address overflow, L1 0x1bf400 on chips 12 and 15)
        zh = ttnn.from_torch(torch.zeros((1, 16 + 4 * kmax), dtype=torch.int32), dtype=ttnn.uint32,
                             layout=ttnn.ROW_MAJOR_LAYOUT)
        for h in (self.hdr, self.hdr_b):
            ttnn.copy_host_to_device_tensor(zh, h, cq_id=0)
        # bucket b holds chunk * (b + 1) rows; the sampler writes the n rows into bucket ceil(n/chunk)-1
        # only, so the host reads back one right-sized tensor
        # SP_KPC_ONE=1 (single D2H): every bucket starts with HR rows holding a copy of the keypoint
        # header (kp_compact2 writes it into the bucket it picks: max(host speculation, rows needed)),
        # so the host reads header + descriptor rows of a request in one transfer
        self.fused = os.environ.get("SP_SAMPLE_FUSED", "1") == "1"
        self.one_d2h = self.fused and os.environ.get("SP_KPC_ONE", "1") == "1"
        if self.one_d2h:
            chunk = int(os.environ.get("SP_KPC_CHUNK", "32"))
        self.chunk = chunk
        self.hr = (4 * (16 + 4 * kmax) + 4 * channels - 1) // (4 * channels) if self.one_d2h else 0
        #: speculative bucket index (word 2 of the NMS parameter tensor), maintained by the host
        self.spec = 0
        assert kmax % chunk == 0 and chunk % 32 == 0
        nb = kmax // chunk
        #: SP_KPC_SPLIT=1 (single-D2H mode with the direct sampler): bucket b of the main trace is TWO tensors, the head
        #: chunks_a[b] (HR header rows + the first split_rows[b] descriptor rows) and the tail tails_a[b] (the remaining
        #: rows of the bucket; none for b = 0), read back by two non-blocking D2H so that the host decodes and
        #: L2-normalises the head rows while the tail is in flight (same bytes, same values)
        self.split = (self.one_d2h and os.environ.get("SP_KPC_SPLIT", "1") == "1"
                      and os.environ.get("SP_SF_PIPE", "1") == "1")
        frac = float(os.environ.get("SP_KPC_SPLIT_FRAC", "0.65"))
        if self.split:
            self.split_rows = [chunk * (b + 1) if b == 0 else chunk * min(b, max(1, round(frac * (b + 1)))) for b in range(nb)]
            self.chunks_a = [alloc([self.hr + self.split_rows[b], channels], ttnn.float32) for b in range(nb)]
            self.tails_a = [alloc([max(chunk, chunk * (b + 1) - self.split_rows[b]), channels], ttnn.float32) for b in range(nb)]
        else:
            self.split_rows = None
            self.chunks_a = [alloc([self.hr + chunk * (b + 1), channels], ttnn.float32) for b in range(nb)]
            self.tails_a = []
        self.chunks_b = [alloc([self.hr + chunk * (b + 1), channels], ttnn.float32) for b in range(kmax // chunk)]
        self.tables = t = SampleTables(height, width)
        tx = torch.stack([t.x0, t.x1], 1).to(torch.int32)
        ty = torch.stack([t.y0 * t.wc, t.y1 * t.wc], 1).to(torch.int32)
        tab = torch.cat([tx.reshape(-1), ty.reshape(-1)])[None]
        self.tab = ttnn.from_torch(tab, dtype=ttnn.uint32, layout=RM, device=device, memory_config=DR)
        self.wtab = ttnn.from_torch(t.weight_table(), dtype=ttnn.float32, layout=RM, device=device, memory_config=DR)
        self.l1 = ttnn.L1_MEMORY_CONFIG
        self._src_compact = _src("kp_compact.cpp")
        self._src_compact2 = _src("kp_compact2.cpp")
        self._src_compact3 = _src("kp_compact3.cpp")
        self._src_gather = _src("sample_gather.cpp")
        self._src_untilize = _src("untilize_chunks.cpp")
        self._src_f = [_src(f"sample_fused_{k}.cpp") for k in ("reader", "compute", "writer")]
        self._src_p = [_src(f"sample_pipe_{k}.cpp") for k in ("reader", "compute", "writer")]
        kpl = _src("sample_kplist.hpp.inc")
        self._src_d = [kpl + _src("sample_direct_reader.cpp"), self._src_p[1], kpl + _src("sample_direct_writer.cpp")]
        g = device.compute_with_storage_grid_size()
        self.ncores = g.x * g.y
        self.all_cores = ttnn.corerange_to_cores(ttnn.num_cores_to_corerangeset(self.ncores, g, row_wise=True), row_wise=True)
        self.tid_b = None
        self.d = None

    def _spread_program(self, units, common, unit_args, src, ct, size, io, extra=()):
        """Units spread round-robin over every core and both data-movement RISCs (unit u -> core
        u % ncores, RISC (u // ncores) % 2, further units appended), so that the active low-index
        units (tile rows < n) land on distinct cores."""
        slots = {}
        for idx, u in enumerate(units):
            key = (idx % self.ncores, (idx // self.ncores) % 2)
            slots.setdefault(key, []).append(u)
        ks = []
        for proc in (0, 1):
            rt = ttnn.RuntimeArgs()
            used = []
            for ci, core in enumerate(self.all_cores):
                us = slots.get((ci, proc))
                if not us:
                    continue
                args = list(common) + [len(us)] + list(extra)
                for u in us:
                    args += unit_args(u)
                rt[core.x][core.y] = args
                used.append(core)
            if not used:
                continue
            cfg = ttnn.ReaderConfigDescriptor() if proc == 0 else ttnn.WriterConfigDescriptor()
            ks.append(
                ttnn.KernelDescriptor(
                    kernel_source=src,
                    source_type=ttnn.KernelDescriptor.SourceType.SOURCE_CODE,
                    core_ranges=ttnn.CoreRangeSet([ttnn.CoreRange(c, c) for c in used]),
                    compile_time_args=[proc] + ct,
                    runtime_args=rt,
                    config=cfg,
                )
            )
        grid = ttnn.CoreRangeSet([ttnn.CoreRange(c, c) for c in self.all_cores[: min(len(units), self.ncores)]])
        cbs = []
        for i in (0, 1):
            fmt = ttnn.CBFormatDescriptor(buffer_index=i, data_format=ttnn.bfloat16, page_size=size)
            cbs.append(ttnn.CBDescriptor(total_size=size, core_ranges=grid, format_descriptors=[fmt]))
        ttnn.generic_op(io, ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))

    def compact(self, cand: ttnn.Tensor):
        if cand.is_sharded() and self.nms.rec_mode and self.nms.rec is not None:
            self._compact3(cand, self.nms.rec)
            ttnn.deallocate(self.nms.rec)
            self.nms.rec = None
            return
        if cand.is_sharded():
            return self._compact2(cand)
        if self.one_d2h:
            raise RuntimeError("SP_KPC_ONE needs the L1-sharded candidate slots (kp_compact2)")
        nk = self.nms
        core = ttnn.CoreCoord(0, 0)
        grid = ttnn.CoreRangeSet([ttnn.CoreRange(core, core)])
        io = [cand, self.hdr, self.tab]
        acc = []
        for t in io:
            acc += ttnn.TensorAccessorArgs(t).get_compile_time_args()
        rt = ttnn.RuntimeArgs()
        rt[0][0] = [t.buffer_address() for t in io]
        local = 1 if cand.is_sharded() else 0  # slots already in this core's L1
        ct = [0, nk.nslots, nk.cap, self.W, self.kmax, nk.rows, nk._split, self.H, local] + acc
        k = ttnn.KernelDescriptor(
            kernel_source=self._src_compact,
            source_type=ttnn.KernelDescriptor.SourceType.SOURCE_CODE,
            core_ranges=grid,
            compile_time_args=ct,
            runtime_args=rt,
            config=ttnn.ReaderConfigDescriptor(),
        )
        size = nk.nslots * (nk.cap + 1) * 4 + 64 * ((2 * (self.W + self.H) * 4 + 63) // 64) + (16 + 4 * self.kmax) * 4 + 64
        fmt = ttnn.CBFormatDescriptor(buffer_index=0, data_format=ttnn.bfloat16, page_size=size)
        cb = ttnn.CBDescriptor(total_size=size, core_ranges=grid, format_descriptors=[fmt])
        ttnn.generic_op(io, ttnn.ProgramDescriptor(kernels=[k], semaphores=[], cbs=[cb]))

    def _compact2(self, cand: ttnn.Tensor):
        """kp_compact2.cpp: both RISCs of core (0, 0), candidate slots read in place from its L1."""
        nk = self.nms
        core = ttnn.CoreCoord(0, 0)
        grid = ttnn.CoreRangeSet([ttnn.CoreRange(core, core)])
        acc = ttnn.TensorAccessorArgs(self.hdr).get_compile_time_args() + ttnn.TensorAccessorArgs(self.tab).get_compile_time_args()
        args = [cand.buffer_address(), self.hdr.buffer_address(), self.tab.buffer_address()]
        defs = []
        io = [cand, self.hdr, self.tab]
        if self.one_d2h:
            acc += ttnn.TensorAccessorArgs(nk.params).get_compile_time_args() + ttnn.TensorAccessorArgs(self.chunks_a[0]).get_compile_time_args()
            args += [nk.params.buffer_address()] + [t.buffer_address() for t in self.chunks_a]
            defs = [("KPC_HR", str(self.hr)), ("KPC_BSTEP", str(self.chunk)), ("KPC_C", str(self.c))]
            io += [nk.params] + list(self.chunks_a)
        ks = []
        for proc, cfg in ((0, ttnn.ReaderConfigDescriptor()), (1, ttnn.WriterConfigDescriptor())):
            rt = ttnn.RuntimeArgs()
            rt[0][0] = list(args)
            ct = [0, nk.nslots, nk.cap, self.W, self.kmax, nk.rows, nk._split, self.H, proc] + acc
            ks.append(ttnn.KernelDescriptor(kernel_source=self._src_compact2, source_type=ttnn.KernelDescriptor.SourceType.SOURCE_CODE,
                                            core_ranges=grid, compile_time_args=ct, runtime_args=rt, config=cfg, defines=defs))
        tab_region = ((2 * self.W + 2 * self.H) * 4 + 63) // 64 * 64
        size = 2 * tab_region + (16 + 4 * self.kmax) * 4 + 64
        fmt = ttnn.CBFormatDescriptor(buffer_index=0, data_format=ttnn.bfloat16, page_size=size)
        cb = ttnn.CBDescriptor(total_size=size, core_ranges=grid, format_descriptors=[fmt])
        sem = ttnn.SemaphoreDescriptor(id=0, core_ranges=grid, initial_value=0)
        ttnn.generic_op(io, ttnn.ProgramDescriptor(kernels=ks, semaphores=[sem], cbs=[cb]))

    def _compact3(self, cand: ttnn.Tensor, rec: ttnn.Tensor):
        """kp_compact3.cpp: both RISCs of core (0, 0); the header entries are gathered by DMA from the
        unfold's per-slot records (one NoC read per non-empty slot)."""
        nk = self.nms
        core = ttnn.CoreCoord(0, 0)
        grid = ttnn.CoreRangeSet([ttnn.CoreRange(core, core)])
        acc = ttnn.TensorAccessorArgs(self.hdr).get_compile_time_args()
        args = [cand.buffer_address(), self.hdr.buffer_address(), rec.buffer_address()]
        defs = []
        io = [cand, self.hdr, rec]
        if self.one_d2h:
            acc += ttnn.TensorAccessorArgs(nk.params).get_compile_time_args() + ttnn.TensorAccessorArgs(self.chunks_a[0]).get_compile_time_args()
            args += [nk.params.buffer_address()] + [t.buffer_address() for t in self.chunks_a]
            defs = [("KPC_HR", str(self.hr)), ("KPC_BSTEP", str(self.chunk)), ("KPC_C", str(self.c))]
            io += [nk.params] + list(self.chunks_a)
        for c in nk.cores:
            w = self.device.worker_core_from_logical_core(c)
            args.append((w.x << 16) | w.y)
        ks = []
        for proc, cfg in ((0, ttnn.ReaderConfigDescriptor()), (1, ttnn.WriterConfigDescriptor())):
            rt = ttnn.RuntimeArgs()
            rt[0][0] = list(args)
            ct = [0, nk.nslots, nk.cap, self.kmax, proc] + acc
            ks.append(ttnn.KernelDescriptor(kernel_source=self._src_compact3, source_type=ttnn.KernelDescriptor.SourceType.SOURCE_CODE,
                                            core_ranges=grid, compile_time_args=ct, runtime_args=rt, config=cfg, defines=defs))
        size = (16 + 4 * self.kmax) * 4 + 128  # header image + parameter page + exchange words
        fmt = ttnn.CBFormatDescriptor(buffer_index=0, data_format=ttnn.bfloat16, page_size=size)
        cb = ttnn.CBDescriptor(total_size=size, core_ranges=grid, format_descriptors=[fmt])
        sems = [ttnn.SemaphoreDescriptor(id=i, core_ranges=grid, initial_value=0) for i in (0, 1)]
        ttnn.generic_op(io, ttnn.ProgramDescriptor(kernels=ks, semaphores=sems, cbs=[cb]))

    def sample_fused(self, d_rm: ttnn.Tensor, hdr: ttnn.Tensor, chunks):
        """Gather + ((G0*W0 + G1*W1) + G2*W2) + G3*W3 (fp32 SFPU) + bucket write, ONE program
        (sample_fused_{reader,compute,writer}.cpp). Unit u = (tile row u // NQ, channel chunk
        u % NQ of CPU channels); data stay row-major in 2 KB / 4 KB "pseudo tiles" (elementwise math
        is order-agnostic). Units are spread round-robin over all cores so the active (low) tile rows
        land on distinct cores. Bit-identical to :meth:`sample_unfused`."""
        kv, c = self.kmax, self.c
        cpu = int(os.environ.get("SP_SF_CPU", "128"))
        nq, kt = c // cpu, 32 * cpu // 1024
        units = list(range((kv // 32) * nq))
        per_core = {}
        for i, u in enumerate(units):
            per_core.setdefault(i % self.ncores, []).append(u)
        grid = self._unit_grid(per_core)
        CB_G, CB_W, CB_O, CB_S, CB_S2 = 0, 1, 2, 3, 4
        rt_r, rt_c, rt_w = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
        for i in sorted(per_core):
            cc, us = self.all_cores[i], per_core[i]
            rt_r[cc.x][cc.y] = [d_rm.buffer_address(), hdr.buffer_address(), len(us)] + us + [self.wtab.buffer_address()]
            rt_c[cc.x][cc.y] = [len(us)]
            rt_w[cc.x][cc.y] = [hdr.buffer_address(), self.wtab.buffer_address(), len(us)] + [t.buffer_address() for t in chunks] + us
        acc_r = ttnn.TensorAccessorArgs(d_rm).get_compile_time_args() + ttnn.TensorAccessorArgs(hdr).get_compile_time_args()
        # SP_SF_SPLIT=k (default 16): with one unit per core the reader fills the tap-weight pages of
        # keypoints 0..k-1, the writer the rest (the fill was ~14 us of scalar stores on the writer)
        split = int(os.environ.get("SP_SF_SPLIT", "16"))
        # the split applies to the FIRST unit of every core (CB_W slot 0, free at kernel start); further units
        # (tile rows >= ncores / NQ, i.e. only more than ~960 keypoints with NQ = 4) are filled by the writer alone
        one_unit = split <= 32
        if split > 0 and one_unit:
            acc_r += ttnn.TensorAccessorArgs(self.wtab).get_compile_time_args()
        acc_w = []
        for t in (hdr, self.wtab, chunks[0]):
            acc_w += ttnn.TensorAccessorArgs(t).get_compile_time_args()
        ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True)
        utd = [ttnn.UnpackToDestMode.Default] * 64
        utd[CB_W] = ttnn.UnpackToDestMode.UnpackToDestFp32
        ccfg.unpack_to_dest_mode = type(ccfg.unpack_to_dest_mode)(utd)
        src_r, src_c, src_w = self._src_f
        SC = ttnn.KernelDescriptor.SourceType.SOURCE_CODE
        defs = [(d, "1") for d in os.environ.get("SP_SF_DEFS", "").split(",") if d]
        if split > 0 and one_unit:
            defs += [("SF_SPLIT", str(split)), ("SF_W", str(self.W)), ("SF_CBW", str(CB_W))]
        if self.one_d2h:
            defs += [("SF_HR", str(self.hr))]
        ks = [
            ttnn.KernelDescriptor(kernel_source=src_r, source_type=SC, core_ranges=grid,
                                  compile_time_args=[CB_G, CB_S, c, kv, cpu] + acc_r,
                                  runtime_args=rt_r, config=ttnn.ReaderConfigDescriptor(), defines=defs),
            ttnn.KernelDescriptor(kernel_source=src_c, source_type=SC, core_ranges=grid,
                                  compile_time_args=[CB_G, CB_W, CB_O, kt], runtime_args=rt_c, config=ccfg, defines=defs),
            ttnn.KernelDescriptor(kernel_source=src_w, source_type=SC, core_ranges=grid,
                                  compile_time_args=[CB_W, CB_O, CB_S2, c, kv, self.chunk, cpu, self.W] + acc_w,
                                  runtime_args=rt_w, config=ttnn.WriterConfigDescriptor(), defines=defs),
        ]

        def cb(idx, fmt, page, npages):
            f = ttnn.CBFormatDescriptor(buffer_index=idx, data_format=fmt, page_size=page)
            return ttnn.CBDescriptor(total_size=page * npages, core_ranges=grid, format_descriptors=[f])

        cbs = [
            cb(CB_G, ttnn.bfloat16, 2048, 8 * kt),
            cb(CB_W, ttnn.float32, 4096, 8 * kt),
            cb(CB_O, ttnn.float32, 4096, 2),
            cb(CB_S, ttnn.bfloat16, 64 + 512 + 32 * 64, 1),
            cb(CB_S2, ttnn.bfloat16, 64 + 512 + 32 * 64, 1),
        ]
        sems = [ttnn.SemaphoreDescriptor(id=0, core_ranges=grid, initial_value=0)] if (split > 0 and one_unit) else []
        ttnn.generic_op([d_rm, hdr, self.wtab] + list(chunks), ttnn.ProgramDescriptor(kernels=ks, semaphores=sems, cbs=cbs))

    def _unit_grid(self, per_core):
        """Core range set of the cores 0..n-1 (row-wise) that hold units: merged rectangles (SP_SF_RECT=1), so the
        dispatcher multicasts the program to a few ranges instead of writing it to n single-core ranges one by one
        (measured: 120 single-core ranges -> 17.5 us launch gap after the NMS op)."""
        idx = sorted(per_core)
        if os.environ.get("SP_SF_RECT", "1") == "1" and idx == list(range(len(idx))):
            g = self.device.compute_with_storage_grid_size()
            return ttnn.num_cores_to_corerangeset(len(idx), g, row_wise=True)
        return ttnn.CoreRangeSet([ttnn.CoreRange(self.all_cores[i], self.all_cores[i]) for i in idx])

    def sample_direct(self, d_rm: ttnn.Tensor, chunks):
        """SP_SF_DIRECT=1 (main trace): :meth:`sample_pipe` that builds its keypoint list itself from the NMS unfold's
        dense slot counts and records (sample_direct_*.cpp + sample_kplist.hpp.inc) and writes the keypoint header into
        the bucket like kp_compact3.cpp: the compaction op is gone. Same list, same sampling math (bit-identical)."""
        nk = self.nms
        kv, c = self.kmax, self.c
        cpu = 128
        nq = c // cpu
        split = self.split and chunks is self.chunks_a
        #: SP_SF_KPU: keypoints per unit (32 = one tile row; 16 = half a tile row: twice the units / cores, half the
        #: gathers per core)
        kpu = int(os.environ.get("SP_SF_KPU", "16"))
        assert kpu in (16, 32)
        units = list(range((kv // kpu) * nq))
        per_core = {}
        for i, u in enumerate(units):
            per_core.setdefault(i % self.ncores, []).append(u)
        grid = self._unit_grid(per_core)
        CB_G, CB_W, CB_O, CB_S, CB_S2, CB_W0 = 0, 1, 2, 3, 4, 5
        cnt, rec = nk.counts, nk.rec
        rt_r, rt_c, rt_w = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
        for i in sorted(per_core):
            cc, us = self.all_cores[i], per_core[i]
            rt_r[cc.x][cc.y] = [d_rm.buffer_address(), cnt.buffer_address(), len(us)] + us + [self.wtab.buffer_address()]
            rt_c[cc.x][cc.y] = [len(us)]
            rt_w[cc.x][cc.y] = [cnt.buffer_address(), self.wtab.buffer_address(), len(us)] + us
        common = [rec.buffer_address()]
        for core in nk.cores:
            w = self.device.worker_core_from_logical_core(core)
            common.append((w.x << 16) | w.y)
        comb = len(common)  # writer's common args from here: bucket addresses, params (, tails, split rows)
        common_w = (common + [t.buffer_address() for t in chunks] + [nk.params.buffer_address()]
                    + ([t.buffer_address() for t in self.tails_a] + list(self.split_rows) if split else []))
        acc_r = []
        for t in (d_rm, cnt, self.wtab):
            acc_r += ttnn.TensorAccessorArgs(t).get_compile_time_args()
        acc_w = []
        for t in (cnt, self.wtab, chunks[0], nk.params):
            acc_w += ttnn.TensorAccessorArgs(t).get_compile_time_args()
        ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True)
        utd = [ttnn.UnpackToDestMode.Default] * 64
        utd[CB_W] = ttnn.UnpackToDestMode.UnpackToDestFp32
        utd[CB_W0] = ttnn.UnpackToDestMode.UnpackToDestFp32
        ccfg.unpack_to_dest_mode = type(ccfg.unpack_to_dest_mode)(utd)
        src_r, src_c, src_w = self._src_d
        SC = ttnn.KernelDescriptor.SourceType.SOURCE_CODE
        defs = [(d, "1") for d in os.environ.get("SP_SF_DEFS", "").split(",") if d]
        if os.environ.get("SP_SF_WC16", "1") == "1":
            defs.append(("SF_WC16", "1"))  # 16-word compact weight blocks, broadcast by SFPTRANSP
        if split:
            defs.append(("SF_TAIL", "1"))  # rows >= split_rows[b] go to tails_a[b]
        defs.append(("SF_KT", str(kpu // 8)))
        defs.append(("SF_COMB", str(comb)))
        defs += [("SF_W", str(self.W)), ("SF_CBW0", str(CB_W0)), ("SF_HR", str(self.hr)), ("SF_NSLOT", str(nk.nslots)),
                 ("SF_CAP", str(nk.cap)), ("SF_CNT_PAGE", str(nk.CNT_PAGE))]
        kr = ttnn.KernelDescriptor(kernel_source=src_r, source_type=SC, core_ranges=grid,
                                   compile_time_args=[CB_G, CB_S, c, kv, cpu] + acc_r,
                                   runtime_args=rt_r, config=ttnn.ReaderConfigDescriptor(), defines=defs)
        kr.common_runtime_args = common
        kw_ = ttnn.KernelDescriptor(kernel_source=src_w, source_type=SC, core_ranges=grid,
                                    compile_time_args=[CB_W, CB_O, CB_S2, c, kv, self.chunk, cpu, self.W] + acc_w,
                                    runtime_args=rt_w, config=ttnn.WriterConfigDescriptor(), defines=defs)
        kw_.common_runtime_args = common_w
        ks = [kr,
              ttnn.KernelDescriptor(kernel_source=src_c, source_type=SC, core_ranges=grid,
                                    compile_time_args=[CB_G, CB_W, CB_O, 4, CB_W0], runtime_args=rt_c, config=ccfg, defines=defs),
              kw_]

        def cb(idx, fmt, page, npages):
            f = ttnn.CBFormatDescriptor(buffer_index=idx, data_format=fmt, page_size=page)
            return ttnn.CBDescriptor(total_size=page * npages, core_ranges=grid, format_descriptors=[f])

        ns = ((nk.nslots * 16 + nk.CNT_PAGE - 1) // nk.CNT_PAGE) * nk.CNT_PAGE
        cbs = [
            cb(CB_G, ttnn.bfloat16, 2048, 32),
            cb(CB_W, ttnn.float32, 4096, 8),
            cb(CB_O, ttnn.float32, 4096, 2),
            cb(CB_S, ttnn.bfloat16, ns + 512 + 8 * 64, 1),
            cb(CB_S2, ttnn.bfloat16, ns + 128 + 512 + 32 * 64, 1),
            cb(CB_W0, ttnn.float32, 4096, 2),
        ]
        io = [d_rm, cnt, rec, self.wtab, nk.params] + list(chunks) + (list(self.tails_a) if split else [])
        ttnn.generic_op(io, ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))

    def sample_pipe(self, d_rm: ttnn.Tensor, hdr: ttnn.Tensor, chunks):
        """SP_SF_PIPE=1: :meth:`sample_fused` with per-keypoint-group pipelining (sample_pipe_*.cpp): the reader pushes
        the gathered pages per group of 8 keypoints, the writer fills compact tap-weight pages (64 words per keypoint and
        tap, replicated over the 128 channels by the SFPU) for groups 1..3 and the reader for group 0 after its gathers,
        the compute kernel runs group by group as soon as both are in (order 1, 2, 3, 0). Same math, bit-identical."""
        kv, c = self.kmax, self.c
        cpu = 128
        nq = c // cpu
        units = list(range((kv // 32) * nq))
        per_core = {}
        for i, u in enumerate(units):
            per_core.setdefault(i % self.ncores, []).append(u)
        grid = self._unit_grid(per_core)
        CB_G, CB_W, CB_O, CB_S, CB_S2, CB_W0 = 0, 1, 2, 3, 4, 5
        rt_r, rt_c, rt_w = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
        for i in sorted(per_core):
            cc, us = self.all_cores[i], per_core[i]
            rt_r[cc.x][cc.y] = [d_rm.buffer_address(), hdr.buffer_address(), len(us)] + us + [self.wtab.buffer_address()]
            rt_c[cc.x][cc.y] = [len(us)]
            rt_w[cc.x][cc.y] = [hdr.buffer_address(), self.wtab.buffer_address(), len(us)] + [t.buffer_address() for t in chunks] + us
        acc_r = []
        for t in (d_rm, hdr, self.wtab):
            acc_r += ttnn.TensorAccessorArgs(t).get_compile_time_args()
        acc_w = []
        for t in (hdr, self.wtab, chunks[0]):
            acc_w += ttnn.TensorAccessorArgs(t).get_compile_time_args()
        ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True)
        utd = [ttnn.UnpackToDestMode.Default] * 64
        utd[CB_W] = ttnn.UnpackToDestMode.UnpackToDestFp32
        utd[CB_W0] = ttnn.UnpackToDestMode.UnpackToDestFp32
        ccfg.unpack_to_dest_mode = type(ccfg.unpack_to_dest_mode)(utd)
        src_r, src_c, src_w = self._src_p
        SC = ttnn.KernelDescriptor.SourceType.SOURCE_CODE
        defs = [(d, "1") for d in os.environ.get("SP_SF_DEFS", "").split(",") if d]
        if os.environ.get("SP_SF_WC16", "1") == "1":
            defs.append(("SF_WC16", "1"))  # 16-word compact weight blocks, broadcast by SFPTRANSP
        defs += [("SF_W", str(self.W)), ("SF_CBW0", str(CB_W0))]
        if self.one_d2h:
            defs += [("SF_HR", str(self.hr))]
        ks = [
            ttnn.KernelDescriptor(kernel_source=src_r, source_type=SC, core_ranges=grid,
                                  compile_time_args=[CB_G, CB_S, c, kv, cpu] + acc_r,
                                  runtime_args=rt_r, config=ttnn.ReaderConfigDescriptor(), defines=defs),
            ttnn.KernelDescriptor(kernel_source=src_c, source_type=SC, core_ranges=grid,
                                  compile_time_args=[CB_G, CB_W, CB_O, 4, CB_W0], runtime_args=rt_c, config=ccfg, defines=defs),
            ttnn.KernelDescriptor(kernel_source=src_w, source_type=SC, core_ranges=grid,
                                  compile_time_args=[CB_W, CB_O, CB_S2, c, kv, self.chunk, cpu, self.W] + acc_w,
                                  runtime_args=rt_w, config=ttnn.WriterConfigDescriptor(), defines=defs),
        ]

        def cb(idx, fmt, page, npages):
            f = ttnn.CBFormatDescriptor(buffer_index=idx, data_format=fmt, page_size=page)
            return ttnn.CBDescriptor(total_size=page * npages, core_ranges=grid, format_descriptors=[f])

        cbs = [
            cb(CB_G, ttnn.bfloat16, 2048, 32),
            cb(CB_W, ttnn.float32, 4096, 8),
            cb(CB_O, ttnn.float32, 4096, 2),
            cb(CB_S, ttnn.bfloat16, 64 + 512 + 8 * 64, 1),
            cb(CB_S2, ttnn.bfloat16, 64 + 512 + 32 * 64, 1),
            cb(CB_W0, ttnn.float32, 4096, 2),
        ]
        ttnn.generic_op([d_rm, hdr, self.wtab] + list(chunks), ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))

    def sample(self, d_rm: ttnn.Tensor, hdr: ttnn.Tensor, chunks):
        if self.fused and os.environ.get("SP_SF_PIPE", "1") == "1" and self.c == 256:
            return self.sample_pipe(d_rm, hdr, chunks)
        if self.fused:
            return self.sample_fused(d_rm, hdr, chunks)
        return self.sample_unfused(d_rm, hdr, chunks)

    def sample_unfused(self, d_rm: ttnn.Tensor, hdr: ttnn.Tensor, chunks):
        TILE = ttnn.TILE_LAYOUT
        kv, c = self.kmax, self.c
        gs = [ttnn.allocate_tensor_on_device(ttnn.Shape([kv, c]), ttnn.bfloat16, TILE, self.device, self.l1) for _ in range(4)]
        ws = [ttnn.allocate_tensor_on_device(ttnn.Shape([kv, 1]), ttnn.float32, TILE, self.device, self.l1) for _ in range(4)]
        acc = []
        for t in (d_rm, hdr, self.wtab, gs[0], ws[0]):
            acc += ttnn.TensorAccessorArgs(t).get_compile_time_args()
        # unit = (tr << 3) | (h << 2) | t, ordered by tile row
        units = [(tr << 3) | (h << 2) | t for tr in range(kv // 32) for h in (0, 1) for t in range(4)]
        ch = c // 2
        self._spread_program(
            units,
            [d_rm.buffer_address(), hdr.buffer_address(), self.wtab.buffer_address()],
            lambda u: [u, gs[u & 3].buffer_address(), ws[u & 3].buffer_address()],
            self._src_gather,
            [c, kv, self.W, self.H] + acc,
            64 + 512 + 32 * 64 + 32 * ch * 2 + (ch // 32) * 2048 + 4096,
            [d_rm, hdr, self.wtab] + gs + ws,
        )
        total = None
        for t in range(4):
            p = ttnn.multiply(gs[t], ws[t], dtype=ttnn.float32, memory_config=self.l1)
            ttnn.deallocate(gs[t])
            ttnn.deallocate(ws[t])
            if total is None:
                total = p
            else:
                nxt = ttnn.add(total, p, memory_config=self.l1)
                ttnn.deallocate(total)
                ttnn.deallocate(p)
                total = nxt
        acc = []
        for t in (total, hdr, chunks[0]):
            acc += ttnn.TensorAccessorArgs(t).get_compile_time_args()
        tc = c // 32
        units = list(range((kv // 32) * tc))
        self._spread_program(
            units,
            [total.buffer_address(), hdr.buffer_address()],
            lambda u: [u],
            self._src_untilize,
            [c, kv, self.chunk] + acc,
            64 + 4096,
            [total, hdr] + list(chunks),
            extra=[t.buffer_address() for t in chunks],
        )
        ttnn.deallocate(total)

    def run_main(self, cand: ttnn.Tensor, d_rm: ttnn.Tensor):
        """Main-trace part: keypoint list from the candidates, then sampling into chunks_a."""
        self.d = d_rm
        self.cand_ref = cand
        if self.nms.cnt_mode and self.one_d2h and self.fused and self.c == 256 and os.environ.get("SP_SF_PIPE", "1") == "1":
            # SP_SF_DIRECT: no compaction op; the sampler gathers the list and writes the header into the bucket
            self.sample_direct(d_rm, self.chunks_a)
            ttnn.deallocate(self.nms.counts)
            ttnn.deallocate(self.nms.rec)
            self.nms.counts = self.nms.rec = None
            return
        self.compact(cand)
        self.sample(d_rm, self.hdr, self.chunks_a)

    def capture_b(self, cq_id: int = 0):
        """Second trace (top-k fallback): sampling at the host-uploaded hdr_b keypoints into chunks_b."""
        self.sample(self.d, self.hdr_b, self.chunks_b)  # compile
        ttnn.synchronize_device(self.device)
        self.tid_b = ttnn.begin_trace_capture(self.device, cq_id=cq_id)
        self.sample(self.d, self.hdr_b, self.chunks_b)
        ttnn.end_trace_capture(self.device, self.tid_b, cq_id=cq_id)
        ttnn.synchronize_device(self.device)

    def release(self):
        if self.tid_b is not None:
            ttnn.release_trace(self.device, self.tid_b)
            self.tid_b = None
        for t in [self.hdr, self.hdr_b, self.tab, self.wtab] + self.chunks_a + self.chunks_b + self.tails_a:
            ttnn.deallocate(t)
        self.chunks_a, self.chunks_b, self.tails_a = [], [], []