code/tt_diffusion_planner/ttaw/ops/conv.py

# SPDX-License-Identifier: Apache-2.0
"""C17: conv2d builders for NHWC feature maps (PLAN.md section 1.2 row C17; probes P5 and P14, common/PROBES.md).

A conv layer of a port is built once from host weights and called on device feature maps, eagerly for the warm-up
and then inside trace captures. What every builder here guarantees:

- **Weights prepared once.** The first call hands the host weights to ``ttnn.conv2d`` (which prepares them for the
  chosen sharding and moves them to the device, a host write); the prepared device tensors are kept per input spec
  and reused, so a traced call never writes from the host ("Writes are not supported during trace capture",
  RP section 1.4). Call each layer once eagerly before capturing (``TraceRunner`` warm-up does).
- **Explicit precision.** ``compute_config`` comes from :mod:`..precision` (default :data:`CONV_PRECISION`: HiFi4 +
  fp32 accumulation + packer L1 accumulation, exact SFPU activations; never ttnn's implicit defaults; without the L1
  accumulation the partial sums of wide convs round to bf16 between inner blocks). Weights default to bf16
  (``Precision.weights`` selects bfp8); ``output_dtype`` may be float32 (e.g. heads whose logits feed thresholds).
- **Fused bias + activation** in the conv epilogue: ``relu`` (packer), ``relu6`` / ``silu`` / ``gelu`` /
  ``sigmoid`` (SFPU). RELU6 is verified on the device (P5: output in [0, 6], PCC >= 0.99997). HARDSWISH is refused:
  the fused form is silently skipped on this tree (P5; apply ``ttnn.hardswish`` as a separate op).
- **Grid from the device.** Nothing here hard-codes 12x10 / 11x10: ttnn's auto-sharding reads
  ``device.compute_with_storage_grid_size()``; explicit shard overrides go through ``conv_config``.
- **Memory placement.** ``slicing="auto"`` keeps ttnn's rule: a DRAM input runs the DRAM-sliced path with automatic
  slice counts (P14: every accepted config captures and replays bit-exactly; explicit counts are often rejected, so
  never hard-code them) and writes a DRAM-interleaved output; an L1 input runs the L1 path. 1x1 stride-1 convs are
  matmuls and always run the L1 path; ``output_memory_config`` (default ``"dram"``) brings their output back to DRAM
  so a graph can keep every activation in DRAM between ops (the robust layout of a first port; L1 residency is an
  optimization). Feed large-kernel convs TILE tensors (``[1, 1, N*H*W, C]``, what a producing conv writes): from
  row-major DRAM they cost 2x (P5: 926 vs 482 us for YOLOX's 12 large-kernel convs).

Feature maps travel as :class:`FeatureMap` (device tensor + logical N, H, W, C), because ttnn's flattened
``[1, 1, N*H*W, C]`` layout no longer carries the spatial size.

Builders: :class:`Conv2d` (any kernel / stride / padding / dilation / groups), :class:`KSplitConv`
(``conv(concat(x_1..x_n))`` as ``sum_i conv_i(x_i)`` without building the concat; P14: 1.7-4.1x faster and more
accurate for CenterPoint / PointPainting / FRNet's wide concats), :class:`ConvTranspose2d` for kernel == stride
(P5: k = 2 -> ``ttnn.conv_transpose2d``; k >= 3 -> ``ttnn.linear`` + depth-to-space, 5x faster at k = 4), and the
host helpers :func:`merge_sibling_convs` (one conv for sibling heads that read the same input) /
:func:`split_k_weights`. Glue: :func:`concat_channels`, :func:`split_channels`, :func:`residual_add`,
:func:`apply_activation`, :func:`fmap_from_numpy` / :func:`fmap_to_numpy` (NCHW <-> device NHWC).
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

import numpy as np

from ..precision import Precision
from ..tensors import dtype_name, round_to_bf16, ttnn_dtype

__all__ = [
    "CONV_PRECISION",
    "FUSED_ACTIVATIONS",
    "BINARY_FUSED_ACTIVATIONS",
    "FeatureMap",
    "Conv2d",
    "KSplitConv",
    "ConvTranspose2d",
    "normalize_activation",
    "unary_with_param",
    "apply_activation",
    "conv_padding",
    "conv_out_hw",
    "merge_sibling_convs",
    "split_k_weights",
    "concat_channels",
    "split_channels",
    "residual_add",
    "fmap_from_numpy",
    "fmap_to_numpy",
    "is_dram",
]

# Default precision of every builder here: HiFi4 + fp32 accumulation + packer L1 accumulation. Without
# ``packer_l1_acc`` a conv whose inner dim spans several blocks keeps its partial sums between blocks in the OUTPUT
# dtype (bf16) (conv2d_op_program_factory_common.cpp:175-178: partials CB = output dtype unless packer_l1_acc), which
# costs accuracy on wide convs; with it the partials stay fp32 (C17 device test, logs/ttaw/c17_c18_device_results.json).
CONV_PRECISION = "HiFi4+fp32+l1acc"

# Activations a conv epilogue applies correctly on this tree (P5 / P14 on the device). HARDSWISH is NOT among them:
# a fused HARDSWISH compiles to nothing (unary_op_utils.cpp:964-965, probe P5) and returns the un-activated result.
FUSED_ACTIVATIONS = ("relu", "relu6", "silu", "gelu", "sigmoid")
# Activations fused into ``ttnn.add(..., activations=[...])`` verified on the device (P5: RELU6 on the clip6 add).
BINARY_FUSED_ACTIVATIONS = ("relu", "relu6")
_REFUSED = {
    "hardswish": "a fused HARDSWISH is silently skipped by ttnn on this tree (probe P5): use activation=None and "
                 "ttnn.hardswish / ttnn.unary_chain([HARDSWISH]) as a separate op",
}

ActivationSpec = Optional[str]


def normalize_activation(activation: Any) -> ActivationSpec:
    """``None`` / ``"none"`` / ``"linear"`` / ``"identity"`` -> None; a fused activation name -> its lower-case name.
    Raises ``ValueError`` for HARDSWISH (silently skipped when fused, P5) and for unknown names."""
    if activation is None:
        return None
    name = str(activation).strip().lower()
    if name in ("", "none", "linear", "identity"):
        return None
    if name in _REFUSED:
        raise ValueError(f"activation {activation!r}: {_REFUSED[name]}")
    if name not in FUSED_ACTIVATIONS:
        raise ValueError(f"activation {activation!r}: expected None or one of {FUSED_ACTIVATIONS}")
    return name


def unary_with_param(activation: Any):
    """The ``ttnn.UnaryWithParam`` of a fused activation (None for no activation). GELU is the exact (erf) form."""
    import ttnn

    name = normalize_activation(activation)
    if name is None:
        return None
    if name == "gelu":
        return ttnn.UnaryWithParam(ttnn.UnaryOpType.GELU, 0.0)
    return ttnn.UnaryWithParam(getattr(ttnn.UnaryOpType, name.upper()))


def apply_activation(tensor: Any, activation: Any):
    """The activation as a separate (unfused) ttnn op; identity for None."""
    import ttnn

    name = normalize_activation(activation)
    if name is None:
        return tensor
    if name == "relu6":
        return ttnn.relu6(tensor)
    if name == "gelu":
        return ttnn.gelu(tensor, fast_and_approximate_mode=False)
    return getattr(ttnn, name)(tensor)


def _pair(value: Union[int, Sequence[int]], what: str) -> Tuple[int, int]:
    if isinstance(value, (int, np.integer)):
        return int(value), int(value)
    v = tuple(int(x) for x in value)
    if len(v) != 2:
        raise ValueError(f"{what} must be an int or a pair, got {value!r}")
    return v


def conv_padding(kernel: Tuple[int, int], padding: Any = "same", dilation: Tuple[int, int] = (1, 1)
                 ) -> Tuple[int, int, int, int]:
    """Padding as (top, bottom, left, right). ``"same"``: ``dilation * (k // 2)`` on every side (odd kernels; the
    "same" padding of stride-1 convs in every model here). An int, a pair (h, w) or a 4-tuple pass through."""
    if isinstance(padding, str):
        if padding != "same":
            raise ValueError(f"padding {padding!r}: expected 'same', an int, a pair or (top, bottom, left, right)")
        if kernel[0] % 2 == 0 or kernel[1] % 2 == 0:
            raise ValueError(f"'same' padding needs odd kernels, got {kernel}")
        ph, pw = dilation[0] * (kernel[0] // 2), dilation[1] * (kernel[1] // 2)
        return ph, ph, pw, pw
    if isinstance(padding, (int, np.integer)):
        p = int(padding)
        return p, p, p, p
    p = tuple(int(x) for x in padding)
    if len(p) == 2:
        return p[0], p[0], p[1], p[1]
    if len(p) == 4:
        return p  # type: ignore[return-value]
    raise ValueError(f"padding {padding!r}: expected 'same', an int, a pair or (top, bottom, left, right)")


def conv_out_hw(hw: Tuple[int, int], kernel: Tuple[int, int], stride: Tuple[int, int],
                padding: Tuple[int, int, int, int], dilation: Tuple[int, int] = (1, 1)) -> Tuple[int, int]:
    """Output size of a convolution (floor), as torch / ttnn compute it."""
    h = (hw[0] + padding[0] + padding[1] - dilation[0] * (kernel[0] - 1) - 1) // stride[0] + 1
    w = (hw[1] + padding[2] + padding[3] - dilation[1] * (kernel[1] - 1) - 1) // stride[1] + 1
    if h <= 0 or w <= 0:
        raise ValueError(f"conv of a {hw} input with kernel {kernel}, stride {stride}, padding {padding} is empty")
    return h, w


def is_dram(tensor_or_config: Any) -> bool:
    """True when a device tensor (or a memory config) lives in DRAM."""
    import ttnn

    mc = tensor_or_config.memory_config() if hasattr(tensor_or_config, "memory_config") else tensor_or_config
    bt = getattr(mc, "buffer_type", None)
    if bt is not None and hasattr(ttnn, "BufferType"):
        return bt == ttnn.BufferType.DRAM
    return mc == ttnn.DRAM_MEMORY_CONFIG


def _memory_config(spec: Any):
    """``"dram"`` / ``"l1"`` / None / a ttnn.MemoryConfig -> a memory config (or None)."""
    import ttnn

    if spec is None:
        return None
    if isinstance(spec, str):
        key = spec.lower()
        if key == "dram":
            return ttnn.DRAM_MEMORY_CONFIG
        if key == "l1":
            return ttnn.L1_MEMORY_CONFIG
        raise ValueError(f"memory config {spec!r}: expected 'dram', 'l1', None or a ttnn.MemoryConfig")
    return spec


def _layout(spec: Any):
    import ttnn

    if isinstance(spec, str):
        key = spec.lower().replace("_layout", "")
        if key == "tile":
            return ttnn.TILE_LAYOUT
        if key in ("row_major", "rm"):
            return ttnn.ROW_MAJOR_LAYOUT
        raise ValueError(f"layout {spec!r}: expected 'tile' or 'row_major'")
    return spec


def _layout_name(spec: Any) -> str:
    """``"tile"`` / ``"row_major"`` of a layout spec (string or ttnn layout), without importing ttnn for strings."""
    if isinstance(spec, str):
        key = spec.lower().replace("_layout", "")
        return "row_major" if key in ("row_major", "rm") else key
    return "tile" if "TILE" in str(spec).upper() else "row_major"


def _as_float32(a: Any, what: str) -> np.ndarray:
    if hasattr(a, "detach"):
        a = a.detach().cpu().float().numpy()
    arr = np.asarray(a, dtype=np.float32)
    if not np.all(np.isfinite(arr)):
        raise ValueError(f"{what} holds non-finite values")
    return np.ascontiguousarray(arr)


def _host_tensor(a: np.ndarray, dtype: str = "bfloat16"):
    """A ROW_MAJOR ttnn *host* tensor (no device involved: only TILE host tensors open the device, probe G4-X).
    The data are copied (ONNX initializers are read-only arrays, which torch warns about)."""
    import torch
    import ttnn

    return ttnn.from_torch(torch.from_numpy(np.array(a, dtype=np.float32, copy=True)), dtype=ttnn_dtype(dtype),
                           layout=ttnn.ROW_MAJOR_LAYOUT)


# --------------------------------------------------------------------------------------------- feature maps

@dataclass(frozen=True)
class FeatureMap:
    """A device feature map: ``tensor`` holds N*H*W rows of C channels (ttnn's flattened ``[1, 1, N*H*W, C]``,
    TILE or ROW_MAJOR; a ``[N, H, W, C]`` ROW_MAJOR tensor is accepted too), plus the logical sizes ttnn's layout
    does not keep."""

    tensor: Any
    batch: int
    height: int
    width: int
    channels: int

    @property
    def rows(self) -> int:
        return self.batch * self.height * self.width

    @property
    def hw(self) -> Tuple[int, int]:
        return self.height, self.width

    def with_tensor(self, tensor: Any, channels: Optional[int] = None) -> "FeatureMap":
        """The same spatial map with another tensor (and channel count)."""
        return FeatureMap(tensor, self.batch, self.height, self.width,
                          self.channels if channels is None else int(channels))

    def deallocate(self) -> None:
        """Free the tensor's device memory (``force=False``: a view's base is left to its owner, probe P2)."""
        import ttnn

        if self.tensor is not None and self.tensor.is_allocated():
            ttnn.deallocate(self.tensor, False)


def fmap_from_numpy(nchw: Any, device: Any, *, dtype: Any = "bfloat16", layout: Any = "tile",
                    memory_config: Any = "dram") -> FeatureMap:
    """Upload an NCHW array as a device feature map ``[1, 1, N*H*W, C]`` (TILE by default). A TILE host tensor
    touches the device: run under ``bin/devrun`` (probe G4-X)."""
    import torch
    import ttnn

    a = _as_float32(nchw, "feature map") if np.asarray(nchw).dtype != np.uint8 else np.asarray(nchw)
    if a.ndim != 4:
        raise ValueError(f"expected NCHW, got shape {a.shape}")
    n, c, h, w = a.shape
    flat = np.ascontiguousarray(a.transpose(0, 2, 3, 1).reshape(1, 1, n * h * w, c))
    t = ttnn.from_torch(torch.from_numpy(flat), dtype=ttnn_dtype(dtype), layout=_layout(layout), device=device,
                        memory_config=_memory_config(memory_config))
    return FeatureMap(t, n, h, w, c)


def fmap_to_numpy(fm: FeatureMap, dtype: Any = np.float32) -> np.ndarray:
    """Read a device feature map back as an NCHW numpy array (blocking; never inside a capture)."""
    import ttnn

    a = ttnn.to_torch(fm.tensor)
    a = a.float().numpy() if a.is_floating_point() else a.numpy()
    a = a.reshape(-1, a.shape[-1])[: fm.rows, : fm.channels]
    return np.ascontiguousarray(a.reshape(fm.batch, fm.height, fm.width, fm.channels).transpose(0, 3, 1, 2),
                                dtype=dtype)


# ------------------------------------------------------------------------------------------------- Conv2d

class Conv2d:
    """``ttnn.conv2d`` with weights prepared by the first call and reused afterwards (see the module docstring).

    Args:
        weight: float array ``[Cout, Cin / groups, kh, kw]`` (OIHW, numpy or torch; BN already folded).
        bias: float array ``[Cout]`` or None.
        stride, dilation: int or pair. ``padding``: ``"same"`` (default), int, pair or (top, bottom, left, right).
        groups: 1 for dense convs; ``groups == Cin`` is a depthwise conv (correct but expanded to dense weights by
            ttnn: P5, 311 MiB for SceneSeg).
        activation: None or one of :data:`FUSED_ACTIVATIONS` (fused in the epilogue).
        precision: a :class:`..precision.Precision` or its spec (default :data:`CONV_PRECISION` = HiFi4 + fp32
            accumulation + packer L1 accumulation; ``Precision.weights`` is the device weight dtype: bf16 or bfp8).
        output_dtype: dtype of the output (``"bfloat16"``; ``"float32"`` for logits that feed thresholds).
        output_layout: ``"tile"`` (default) or ``"row_major"``.
        slicing: ``"auto"`` (ttnn's rule: DRAM input -> automatic DRAM slicing, L1 input -> L1 path), ``"l1"``
            (``Conv2dL1FullSliceConfig``: the op moves the whole input to L1) or ``("height" | "width", n)``
            (explicit DRAM slices; P14 found most explicit counts rejected or slower: diagnostics only).
        output_memory_config: where an L1-path output (1x1 matmul convs, L1 inputs) is placed: ``"dram"``
            (default), ``"l1"``, None (the op's sharded output) or a ttnn.MemoryConfig. DRAM-sliced convs always
            write DRAM interleaved; ``"l1"`` then adds a move.
        conv_config: extra ``ttnn.Conv2dConfig`` fields (``act_block_h_override``, ``shard_layout``,
            ``enable_act_double_buffer``, ``deallocate_activation``, ...): optimization knobs.
        weight_terms: 1 (default) or 2. With 2 the weight and the bias are split into two bf16 terms
            (``hi = bf16(w)``, ``lo = bf16(w - hi)``: about 16 significant bits) run as two convs into float32
            outputs, summed in fp32 with the activation fused into the add, then cast to ``output_dtype``: the conv
            then carries ~fp32 weights and bias for twice the MACs and 3 more programs. The device's own float32
            weight path is no substitute (1.5x less error than bf16, vs 3.5x for the split; YOLOX port,
            logs/yolox/m3_weight_precision_probe.log). Use it where the bf16 rounding of weights and biases, which
            does not average out over a deep network (a per-channel bias error is the same at every pixel),
            decides an output: YOLOX's 16-class mask on public frames (97.5 % -> 99.3 % agreement in emulation).
        name: label for messages and :meth:`describe`.
    """

    def __init__(self, weight: Any, bias: Any = None, *, stride: Any = 1, padding: Any = "same",
                 dilation: Any = 1, groups: int = 1, activation: Any = None, precision: Any = CONV_PRECISION,
                 output_dtype: Any = "bfloat16", output_layout: Any = "tile", slicing: Any = "auto",
                 output_memory_config: Any = "dram", conv_config: Optional[Dict[str, Any]] = None,
                 weight_terms: int = 1, name: str = "conv"):
        w = _as_float32(weight, f"{name}: weight")
        if w.ndim != 4:
            raise ValueError(f"{name}: weight must be [Cout, Cin/groups, kh, kw], got {w.shape}")
        self.name = name
        self.groups = int(groups)
        self.out_channels, cin_g, kh, kw = (int(s) for s in w.shape)
        self.in_channels = cin_g * self.groups
        if self.groups < 1 or self.out_channels % self.groups:
            raise ValueError(f"{name}: groups={groups} does not divide Cout={self.out_channels}")
        self.kernel = (kh, kw)
        self.stride = _pair(stride, "stride")
        self.dilation = _pair(dilation, "dilation")
        self.padding = conv_padding(self.kernel, padding, self.dilation)
        self.activation = normalize_activation(activation)
        self.precision = Precision.parse(precision)
        self.output_dtype = dtype_name(output_dtype)
        self.output_layout = output_layout
        if not (slicing == "auto" or slicing == "l1" or (isinstance(slicing, (tuple, list)) and len(slicing) == 2
                                                         and slicing[0] in ("height", "width"))):
            raise ValueError(f"{name}: slicing {slicing!r}: expected 'auto', 'l1' or ('height' | 'width', n)")
        self.slicing = slicing if isinstance(slicing, str) else (str(slicing[0]), int(slicing[1]))
        self.output_memory_config = output_memory_config
        self.conv_config_overrides = dict(conv_config or {})
        if "activation" in self.conv_config_overrides or "weights_dtype" in self.conv_config_overrides:
            raise ValueError(f"{name}: pass activation / weights dtype as arguments, not in conv_config")
        b = np.zeros(self.out_channels, np.float32) if bias is None else _as_float32(bias, f"{name}: bias")
        if b.shape != (self.out_channels,):
            raise ValueError(f"{name}: bias must have shape ({self.out_channels},), got {b.shape}")
        if int(weight_terms) not in (1, 2):
            raise ValueError(f"{name}: weight_terms must be 1 or 2, got {weight_terms!r}")
        self.weight_terms = int(weight_terms)
        if self.weight_terms == 2 and str(_layout_name(output_layout)) != "tile":
            raise ValueError(f"{name}: weight_terms=2 writes TILE outputs (the fp32 sum is cast in TILE)")
        self._weight = w
        self._bias = b
        self._host: Optional[List[Tuple[Any, Any]]] = None
        self._prepared: Dict[Tuple, Tuple[Any, Any]] = {}

    # ---- geometry
    @property
    def is_matmul(self) -> bool:
        """1x1, stride 1, no padding, no dilation, dense: ttnn runs it as a matmul on the L1 path."""
        return (self.kernel == (1, 1) and self.stride == (1, 1) and self.padding == (0, 0, 0, 0)
                and self.dilation == (1, 1) and self.groups == 1)

    def out_hw(self, height: int, width: int) -> Tuple[int, int]:
        return conv_out_hw((height, width), self.kernel, self.stride, self.padding, self.dilation)

    @property
    def weight(self) -> np.ndarray:
        """The float32 host weight (OIHW)."""
        return self._weight

    @property
    def bias(self) -> np.ndarray:
        return self._bias

    # ---- ttnn arguments
    def host_terms(self) -> List[Tuple[np.ndarray, np.ndarray]]:
        """The float32 (weight, bias) of each term: one term, or ``hi = bf16(x)`` and ``lo = bf16(x - hi)`` (both
        exactly representable in bf16, so the bf16 upload is exact)."""
        if self.weight_terms == 1:
            return [(self._weight, self._bias)]
        hi_w, hi_b = round_to_bf16(self._weight), round_to_bf16(self._bias)
        return [(hi_w, hi_b), (round_to_bf16(self._weight - hi_w), round_to_bf16(self._bias - hi_b))]

    def _host_tensors(self, term: int = 0) -> Tuple[Any, Any]:
        if self._host is None:
            # host tensors in the device weight dtype's host precision: bf16 for bf16 / bfp8 weights (ttnn rounds
            # to bfp8 itself), float32 for float32 weights (a bf16 host tensor would round them first)
            host_dtype = "float32" if self.precision.weights == "float32" else "bfloat16"
            self._host = [(_host_tensor(w, host_dtype), _host_tensor(b.reshape(1, 1, 1, -1), host_dtype))
                          for w, b in self.host_terms()]
        return self._host[term]

    def _conv_config(self, fused: bool = True):
        import ttnn

        kw = dict(self.conv_config_overrides)
        kw["weights_dtype"] = ttnn_dtype(self.precision.weights)
        kw["activation"] = unary_with_param(self.activation) if fused else None
        kw.setdefault("output_layout", _layout(self.output_layout))
        return ttnn.Conv2dConfig(**kw)

    def _slice_config(self):
        import ttnn

        if self.slicing == "auto":
            return None
        if self.slicing == "l1":
            return ttnn.Conv2dL1FullSliceConfig
        kind, num = self.slicing
        slice_type = ttnn.Conv2dDRAMSliceHeight if kind == "height" else ttnn.Conv2dDRAMSliceWidth
        return ttnn.Conv2dSliceConfig(slice_type=slice_type, num_slices=num)

    def _spec_key(self, fm: FeatureMap) -> Tuple:
        t = fm.tensor
        return (fm.batch, fm.height, fm.width, str(t.dtype), str(t.layout), str(t.memory_config()),
                tuple(int(s) for s in t.shape))

    def takes_l1_path(self, x: Any) -> bool:
        """Whether a call on ``x`` (a device tensor) runs ttnn's L1 path (else the DRAM-sliced path)."""
        if self.slicing == "l1" or self.is_matmul:
            return True
        return self.slicing == "auto" and not is_dram(x)

    def __call__(self, x: Union[FeatureMap, Any], *, batch: Optional[int] = None, height: Optional[int] = None,
                 width: Optional[int] = None) -> FeatureMap:
        """Run the conv on a :class:`FeatureMap` (or a raw device tensor plus ``batch`` / ``height`` / ``width``)
        and return the output :class:`FeatureMap` (``[1, 1, N*OH*OW, Cout]``)."""
        import ttnn

        fm = x if isinstance(x, FeatureMap) else FeatureMap(x, int(batch), int(height), int(width), self.in_channels)
        if fm.channels != self.in_channels:
            raise ValueError(f"{self.name}: input has {fm.channels} channels, the conv takes {self.in_channels}")
        key = self._spec_key(fm)
        l1_path = self.takes_l1_path(fm.tensor)
        out_mc = _memory_config(self.output_memory_config)
        split = self.weight_terms == 2
        outs = []
        for term in range(self.weight_terms):
            tkey = key + (term,)
            weight, bias = self._prepared.get(tkey) or self._host_tensors(term)
            out, (oh, ow), (weight, bias) = ttnn.conv2d(
                input_tensor=fm.tensor, weight_tensor=weight, device=fm.tensor.device(),
                in_channels=self.in_channels, out_channels=self.out_channels, batch_size=fm.batch,
                input_height=fm.height, input_width=fm.width, kernel_size=self.kernel, stride=self.stride,
                padding=self.padding, dilation=self.dilation, groups=self.groups, bias_tensor=bias,
                conv_config=self._conv_config(fused=not split), compute_config=self.precision.compute_kernel_config(),
                slice_config=self._slice_config(), memory_config=out_mc if l1_path else None,
                dtype=ttnn.float32 if split else ttnn_dtype(self.output_dtype), return_output_dim=True,
                return_weights_and_bias=True)
            self._prepared[tkey] = (weight, bias)
            outs.append(out)
        if (int(oh), int(ow)) != self.out_hw(fm.height, fm.width):
            raise RuntimeError(f"{self.name}: ttnn returned a {oh}x{ow} output, expected "
                               f"{self.out_hw(fm.height, fm.width)}")
        out = outs[0]
        if split:
            fuse = self.activation in BINARY_FUSED_ACTIVATIONS
            out = ttnn.add(outs[0], outs[1], activations=[unary_with_param(self.activation)] if fuse else [])
            for t in outs:
                ttnn.deallocate(t, False)
            if not fuse:
                out = apply_activation(out, self.activation)
            if self.output_dtype != "float32":
                out = ttnn.typecast(out, ttnn_dtype(self.output_dtype))
        if not l1_path and out_mc is not None and not is_dram(out_mc):
            out = ttnn.to_memory_config(out, out_mc)
        return FeatureMap(out, fm.batch, int(oh), int(ow), self.out_channels)

    @property
    def prepared_specs(self) -> int:
        """Number of input specs this layer holds prepared device weights for."""
        return len({k[:-1] for k in self._prepared})

    def release(self) -> None:
        """Free the prepared device weights (the host weights stay; the next call prepares them again)."""
        import ttnn

        for w, b in self._prepared.values():
            for t in (w, b):
                if t is not None and t.is_allocated():
                    ttnn.deallocate(t)
        self._prepared.clear()

    def describe(self) -> Dict[str, Any]:
        return {"name": self.name, "cin": self.in_channels, "cout": self.out_channels, "kernel": list(self.kernel),
                "stride": list(self.stride), "padding": list(self.padding), "dilation": list(self.dilation),
                "groups": self.groups, "activation": self.activation, "precision": self.precision.label,
                "output_dtype": self.output_dtype, "slicing": self.slicing if isinstance(self.slicing, str)
                else list(self.slicing), "matmul": self.is_matmul, "weight_terms": self.weight_terms}


# ----------------------------------------------------------------------------------- K-split (concat-free)

def split_k_weights(weight: Any, parts: Sequence[int]) -> List[np.ndarray]:
    """Split a conv weight ``[Cout, sum(parts), kh, kw]`` along its input channels: ``conv(concat(x_i), W) ==
    sum_i conv(x_i, W_i)``."""
    w = _as_float32(weight, "weight")
    if w.ndim != 4 or w.shape[1] != sum(int(p) for p in parts):
        raise ValueError(f"weight {w.shape} does not have sum(parts) = {sum(parts)} input channels")
    out, start = [], 0
    for p in parts:
        out.append(np.ascontiguousarray(w[:, start:start + int(p)]))
        start += int(p)
    return out


class KSplitConv:
    """``conv(concat([x_1, ..., x_n], channels))`` computed as ``sum_i conv_i(x_i)`` without building the concat
    (S:centerpoint:448; P14: CenterPoint's shared conv 3.6 vs 14.8 ms and PCC 0.999992 vs 0.999961).

    The bias rides on the first partial conv, the partial sums are added in order, and the activation is fused into
    the last add (RELU / RELU6, P5) or applied after it. ``accumulate_dtype="float32"`` keeps the partial outputs and
    their sums in fp32 and casts to ``output_dtype`` once (more accurate, more traffic); the default ``"bfloat16"``
    rounds each partial (what P14 measured). Every keyword except ``activation`` / ``output_dtype`` /
    ``accumulate_dtype`` goes to the partial :class:`Conv2d` layers."""

    def __init__(self, weight: Any, bias: Any, parts: Sequence[int], *, activation: Any = None,
                 output_dtype: Any = "bfloat16", accumulate_dtype: Any = "bfloat16", name: str = "ksplit",
                 **conv_kwargs: Any):
        if len(parts) < 2:
            raise ValueError(f"{name}: a K-split needs at least two parts")
        self.name = name
        self.parts = tuple(int(p) for p in parts)
        self.activation = normalize_activation(activation)
        self.output_dtype = dtype_name(output_dtype)
        self.accumulate_dtype = dtype_name(accumulate_dtype)
        weights = split_k_weights(weight, self.parts)
        cout = weights[0].shape[0]
        b = np.zeros(cout, np.float32) if bias is None else _as_float32(bias, f"{name}: bias")
        self.convs = [Conv2d(w, b if i == 0 else None, activation=None, output_dtype=self.accumulate_dtype,
                             name=f"{name}.{i}", **conv_kwargs) for i, w in enumerate(weights)]
        self.out_channels = cout

    def __call__(self, inputs: Sequence[FeatureMap]) -> FeatureMap:
        import ttnn

        if len(inputs) != len(self.convs):
            raise ValueError(f"{self.name}: expected {len(self.convs)} inputs, got {len(inputs)}")
        total: Optional[FeatureMap] = None
        n = len(inputs)
        for i, (conv, x) in enumerate(zip(self.convs, inputs)):
            part = conv(x)
            if total is None:
                total = part
                continue
            last = i == n - 1
            fuse = last and self.activation in BINARY_FUSED_ACTIVATIONS
            acts = [unary_with_param(self.activation)] if fuse else []
            summed = ttnn.add(total.tensor, part.tensor, activations=acts)
            for t in (total, part):
                t.deallocate()
            total = total.with_tensor(summed)
            if last and not fuse:
                total = total.with_tensor(apply_activation(total.tensor, self.activation))
        assert total is not None
        if self.accumulate_dtype != self.output_dtype:
            total = total.with_tensor(ttnn.typecast(total.tensor, ttnn_dtype(self.output_dtype)))
        return total

    def release(self) -> None:
        for c in self.convs:
            c.release()

    def describe(self) -> Dict[str, Any]:
        return {"name": self.name, "parts": list(self.parts), "activation": self.activation,
                "accumulate_dtype": self.accumulate_dtype, "convs": [c.describe() for c in self.convs]}


# ----------------------------------------------------------------------------------- sibling merge / glue

def merge_sibling_convs(weights: Sequence[Any], biases: Sequence[Any]) -> Tuple[np.ndarray, np.ndarray, List[int]]:
    """One conv for sibling convs reading the same input (S:centerpoint:450; S:yolox section 8.5 #1): output
    channels concatenated in order -> ``(weight, bias, out_channels_per_sibling)``; split the output with
    :func:`split_channels`. Exact (each output channel keeps its own weights)."""
    ws = [_as_float32(w, "weight") for w in weights]
    bs = [np.zeros(w.shape[0], np.float32) if b is None else _as_float32(b, "bias") for w, b in zip(ws, biases)]
    shapes = {w.shape[1:] for w in ws}
    if len(shapes) != 1:
        raise ValueError(f"sibling convs must share Cin and kernel, got {[w.shape for w in ws]}")
    return (np.ascontiguousarray(np.concatenate(ws, axis=0)), np.ascontiguousarray(np.concatenate(bs)),
            [int(w.shape[0]) for w in ws])


def concat_channels(maps: Sequence[FeatureMap], *, memory_config: Any = "dram") -> FeatureMap:
    """Concatenate feature maps of one spatial size along channels (``ttnn.concat(dim=3)``); every input must have
    the same layout and dtype. TILE inputs with channel counts that are multiples of 32 avoid any re-tiling."""
    import ttnn

    if not maps:
        raise ValueError("concat_channels needs at least one map")
    first = maps[0]
    for m in maps[1:]:
        if (m.batch, m.height, m.width) != (first.batch, first.height, first.width):
            raise ValueError(f"concat of maps of different sizes: {[(m.batch, m.height, m.width) for m in maps]}")
    if len(maps) == 1:
        return first
    t = ttnn.concat([m.tensor for m in maps], dim=3, memory_config=_memory_config(memory_config))
    return first.with_tensor(t, sum(m.channels for m in maps))


def split_channels(fm: FeatureMap, sizes: Sequence[int]) -> List[FeatureMap]:
    """Slice a map's channels into consecutive groups of ``sizes`` (``ttnn.slice`` on the last dim)."""
    import ttnn

    if sum(int(s) for s in sizes) != fm.channels:
        raise ValueError(f"split sizes {list(sizes)} do not add up to {fm.channels} channels")
    shape = [int(s) for s in fm.tensor.shape]
    out, start = [], 0
    for s in sizes:
        begin = [0] * len(shape)
        end = list(shape)
        begin[-1], end[-1] = start, start + int(s)
        out.append(fm.with_tensor(ttnn.slice(fm.tensor, begin, end), int(s)))
        start += int(s)
    return out


def residual_add(a: FeatureMap, b: FeatureMap, activation: Any = None) -> FeatureMap:
    """``a + b`` (same spec), with RELU / RELU6 fused into the add (P5: the clip6 residual of YOLOX) or another
    activation applied after it."""
    import ttnn

    if (a.rows, a.channels) != (b.rows, b.channels):
        raise ValueError(f"residual add of {a.rows}x{a.channels} and {b.rows}x{b.channels}")
    act = normalize_activation(activation)
    if act in BINARY_FUSED_ACTIVATIONS:
        return a.with_tensor(ttnn.add(a.tensor, b.tensor, activations=[unary_with_param(act)]))
    return a.with_tensor(apply_activation(ttnn.add(a.tensor, b.tensor), act))


# ----------------------------------------------------------------------------- ConvTranspose (kernel == stride)

class ConvTranspose2d:
    """ConvTranspose2d with kernel == stride and no padding (the BEV deblocks of CenterPoint / PointPainting /
    TransFusion and SceneSeg's up-sampling layers), weights prepared by the first call.

    ``method="auto"`` follows probe P5: k = 2 -> ``ttnn.conv_transpose2d`` (CP deblock1 0.64 ms vs 2.08 ms for the
    rewrite), k >= 3 -> ``"linear_d2s"`` = ``ttnn.linear(Cin -> k*k*Cout)`` with the bias and activation fused, then
    depth-to-space (``to_layout(ROW_MAJOR)``, ``reshape [N*h, w, k, k, C]``, ``permute (0, 2, 1, 3, 4)``, ``reshape``;
    CP deblock2 k = 4: 1.97 vs 10.27 ms). The linear output channel order is (i, j, c):
    ``W2[cin, (i*k + j)*Cout + c] = W[cin, c, i, j]``. Output: a ROW_MAJOR ``[1, 1, N*h*k*w*k, Cout]`` map for
    ``linear_d2s`` (``output_layout="tile"`` re-tiles it), the conv's TILE output for ``conv_transpose2d``.

    ``weight``: ``[Cin, Cout, k, k]`` (torch ConvTranspose2d layout); ``bias``: ``[Cout]`` or None."""

    METHODS = ("auto", "conv_transpose2d", "linear_d2s")

    def __init__(self, weight: Any, bias: Any = None, *, stride: Optional[int] = None, activation: Any = None,
                 precision: Any = CONV_PRECISION, output_dtype: Any = "bfloat16", method: str = "auto",
                 output_layout: Any = "tile", memory_config: Any = "dram", name: str = "deconv"):
        w = _as_float32(weight, f"{name}: weight")
        if w.ndim != 4 or w.shape[2] != w.shape[3]:
            raise ValueError(f"{name}: weight must be [Cin, Cout, k, k], got {w.shape}")
        self.name = name
        self.in_channels, self.out_channels, self.k = int(w.shape[0]), int(w.shape[1]), int(w.shape[2])
        if stride is not None and int(stride) != self.k:
            raise ValueError(f"{name}: only kernel == stride is supported (k={self.k}, stride={stride})")
        if method not in self.METHODS:
            raise ValueError(f"{name}: method {method!r}: expected one of {self.METHODS}")
        self.method = ("conv_transpose2d" if self.k == 2 else "linear_d2s") if method == "auto" else method
        self.activation = normalize_activation(activation)
        self.precision = Precision.parse(precision)
        self.output_dtype = dtype_name(output_dtype)
        self.output_layout = output_layout
        self.memory_config = memory_config
        b = np.zeros(self.out_channels, np.float32) if bias is None else _as_float32(bias, f"{name}: bias")
        if b.shape != (self.out_channels,):
            raise ValueError(f"{name}: bias must have shape ({self.out_channels},), got {b.shape}")
        self._weight, self._bias = w, b
        self._prepared: Dict[Tuple, Tuple[Any, Any]] = {}
        self._linear: Optional[Tuple[Any, Any]] = None

    def _spec_key(self, fm: FeatureMap) -> Tuple:
        t = fm.tensor
        return (fm.batch, fm.height, fm.width, str(t.dtype), str(t.layout), str(t.memory_config()))

    def linear_weights(self) -> Tuple[np.ndarray, np.ndarray]:
        """Host form of the ``linear_d2s`` weights: ``W2 [Cin, k*k*Cout]`` in (i, j, c) order, bias repeated."""
        k = self.k
        w2 = self._weight.transpose(0, 2, 3, 1).reshape(self.in_channels, k * k * self.out_channels)
        return np.ascontiguousarray(w2), np.ascontiguousarray(np.tile(self._bias, k * k))

    def __call__(self, x: FeatureMap) -> FeatureMap:
        import ttnn

        if x.channels != self.in_channels:
            raise ValueError(f"{self.name}: input has {x.channels} channels, the layer takes {self.in_channels}")
        k = self.k
        if self.method == "conv_transpose2d":
            key = self._spec_key(x)
            # float32 host tensors for float32 weights, as Conv2d (a bf16 host tensor would round them first)
            host_dtype = "float32" if self.precision.weights == "float32" else "bfloat16"
            weight, bias = self._prepared.get(key) or (_host_tensor(self._weight, host_dtype),
                                                       _host_tensor(self._bias.reshape(1, 1, 1, -1), host_dtype))
            conv_config = ttnn.Conv2dConfig(weights_dtype=ttnn_dtype(self.precision.weights),
                                            activation=unary_with_param(self.activation),
                                            output_layout=_layout(self.output_layout))
            out, (weight, bias) = ttnn.conv_transpose2d(
                input_tensor=x.tensor, weight_tensor=weight, device=x.tensor.device(),
                in_channels=self.in_channels, out_channels=self.out_channels, batch_size=x.batch,
                input_height=x.height, input_width=x.width, kernel_size=(k, k), stride=(k, k), padding=(0, 0),
                output_padding=(0, 0), dilation=(1, 1), groups=1, bias_tensor=bias, conv_config=conv_config,
                compute_config=self.precision.compute_kernel_config(), dtype=ttnn_dtype(self.output_dtype),
                mirror_kernel=True, return_weights_and_bias=True)
            self._prepared[key] = (weight, bias)
            return FeatureMap(out, x.batch, x.height * k, x.width * k, self.out_channels)
        # linear + depth-to-space
        import torch

        device = x.tensor.device()
        if self._linear is None:
            w2, b2 = self.linear_weights()
            wdt = ttnn_dtype(self.precision.weights)
            self._linear = (
                ttnn.from_torch(torch.from_numpy(w2), dtype=wdt, layout=ttnn.TILE_LAYOUT, device=device),
                ttnn.from_torch(torch.from_numpy(b2.reshape(1, -1)), dtype=wdt, layout=ttnn.TILE_LAYOUT,
                                device=device))
        weight, bias = self._linear
        t = x.tensor if x.tensor.layout == ttnn.TILE_LAYOUT else ttnn.to_layout(x.tensor, ttnn.TILE_LAYOUT)
        y = ttnn.linear(t, weight, bias=bias, activation=unary_with_param(self.activation),
                        compute_kernel_config=self.precision.compute_kernel_config(),
                        memory_config=_memory_config(self.memory_config), dtype=ttnn_dtype(self.output_dtype))
        y_rm = ttnn.to_layout(y, ttnn.ROW_MAJOR_LAYOUT)
        ttnn.deallocate(y, False)
        c = self.out_channels
        y5 = ttnn.reshape(y_rm, (x.batch * x.height, x.width, k, k, c))
        y5 = ttnn.permute(y5, (0, 2, 1, 3, 4))
        out = ttnn.reshape(y5, (1, 1, x.batch * x.height * k * x.width * k, c))
        if _layout(self.output_layout) == ttnn.TILE_LAYOUT:
            out = ttnn.to_layout(out, ttnn.TILE_LAYOUT)
        return FeatureMap(out, x.batch, x.height * k, x.width * k, c)

    def release(self) -> None:
        import ttnn

        tensors = [t for pair in self._prepared.values() for t in pair] + list(self._linear or ())
        for t in tensors:
            if t is not None and t.is_allocated():
                ttnn.deallocate(t)
        self._prepared.clear()
        self._linear = None

    def describe(self) -> Dict[str, Any]:
        return {"name": self.name, "cin": self.in_channels, "cout": self.out_channels, "k": self.k,
                "method": self.method, "activation": self.activation, "precision": self.precision.label,
                "output_dtype": self.output_dtype}