code/models/tt/pool_kernels.py

# 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))