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