# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # # SPDX-License-Identifier: Apache-2.0 """On-device YUV 4:2:0 conversion + fast device-to-host into ffmpeg yuv420p planar bytes. Isolated from ``tensor.py`` so the LTX wiring stays decoupled from that module's d2h internals; reuses the shard-extraction primitives it already exposes. """ from __future__ import annotations import os from concurrent.futures import ThreadPoolExecutor import numpy as np import torch import ttnn from .planar_concat import HAS_CPP_PLANAR_CONCAT from .planar_concat import planar_concat_cpp as _planar_concat_cpp_impl from .tensor import _get_inter_host_axis, _host_buffer_to_torch, _to_torch_zero_copy # Persistent host-side reassembly pool — strided uint8 copies release the GIL, so a few threads scale near-linearly _DEFAULT_REASSEMBLE_POOL: ThreadPoolExecutor | None = None _DEFAULT_REASSEMBLE_WORKERS = min(8, os.cpu_count() or 8) def _get_default_reassemble_pool() -> ThreadPoolExecutor: global _DEFAULT_REASSEMBLE_POOL if _DEFAULT_REASSEMBLE_POOL is None: _DEFAULT_REASSEMBLE_POOL = ThreadPoolExecutor( max_workers=_DEFAULT_REASSEMBLE_WORKERS, thread_name_prefix="tt_dit_yuv_reassemble", ) return _DEFAULT_REASSEMBLE_POOL # Persistent output buffer for the C++ planar-concat fast path _PLANAR_OUT_BUF: np.ndarray | None = None _PLANAR_OUT_SHAPE: tuple[int, int] | None = None def _get_planar_out_buf(T: int, row_stride: int) -> np.ndarray: global _PLANAR_OUT_BUF, _PLANAR_OUT_SHAPE shape = (T, row_stride) if _PLANAR_OUT_SHAPE != shape: _PLANAR_OUT_BUF = np.empty(shape, dtype=np.uint8) _PLANAR_OUT_SHAPE = shape return _PLANAR_OUT_BUF def _bt601_yuv_coefficients(): """BT.601 coefficients for input in [-1, 1] -> limited-range uint8 (Y 16-235, CbCr 16-240).""" return ttnn.experimental.yuv_bt601_coefficients() def _yuv_planar_d2h( tt_Y: ttnn.Tensor, tt_Cb: ttnn.Tensor, tt_Cr: ttnn.Tensor, mesh_device: ttnn.MeshDevice, H: int, W: int, T: int, *, out_H: int | None = None, out_W: int | None = None, view=None, pool: ThreadPoolExecutor | None = None, ) -> np.ndarray: """Batched D2H of three YUV ttnn tensors into ffmpeg yuv420p planar uint8. Per-shard input shapes (kernel-native BHWT with C=1): * tt_Y: ``(1, h_per_y, w_per_y, T)`` * tt_Cb/tt_Cr: ``(1, h_per_uv, w_per_uv, T)`` Output: ``(T, H*W + 2*(H/2 * W/2))`` numpy uint8 — per-frame ``[Y plane | Cb plane | Cr plane]`` in row-major. Kicks off all three ``cpu(blocking=False)`` calls before a single ``synchronize_device`` so the reads overlap. Per-shard scatters then take one of two paths: * **C++/AVX2 fast path** (when ``HAS_CPP_PLANAR_CONCAT`` is True and the local mesh is rectangular): one ``planar_concat_cpp`` call into a module-level persistent output buffer. ~2× faster than the Python path on warm calls, but the returned buffer is **reused across calls** — copy out (or feed ffmpeg) before the next call. * **Python fallback**: per-shard ``_write`` tasks on the shared reassembly ThreadPoolExecutor (torch's strided-copy backend). Each scatter is a strided->strided copy; allocates a fresh output per call. Either way, the byte layout is identical. Args: view: Optional mesh device view for multi-host environments. When provided, uses ``host_buffer()`` / ``get_shard()`` with ``view.is_local()`` filtering instead of ``get_device_tensors()``. """ Hu, Wu = H // 2, W // 2 hw = H * W uv = Hu * Wu row_stride = hw + 2 * uv # Logical output dims: when the VAE pads out_H = H if out_H is None else out_H out_W = W if out_W is None else out_W # Async D2H all 3 outputs, single sync — overlaps three D2H reads. host_Y = tt_Y.cpu(blocking=False) host_Cb = tt_Cb.cpu(blocking=False) host_Cr = tt_Cr.cpu(blocking=False) ttnn.synchronize_device(mesh_device) if view is not None: # --- Multi-host: extract local shards via host_buffer/get_shard --- def _extract_local(host_tensor): host_mesh_coords = list(host_tensor.tensor_topology().mesh_coords()) distributed_buf = host_tensor.host_buffer() tt_dtype = host_tensor.dtype padded_shape = list(host_tensor.padded_shape) logical_shape = list(host_tensor.shape) trim = tuple(slice(0, d) for d in logical_shape) coords_and_shards = [] for c in host_mesh_coords: if not view.is_local(c): continue buf = distributed_buf.get_shard(c) if buf is not None: coords_and_shards.append((c, _host_buffer_to_torch(buf, padded_shape, tt_dtype)[trim])) return coords_and_shards Y_coords_shards = _extract_local(host_Y) Cb_coords_shards = _extract_local(host_Cb) Cr_coords_shards = _extract_local(host_Cr) # Remap global mesh coordinates to 0-based local coordinates. all_local_coords = [c for c, _ in Y_coords_shards] local_row_positions = sorted({int(c[0]) for c in all_local_coords}) local_col_positions = sorted({int(c[1]) for c in all_local_coords}) row_remap = {pos: i for i, pos in enumerate(local_row_positions)} col_remap = {pos: i for i, pos in enumerate(local_col_positions)} TP_eff = len(local_row_positions) SP_eff = len(local_col_positions) h_per_y, w_per_y = H // TP_eff, W // SP_eff h_per_uv, w_per_uv = Hu // TP_eff, Wu // SP_eff mesh_coords = [(row_remap[int(c[0])], col_remap[int(c[1])]) for c in all_local_coords] Y_shards = [s for _, s in Y_coords_shards] Cb_shards = [s for _, s in Cb_coords_shards] Cr_shards = [s for _, s in Cr_coords_shards] else: # --- Single-host: extract all shards via get_device_tensors --- TP_eff, SP_eff = tuple(mesh_device.shape) h_per_y, w_per_y = H // TP_eff, W // SP_eff h_per_uv, w_per_uv = Hu // TP_eff, Wu // SP_eff mesh_coords = list(tt_Y.tensor_topology().mesh_coords()) def _extract(host_tensor): host_shards = ttnn.get_device_tensors(host_tensor) logical_shape = list(host_shards[0].shape) trim = tuple(slice(0, d) for d in logical_shape) return [_to_torch_zero_copy(s)[trim] for s in host_shards] Y_shards = _extract(host_Y) # each (1, h_per_y, w_per_y, T) Cb_shards = _extract(host_Cb) # each (1, h_per_uv, w_per_uv, T) Cr_shards = _extract(host_Cr) # --- C++/AVX2 fast path --------------------------------------------- Drop-in replacement for the torch_threaded if HAS_CPP_PLANAR_CONCAT and len(mesh_coords) == TP_eff * SP_eff: triples = sorted( zip(mesh_coords, Y_shards, Cb_shards, Cr_shards), key=lambda t: (int(t[0][0]), int(t[0][1])), ) out_Hu, out_Wu = out_H // 2, out_W // 2 out_row = out_H * out_W + 2 * out_Hu * out_Wu out = _get_planar_out_buf(T, out_row) return _planar_concat_cpp_impl( [t[1] for t in triples], [t[2] for t in triples], [t[3] for t in triples], "CHWT", (TP_eff, SP_eff), out=out, out_H=out_H, out_W=out_W, ) # --- Python fallback (torch_threaded scatter) ------------------------ Assemble directly into the logical-sized out_Hu, out_Wu = out_H // 2, out_W // 2 out_hw, out_uv = out_H * out_W, out_Hu * out_Wu out_row = out_hw + 2 * out_uv out = np.empty((T, out_row), dtype=np.uint8) out_t = torch.from_numpy(out) y_view = out_t.as_strided((T, out_H, out_W), (out_row, out_W, 1), 0) u_view = out_t.as_strided((T, out_Hu, out_Wu), (out_row, out_Wu, 1), out_hw) v_view = out_t.as_strided((T, out_Hu, out_Wu), (out_row, out_Wu, 1), out_hw + out_uv) if pool is None: pool = _get_default_reassemble_pool() def _write(view, shard, r, c, h_per, w_per, bound_h, bound_w): r0, c0 = r * h_per, c * w_per vh = min(h_per, bound_h - r0) vw = min(w_per, bound_w - c0) if vh <= 0 or vw <= 0: return # shard lies entirely in the padded tail # shard (1, h_per, w_per, T) -> squeeze(0).permute(2, 0, 1) -> (T, h_per, w_per). src = shard.squeeze(0).permute(2, 0, 1)[:, :vh, :vw] view[:, r0 : r0 + vh, c0 : c0 + vw].copy_(src) futures = [] for coord, shard in zip(mesh_coords, Y_shards): r, c = int(coord[0]), int(coord[1]) futures.append(pool.submit(_write, y_view, shard, r, c, h_per_y, w_per_y, out_H, out_W)) for coord, shard in zip(mesh_coords, Cb_shards): r, c = int(coord[0]), int(coord[1]) futures.append(pool.submit(_write, u_view, shard, r, c, h_per_uv, w_per_uv, out_Hu, out_Wu)) for coord, shard in zip(mesh_coords, Cr_shards): r, c = int(coord[0]), int(coord[1]) futures.append(pool.submit(_write, v_view, shard, r, c, h_per_uv, w_per_uv, out_Hu, out_Wu)) for f in futures: f.result() return out def fast_device_to_host_yuv( tt_video_BCTHW: ttnn.Tensor, mesh_device: ttnn.MeshDevice, *, ccl_manager=None, root: int | None = None, coefficients=None, pool: ThreadPoolExecutor | None = None, debug: bool = False, logical_h: int | None = None, logical_w: int | None = None, ) -> np.ndarray | None: """On-device YUV 4:2:0 conversion + batched D2H + planar uint8 concat. Takes a sharded BCTHW bf16 row-major tensor with values in ``[-1, 1]`` — typically the output of the Wan VAE — and returns a single numpy uint8 array in ffmpeg ``AV_PIX_FMT_YUV420P`` layout. On a single-host system, reads all per-device shards concurrently with async DMA, converts each shard's YUV planes, and scatters into the planar output on host. On a multi-host (distributed) system, uses a hybrid approach: an on-device all_gather for the inter-host axis only, then re-shards with repeat + mesh_partition so each local device holds unique data, runs the YUV conversion on the gathered data, and performs fast async DMA + planar concat for local shards only. Pipeline: 1. (Multi-host only) ``all_gather`` + ``repeat`` + ``mesh_partition`` on the inter-host axis of the bf16 BCTHW input. 2. Permute BCTHW -> BCHWT (T moves to the last position) and reshape to drop the B=1 dim, landing in CHWT — the layout the YUV kernel expects. 3. ``ttnn.experimental.rgb_to_yuv`` runs on each device's local shard, producing 3 uint8 outputs (Y full-res, Cb/Cr 4:2:0 subsampled). 4. Async ``cpu(blocking=False)`` on all three outputs followed by a single ``synchronize_device`` so the D2H reads overlap. 5. Per-shard scatters are dispatched across the shared reassembly thread pool using torch's strided-copy backend. Output shape: ``(T, H*W + 2*(H/2 * W/2))`` numpy uint8 — one row per frame, ``[Y plane | Cb plane | Cr plane]`` in row-major. Args: tt_video_BCTHW: Sharded ttnn tensor with shape ``(1, 3, T, H, W)``, bfloat16, ROW_MAJOR_LAYOUT, sharded ``{axis 0: dim 3 (H), axis 1: dim 4 (W)}``. Values must lie in ``[-1, 1]`` — the YUV kernel's expected range. mesh_device: The mesh device. ccl_manager: Optional :class:`CCLManager` instance. Required for multi-host environments where only local devices are accessible. root: If set, only the host with this MPI rank performs the D2H transfer and returns the assembled array; all other ranks return ``None``. If ``None`` (default), all ranks perform D2H. coefficients: ``ttnn.experimental.YUVCoefficients`` to use for the per-channel weights and offsets. Defaults to BT.601 limited range. pool: Optional ``ThreadPoolExecutor`` for the host-side reassembly. If ``None``, the module-level lazy default pool is used. debug: If ``True``, print diagnostic shape information. logical_h: Optional logical (un-padded) height of the output. When the VAE pads ``H`` to a coarser size, pass the true logical height here and the function will trim the bottom rows of each plane in the planar buffer on host (Y -> top ``logical_h`` rows; Cb/Cr -> top ``logical_h/2`` rows). Must be even and ``<= H``. Defaults to ``None`` (no trim). logical_w: Optional logical (un-padded) width of the output, trimming the right columns of each plane exactly as ``logical_h`` trims the bottom rows. Must be even and ``<= W``. Defaults to ``None``. Returns: ``np.ndarray`` of shape ``(T, H'*W' + 2*(H'/2 * W'/2))``, dtype uint8, where ``H'``/``W'`` are ``logical_h``/``logical_w`` when set and ``H``/``W`` otherwise. Returns ``None`` for non-root ranks when ``root`` is set. Raises: AssertionError: if ``B != 1``, ``C != 3``, or H/W are not even. ValueError: if ``logical_h``/``logical_w`` is non-positive, odd, or exceeds ``H``/``W``. """ if coefficients is None: coefficients = _bt601_yuv_coefficients() # NOTE: ttnn ``.shape`` on a multi-device sharded tensor returns the per-shard (local) shape mesh_shape = tuple(mesh_device.shape) B, C, T, h_per, w_per = tt_video_BCTHW.shape assert B == 1, f"fast_device_to_host_yuv requires B=1, got {B}" assert C == 3, f"fast_device_to_host_yuv requires C=3 (RGB), got {C}" TP, SP = mesh_shape H, W = h_per * TP, w_per * SP # The shards only ever cover HxW, so an oversized crop leaves the tail rows/columns # of the planar buffer untouched; an odd one halves to a chroma plane that no longer # matches its luma plane. Both surface far downstream, as corrupt video or an opaque # reshape failure. for name, value, bound_name, bound in ( ("logical_h", logical_h, "H", H), ("logical_w", logical_w, "W", W), ): if value is None: continue if value <= 0 or value % 2 != 0: msg = f"{name} must be positive and even for 4:2:0, got {value}" raise ValueError(msg) if value > bound: msg = f"{name}={value} exceeds the assembled {bound_name}={bound}" raise ValueError(msg) if debug: print(f" [yuv-d2h] input per-shard: {list(tt_video_BCTHW.shape)}") print(f" [yuv-d2h] global H={H}, W={W}, T={T} (mesh TP={TP}, SP={SP})") # Sharding convention: axis 0 -> dim 3 (H), axis 1 -> dim 4 (W). concat_dims: list[int | None] = [3, 4] # --- Multi-host: hybrid on-device collective + fast local DMA ----------- d2h_view = None if ttnn.using_distributed_env(): if ccl_manager is None: msg = "fast_device_to_host_yuv requires ccl_manager in a distributed (multi-host) environment" raise ValueError(msg) d2h_view = mesh_device.get_view() rank = int(ttnn.distributed_context_get_rank()) inter_host_axis = _get_inter_host_axis(mesh_device, d2h_view, mesh_shape) inter_dim = concat_dims[inter_host_axis] if inter_dim is not None and mesh_shape[inter_host_axis] > 1: # Move the gather dim out of the tile dims (last two) to position 2 (dim=-3) if inter_dim == 4: pre_dims = (0, 1, 4, 2, 3) # BCTHW -> BCWTH post_dims = (0, 1, 3, 4, 2) # BCWTH -> BCTHW else: # inter_dim == 3 pre_dims = (0, 1, 3, 2, 4) # BCTHW -> BCHTW post_dims = (0, 1, 3, 2, 4) # BCHTW -> BCTHW (same swap) ag_dim = 2 tt_video_BCTHW = ttnn.permute(tt_video_BCTHW, pre_dims) tt_video_BCTHW = ttnn.to_layout(tt_video_BCTHW, ttnn.TILE_LAYOUT) tt_video_BCTHW = ccl_manager.all_gather( tt_video_BCTHW, dim=ag_dim, mesh_axis=inter_host_axis, use_hyperparams=True, use_persistent_buffer=True, ) # Drop back to ROW_MAJOR before repeat/mesh_partition so ttnn.repeat doesn't wrap itself in an Untilize → tt_video_BCTHW = ttnn.to_layout(tt_video_BCTHW, ttnn.ROW_MAJOR_LAYOUT) n_hosts = int(ttnn.distributed_context_get_size()) if n_hosts > 1: repeat_dims = [1] * len(tt_video_BCTHW.shape) repeat_dims[ag_dim] = n_hosts tt_video_BCTHW = ttnn.repeat(tt_video_BCTHW, repeat_dims) tt_video_BCTHW = ttnn.mesh_partition(tt_video_BCTHW, dim=ag_dim, cluster_axis=inter_host_axis) tt_video_BCTHW = ttnn.permute(tt_video_BCTHW, post_dims) # Recompute per-shard dims after gather — they may have grown. B, C, T, h_per, w_per = tt_video_BCTHW.shape if debug: print(f" [yuv-d2h] (distributed) post-gather per-shard: {list(tt_video_BCTHW.shape)}") if root is not None and rank != root: return None assert ( h_per % 2 == 0 and w_per % 2 == 0 ), f"per-shard H and W must be even for 4:2:0 (got h_per={h_per}, w_per={w_per})" # 1. Reorder BCTHW -> CHWT for the YUV kernel. Shapes here are per-shard. tt_BCHWT = ttnn.permute(tt_video_BCTHW, (0, 1, 3, 4, 2)) if debug: print(f" [yuv-d2h] after permute(0,1,3,4,2) per-shard: {list(tt_BCHWT.shape)}") tt_CHWT = ttnn.reshape(tt_BCHWT, (C, h_per, w_per, T)) if debug: print(f" [yuv-d2h] after reshape to (C,h_per,w_per,T) per-shard: {list(tt_CHWT.shape)}") # 2. On-device YUV 4:2:0 -> 3 uint8 tensors. tt_Y, tt_Cb, tt_Cr = ttnn.experimental.rgb_to_yuv(tt_CHWT, coefficients=coefficients) if debug: print(f" [yuv-d2h] yuv outputs per-shard:") print(f" [yuv-d2h] Y : {list(tt_Y.shape)}") print(f" [yuv-d2h] Cb: {list(tt_Cb.shape)}") print(f" [yuv-d2h] Cr: {list(tt_Cr.shape)}") # 3+4 new_H = logical_h if logical_h is not None else H new_W = logical_w if logical_w is not None else W out = _yuv_planar_d2h(tt_Y, tt_Cb, tt_Cr, mesh_device, H, W, T, out_H=new_H, out_W=new_W, view=d2h_view, pool=pool) return out