File size: 5,522 Bytes
629b123 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | # SPDX-FileCopyrightText: (c) 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
"""Python wrapper around the C++ planar concat extension.
The C++ source lives in ``models/tt_dit/utils/cpp/`` and produces
``cpp/build/_planar_concat<abi>.so`` when built via ``cpp/build.sh``. This
module loads it via ``importlib.util`` so callers don't need to put the
build directory on ``sys.path``.
If the ``.so`` isn't built (or fails to import), :data:`HAS_CPP_PLANAR_CONCAT`
is False and :func:`planar_concat_cpp` is None. The benchmark harness in
``test_fast_device_to_host.py`` already has a ``HAS_*`` skip pattern that
this slots into.
"""
from __future__ import annotations
import importlib.util
import os
from typing import TYPE_CHECKING, Sequence
import numpy as np
if TYPE_CHECKING:
pass
_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
_BUILD_DIR = os.path.join(_THIS_DIR, "cpp", "build")
def _try_load_extension():
if not os.path.isdir(_BUILD_DIR):
return None
candidates = [f for f in os.listdir(_BUILD_DIR) if f.startswith("_planar_concat") and f.endswith(".so")]
if not candidates:
return None
so_path = os.path.join(_BUILD_DIR, candidates[0])
spec = importlib.util.spec_from_file_location("_planar_concat", so_path)
if spec is None or spec.loader is None:
return None
mod = importlib.util.module_from_spec(spec)
try:
spec.loader.exec_module(mod)
except Exception:
return None
return mod
_ext = _try_load_extension()
HAS_CPP_PLANAR_CONCAT: bool = _ext is not None
def _to_numpy(shard) -> np.ndarray:
"""Convert a torch tensor to a contiguous numpy uint8 array (zero-copy when possible)."""
if isinstance(shard, np.ndarray):
if not shard.flags.c_contiguous:
shard = np.ascontiguousarray(shard)
return shard
# torch.Tensor or anything else with .numpy()
if hasattr(shard, "contiguous") and hasattr(shard, "numpy"):
t = shard.contiguous()
return t.numpy()
raise TypeError(f"unsupported shard type: {type(shard)!r}")
if HAS_CPP_PLANAR_CONCAT:
def planar_concat_cpp(
y_shards: Sequence,
u_shards: Sequence,
v_shards: Sequence,
dim_order: str,
mesh_shape: tuple[int, int] = (4, 8),
out: np.ndarray | None = None,
out_H: int | None = None,
out_W: int | None = None,
) -> np.ndarray:
"""Vectorized YUV 4:2:0 planar concat — C++/AVX2 implementation.
Equivalent in output to the ``planar_concat_torch_threaded``
reference in ``test_fast_device_to_host.py`` but runs in C++ with a
persistent ``std::thread`` pool and AVX2 byte-tile transposes for
the CHWT path.
Args:
y_shards, u_shards, v_shards: lists of ``TP*SP`` per-shard
tensors/arrays. Per-shard shape:
* CHWT: ``(1, h_per, w_per, T)`` uint8 contiguous.
* CTHW: ``(1, T, h_per, w_per)`` uint8 contiguous.
UV shards must have ``h_per/2`` and ``w_per/2`` of the Y
shards' dims (4:2:0 subsampling).
dim_order: ``"CHWT"`` or ``"CTHW"``.
mesh_shape: ``(TP, SP)``. Defaults to ``(4, 8)``.
out: Optional pre-allocated output buffer of shape
``(T, H*W + 2*(H/2 * W/2))`` uint8. When ``None``, a fresh
``np.empty`` is allocated each call (matches the convention
of the other variants in the test harness). Reusing a
buffer across calls eliminates ~50 ms of first-touch page
faults per call on systems without THP=always.
Returns:
``np.ndarray`` of shape ``(T, H*W + 2*(H/2 * W/2))``, dtype
``uint8``. Per-frame layout: ``[Y plane | Cb plane | Cr plane]``.
"""
y_np = [_to_numpy(s) for s in y_shards]
u_np = [_to_numpy(s) for s in u_shards]
v_np = [_to_numpy(s) for s in v_shards]
TP, SP = int(mesh_shape[0]), int(mesh_shape[1])
if dim_order == "CHWT":
_, h_per, w_per, T = y_np[0].shape
elif dim_order == "CTHW":
_, T, h_per, w_per = y_np[0].shape
else:
raise ValueError(f"dim_order must be 'CHWT' or 'CTHW', got {dim_order!r}")
H, W = h_per * TP, w_per * SP
# Logical (cropped) output: when the VAE pads a global tail, write only the valid frame.
oH = H if out_H is None else out_H
oW = W if out_W is None else out_W
oHu, oWu = oH // 2, oW // 2
row_stride = oH * oW + 2 * oHu * oWu
if out is None:
out = np.empty((T, row_stride), dtype=np.uint8)
elif out.shape != (T, row_stride) or out.dtype != np.uint8 or not out.flags.c_contiguous:
raise ValueError(
f"out must be C-contiguous uint8 with shape ({T}, {row_stride}); "
f"got shape {out.shape} dtype {out.dtype} c_contig={out.flags.c_contiguous}"
)
_ext.planar_concat(y_np, u_np, v_np, dim_order, mesh_shape, out, oH, oW)
return out
def set_thread_pool_size(n_threads: int) -> None:
"""Set the C++ thread pool size. Must be called BEFORE first scatter."""
_ext.set_thread_pool_size(int(n_threads))
else:
planar_concat_cpp = None # type: ignore[assignment]
def set_thread_pool_size(n_threads: int) -> None: # noqa: D401
"""No-op fallback when the C++ extension isn't built."""
|