# 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("> 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 = [], [], []