File size: 8,527 Bytes
b9db84c d18e0a7 b9db84c | 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 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | # SPDX-License-Identifier: Apache-2.0
"""Fused host input path: RGB uint8 image -> the window-order fp32 im2col rows of the patch-embed matmul, in one
threaded pass, written straight into the persistent host buffers that are uploaded to the device.
Reference path (``app._to_pixel_values`` + ``fused_tables.im2col_window_order``):
x = u8 (CHW, fp32) / 255; x = F.interpolate(x, (1008, 1008), bilinear, align_corners=False); x = (x - 0.5) / 0.5
im2col: [nW = 9, 576 tokens, 608 cols], cols (c, kh, kw), 588 .. 607 zero
This module computes the same values with the same fp32 arithmetic as torch's CPU bilinear kernel (source index
``fma(scale, dst + 0.5, -0.5)``, per axis ``fma(t0, w0, t1 * w1)``, width first; the tables of ``maskpost._axis_table``)
and writes each value to its im2col position directly (no fp32 CHW image, no interpolate output, no im2col copy).
The work is split by window row ``band`` (output pixel rows 336 b .. 336 b + 335 = the 3 windows of one window
row = rows 1728 b .. 1728 b + 1727 of the im2col matrix), so band b can be uploaded while band b + 1 is computed.
``self_check`` compares the result bit for bit with the reference path on several image sizes; the caller falls back
to the reference path when it fails (another torch build) or when numba is missing.
"""
from __future__ import annotations
import numpy as np
try:
import numba
from numba import njit, prange, types
from numba.extending import intrinsic
except ImportError: # optional dependency
numba = None
from tt_sam3.server.maskpost import _axis_table
PATCH = 14
WIN = 24 # tokens per window side
OUT = 1008 # model input side
BAND_ROWS = PATCH * WIN # 336 output pixel rows per window row
N_BANDS = OUT // BAND_ROWS # 3
K = 3 * PATCH * PATCH # 588
K_PAD = 608
_F32 = np.float32
# u8 / 255 in fp32 (correctly rounded division, == torch's float32 x.div_(255.0) for all 256 values)
_LUT = (np.arange(256, dtype=_F32) / _F32(255.0)).astype(_F32)
if numba is not None:
@intrinsic
def _fmaf(typingctx, a, b, c): # noqa: ARG001 - numba intrinsic signature
sig = types.float32(types.float32, types.float32, types.float32)
def codegen(context, builder, signature, args): # noqa: ARG001
return builder.fma(args[0], args[1], args[2])
return sig, codegen
@njit(parallel=True, nogil=True, cache=True)
def _band_kernel(src, pstride, lut, h0, h1, hl0, hl1, w0, w1, wl0, wl1, band, out):
"""src: (H, W * pstride) uint8 rows (RGB in the first 3 bytes of each pixel); out: (3, 576, 608) fp32 band."""
half = np.float32(0.5)
for r in prange(BAND_ROWS):
y = band * BAND_ROWS + r
a0 = h0[y]
a1 = h1[y]
la0 = hl0[y]
la1 = hl1[y]
ti = r // PATCH # token row inside the window
kh = r - ti * PATCH
row0 = src[a0]
row1 = src[a1]
for x in range(OUT):
b0 = w0[x] * pstride
b1 = w1[x] * pstride
lb0 = wl0[x]
lb1 = wl1[x]
tx = x // PATCH
kw = x - tx * PATCH
ww = tx // WIN
tok = ti * WIN + (tx - ww * WIN)
for c in range(3):
# width first (two source rows), then height: fma(t0, w0, t1 * w1) per axis
t0 = _fmaf(lut[row0[b0 + c]], lb0, lut[row0[b1 + c]] * lb1)
t1 = _fmaf(lut[row1[b0 + c]], lb0, lut[row1[b1 + c]] * lb1)
v = _fmaf(t0, la0, t1 * la1)
out[ww, tok, c * (PATCH * PATCH) + kh * PATCH + kw] = (v - half) / half
@njit(parallel=True, nogil=True, cache=True)
def _band_kernel_identity(src, pstride, lut, band, out):
"""1008 x 1008 input: no interpolation (``_to_pixel_values`` skips F.interpolate)."""
half = np.float32(0.5)
for r in prange(BAND_ROWS):
y = band * BAND_ROWS + r
ti = r // PATCH
kh = r - ti * PATCH
row = src[y]
for x in range(OUT):
tx = x // PATCH
kw = x - tx * PATCH
ww = tx // WIN
tok = ti * WIN + (tx - ww * WIN)
for c in range(3):
out[ww, tok, c * (PATCH * PATCH) + kh * PATCH + kw] = (lut[row[x * pstride + c]] - half) / half
_NUM_THREADS = [8]
def configure(num_threads: int) -> None:
"""Thread count of the band kernels (numba's setting is per calling thread; ``BandWriter.write`` applies it)."""
if numba is not None:
_NUM_THREADS[0] = max(1, min(int(num_threads), numba.config.NUMBA_NUM_THREADS))
def available() -> bool:
return numba is not None
def _rows_view(rgb: np.ndarray):
"""(H, W, 3|4) uint8 (possibly the RGB view of an RGBA buffer) -> ((H, W * pstride) uint8 rows, pstride)."""
if rgb.dtype != np.uint8 or rgb.ndim != 3 or rgb.shape[2] < 3:
raise ValueError("expected (H, W, 3) uint8")
h, w = rgb.shape[:2]
s0, s1, s2 = rgb.strides
if s2 == 1 and s1 in (3, 4) and s0 >= w * s1:
# rows of the decoded buffer as they are (RGBA pixels, PNG filter-byte row padding): a strided 2-D view that
# covers up to the last pixel's blue byte (never past the buffer)
return np.lib.stride_tricks.as_strided(rgb, ((h, (w - 1) * s1 + 3)), (s0, 1), writeable=False), s1
a = np.ascontiguousarray(rgb[..., :3])
return a.reshape(h, w * 3), 3
class BandWriter:
"""Prepared per image: ``write(band, out)`` fills one (3, 576, 608) band of the window-order im2col matrix."""
def __init__(self, rgb: np.ndarray):
self.rows, self.pstride = _rows_view(rgb)
h, w = rgb.shape[:2]
self.identity = (h, w) == (OUT, OUT)
if not self.identity:
self.ht = _axis_table(h, OUT)
self.wt = _axis_table(w, OUT)
def write(self, band: int, out: np.ndarray) -> None:
if out.shape != (N_BANDS, WIN * WIN, K_PAD) or out.dtype != _F32:
raise ValueError("band buffer must be fp32 (3, 576, 608)")
if numba.get_num_threads() != _NUM_THREADS[0]:
numba.set_num_threads(_NUM_THREADS[0])
if self.identity:
_band_kernel_identity(self.rows, self.pstride, _LUT, band, out)
else:
h0, h1, hl0, hl1 = self.ht
w0, w1, wl0, wl1 = self.wt
_band_kernel(self.rows, self.pstride, _LUT, h0, h1, hl0, hl1, w0, w1, wl0, wl1, band, out)
def reference_bands(rgb: np.ndarray) -> np.ndarray:
"""The served reference path (``_to_pixel_values`` + window-order im2col) -> (9, 576, 608) fp32."""
import torch
from tt_sam3.imaging import DecodedImage, to_pixel_values as _to_pixel_values
from tt_sam3.tt.fused_tables import im2col_window_order
px = _to_pixel_values(DecodedImage(rgb))
out = torch.zeros(N_BANDS * N_BANDS, WIN * WIN, K_PAD)
return im2col_window_order(px, PATCH, WIN, out=out, k_pad=K_PAD, window_order=True).numpy()
def fused_bands(rgb: np.ndarray) -> np.ndarray:
out = np.zeros((N_BANDS * N_BANDS, WIN * WIN, K_PAD), dtype=_F32)
bw = BandWriter(rgb)
for b in range(N_BANDS):
bw.write(b, out[N_BANDS * b : N_BANDS * (b + 1)])
return out
def self_check(seed: int = 0) -> bool:
"""Bit-for-bit comparison with the reference path (values incl. the zero pad columns) on random images of several
sizes (down / up / identity / mixed scales, RGB and the RGB view of RGBA)."""
if numba is None:
return False
rng = np.random.default_rng(seed)
try: # numba JIT / cache errors (e.g. no writable cache location in a container) -> the reference path
for h, w, ch in ((1080, 1920, 4), (1008, 1008, 3), (720, 1280, 3), (1500, 1100, 4), (333, 517, 3), (1008, 2000, 3)):
a = rng.integers(0, 256, size=(h, w, ch), dtype=np.uint8)
rgb = a[..., :3] if ch == 4 else a
if not np.array_equal(fused_bands(rgb).view(np.uint32), reference_bands(rgb).view(np.uint32)):
return False
# the PNG fast path's layout: RGBA rows W * 4 + 1 bytes apart (filter byte)
rows = rng.integers(0, 256, size=(h, w * 4 + 1), dtype=np.uint8)
rgb = rows[:, 1:].reshape(h, w, 4)[..., :3]
if not np.array_equal(fused_bands(rgb).view(np.uint32), reference_bands(rgb).view(np.uint32)):
return False
except Exception: # noqa: BLE001
return False
return True
|