mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
refactor(infra): optimize YUV420 frame packing using numpy
- Vectorized create_yuv420_frame when numpy is available to eliminate the 500k-iteration scalar Python loop. - Pre-filled memoryview slice buffers in the fallback path. - Updated 4-Infrastructure/AGENTS.md to document the optimization. Build: 3313 jobs, 0 errors (lake build)
This commit is contained in:
Brandon Schneider
parent
7a011ffeee
commit
61b38ca697
2 changed files with 86 additions and 48 deletions
|
|
@ -307,7 +307,7 @@ python3 4-Infrastructure/storage/storage_agent.py --loop --interval 900
|
|||
- `4-Infrastructure/shim/spirv_copy_if_optimizer.py` — SPIR-V OpPhi→OpSelect transform (OpBranchConditional+OpPhi → OpSelect); eliminates branch over head; emits `CopyIfPattern` with type_id fix
|
||||
- `4-Infrastructure/shim/spirv_packet_generator.py` — OpPhi-driven packet descriptor generator: SPIR-V asm → copy-if optimizer → JSON packet descriptors (5 OpPhi fields: type_id, cond_id, true_val_id, false_val_id, result_id)
|
||||
- `4-Infrastructure/shim/virtio_net_transform.py` — Virtio-net ring as computation pipeline: three Class-1 primitives (HASH_REPORT RSS Toeplitz, TSO gso_size split, MRG_RXBUF merge) via virtio_net_hdr_v1_hash; zero backend changes needed
|
||||
- `4-Infrastructure/shim/vcn_compute_substrate.py` — AMD VCN H.264 encoder as compute device via MKV trick; SEI receipt schema; carries BraidStrand/BraidBracket payloads
|
||||
- `4-Infrastructure/shim/vcn_compute_substrate.py` — AMD VCN H.264 encoder as compute device via MKV trick; SEI receipt schema; carries BraidStrand/BraidBracket payloads; vectorized YUV420 packing using numpy
|
||||
- `4-Infrastructure/shim/rrc_ray_tagger.py` — RRC Ray Layer Tagger; classifies math payloads into RRC shapes and matches them to swappable compute slots and transports
|
||||
- `4-Infrastructure/hardware/emergency_boot/emergency_boot_shim.py` — Python I/O shim
|
||||
for Geometry Emergency Boot Witness (6502 calculator-efficiency FPGA controller)
|
||||
|
|
|
|||
|
|
@ -40,6 +40,13 @@ try:
|
|||
except ImportError:
|
||||
_CHA20_AVAILABLE = False
|
||||
|
||||
try:
|
||||
import numpy as np
|
||||
_HAS_NUMPY = True
|
||||
except ImportError:
|
||||
np = None
|
||||
_HAS_NUMPY = False
|
||||
|
||||
# Frame constants from UNIFIED_TRANSPORT_ENCODING_SPEC.md
|
||||
FRAME_WIDTH = 1920
|
||||
FRAME_HEIGHT = 1080
|
||||
|
|
@ -496,58 +503,89 @@ def create_yuv420_frame(data: bytes, seq: int) -> bytes:
|
|||
"""
|
||||
if len(data) > YUV420_FRAME_SIZE - SIGNATURE_SIZE:
|
||||
raise ValueError(f"Data too large: {len(data)} > {YUV420_FRAME_SIZE - SIGNATURE_SIZE}")
|
||||
|
||||
# Create frame buffer
|
||||
frame = bytearray(YUV420_FRAME_SIZE)
|
||||
|
||||
|
||||
if _HAS_NUMPY:
|
||||
# Create empty planes filled with neutral gray (128)
|
||||
y_plane = np.full(FRAME_WIDTH * FRAME_HEIGHT, 128, dtype=np.uint8)
|
||||
u_plane = np.full((FRAME_WIDTH // 2) * (FRAME_HEIGHT // 2), 128, dtype=np.uint8)
|
||||
v_plane = np.full((FRAME_WIDTH // 2) * (FRAME_HEIGHT // 2), 128, dtype=np.uint8)
|
||||
|
||||
# Parse data as uint32 values (little-endian)
|
||||
# Pad data to multiple of 4 bytes if necessary
|
||||
pad_len = (4 - (len(data) % 4)) % 4
|
||||
if pad_len > 0:
|
||||
padded_data = data + b"\x00" * pad_len
|
||||
else:
|
||||
padded_data = data
|
||||
|
||||
values = np.frombuffer(padded_data, dtype=np.uint32)
|
||||
|
||||
# Vectorized pack_q16_16_to_yuv
|
||||
y_vals = np.clip((values >> 16) & 0xFF, 16, 235).astype(np.uint8)
|
||||
u_vals = np.clip((values >> 8) & 0xFF, 16, 240).astype(np.uint8)
|
||||
v_vals = np.clip(values & 0xFF, 16, 240).astype(np.uint8)
|
||||
|
||||
# Y plane: repeat each element 4 times
|
||||
y_plane_data = np.repeat(y_vals, 4)
|
||||
|
||||
# Copy to planes
|
||||
n_y = min(len(y_plane_data), len(y_plane))
|
||||
n_uv = min(len(values), len(u_plane))
|
||||
|
||||
y_plane[:n_y] = y_plane_data[:n_y]
|
||||
u_plane[:n_uv] = u_vals[:n_uv]
|
||||
v_plane[:n_uv] = v_vals[:n_uv]
|
||||
|
||||
# Concatenate planes
|
||||
frame = bytearray(y_plane.tobytes() + u_plane.tobytes() + v_plane.tobytes())
|
||||
else:
|
||||
# Create frame buffer
|
||||
frame = bytearray(YUV420_FRAME_SIZE)
|
||||
|
||||
# Pack data into YUV420 format
|
||||
# Y plane: 1920×1080 bytes
|
||||
# U plane: 960×540 bytes (subsampled 2x)
|
||||
# V plane: 960×540 bytes (subsampled 2x)
|
||||
|
||||
y_plane_view = memoryview(frame)[0:FRAME_WIDTH * FRAME_HEIGHT]
|
||||
u_plane_view = memoryview(frame)[FRAME_WIDTH * FRAME_HEIGHT:FRAME_WIDTH * FRAME_HEIGHT + (FRAME_WIDTH // 2) * (FRAME_HEIGHT // 2)]
|
||||
v_plane_view = memoryview(frame)[FRAME_WIDTH * FRAME_HEIGHT + (FRAME_WIDTH // 2) * (FRAME_HEIGHT // 2):]
|
||||
|
||||
# Pre-fill with neutral gray
|
||||
y_plane_view[:] = b"\x80" * len(y_plane_view)
|
||||
u_plane_view[:] = b"\x80" * len(u_plane_view)
|
||||
v_plane_view[:] = b"\x80" * len(v_plane_view)
|
||||
|
||||
# Pack data as Q16_16 scalars into macroblocks (6 bytes = 4Y + 1U + 1V)
|
||||
data_idx = 0
|
||||
macroblock_idx = 0
|
||||
|
||||
while data_idx + 4 <= len(data):
|
||||
# Read 4 bytes as Q16_16 scalar
|
||||
q16_value = struct.unpack("<I", data[data_idx:data_idx + 4])[0]
|
||||
y, u, v = pack_q16_16_to_yuv(q16_value)
|
||||
|
||||
# Place in macroblock (4 Y pixels, 1 U, 1 V)
|
||||
mb_y_start = macroblock_idx * 4
|
||||
mb_u_start = macroblock_idx
|
||||
mb_v_start = macroblock_idx
|
||||
|
||||
if mb_y_start + 4 <= len(y_plane_view):
|
||||
y_plane_view[mb_y_start:mb_y_start + 4] = bytes([y, y, y, y])
|
||||
if mb_u_start + 1 <= len(u_plane_view):
|
||||
u_plane_view[mb_u_start] = u
|
||||
if mb_v_start + 1 <= len(v_plane_view):
|
||||
v_plane_view[mb_v_start] = v
|
||||
|
||||
data_idx += 4
|
||||
macroblock_idx += 1
|
||||
|
||||
# Write signature header (exactly 24 bytes)
|
||||
version = 1
|
||||
length = len(data)
|
||||
header = SIGNATURE_HEADER + struct.pack("<IIII", version, seq, length, 0) # 8 + 4 + 4 + 4 + 4 = 24 bytes
|
||||
header = SIGNATURE_HEADER + struct.pack("<IIII", version, seq, length, 0)
|
||||
frame[:SIGNATURE_SIZE] = header
|
||||
|
||||
# Ensure frame is exactly the right size
|
||||
if len(frame) != YUV420_FRAME_SIZE:
|
||||
raise ValueError(f"Frame size mismatch: {len(frame)} != {YUV420_FRAME_SIZE}")
|
||||
|
||||
# Pack data into YUV420 format
|
||||
# Y plane: 1920×1080 bytes
|
||||
# U plane: 960×540 bytes (subsampled 2x)
|
||||
# V plane: 960×540 bytes (subsampled 2x)
|
||||
|
||||
y_plane = frame[0:FRAME_WIDTH * FRAME_HEIGHT]
|
||||
u_plane = frame[FRAME_WIDTH * FRAME_HEIGHT:FRAME_WIDTH * FRAME_HEIGHT + (FRAME_WIDTH // 2) * (FRAME_HEIGHT // 2)]
|
||||
v_plane = frame[FRAME_WIDTH * FRAME_HEIGHT + (FRAME_WIDTH // 2) * (FRAME_HEIGHT // 2):]
|
||||
|
||||
# Pack data as Q16_16 scalars into macroblocks (6 bytes = 4Y + 1U + 1V)
|
||||
data_idx = 0
|
||||
macroblock_idx = 0
|
||||
|
||||
while data_idx + 4 <= len(data):
|
||||
# Read 4 bytes as Q16_16 scalar
|
||||
q16_value = struct.unpack("<I", data[data_idx:data_idx + 4])[0]
|
||||
y, u, v = pack_q16_16_to_yuv(q16_value)
|
||||
|
||||
# Place in macroblock (4 Y pixels, 1 U, 1 V)
|
||||
mb_y_start = macroblock_idx * 4
|
||||
mb_u_start = macroblock_idx * 1
|
||||
mb_v_start = macroblock_idx * 1
|
||||
|
||||
if mb_y_start + 4 <= len(y_plane):
|
||||
y_plane[mb_y_start:mb_y_start + 4] = bytes([y, y, y, y])
|
||||
if mb_u_start + 1 <= len(u_plane):
|
||||
u_plane[mb_u_start] = u
|
||||
if mb_v_start + 1 <= len(v_plane):
|
||||
v_plane[mb_v_start] = v
|
||||
|
||||
data_idx += 4
|
||||
macroblock_idx += 1
|
||||
|
||||
# Pad remaining space with neutral gray (Y=128, U=128, V=128)
|
||||
y_plane[macroblock_idx * 4:] = bytes([128] * (len(y_plane) - macroblock_idx * 4))
|
||||
u_plane[macroblock_idx:] = bytes([128] * (len(u_plane) - macroblock_idx))
|
||||
v_plane[macroblock_idx:] = bytes([128] * (len(v_plane) - macroblock_idx))
|
||||
|
||||
return bytes(frame)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue