code/models/tt_dit/utils/yuv_d2h.py · sunandbean/diffrhythm-1.2-p100a at main

diffrhythm-1.2-p100a / code /models /tt_dit /utils /yuv_d2h.py

Add files using upload-large-folder tool

629b123 verified 11 days ago

18.6 kB

# 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

	# 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, HW + 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, HW + 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