# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 """2x2 / stride-2 max pool as ONE custom generic_op (no halo), for the 64-channel encoder blocks. Input: L1 height-sharded ROW_MAJOR bf16 [1, 1, N*H*W, 64] whose shards hold whole, even numbers of image rows (block 0: 4 rows of 640 px per core on 120 cores; block 1: 2 rows of 320). Output: L1 height-sharded TILE bf16 [1, 1, N*H/2*W/2, 64] on the same cores (shard = the pooled rows). The compute kernel tilizes 64-pixel row chunks read as 32 "pixel pairs" x 128 channels (even/odd pixels land in different tiles) and takes the SFPU max of the 4 tiles of each 2x2 window: exact. """ 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_pool") def _src(name): with open(os.path.join(_KDIR, name)) as f: return f.read() class Pool2x2: C = 64 def __init__(self, device): self.device = device self._src = [_src(f"pool2x2_{k}.cpp") for k in ("reader", "compute", "writer")] def supports(self, x: ttnn.Tensor, h: int, w: int, c: int) -> bool: if c != self.C or w % 64 or h % 2 or not x.is_sharded() or x.layout != ttnn.ROW_MAJOR_LAYOUT: return False mc = x.memory_config() if mc.memory_layout != ttnn.TensorMemoryLayout.HEIGHT_SHARDED or mc.buffer_type != ttnn.BufferType.L1: return False rows = self._rows(x) n = x.shape[-2] ncores = mc.shard_spec.grid.num_cores() return x.shape[-1] == c and rows % (2 * w) == 0 and rows * ncores == n def _rows(self, x: ttnn.Tensor) -> int: """Pixels (C-channel rows) per shard; also right for a zero-copy view of a cell-conv output whose shard spec still says [rows / cell, cell * C].""" sh = x.memory_config().shard_spec.shape return sh[0] * sh[1] // self.C def __call__(self, x: ttnn.Tensor, h: int, w: int) -> ttnn.Tensor: mc = x.memory_config() grid = mc.shard_spec.grid rows = self._rows(x) rows_out_img = rows // (2 * w) # output image rows per core ch = w // 64 # 64-pixel chunks per image row n_out = x.shape[-2] // 4 out_rows = rows // 4 omc = ttnn.MemoryConfig( ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, ttnn.ShardSpec(grid, [out_rows, self.C], mc.shard_spec.orientation), ) out = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, n_out, self.C]), ttnn.bfloat16, ttnn.TILE_LAYOUT, self.device, omc) CB_IN, CB_T, CB_OUT = 0, 1, 2 n_in_pages = rows * self.C * 2 // 8192 n_out_pages = out_rows * self.C * 2 // 2048 cb_in = ttnn.cb_descriptor_from_sharded_tensor(CB_IN, x) cb_in.format_descriptors = [ttnn.CBFormatDescriptor(buffer_index=CB_IN, data_format=ttnn.bfloat16, page_size=8192)] cb_out = ttnn.cb_descriptor_from_sharded_tensor(CB_OUT, out) cb_out.format_descriptors = [ttnn.CBFormatDescriptor(buffer_index=CB_OUT, data_format=ttnn.bfloat16, page_size=2048)] f = ttnn.CBFormatDescriptor(buffer_index=CB_T, data_format=ttnn.bfloat16, page_size=2048) cb_t = ttnn.CBDescriptor(total_size=2048 * 8 * ch, core_ranges=grid, format_descriptors=[f]) SC = ttnn.KernelDescriptor.SourceType.SOURCE_CODE r, c, wr = self._src ks = [ ttnn.KernelDescriptor(kernel_source=r, source_type=SC, core_ranges=grid, compile_time_args=[CB_IN, n_in_pages], runtime_args=[], config=ttnn.ReaderConfigDescriptor()), ttnn.KernelDescriptor(kernel_source=c, source_type=SC, core_ranges=grid, compile_time_args=[CB_IN, CB_T, CB_OUT, rows_out_img, ch], runtime_args=[], config=ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4)), ttnn.KernelDescriptor(kernel_source=wr, source_type=SC, core_ranges=grid, compile_time_args=[CB_OUT, n_out_pages], runtime_args=[], config=ttnn.WriterConfigDescriptor()), ] return ttnn.generic_op([x, out], ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=[cb_in, cb_t, cb_out])) class U8ToBf16: """uint8 image upload -> the bf16 cell input of the ``l1`` encoder, ONE data-movement generic_op (kernels/sp_input/u8_to_bf16.cpp): table lookup bf16(fp32(u) / 255) on both RISCs of every core, local L1 only. Bit-identical to the host fp32 /255 + bf16 cast (fused_host.U8_TO_BF16).""" def __init__(self, device): import torch from .fused_host import U8_TO_BF16 self.device = device kdir = os.path.join(os.path.dirname(_KDIR), "sp_input") with open(os.path.join(kdir, "u8_to_bf16.cpp")) as f: self._src = f.read() bits = (U8_TO_BF16.view(torch.int16).to(torch.int32) & 0xFFFF).tolist() self._src = "#define SP_U8_LUT " + ",".join(f"0x{b:04x}" for b in bits) + "\n" + self._src def __call__(self, x: ttnn.Tensor, out_mc: ttnn.MemoryConfig, out_shape) -> ttnn.Tensor: out = ttnn.allocate_tensor_on_device(ttnn.Shape(list(out_shape)), ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT, self.device, out_mc) grid = out_mc.shard_spec.grid sh = out_mc.shard_spec.shape n_px = sh[0] * sh[1] xs = x.memory_config().shard_spec if xs.grid != grid or xs.shape[0] * xs.shape[1] != n_px or n_px % 8: raise ValueError("u8 input and bf16 cell input must split the pixels identically") CB_IN, CB_OUT = 0, 1 cb_in = ttnn.cb_descriptor_from_sharded_tensor(CB_IN, x) cb_in.format_descriptors = [ttnn.CBFormatDescriptor(buffer_index=CB_IN, data_format=ttnn.bfloat16, page_size=n_px)] cb_out = ttnn.cb_descriptor_from_sharded_tensor(CB_OUT, out) cb_out.format_descriptors = [ttnn.CBFormatDescriptor(buffer_index=CB_OUT, data_format=ttnn.bfloat16, page_size=2 * n_px)] SC = ttnn.KernelDescriptor.SourceType.SOURCE_CODE ks = [] for proc, cfg in ((0, ttnn.ReaderConfigDescriptor()), (1, ttnn.WriterConfigDescriptor())): ks.append(ttnn.KernelDescriptor(kernel_source=self._src, source_type=SC, core_ranges=grid, compile_time_args=[CB_IN, CB_OUT, n_px, proc], runtime_args=[], config=cfg)) return ttnn.generic_op([x, out], ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=[cb_in, cb_out])) class Pool2x2Rows: """2x2 / stride-2 max pool for an activation sharded as ONE image row per core (block 2: 160 px x 128 ch on 120 cores), ONE generic_op instead of halo + Move + max_pool2d (kernels/sp_pool/pool_rows_{reader,compute}.cpp). Input: L1 height-sharded ROW_MAJOR bf16 [1, 1, H*W, C], shard [W, C] on H cores; output: L1 height-sharded ROW_MAJOR bf16 [1, 1, H/2*W/2, C], shard [W/4, C] on the same cores (what ttnn.max_pool2d produces): output core c holds half c % 2 of output row c // 2. Each RISC gathers the even / odd pixels of one of the two input rows (NoC reads, one row is local), the compute kernel takes the SFPU max of the four 2 KB pseudo tiles: exact.""" def __init__(self, device): self.device = device self._r = _src("pool_rows_reader.cpp") self._c = _src("pool_rows_compute.cpp") def supports(self, x: ttnn.Tensor, h: int, w: int, c: int) -> bool: if not x.is_sharded() or x.layout != ttnn.ROW_MAJOR_LAYOUT or x.dtype != ttnn.bfloat16 or h % 2 or w % 4: return False mc = x.memory_config() if mc.memory_layout != ttnn.TensorMemoryLayout.HEIGHT_SHARDED or mc.buffer_type != ttnn.BufferType.L1: return False sh = mc.shard_spec.shape return (x.shape[-1] == c and sh[1] == c and sh[0] == w and mc.shard_spec.grid.num_cores() == h and (w // 4 * c * 2) % 2048 == 0 and 2048 % (c * 2) == 0 and x.shape[-2] == h * w) def __call__(self, x: ttnn.Tensor, h: int, w: int) -> ttnn.Tensor: mc = x.memory_config() grid = mc.shard_spec.grid c = x.shape[-1] npx = w // 4 omc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, ttnn.ShardSpec(grid, [npx, c], mc.shard_spec.orientation)) out = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, h * w // 4, c]), ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT, self.device, omc) cores = ttnn.corerange_to_cores(grid, row_wise=mc.shard_spec.orientation == ttnn.ShardOrientation.ROW_MAJOR) pages = npx * c * 2 // 2048 CB_E0, CB_O0, CB_E1, CB_O1, CB_OUT = 0, 1, 2, 3, 4 cbs = [] for i in (CB_E0, CB_O0, CB_E1, CB_O1): f = ttnn.CBFormatDescriptor(buffer_index=i, data_format=ttnn.bfloat16, page_size=2048) cbs.append(ttnn.CBDescriptor(total_size=2048 * min(pages, 4), core_ranges=grid, format_descriptors=[f])) cb_out = ttnn.cb_descriptor_from_sharded_tensor(CB_OUT, out) cb_out.format_descriptors = [ttnn.CBFormatDescriptor(buffer_index=CB_OUT, data_format=ttnn.bfloat16, page_size=2048)] cbs.append(cb_out) addr = x.buffer_address() ks = [] SC = ttnn.KernelDescriptor.SourceType.SOURCE_CODE for proc, cfg, ce, co in ((0, ttnn.ReaderConfigDescriptor(), CB_E0, CB_O0), (1, ttnn.WriterConfigDescriptor(), CB_E1, CB_O1)): rt = ttnn.RuntimeArgs() for i, core in enumerate(cores): hh = i % 2 src = cores[i - hh + proc] nc = self.device.worker_core_from_logical_core(src) rt[core.x][core.y] = [nc.x, nc.y, addr, hh * (w // 2)] ks.append(ttnn.KernelDescriptor(kernel_source=self._r, source_type=SC, core_ranges=grid, compile_time_args=[ce, co, npx, c * 2], runtime_args=rt, config=cfg)) ks.append(ttnn.KernelDescriptor(kernel_source=self._c, source_type=SC, core_ranges=grid, compile_time_args=[CB_E0, CB_O0, CB_E1, CB_O1, CB_OUT, pages], runtime_args=[], config=ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4))) return ttnn.generic_op([x, out], ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))