diff --git a/.gitignore b/.gitignore index c23e9a40..68404c62 100644 --- a/.gitignore +++ b/.gitignore @@ -12,3 +12,5 @@ __pycache__/ # macOS .DS_Store + +__pycache__/ diff --git a/docs/images/daqiri_rx_path.svg b/docs/images/daqiri_rx_path.svg new file mode 100644 index 00000000..75bff4ee --- /dev/null +++ b/docs/images/daqiri_rx_path.svg @@ -0,0 +1,351 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + DAQIRI RX Path — Pointer Lifecycle & Multiple Receive Modes + + + DPDK backend · ConnectX-6/7 · GPUDirect via peermem or DMA-BUF (or host_pinned on integrated GPUs) + + + + + + + + Legend + + + packet data (DMA, zero-copy) + + + descriptor / handle exchange (rte_mbuf*, BurstParams*) + + + free-return path (caller must close) + + + setup-time wiring + + + + + + Wire — mixed UDP flows + + + udp_dst=4000 + udp_dst=4001 + udp_dst=4002 + flex_item match + udp_dst=5000 (batched) + udp_dst=5001 (batched) + + + + + + + + CONNECTX-6 / -7 NIC + Flow steering engine (rte_flow) + Matches on UDP src/dst, IPv4 length, or flex-item bytes after UDP header. Action = queue.id. + + + + rule → queue 0 (HDS, 2 segments) + + + rule → queue 1 (GPU-only, batched) + + + + + + + + + + + + Queue 0 — Header-Data Split + 2 memory_regions → segment 0 (CPU hdr) + segment 1 (GPU payload) + + + + + SEGMENT 0 — CPU MEMPOOL + kind: huge (hugepages) + buf_size ≈ 64 B (Eth+IPv4+UDP) + num_bufs ≥ 3× batch_size + Pre-allocated. NIC DMAs into + these slots — never copied. + + + + + + SEGMENT 1 — GPU MEMPOOL + kind: device (GPU VRAM) + DMA-mapped via peermem/DMA-BUF + buf_size = MTU − header + GPUDirect: NIC → GPU VRAM, + no CPU bounce. + + + + + + NIC DESCRIPTOR RING (per-queue, SPSC) + rte_mbuf * — one descriptor pair per packet (seg-0 + seg-1 mbuf chain) + Filled by NIC DMA · drained by rte_eth_rx_burst() in worker + + + + + + RX worker thread + pinned to rx.queue.cpu_core (e.g. 5) + + + + + + HOST RX RING (lock-free SPSC) + BurstParams * handles (allocated from rx_meta_buffers pool) + Worker enqueues · user dequeues via get_rx_burst() + + + + + + USER READS (HDS) + get_rx_burst(&burst, port, /*q=*/0) + N = get_num_packets(burst) + hdr_ptr = get_segment_packet_ptr(burst, /*seg=*/0, i) + pay_ptr = get_segment_packet_ptr(burst, /*seg=*/1, i) + + + + + + + + + Queue 1 — Batched GPU-only + 1 memory_region → segment 0 (full packet in GPU VRAM) + + + + + SEGMENT 0 — GPU MEMPOOL + kind: device (GPU VRAM) + buf_size = MTU (full packet, headers + payload) + num_bufs ≥ 3× batch_size (auto-bumped if <1.5× ring) + DMA-mapped via peermem/DMA-BUF; full packet lands + in VRAM in one shot — no segmentation. + + + + + + NIC DESCRIPTOR RING (per-queue, SPSC) + rte_mbuf * — one descriptor per packet (single segment) + Filled by NIC DMA · drained by rte_eth_rx_burst() + + + + + + RX worker thread + pinned to rx.queue.cpu_core (e.g. 6) + + + + + + HOST RX RING (lock-free SPSC) + BurstParams * handles (from rx_meta_buffers pool) + Worker enqueues · user dequeues via get_rx_burst() + + + + + + USER READS (BATCHED) + get_rx_burst(&burst, port, /*q=*/1) + N = get_num_packets(burst) + ptr = get_packet_ptr(burst, i) // single segment + len = get_packet_length(burst, i) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + User Application + Owns the BurstParams* between get_rx_burst() and the matching free. Reads from packet buffers in place — zero-copy. + + + + HDS PATH (queue 0) + hdr = get_segment_packet_ptr(burst, 0, i); pay = get_segment_packet_ptr(burst, 1, i); + + + + + BATCHED PATH (queue 1) + ptr = get_packet_ptr(burst, i); flow_id = get_packet_flow_id(burst, i); + + + + + + + + + Buffer return path — packets first, then burst (MUST NOT LEAK) + + free_rx_burst() releases only the BurstParams handle. It does NOT free packet buffers — those must be returned to the mempool first. + + + Skipping the per-packet step drains the packet mempool → NO_FREE_PACKET_BUFFERS → NIC drops new traffic. + + + + + + A · ONE-SHOT (RECOMMENDED) + Frees every packet (all segments) + the burst + handle in a single call. + + free_all_packets_and_burst_rx(burst); + // works for HDS & batched, any segment count + + + + + + B · PER-SEGMENT (HDS, FINE-GRAIN) + Free each segment's packet pool, then release + the burst handle. + + free_all_segment_packets(burst, 0); + free_all_segment_packets(burst, 1); free_rx_burst(burst); + + + + + + C · PER-PACKET (SELECTIVE) + Use when only some packets are kept (e.g. forwarded + to another stage). Then close the burst. + + for (i …) free_packet(burst, i); + free_rx_burst(burst); + + + + + Where buffers actually go back: + + • free_packet* / free_all_segment_packets → returns rte_mbuf to the segment's packet mempool (CPU hugepage / GPU device). + + + • free_rx_burst → returns the BurstParams handle to the rx_meta_buffers pool (separate from the packet pool). + + + • free_all_packets_and_burst_rx = both of the above, in order. + + + + + + + SIZING / FAILURE NOTES + + Rule of thumb: num_bufs ≥ 3× batch_size; DPDK auto-bumps to 3× ring (default 8192 → 24576) and logs WARN if too small (deadlock guard). + + + Symptoms of a missed free: NO_FREE_PACKET_BUFFERS / NO_FREE_BURST_BUFFERS, then NIC-level RX drops in imissed. + + + + + + + + + + + + packets back to mempool (seg 0/1) + + + + packets back to mempool + + + + + + Notes + + ① Number of memory_regions per queue picks the receive mode (1 = batched / CPU-only, 2 = header-data split). See docs/configuration.md §Memory Regions. + + + ② Same flow rule schema serves both columns: rx.flows[*].action.id chooses the queue; mode is determined by that queue's memory_regions list. + + + ③ On integrated GPUs (e.g. NVIDIA GB10 / DGX Spark) substitute kind: host_pinned for device — the NIC can't peer-DMA into iGPU VRAM. + + + ④ Sources: include/daqiri/common.h (API), include/daqiri/types.h (MemoryKind), src/managers/dpdk/daqiri_dpdk_mgr.cpp (rte_flow), docs/configuration.md. + + + diff --git a/docs/images/packet_diagrams/Makefile b/docs/images/packet_diagrams/Makefile new file mode 100644 index 00000000..d38729a7 --- /dev/null +++ b/docs/images/packet_diagrams/Makefile @@ -0,0 +1,38 @@ +PYTHON ?= python3 +REPO_ROOT := $(shell git -C "$(CURDIR)" rev-parse --show-toplevel 2>/dev/null) +GENERATED_PATHS := \ + docs/images/packet_diagrams/hds \ + docs/images/packet_diagrams/flow_steering \ + docs/images/packet_diagrams/reorder \ + docs/images/packet_diagrams/reorder_quantize + +.PHONY: all hds flow-steering reorder check clean + +all: hds flow-steering reorder + +hds: + $(PYTHON) hds_animation.py + +flow-steering: + $(PYTHON) flow_steering_animation.py + +reorder: + $(PYTHON) reorder_animation.py + +check: all + @if [ -z "$(REPO_ROOT)" ]; then \ + echo "error: packet diagram reproducibility check requires a git checkout" >&2; \ + exit 1; \ + fi + @if [ -n "$$(git -C "$(REPO_ROOT)" status --porcelain -- $(GENERATED_PATHS))" ]; then \ + git -C "$(REPO_ROOT)" status --short -- $(GENERATED_PATHS); \ + echo "error: packet diagram outputs are stale; rerun make -C docs/images/packet_diagrams and commit the results" >&2; \ + exit 1; \ + fi + @echo "Packet diagram outputs are reproducible." + +clean: + $(RM) hds/*.webp hds/*-poster.png + $(RM) flow_steering/*.webp flow_steering/*-poster.png + $(RM) reorder/*.webp reorder/*-poster.png + $(RM) reorder_quantize/*.webp reorder_quantize/*-poster.png diff --git a/docs/images/packet_diagrams/README.md b/docs/images/packet_diagrams/README.md new file mode 100644 index 00000000..765ce7e1 --- /dev/null +++ b/docs/images/packet_diagrams/README.md @@ -0,0 +1,63 @@ +# Packet Diagram Animations + +These scripts generate the animated packet-path diagrams used by the DAQIRI docs. + +## Requirements + +- Python 3 +- Pillow 11.1.x with animated WebP support +- DejaVu Sans or Liberation Sans fonts installed in a standard system font path + +Check the encoder support with: + +```bash +python3 - <<'PY' +from PIL import Image, features +print("Pillow", Image.__version__) +print("webp", features.check("webp")) +print("webp_anim", features.check("webp_anim")) +PY +``` + +## Regenerate + +From the repository root: + +```bash +make -C docs/images/packet_diagrams +``` + +Before committing, run: + +```bash +make -C docs/images/packet_diagrams check +``` + +The check regenerates the assets and fails if any generated output directory +has uncommitted changes afterward. + +Or run one generator directly: + +```bash +python3 docs/images/packet_diagrams/hds_animation.py +python3 docs/images/packet_diagrams/flow_steering_animation.py +python3 docs/images/packet_diagrams/reorder_animation.py +``` + +Each generator writes animated WebP files and PNG posters into its adjacent output directory. The reorder generator emits both the reorder-only and reorder-plus-convert variants. + +Animated WebP encoders may coalesce identical consecutive frames, so a decoded +WebP frame count can be lower than the script's logical frame count. Use +`make check` to verify reproducibility rather than comparing decoded frame +counts directly. + +## Size Notes + +`anim_common.py` currently uses `WEBP_METHOD = 0` to keep local regeneration +fast. Raising the method is still available when smaller artifacts matter more +than encoder time; expect roughly another 20-35% size reduction depending on the +animation. Run `make check` after changing it. + +The PR history before the WebP conversion contains large GIF blobs across the +animation commits. Squash before merging so those transient GIF binaries do not +remain in the main branch history. diff --git a/docs/images/packet_diagrams/anim_common.py b/docs/images/packet_diagrams/anim_common.py new file mode 100644 index 00000000..5d2d923d --- /dev/null +++ b/docs/images/packet_diagrams/anim_common.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from pathlib import Path +from collections.abc import Sequence + +from PIL import Image, ImageFont, features + + +WEBP_QUALITY = 80 +WEBP_METHOD = 0 + + +def font(size: int, *, scale: int = 2, bold: bool = False, mono: bool = False) -> ImageFont.FreeTypeFont: + if mono: + candidates = [ + "/usr/share/fonts/truetype/dejavu/DejaVuSansMono-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf", + "/usr/share/fonts/truetype/noto/NotoSansMono-Bold.ttf" if bold else "/usr/share/fonts/truetype/noto/NotoSansMono-Regular.ttf", + "C:/Windows/Fonts/consolab.ttf" if bold else "C:/Windows/Fonts/consola.ttf", + ] + elif bold: + candidates = [ + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", + "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf", + "C:/Windows/Fonts/segoeuib.ttf", + ] + else: + candidates = [ + "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf", + "C:/Windows/Fonts/segoeui.ttf", + ] + for path in candidates: + if Path(path).exists(): + return ImageFont.truetype(path, size * scale) + raise FileNotFoundError(f"missing required font; tried: {', '.join(candidates)}") + + +def font_scheme(scale: int = 2) -> dict[str, ImageFont.FreeTypeFont]: + return { + "label": font(18, scale=scale, bold=True), + "small": font(14, scale=scale), + "tiny": font(12, scale=scale), + "chip": font(21, scale=scale, bold=True), + "badge": font(15, scale=scale, bold=True), + "nic_center": font(20, scale=scale, bold=True), + "packet": font(18, scale=scale, bold=True), + "queue": font(18, scale=scale, bold=True), + "slot": font(18, scale=scale, bold=True), + } + + +def require_webp_animation_support() -> None: + if not features.check("webp") or not features.check("webp_anim"): + raise RuntimeError("Pillow must be built with animated WebP support to render packet diagrams") + + +def save_webp_animation(frames: Sequence[Image.Image], path: Path, duration_ms: int) -> None: + if not frames: + raise ValueError("cannot save animation with no frames") + require_webp_animation_support() + + rgba_frames = [frame.convert("RGBA") for frame in frames] + tmp_path = path.with_name(f".{path.name}.tmp") + if tmp_path.exists(): + tmp_path.unlink() + # libwebp may coalesce identical consecutive frames; verify outputs with + # Makefile's reproducibility check rather than decoded frame counts. + rgba_frames[0].save( + tmp_path, + format="WEBP", + save_all=True, + append_images=rgba_frames[1:], + duration=duration_ms, + loop=0, + lossless=False, + quality=WEBP_QUALITY, + method=WEBP_METHOD, + ) + tmp_path.replace(path) diff --git a/docs/images/packet_diagrams/flow_steering/flow-steering-light-poster.png b/docs/images/packet_diagrams/flow_steering/flow-steering-light-poster.png new file mode 100644 index 00000000..eba384d9 Binary files /dev/null and b/docs/images/packet_diagrams/flow_steering/flow-steering-light-poster.png differ diff --git a/docs/images/packet_diagrams/flow_steering/flow-steering-light.webp b/docs/images/packet_diagrams/flow_steering/flow-steering-light.webp new file mode 100644 index 00000000..062ccc1a Binary files /dev/null and b/docs/images/packet_diagrams/flow_steering/flow-steering-light.webp differ diff --git a/docs/images/packet_diagrams/flow_steering/flow-steering-poster.png b/docs/images/packet_diagrams/flow_steering/flow-steering-poster.png new file mode 100644 index 00000000..3b733a6a Binary files /dev/null and b/docs/images/packet_diagrams/flow_steering/flow-steering-poster.png differ diff --git a/docs/images/packet_diagrams/flow_steering/flow-steering.webp b/docs/images/packet_diagrams/flow_steering/flow-steering.webp new file mode 100644 index 00000000..b8e82bad Binary files /dev/null and b/docs/images/packet_diagrams/flow_steering/flow-steering.webp differ diff --git a/docs/images/packet_diagrams/flow_steering_animation.py b/docs/images/packet_diagrams/flow_steering_animation.py new file mode 100644 index 00000000..8def8884 --- /dev/null +++ b/docs/images/packet_diagrams/flow_steering_animation.py @@ -0,0 +1,867 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path + +from PIL import Image, ImageDraw + +from anim_common import font_scheme, save_webp_animation +from incoming_wire import center_wire_y, draw_rx_wire, draw_rx_wire_label, path_point_and_tangent, wire_t_end + + +ROOT = Path(__file__).resolve().parent +OUTPUT_DIR = ROOT / "flow_steering" + +WIDTH = 1180 +HEIGHT = 660 +SCALE = 2 +DURATION_MS = 40 + +QUEUE_TRAVEL_FRAMES = 76 +ARRIVAL_GAP = 14 + +LAYOUT_SHIFT_Y = 0 + +NIC_RECT = (300, 250, 500, 410) +HOST_RECT = (750, 86, 1120, 282) +GPU_RECT = (750, 360, 1120, 620) +HOST_ROW = (774, 170, 1096, 238) +GPU_QUEUE_AREA = (774, 430, 1096, 596) + +WIRE_Y = 330 +KERNEL_WIDTH = 130 +HOST_SLOT_COUNT = 3 +QUEUE_COUNT = 3 +QUEUE_SLOT_COUNT = 2 +WIRE_PACKET_WIDTH = 86 +WIRE_PACKET_HEIGHT = 38 +HOST_PACKET_HEIGHT = 38 +QUEUE_PACKET_HEIGHT = 32 +FLOW_RULES = { + 4096: 0, + 4097: 1, + 4098: 2, +} + +TRANSPARENT = (0, 0, 0, 0) + +ACCENTS = { + "nvidia": "#76b900", + "host": "#9b8cff", + "gpu": "#59d4ff", + "kernel": "#3b82f6", + "host_pkt": "#9b8cff", + "gpu_pkt": "#59d4ff", + "queue_0": "#59d4ff", + "queue_1": "#78e08f", + "queue_2": "#ffcf5a", + "ink": "#07111f", +} + +THEMES: dict[str, dict[str, str | None]] = { + "default": { + **ACCENTS, + "bg": None, + "panel": "#0d1b2e", + "panel_2": "#101f34", + "line": "#35516f", + "text": "#f6fbff", + "muted": "#afc0cf", + "wire": "#6d7e92", + "shadow": "#03101d", + "row_bg": "#061220", + "arrow": "#ffffff", + "stroke": "#ffffff", + "bar_glow": "#ffffff", + "kernel_fill": "#0f2744", + "bypass": "#64748b", + "match_ok": "#76b900", + "match_no": "#ef4444", + }, + "light": { + **ACCENTS, + "bg": "#ffffff", + "panel": "#f5f5f5", + "panel_2": "#eeeeee", + "line": "#1a1a1a", + "text": "#1a1a1a", + "muted": "#404040", + "wire": "#404040", + "shadow": "#bdbdbd", + "row_bg": "#e8e8e8", + "arrow": "#1a1a1a", + "stroke": "#1a1a1a", + "bar_glow": "#cccccc", + "kernel_fill": "#f0f4ff", + "bypass": "#64748b", + "match_ok": "#15803d", + "match_no": "#dc2626", + }, +} + +THEME_OUTPUT_SUFFIX = { + "default": "", + "light": "-light", +} + +COLORS = THEMES["default"] + + +@dataclass(frozen=True) +class PacketSpec: + num: int + udp_port: int + queue_id: int | None + start: int + decide: int + end: int + + @property + def matched(self) -> bool: + return self.queue_id is not None + + +def shift_y(y: float) -> float: + return y + LAYOUT_SHIFT_Y + + +def shift_rect(rect: tuple[float, float, float, float]) -> tuple[float, float, float, float]: + x1, y1, x2, y2 = rect + return x1, y1 + LAYOUT_SHIFT_Y, x2, y2 + LAYOUT_SHIFT_Y + + +NIC = shift_rect(NIC_RECT) +HOST = shift_rect(HOST_RECT) +GPU = shift_rect(GPU_RECT) +HOST_ROW_R = shift_rect(HOST_ROW) +GPU_QUEUE_AREA_R = shift_rect(GPU_QUEUE_AREA) +WIRE = shift_y(WIRE_Y) +KERNEL_CX = NIC_RECT[2] + 102 +KERNEL = shift_rect((KERNEL_CX - KERNEL_WIDTH / 2, 22, KERNEL_CX + KERNEL_WIDTH / 2, 92)) +NIC_CX = (NIC[0] + NIC[2]) / 2 + + +def rgb(hex_color: str) -> tuple[int, int, int]: + hex_color = hex_color.strip("#") + return tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4)) + + +def rgba(hex_color: str, alpha: int = 255) -> tuple[int, int, int, int]: + r, g, b = rgb(hex_color) + return r, g, b, alpha + + +FONTS = font_scheme(SCALE) + + +def s(value: float) -> int: + return int(round(value * SCALE)) + + +def pt(point: tuple[float, float]) -> tuple[int, int]: + return s(point[0]), s(point[1]) + + +def box(rect: tuple[float, float, float, float]) -> tuple[int, int, int, int]: + return tuple(s(v) for v in rect) + + +def text_size(draw: ImageDraw.ImageDraw, text: str, face: ImageFont.FreeTypeFont) -> tuple[int, int]: + bbox = draw.textbbox((0, 0), text, font=face) + return bbox[2] - bbox[0], bbox[3] - bbox[1] + + +def draw_text( + draw: ImageDraw.ImageDraw, + xy: tuple[float, float], + text: str, + face: ImageFont.FreeTypeFont, + fill: str | tuple[int, int, int, int] = COLORS["text"], + anchor: str = "la", +) -> None: + draw.text(pt(xy), text, font=face, fill=fill, anchor=anchor) + + +def centered_text( + draw: ImageDraw.ImageDraw, + rect: tuple[float, float, float, float], + text: str, + face: ImageFont.FreeTypeFont, + fill: str | tuple[int, int, int, int] = COLORS["text"], +) -> None: + x1, y1, x2, y2 = rect + draw.text((s((x1 + x2) / 2), s((y1 + y2) / 2)), text, font=face, fill=fill, anchor="mm") + + +def centered_multiline_text( + draw: ImageDraw.ImageDraw, + rect: tuple[float, float, float, float], + text: str, + face: ImageFont.FreeTypeFont, + fill: str | tuple[int, int, int, int] = COLORS["text"], +) -> None: + x1, y1, x2, y2 = rect + draw.multiline_text( + (s((x1 + x2) / 2), s((y1 + y2) / 2)), + text, + font=face, + fill=fill, + anchor="mm", + align="center", + spacing=s(4), + ) + + +def lerp(a: float, b: float, t: float) -> float: + return a + (b - a) * t + + +def clamp(value: float, lo: float = 0.0, hi: float = 1.0) -> float: + return max(lo, min(hi, value)) + + +def progress_linear(frame: int, start: int, end: int) -> float: + return clamp((frame - start) / max(1, end - start)) + + +def rectangular_path(waypoints: tuple[tuple[float, float], ...], steps: int = 32) -> list[tuple[float, float]]: + if len(waypoints) < 2: + return list(waypoints) + path: list[tuple[float, float]] = [] + for i in range(len(waypoints) - 1): + x0, y0 = waypoints[i] + x1, y1 = waypoints[i + 1] + for step in range(steps + 1): + if i > 0 and step == 0: + continue + t = step / steps + path.append((lerp(x0, x1, t), lerp(y0, y1, t))) + return path + + +def draw_polyline( + draw: ImageDraw.ImageDraw, + points: list[tuple[float, float]], + color: str | tuple[int, int, int, int], + width: int = 4, +) -> None: + draw.line([pt(p) for p in points], fill=color, width=s(width)) + + +def draw_dashed_polyline( + draw: ImageDraw.ImageDraw, + points: list[tuple[float, float]], + color: str | tuple[int, int, int, int], + width: int = 3, + dash: float = 10, + gap: float = 14, +) -> None: + for i in range(len(points) - 1): + x0, y0 = points[i] + x1, y1 = points[i + 1] + seg_len = math.hypot(x1 - x0, y1 - y0) + if seg_len <= 0: + continue + ux, uy = (x1 - x0) / seg_len, (y1 - y0) / seg_len + pos = 0.0 + draw_on = True + while pos < seg_len: + step = dash if draw_on else gap + end = min(seg_len, pos + step) + if draw_on: + draw.line( + (pt((x0 + ux * pos, y0 + uy * pos)), pt((x0 + ux * end, y0 + uy * end))), + fill=color, + width=s(width), + ) + pos = end + draw_on = not draw_on + + +def draw_arrow_head( + draw: ImageDraw.ImageDraw, + end: tuple[float, float], + prev: tuple[float, float], + color: str | tuple[int, int, int, int], + size: float = 12, +) -> None: + angle = math.atan2(end[1] - prev[1], end[0] - prev[0]) + wing = math.radians(28) + p1 = (end[0] - size * math.cos(angle - wing), end[1] - size * math.sin(angle - wing)) + p2 = (end[0] - size * math.cos(angle + wing), end[1] - size * math.sin(angle + wing)) + draw.polygon([pt(end), pt(p1), pt(p2)], fill=color) + + +def draw_arrow( + draw: ImageDraw.ImageDraw, + points: list[tuple[float, float]], + color: str | tuple[int, int, int, int], + width: int = 4, + arrow_size: float = 13, +) -> None: + draw_polyline(draw, points, color, width) + if len(points) >= 2: + prev_idx = max(0, len(points) - 5) + draw_arrow_head(draw, points[-1], points[prev_idx], color, arrow_size) + + +def draw_glow_line( + base: Image.Image, + points: list[tuple[float, float]], + color: str, + alpha: int, + width: int, +) -> None: + layer = Image.new("RGBA", base.size, (0, 0, 0, 0)) + d = ImageDraw.Draw(layer) + draw_polyline(d, points, rgba(color, alpha), width) + base.alpha_composite(layer) + + +def draw_path_tag( + draw: ImageDraw.ImageDraw, + xy: tuple[float, float], + text: str, + color: str, +) -> None: + tx, ty = xy + tw, th = text_size(draw, text, FONTS["tiny"]) + pad_x = 8 + pad_y = 5 + tag = ( + tx - tw / (2 * SCALE) - pad_x, + ty - th / (2 * SCALE) - pad_y, + tx + tw / (2 * SCALE) + pad_x, + ty + th / (2 * SCALE) + pad_y, + ) + draw.rounded_rectangle(box(tag), radius=s(6), fill=rgba(COLORS["panel"], 235), outline=rgba(color, 180), width=s(1)) + centered_text(draw, tag, text, FONTS["tiny"], fill=COLORS["text"]) + + +def panel_left_center(rect: tuple[float, float, float, float]) -> tuple[float, float]: + x1, y1, x2, y2 = rect + return x1, (y1 + y2) / 2 + + +def panel_top_center(rect: tuple[float, float, float, float]) -> tuple[float, float]: + x1, y1, x2, _ = rect + return (x1 + x2) / 2, y1 + + +def flow_queue_rows() -> list[tuple[float, float, float, float]]: + x1, y1, x2, y2 = GPU_QUEUE_AREA_R + gap = 8 + row_h = (y2 - y1 - gap * (QUEUE_COUNT - 1)) / QUEUE_COUNT + rows: list[tuple[float, float, float, float]] = [] + y = y1 + for _ in range(QUEUE_COUNT): + rows.append((x1, y, x2, y + row_h)) + y += row_h + gap + return rows + + +QUEUE_ROWS = flow_queue_rows() + + +def queue_center_y(queue_id: int) -> float: + row = QUEUE_ROWS[queue_id] + return (row[1] + row[3]) / 2 + + +def queue_route_path(queue_id: int) -> list[tuple[float, float]]: + qy = queue_center_y(queue_id) + tap_x = NIC[2] + 54 + (QUEUE_COUNT - 1 - queue_id) * 44 + return rectangular_path( + ( + (NIC[2], WIRE), + (tap_x, WIRE), + (tap_x, qy), + (GPU_QUEUE_AREA_R[0], qy), + ) + ) + + +def build_paths() -> dict[str, list[tuple[float, float]]]: + host_y = (HOST_ROW_R[1] + HOST_ROW_R[3]) / 2 + bus_y = (KERNEL[1] + KERNEL[3]) / 2 + kernel_in = (KERNEL[0], bus_y) + kernel_out = (KERNEL[2], bus_y) + nic_top = (NIC_CX, NIC[1]) + host_entry = panel_top_center(HOST) + + host_nic_to_kernel = rectangular_path((nic_top, (nic_top[0], bus_y), kernel_in)) + host_kernel_to_queue = rectangular_path((kernel_out, (host_entry[0], bus_y), host_entry)) + host_direct = host_nic_to_kernel + host_kernel_to_queue[1:] + + wire = rectangular_path(((48, WIRE), (NIC[0], WIRE))) + host_tail = rectangular_path( + ( + (NIC_CX, WIRE), + nic_top, + (nic_top[0], bus_y), + kernel_in, + kernel_out, + (host_entry[0], bus_y), + host_entry, + (host_entry[0], host_y), + (HOST_ROW_R[0], host_y), + ) + ) + paths = { + "wire": wire, + "host": wire + host_tail[1:], + "host_nic_to_kernel": host_nic_to_kernel, + "host_kernel_to_queue": host_kernel_to_queue, + "host_direct": host_direct, + } + for queue_id in range(QUEUE_COUNT): + route = queue_route_path(queue_id) + paths[f"queue{queue_id}_route"] = route + paths[f"queue{queue_id}"] = wire + route[1:] + return paths + + +PATHS = build_paths() + + +def path_total(points: list[tuple[float, float]]) -> float: + total = 0.0 + for i in range(1, len(points)): + x0, y0 = points[i - 1] + x1, y1 = points[i] + total += math.hypot(x1 - x0, y1 - y0) + return total + + +def make_packets() -> tuple[PacketSpec, ...]: + queue_len = max(path_total(PATHS[f"queue{queue_id}"]) for queue_id in range(QUEUE_COUNT)) + host_len = path_total(PATHS["host"]) + wire_len = path_total(PATHS["wire"]) + px_per_frame = queue_len / QUEUE_TRAVEL_FRAMES + wire_frames = max(1, int(math.ceil(wire_len / px_per_frame))) + + def travel_frames(queue_id: int | None) -> int: + length = host_len if queue_id is None else path_total(PATHS[f"queue{queue_id}"]) + return max(1, int(math.ceil(length / px_per_frame))) + + routing = ( + (1, 4096), + (2, 4097), + (3, 5000), + (4, 4098), + (5, 4096), + (6, 4097), + (7, 5001), + (8, 4098), + (9, 5002), + ) + packets: list[PacketSpec] = [] + cursor = 8 + for num, udp_port in routing: + queue_id = FLOW_RULES.get(udp_port) + duration = travel_frames(queue_id) + packets.append(PacketSpec(num, udp_port, queue_id, cursor, cursor + wire_frames, cursor + duration)) + cursor += ARRIVAL_GAP + return tuple(packets) + + +PACKETS = make_packets() +FRAMES = PACKETS[-1].end + 24 + + +def decision_packet(frame: int) -> PacketSpec | None: + for pkt in reversed(PACKETS): + if pkt.decide <= frame < pkt.decide + 18: + return pkt + return None + + +def decision_pulse(frame: int, pkt: PacketSpec) -> float: + return max(progress_linear(frame, pkt.decide, pkt.decide + 8), 1.0 - progress_linear(frame, pkt.decide + 10, pkt.decide + 18)) + + +def queued_packets(frame: int) -> tuple[list[int], list[list[int]]]: + host_nums: list[int] = [] + queue_nums: list[list[int]] = [[] for _ in range(QUEUE_COUNT)] + for pkt in PACKETS: + if frame >= pkt.end: + if pkt.queue_id is None: + host_nums.append(pkt.num) + else: + queue_nums[pkt.queue_id].append(pkt.num) + return host_nums, queue_nums + + +def draw_whole_packet( + draw: ImageDraw.ImageDraw, + cx: float, + cy: float, + width: float, + color: str, + alpha: int = 255, + *, + label: str | None = None, + height: float = HOST_PACKET_HEIGHT, +) -> None: + h = height + x1, y1 = cx - width / 2, cy - h / 2 + x2, y2 = cx + width / 2, cy + h / 2 + draw.rounded_rectangle( + box((x1, y1, x2, y2)), + radius=s(6), + fill=rgba(color, alpha), + outline=rgba(str(COLORS["stroke"]), 100 * alpha // 255), + width=s(2), + ) + if label: + centered_text(draw, (x1, y1, x2, y2), label, FONTS["packet"], fill=rgba(COLORS["ink"], alpha)) + + +def queue_color(queue_id: int) -> str: + return str(COLORS[f"queue_{queue_id}"]) + + +def packet_by_num(num: int) -> PacketSpec | None: + for pkt in PACKETS: + if pkt.num == num: + return pkt + return None + + +def packet_color(num: int) -> str: + pkt = packet_by_num(num) + if pkt is None or pkt.queue_id is None: + return str(COLORS["host_pkt"]) + return queue_color(pkt.queue_id) + + +def packet_label(num: int) -> str: + return f"PKT {num}" + + +def queue_chip_layout(rect: tuple[float, float, float, float], count: int) -> list[tuple[float, float, float]]: + x1, y1, x2, y2 = rect + if count <= 0: + return [] + pad = 8 + gap = 8 + inner_w = x2 - x1 - 2 * pad + chip_w = (inner_w - gap * (count - 1)) / count + cy = (y1 + y2) / 2 + slots: list[tuple[float, float, float]] = [] + x = x1 + pad + for _ in range(count): + slots.append((x + chip_w / 2, cy, chip_w)) + x += chip_w + gap + return slots + + +def host_slot_index(pkt: PacketSpec) -> int: + return sum(1 for other in PACKETS if other.queue_id is None and other.end <= pkt.end) - 1 + + +def queue_slot_index(pkt: PacketSpec) -> int: + if pkt.queue_id is None: + return 0 + return sum(1 for other in PACKETS if other.queue_id == pkt.queue_id and other.end <= pkt.end) - 1 + + +def host_slot_target(pkt: PacketSpec) -> tuple[float, float, float, float]: + slots = queue_chip_layout(HOST_ROW_R, HOST_SLOT_COUNT) + idx = max(0, min(HOST_SLOT_COUNT - 1, host_slot_index(pkt))) + cx, cy, chip_w = slots[idx] + return cx, cy, chip_w - 4, HOST_PACKET_HEIGHT + + +def flow_queue_slots_rect(row: tuple[float, float, float, float]) -> tuple[float, float, float, float]: + x1, y1, x2, y2 = row + return x1 + 112, y1 + 3, x2 - 8, y2 - 3 + + +def queue_slot_target(pkt: PacketSpec) -> tuple[float, float, float, float]: + queue_id = 0 if pkt.queue_id is None else pkt.queue_id + slots = queue_chip_layout(flow_queue_slots_rect(QUEUE_ROWS[queue_id]), QUEUE_SLOT_COUNT) + idx = max(0, min(QUEUE_SLOT_COUNT - 1, queue_slot_index(pkt))) + cx, cy, chip_w = slots[idx] + return cx, cy, chip_w - 4, QUEUE_PACKET_HEIGHT + + +def packet_target(pkt: PacketSpec) -> tuple[float, float, float, float]: + return host_slot_target(pkt) if pkt.queue_id is None else queue_slot_target(pkt) + + +def packet_memory_entry(pkt: PacketSpec) -> tuple[float, float]: + if pkt.queue_id is None: + return panel_top_center(HOST) + return PATHS[f"queue{pkt.queue_id}_route"][-1] + + +def packet_draw_dimensions(pkt: PacketSpec, cx: float, cy: float) -> tuple[float, float]: + target_x, target_y, target_w, target_h = packet_target(pkt) + entry_x, entry_y = packet_memory_entry(pkt) + if pkt.queue_id is None: + in_memory = HOST[0] <= cx <= HOST[2] and HOST[1] <= cy <= HOST[3] + else: + in_memory = GPU[0] <= cx <= GPU[2] and GPU[1] <= cy <= GPU[3] + if not in_memory: + return WIRE_PACKET_WIDTH, WIRE_PACKET_HEIGHT + + total = math.hypot(target_x - entry_x, target_y - entry_y) + covered = math.hypot(cx - entry_x, cy - entry_y) + t = clamp(covered / total if total else 1.0) + return lerp(WIRE_PACKET_WIDTH, target_w, t), lerp(WIRE_PACKET_HEIGHT, target_h, t) + + +def packet_path(pkt: PacketSpec) -> list[tuple[float, float]]: + target_x, target_y, _, _ = packet_target(pkt) + wire = PATHS["wire"] + if pkt.queue_id is None: + bus_y = (KERNEL[1] + KERNEL[3]) / 2 + nic_top = (NIC_CX, NIC[1]) + host_entry = packet_memory_entry(pkt) + tail = rectangular_path( + ( + (NIC_CX, WIRE), + nic_top, + (nic_top[0], bus_y), + (KERNEL[0], bus_y), + (KERNEL[2], bus_y), + (host_entry[0], bus_y), + host_entry, + (target_x, target_y), + ) + ) + return wire + tail[1:] + + queue_path = PATHS[f"queue{pkt.queue_id}"] + branch = rectangular_path((queue_path[-1], (target_x, target_y))) + return queue_path + branch[1:] + + +def draw_queue_row( + draw: ImageDraw.ImageDraw, + rect: tuple[float, float, float, float], + packet_nums: list[int], + accent: str, + slot_count: int, +) -> None: + x1, y1, x2, y2 = rect + draw.rounded_rectangle(box(rect), radius=s(8), fill=rgba(COLORS["row_bg"], 235), outline=rgba(accent, 145), width=s(2)) + slots = queue_chip_layout(rect, slot_count) + for idx, (cx, cy, chip_w) in enumerate(slots): + slot_rect = (cx - chip_w / 2, cy - HOST_PACKET_HEIGHT / 2, cx + chip_w / 2, cy + HOST_PACKET_HEIGHT / 2) + draw.rounded_rectangle( + box(slot_rect), + radius=s(7), + fill=rgba(COLORS["panel"], 165), + outline=rgba(accent, 135), + width=s(1), + ) + if idx >= len(packet_nums): + continue + num = packet_nums[idx] + color = packet_color(num) + draw_whole_packet(draw, cx, cy, chip_w - 4, color, label=packet_label(num), height=HOST_PACKET_HEIGHT) + + +def draw_flow_queue_row( + draw: ImageDraw.ImageDraw, + rect: tuple[float, float, float, float], + queue_id: int, + packet_nums: list[int], +) -> None: + x1, y1, x2, y2 = rect + accent = queue_color(queue_id) + draw.rounded_rectangle(box(rect), radius=s(8), fill=rgba(COLORS["row_bg"], 235), outline=rgba(accent, 145), width=s(2)) + label_rect = (x1 + 8, y1 + 3, x1 + 104, y2 - 3) + centered_text(draw, label_rect, f"queue {queue_id + 1}", FONTS["slot"], fill=COLORS["muted"]) + slots_rect = flow_queue_slots_rect(rect) + slots = queue_chip_layout(slots_rect, QUEUE_SLOT_COUNT) + for idx, (cx, cy, chip_w) in enumerate(slots): + slot_rect = (cx - chip_w / 2, cy - QUEUE_PACKET_HEIGHT / 2, cx + chip_w / 2, cy + QUEUE_PACKET_HEIGHT / 2) + draw.rounded_rectangle( + box(slot_rect), + radius=s(6), + fill=rgba(COLORS["panel"], 165), + outline=rgba(accent, 135), + width=s(1), + ) + if idx >= len(packet_nums): + continue + num = packet_nums[idx] + draw_whole_packet(draw, cx, cy, chip_w - 4, accent, label=packet_label(num), height=QUEUE_PACKET_HEIGHT) + + +def draw_host_route(draw: ImageDraw.ImageDraw) -> None: + color = rgba(COLORS["host_pkt"], 255) + draw_arrow(draw, PATHS["host_nic_to_kernel"], color, 3, 12) + draw_arrow(draw, PATHS["host_kernel_to_queue"], color, 3, 12) + + +def draw_kernel_bypass(draw: ImageDraw.ImageDraw) -> None: + x1, y1, x2, y2 = KERNEL + draw.rounded_rectangle(box(KERNEL), radius=s(12), fill=rgba(COLORS["kernel_fill"], 245), outline=rgba(COLORS["kernel"], 210), width=s(2)) + centered_text(draw, (x1, y1 + 8, x2, y2 - 8), "Linux kernel", FONTS["label"], fill=COLORS["text"]) + + +def draw_matched_routes(draw: ImageDraw.ImageDraw) -> None: + for queue_id in range(QUEUE_COUNT): + draw_arrow(draw, PATHS[f"queue{queue_id}_route"], rgba(str(COLORS["arrow"]), 130), 3, 12) + draw_path_tag(draw, (600, 416), "flow match -> queue", COLORS["gpu"]) + + +def nic_status(frame: int) -> tuple[str, str, str, float]: + pkt = decision_packet(frame) + deciding = pkt is not None + pulse = decision_pulse(frame, pkt) if deciding and pkt else 0.0 + if deciding and pkt: + if pkt.queue_id is not None: + return f"{packet_label(pkt.num)}\nqueue {pkt.queue_id + 1}", queue_color(pkt.queue_id), "#ffffff", pulse + return f"{packet_label(pkt.num)}\nto kernel", COLORS["host_pkt"], "#ffffff", pulse + return "flow\nsteering", COLORS["nvidia"], COLORS["ink"], pulse + + +def draw_nic_status_pill(draw: ImageDraw.ImageDraw, frame: int) -> None: + x1, y1, x2, _ = NIC + pill_text, pill_accent, pill_ink, pulse = nic_status(frame) + pill = (x1 + 12, y1 + 38, x2 - 12, y1 + 122) + draw.rounded_rectangle( + box(pill), + radius=s(8), + fill=rgba(pill_accent, 255), + outline=rgba(pill_accent, 220), + width=s(2), + ) + centered_multiline_text(draw, pill, pill_text, FONTS["nic_center"], fill=pill_ink) + + +def draw_nic_chip_base(draw: ImageDraw.ImageDraw, frame: int) -> None: + x1, y1, x2, y2 = NIC + pkt = decision_packet(frame) + deciding = pkt is not None + pulse = decision_pulse(frame, pkt) if deciding and pkt else 0.0 + accent = COLORS["nvidia"] + + glow_alpha = int(55 + 120 * pulse) if deciding else 55 + draw.rounded_rectangle(box((x1 - 10, y1 - 10, x2 + 10, y2 + 10)), radius=s(24), fill=rgba(accent, 16 + glow_alpha // 3)) + draw.rounded_rectangle(box(NIC), radius=s(18), fill=COLORS["panel"], outline=rgba(accent, 190 + int(65 * pulse)), width=s(3 + int(2 * pulse))) + for i in range(7): + y = lerp(y1 + 18, y2 - 18, i / 6) + draw.line((s(x1 - 14), s(y), s(x1), s(y)), fill=rgba(COLORS["line"], 220), width=s(3)) + draw.line((s(x2), s(y), s(x2 + 14), s(y)), fill=rgba(COLORS["line"], 220), width=s(3)) + for i in range(5): + x = lerp(x1 + 25, x2 - 25, i / 4) + draw.line((s(x), s(y1 - 12), s(x), s(y1)), fill=rgba(COLORS["line"], 220), width=s(3)) + draw.line((s(x), s(y2), s(x), s(y2 + 12)), fill=rgba(COLORS["line"], 220), width=s(3)) + centered_text(draw, (x1 + 18, y1 + 4, x2 - 18, y1 + 38), "NVIDIA NIC", FONTS["chip"], fill=COLORS["text"]) + + +def draw_device_memory(draw: ImageDraw.ImageDraw, frame: int) -> None: + host_nums, queue_nums = queued_packets(frame) + panels = ( + (HOST, "Host memory", "kernel fallback packets", COLORS["host"]), + (GPU, "GPU memory", "RX queues selected by flow rules", COLORS["gpu"]), + ) + for panel_rect, title, subtitle, accent in panels: + x1, y1, _, _ = panel_rect + draw.rounded_rectangle(box(panel_rect), radius=s(18), fill=COLORS["panel_2"], outline=rgba(accent, 170), width=s(3)) + draw_text(draw, (x1 + 22, y1 + 22), title, FONTS["label"], fill=COLORS["text"], anchor="lt") + draw_text(draw, (x1 + 22, y1 + 46), subtitle, FONTS["small"], fill=COLORS["muted"], anchor="lt") + draw_queue_row(draw, HOST_ROW_R, host_nums, COLORS["host_pkt"], HOST_SLOT_COUNT) + for queue_id, row in enumerate(QUEUE_ROWS): + draw_flow_queue_row(draw, row, queue_id, queue_nums[queue_id]) + + +def draw_wires(draw: ImageDraw.ImageDraw, frame: int) -> None: + draw_rx_wire_label(lambda xy, text: draw_text(draw, xy, text, FONTS["label"], fill=COLORS["text"]), WIRE) + draw_rx_wire(draw, frame, 48, NIC[0], WIRE, str(COLORS["wire"]), "#ffcf5a", pt, s, rgba, box) + + +def draw_route_glow(base: Image.Image, frame: int) -> None: + return + + +def draw_flowing_packets(draw: ImageDraw.ImageDraw, frame: int) -> None: + for pkt in PACKETS: + if not (pkt.start <= frame < pkt.end): + continue + path = packet_path(pkt) + color = packet_color(pkt.num) + t = progress_linear(frame, pkt.start, pkt.end) + if 0 <= t < 1: + on_wire = t <= wire_t_end(PATHS["wire"], path) + cx, cy, _ = path_point_and_tangent(path, t) + if on_wire: + cy = center_wire_y(cx, 48, WIRE, frame) + inside_nic = NIC[0] <= cx <= NIC[2] and NIC[1] <= cy <= NIC[3] + label = None if inside_nic else packet_label(pkt.num) + packet_w, packet_h = packet_draw_dimensions(pkt, cx, cy) + draw_whole_packet(draw, cx, cy, packet_w, color, label=label, height=packet_h) + + +def path_lengths(points: list[tuple[float, float]]) -> tuple[list[float], float]: + lengths = [0.0] + for i in range(1, len(points)): + x0, y0 = points[i - 1] + x1, y1 = points[i] + lengths.append(lengths[-1] + math.hypot(x1 - x0, y1 - y0)) + return lengths, lengths[-1] + + +def point_on_path(points: list[tuple[float, float]], t: float) -> tuple[float, float]: + if len(points) < 2: + return points[0] if points else (0.0, 0.0) + t = clamp(t) + lengths, total = path_lengths(points) + if total <= 0: + return points[0] + target = t * total + for i in range(1, len(points)): + if lengths[i] >= target: + seg_len = lengths[i] - lengths[i - 1] + local = 0.0 if seg_len <= 0 else (target - lengths[i - 1]) / seg_len + x0, y0 = points[i - 1] + x1, y1 = points[i] + return lerp(x0, x1, local), lerp(y0, y1, local) + return points[-1] + + +def render_frame(frame: int) -> Image.Image: + bg = COLORS.get("bg") + img = Image.new("RGBA", (WIDTH * SCALE, HEIGHT * SCALE), rgba(str(bg), 255) if bg else TRANSPARENT) + draw = ImageDraw.Draw(img) + draw_wires(draw, frame) + draw_device_memory(draw, frame) + draw_host_route(draw) + draw_kernel_bypass(draw) + draw_nic_chip_base(draw, frame) + draw_matched_routes(draw) + draw_route_glow(img, frame) + draw_flowing_packets(draw, frame) + draw_nic_status_pill(draw, frame) + return img.resize((WIDTH, HEIGHT), Image.Resampling.LANCZOS) + + +def render_theme(theme: str) -> None: + global COLORS + COLORS = THEMES[theme] + suffix = THEME_OUTPUT_SUFFIX[theme] + animation_path = OUTPUT_DIR / f"flow-steering{suffix}.webp" + poster_path = OUTPUT_DIR / f"flow-steering{suffix}-poster.png" + + frames = [render_frame(i) for i in range(FRAMES)] + save_webp_animation(frames, animation_path, DURATION_MS) + render_frame(min(FRAMES - 1, PACKETS[-1].end + 10)).save(poster_path, optimize=True) + print(f"Wrote {animation_path}") + print(f"Wrote {poster_path}") + + +def main() -> None: + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + for theme in THEMES: + render_theme(theme) + + +if __name__ == "__main__": + main() diff --git a/docs/images/packet_diagrams/hds/header-data-split-light-poster.png b/docs/images/packet_diagrams/hds/header-data-split-light-poster.png new file mode 100644 index 00000000..6bba66ab Binary files /dev/null and b/docs/images/packet_diagrams/hds/header-data-split-light-poster.png differ diff --git a/docs/images/packet_diagrams/hds/header-data-split-light.webp b/docs/images/packet_diagrams/hds/header-data-split-light.webp new file mode 100644 index 00000000..0244a7f2 Binary files /dev/null and b/docs/images/packet_diagrams/hds/header-data-split-light.webp differ diff --git a/docs/images/packet_diagrams/hds/header-data-split-poster.png b/docs/images/packet_diagrams/hds/header-data-split-poster.png new file mode 100644 index 00000000..ffd31d8b Binary files /dev/null and b/docs/images/packet_diagrams/hds/header-data-split-poster.png differ diff --git a/docs/images/packet_diagrams/hds/header-data-split.webp b/docs/images/packet_diagrams/hds/header-data-split.webp new file mode 100644 index 00000000..b820999d Binary files /dev/null and b/docs/images/packet_diagrams/hds/header-data-split.webp differ diff --git a/docs/images/packet_diagrams/hds_animation.py b/docs/images/packet_diagrams/hds_animation.py new file mode 100644 index 00000000..4ffa73e9 --- /dev/null +++ b/docs/images/packet_diagrams/hds_animation.py @@ -0,0 +1,818 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path + +from PIL import Image, ImageDraw + +from anim_common import font_scheme, save_webp_animation +from incoming_wire import draw_rx_wire, draw_rx_wire_label + + +ROOT = Path(__file__).resolve().parent +OUTPUT_DIR = ROOT / "hds" + +WIDTH = 1180 +HEIGHT = 660 +SCALE = 2 +DURATION_MS = 40 + +TOTAL_BYTES = 1064 +HEADER_BYTES = 64 +PAYLOAD_BYTES = 1000 + +PACKET_COUNT = 5 +WIRE_Y = 330 +ARRIVAL_GAP = 42 +PIXELS_PER_FRAME = 12.0 +PACKET_BAR_WIDTH = 132 +PACKET_BAR_HEIGHT = 34 +HEADER_SEGMENT_WIDTH = 68 +PAYLOAD_SEGMENT_WIDTH = 190 +PAYLOAD_SEGMENT_HEIGHT = 38 +HEADER_SEGMENT_HEIGHT = PAYLOAD_SEGMENT_HEIGHT +SLOT_GAP = 8 + +LAYOUT_SHIFT_Y = 0 + + +def dy(y: float) -> float: + return y + LAYOUT_SHIFT_Y + + +def shift_rect(rect: tuple[float, float, float, float]) -> tuple[float, float, float, float]: + x1, y1, x2, y2 = rect + return x1, y1 + LAYOUT_SHIFT_Y, x2, y2 + LAYOUT_SHIFT_Y + + +NIC_RECT = shift_rect((300, 250, 500, 410)) +HOST_RECT = shift_rect((744, 74, 1150, 234)) +GPU_RECT = shift_rect((770, 324, 1110, 610)) +KERNEL_WIDTH = 130 +KERNEL_CX = NIC_RECT[2] + 102 +HEADER_SLOT_AREA = shift_rect((756, 128, 1140, 180)) +PAYLOAD_SLOT_AREA = shift_rect((832, 370, 1048, 604)) +KERNEL_BOX = shift_rect((KERNEL_CX - KERNEL_WIDTH / 2, 22, KERNEL_CX + KERNEL_WIDTH / 2, 92)) + +TRANSPARENT = (0, 0, 0, 0) + +ACCENTS = { + "nvidia": "#76b900", + "header": "#ffcf5a", + "payload": "#78e08f", + "host": "#9b8cff", + "gpu": "#76b900", + "kernel": "#64748b", + "ink": "#07111f", +} + +THEMES: dict[str, dict[str, str | None]] = { + "default": { + **ACCENTS, + "bg": None, + "panel": "#0d1b2e", + "panel_2": "#101f34", + "line": "#35516f", + "line_soft": "#20384f", + "text": "#f6fbff", + "muted": "#afc0cf", + "wire": "#6d7e92", + "white": "#ffffff", + "shadow": "#03101d", + "row_bg": "#061220", + "kernel_fill": "#0f2744", + "arrow": "#ffffff", + "stroke": "#ffffff", + "bar_glow": "#ffffff", + "bypass": "#afc0cf", + }, + "light": { + **ACCENTS, + "bg": "#ffffff", + "panel": "#f5f5f5", + "panel_2": "#eeeeee", + "line": "#1a1a1a", + "line_soft": "#cccccc", + "text": "#1a1a1a", + "muted": "#404040", + "wire": "#404040", + "white": "#ffffff", + "shadow": "#bdbdbd", + "row_bg": "#e8e8e8", + "kernel_fill": "#f0f4ff", + "arrow": "#1a1a1a", + "stroke": "#1a1a1a", + "bar_glow": "#cccccc", + "bypass": "#404040", + }, +} + +THEME_OUTPUT_SUFFIX = { + "default": "", + "light": "-light", +} + +COLORS = THEMES["default"] + + +@dataclass(frozen=True) +class PacketSpec: + num: int + start: int + nic_arrive: int + header_start: int + header_end: int + payload_start: int + payload_end: int + + +def make_packets() -> tuple[PacketSpec, ...]: + packets: list[PacketSpec] = [] + cursor = 8 + for num in range(PACKET_COUNT): + header_start = cursor + frames_for_path(segment_wire_path("header")) + nic_arrive = cursor + frames_for_path(rectangular_path((header_wire_start(), (NIC_RECT[0], dy(WIRE_Y))))) + payload_start = header_start + header_end = header_start + frames_for_path(segment_post_split_path("header", num)) + payload_end = payload_start + frames_for_path(segment_post_split_path("payload", num)) + packets.append(PacketSpec(num, cursor, nic_arrive, header_start, header_end, payload_start, payload_end)) + cursor += ARRIVAL_GAP + return tuple(packets) + + +def rgb(hex_color: str) -> tuple[int, int, int]: + hex_color = hex_color.strip("#") + return tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4)) + + +def rgba(hex_color: str, alpha: int = 255) -> tuple[int, int, int, int]: + r, g, b = rgb(hex_color) + return r, g, b, alpha + + +FONTS = font_scheme(SCALE) + + +def s(value: float) -> int: + return int(round(value * SCALE)) + + +def pt(point: tuple[float, float]) -> tuple[int, int]: + return s(point[0]), s(point[1]) + + +def box(rect: tuple[float, float, float, float]) -> tuple[int, int, int, int]: + return tuple(s(v) for v in rect) + + +def text_size(draw: ImageDraw.ImageDraw, text: str, face: ImageFont.FreeTypeFont) -> tuple[int, int]: + bbox = draw.textbbox((0, 0), text, font=face) + return bbox[2] - bbox[0], bbox[3] - bbox[1] + + +def draw_text( + draw: ImageDraw.ImageDraw, + xy: tuple[float, float], + text: str, + face: ImageFont.FreeTypeFont, + fill: str | tuple[int, int, int, int] = COLORS["text"], + anchor: str = "la", +) -> None: + draw.text(pt(xy), text, font=face, fill=fill, anchor=anchor) + + +def centered_text( + draw: ImageDraw.ImageDraw, + rect: tuple[float, float, float, float], + text: str, + face: ImageFont.FreeTypeFont, + fill: str | tuple[int, int, int, int] = COLORS["text"], +) -> None: + x1, y1, x2, y2 = rect + draw.text((s((x1 + x2) / 2), s((y1 + y2) / 2)), text, font=face, fill=fill, anchor="mm") + + +def centered_multiline_text( + draw: ImageDraw.ImageDraw, + rect: tuple[float, float, float, float], + text: str, + face: ImageFont.FreeTypeFont, + fill: str | tuple[int, int, int, int] = COLORS["text"], +) -> None: + x1, y1, x2, y2 = rect + draw.multiline_text( + (s((x1 + x2) / 2), s((y1 + y2) / 2)), + text, + font=face, + fill=fill, + anchor="mm", + align="center", + spacing=s(4), + ) + + +def lerp(a: float, b: float, t: float) -> float: + return a + (b - a) * t + + +def clamp(value: float, lo: float = 0.0, hi: float = 1.0) -> float: + return max(lo, min(hi, value)) + + +def smoothstep(t: float) -> float: + t = clamp(t) + return t * t * (3 - 2 * t) + + +def progress(frame: int, start: int, end: int) -> float: + return smoothstep((frame - start) / max(1, end - start)) + + +def progress_linear(frame: int, start: int, end: int) -> float: + return clamp((frame - start) / max(1, end - start)) + + +def bezier_point( + p0: tuple[float, float], + p1: tuple[float, float], + p2: tuple[float, float], + p3: tuple[float, float], + t: float, +) -> tuple[float, float]: + u = 1 - t + x = u**3 * p0[0] + 3 * u**2 * t * p1[0] + 3 * u * t**2 * p2[0] + t**3 * p3[0] + y = u**3 * p0[1] + 3 * u**2 * t * p1[1] + 3 * u * t**2 * p2[1] + t**3 * p3[1] + return x, y + + +def bezier_path(points: tuple[tuple[float, float], ...], steps: int = 80) -> list[tuple[float, float]]: + p0, p1, p2, p3 = points + return [bezier_point(p0, p1, p2, p3, i / steps) for i in range(steps + 1)] + + +def rectangular_path(waypoints: tuple[tuple[float, float], ...], steps: int = 14) -> list[tuple[float, float]]: + if len(waypoints) < 2: + return list(waypoints) + path: list[tuple[float, float]] = [] + for i in range(len(waypoints) - 1): + x0, y0 = waypoints[i] + x1, y1 = waypoints[i + 1] + for step in range(steps + 1): + if i > 0 and step == 0: + continue + t = step / steps + path.append((lerp(x0, x1, t), lerp(y0, y1, t))) + return path + + +def draw_polyline( + draw: ImageDraw.ImageDraw, + points: list[tuple[float, float]], + color: str | tuple[int, int, int, int], + width: int = 4, +) -> None: + draw.line([pt(p) for p in points], fill=color, width=s(width)) + + +def draw_dotted_segment( + draw: ImageDraw.ImageDraw, + start: tuple[float, float], + end: tuple[float, float], + color: str | tuple[int, int, int, int], + *, + width: int = 2, + dash: float = 7, + gap: float = 5, +) -> None: + x0, y0 = start + x1, y1 = end + length = math.hypot(x1 - x0, y1 - y0) + if length == 0: + return + dx = (x1 - x0) / length + dy = (y1 - y0) / length + pos = 0.0 + while pos < length: + seg_end = min(pos + dash, length) + draw.line( + (pt((x0 + dx * pos, y0 + dy * pos)), pt((x0 + dx * seg_end, y0 + dy * seg_end))), + fill=color, + width=s(width), + ) + pos += dash + gap + + +def draw_dotted_polyline( + draw: ImageDraw.ImageDraw, + points: list[tuple[float, float]], + color: str | tuple[int, int, int, int], + *, + width: int = 2, + dash: float = 8, + gap: float = 6, +) -> None: + if len(points) < 2: + return + drawing = True + remaining = dash + for i in range(len(points) - 1): + x0, y0 = points[i] + x1, y1 = points[i + 1] + seg_len = math.hypot(x1 - x0, y1 - y0) + if seg_len == 0: + continue + dx = (x1 - x0) / seg_len + dy = (y1 - y0) / seg_len + pos = 0.0 + while pos < seg_len: + step = min(remaining, seg_len - pos) + if drawing: + draw.line( + ( + pt((x0 + dx * pos, y0 + dy * pos)), + pt((x0 + dx * (pos + step), y0 + dy * (pos + step))), + ), + fill=color, + width=s(width), + ) + pos += step + remaining -= step + if remaining <= 0: + drawing = not drawing + remaining = dash if drawing else gap + + +def draw_arrow_head( + draw: ImageDraw.ImageDraw, + end: tuple[float, float], + prev: tuple[float, float], + color: str | tuple[int, int, int, int], + size: float = 12, +) -> None: + angle = math.atan2(end[1] - prev[1], end[0] - prev[0]) + wing = math.radians(28) + p1 = (end[0] - size * math.cos(angle - wing), end[1] - size * math.sin(angle - wing)) + p2 = (end[0] - size * math.cos(angle + wing), end[1] - size * math.sin(angle + wing)) + draw.polygon([pt(end), pt(p1), pt(p2)], fill=color) + + +def draw_arrow( + draw: ImageDraw.ImageDraw, + points: list[tuple[float, float]], + color: str | tuple[int, int, int, int], + width: int = 4, + arrow_size: float = 13, +) -> None: + draw_polyline(draw, points, color, width) + if len(points) >= 2: + prev_idx = max(0, len(points) - 5) + draw_arrow_head(draw, points[-1], points[prev_idx], color, arrow_size) + + +def draw_glow_line( + base: Image.Image, + points: list[tuple[float, float]], + color: str, + alpha: int, + width: int, +) -> None: + layer = Image.new("RGBA", base.size, (0, 0, 0, 0)) + d = ImageDraw.Draw(layer) + draw_polyline(d, points, rgba(color, alpha), width) + base.alpha_composite(layer) + + +def path_lengths(points: list[tuple[float, float]]) -> tuple[list[float], float]: + lengths = [0.0] + for i in range(1, len(points)): + x0, y0 = points[i - 1] + x1, y1 = points[i] + lengths.append(lengths[-1] + math.hypot(x1 - x0, y1 - y0)) + return lengths, lengths[-1] + + +def path_total(points: list[tuple[float, float]]) -> float: + return path_lengths(points)[1] + + +def frames_for_path(points: list[tuple[float, float]]) -> int: + return max(1, int(math.ceil(path_total(points) / PIXELS_PER_FRAME))) + + +def point_on_path(points: list[tuple[float, float]], t: float) -> tuple[float, float]: + if len(points) < 2: + return points[0] if points else (0.0, 0.0) + t = clamp(t) + lengths, total = path_lengths(points) + if total <= 0: + return points[0] + target = t * total + for i in range(1, len(points)): + if lengths[i] >= target: + seg_len = lengths[i] - lengths[i - 1] + local = 0.0 if seg_len <= 0 else (target - lengths[i - 1]) / seg_len + x0, y0 = points[i - 1] + x1, y1 = points[i] + return lerp(x0, x1, local), lerp(y0, y1, local) + return points[-1] + + +def rect_at_center(cx: float, cy: float, w: float, h: float) -> tuple[float, float, float, float]: + return cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2 + + +def draw_packet_bar( + draw: ImageDraw.ImageDraw, + x: float, + y: float, + w: float, + h: float, + alpha: int = 255, + label: bool = True, + packet_num: int | None = None, +) -> None: + draw.rounded_rectangle(box((x - 3, y - 3, x + w + 3, y + h + 3)), radius=s(11), fill=rgba(str(COLORS["bar_glow"]), 22 * alpha // 255)) + header_w = min(HEADER_SEGMENT_WIDTH, w * 0.45) + outer = (x, y, x + w, y + h) + payload_rect = (x + header_w, y, x + w, y + h) + header_rect = (x, y, x + header_w, y + h) + draw.rounded_rectangle(box(outer), radius=s(9), fill=rgba(COLORS["payload"], alpha), outline=rgba(COLORS["ink"], alpha), width=s(2)) + draw.rectangle(box(header_rect), fill=rgba(COLORS["header"], alpha)) + draw.line((s(x + header_w), s(y + 2), s(x + header_w), s(y + h - 2)), fill=rgba(COLORS["ink"], alpha), width=s(1)) + if label: + centered_text(draw, payload_rect, "Payload", FONTS["tiny"], fill=rgba(COLORS["ink"], alpha)) + if packet_num is not None: + draw_text(draw, (x + w / 2, y - 5), f"P{packet_num}", FONTS["tiny"], fill=COLORS["text"], anchor="mb") + + +def segment_slot_rect( + area: tuple[float, float, float, float], + idx: int, + width: float, + height: float, + orientation: str, +) -> tuple[float, float, float, float]: + x1, y1, x2, y2 = area + if orientation == "horizontal": + total_w = PACKET_COUNT * width + (PACKET_COUNT - 1) * SLOT_GAP + left = x1 + (x2 - x1 - total_w) / 2 + x = left + idx * (width + SLOT_GAP) + cy = (y1 + y2) / 2 + return rect_at_center(x + width / 2, cy, width, height) + + total_h = PACKET_COUNT * height + (PACKET_COUNT - 1) * SLOT_GAP + top = y1 + (y2 - y1 - total_h) / 2 + y = top + idx * (height + SLOT_GAP) + cx = (x1 + x2) / 2 + return rect_at_center(cx, y + height / 2, width, height) + + +def rect_center(rect: tuple[float, float, float, float]) -> tuple[float, float]: + x1, y1, x2, y2 = rect + return (x1 + x2) / 2, (y1 + y2) / 2 + + +def header_slot_rect(idx: int) -> tuple[float, float, float, float]: + return segment_slot_rect(HEADER_SLOT_AREA, idx, HEADER_SEGMENT_WIDTH, HEADER_SEGMENT_HEIGHT, "horizontal") + + +def payload_slot_rect(idx: int) -> tuple[float, float, float, float]: + return segment_slot_rect(PAYLOAD_SLOT_AREA, idx, PAYLOAD_SEGMENT_WIDTH, PAYLOAD_SEGMENT_HEIGHT, "vertical") + + +def draw_segment_slots( + draw: ImageDraw.ImageDraw, + area: tuple[float, float, float, float], + color: str, + completed_slots: set[int], + width: float, + height: float, + slot_label_prefix: str, + orientation: str, +) -> None: + draw.rounded_rectangle(box(area), radius=s(8), fill=rgba(COLORS["row_bg"], 235), outline=rgba(color, 145), width=s(2)) + + for idx in range(PACKET_COUNT): + slot = segment_slot_rect(area, idx, width, height, orientation) + label = f"{slot_label_prefix} {idx + 1}" + if idx in completed_slots: + draw_packet_segment(draw, rect_center(slot), width, height, color, label) + continue + draw.rounded_rectangle( + box(slot), + radius=s(7), + fill=rgba(COLORS["panel"], 180), + outline=rgba(COLORS["line"], 170), + width=s(1), + ) + centered_text(draw, slot, label, FONTS["slot"], fill=COLORS["muted"]) + + +def draw_packet_segment( + draw: ImageDraw.ImageDraw, + center: tuple[float, float], + width: float, + height: float, + color: str, + label: str, + alpha: int = 255, +) -> None: + cx, cy = center + rect = rect_at_center(cx, cy, width, height) + draw.rounded_rectangle(box(rect), radius=s(9), fill=rgba(color, alpha), outline=rgba(str(COLORS["stroke"]), 70 * alpha // 255), width=s(1)) + draw_text(draw, (cx, cy), label, FONTS["packet"], fill=rgba(COLORS["ink"], alpha), anchor="mm") + + +def draw_path_tag( + draw: ImageDraw.ImageDraw, + xy: tuple[float, float], + text: str, + color: str, + *, + anchor: str = "mm", +) -> None: + tx, ty = xy + tw, th = text_size(draw, text, FONTS["tiny"]) + pad_x = 8 + pad_y = 5 + x1 = tx - tw / (2 * SCALE) - pad_x if anchor == "mm" else tx + y1 = ty - th / (2 * SCALE) - pad_y if anchor == "mm" else ty + tag = (x1, y1, x1 + tw / SCALE + pad_x * 2, y1 + th / SCALE + pad_y * 2) + draw.rounded_rectangle(box(tag), radius=s(6), fill=rgba(COLORS["panel"], 235), outline=rgba(color, 180), width=s(1)) + centered_text(draw, tag, text, FONTS["tiny"], fill=COLORS["text"]) + + +def draw_nic_chip_base( + draw: ImageDraw.ImageDraw, + rect: tuple[float, float, float, float], + accent: str, + pulse: float = 0.0, +) -> None: + x1, y1, x2, y2 = rect + glow_alpha = int(55 + 95 * pulse) + draw.rounded_rectangle(box((x1 - 10, y1 - 10, x2 + 10, y2 + 10)), radius=s(24), fill=rgba(accent, 16 + glow_alpha // 4)) + draw.rounded_rectangle(box(rect), radius=s(18), fill=COLORS["panel"], outline=rgba(accent, 190), width=s(3)) + for i in range(7): + y = lerp(y1 + 18, y2 - 18, i / 6) + draw.line((s(x1 - 14), s(y), s(x1), s(y)), fill=rgba(COLORS["line"], 220), width=s(3)) + draw.line((s(x2), s(y), s(x2 + 14), s(y)), fill=rgba(COLORS["line"], 220), width=s(3)) + for i in range(5): + x = lerp(x1 + 25, x2 - 25, i / 4) + draw.line((s(x), s(y1 - 12), s(x), s(y1)), fill=rgba(COLORS["line"], 220), width=s(3)) + draw.line((s(x), s(y2), s(x), s(y2 + 12)), fill=rgba(COLORS["line"], 220), width=s(3)) + centered_text(draw, (x1 + 18, y1 + 4, x2 - 18, y1 + 38), "NVIDIA NIC", FONTS["chip"], fill=COLORS["text"]) + + +def draw_nic_engine_box(draw: ImageDraw.ImageDraw) -> None: + x1, y1, x2, _ = NIC_RECT + engine = (x1 + 12, y1 + 38, x2 - 12, y1 + 122) + draw.rounded_rectangle(box(engine), radius=s(9), fill=rgba(COLORS["nvidia"], 245), outline=rgba(COLORS["nvidia"], 220), width=s(2)) + centered_multiline_text(draw, engine, "buffer splitting\nengine", FONTS["nic_center"], fill=COLORS["ink"]) + + +def draw_wire(draw: ImageDraw.ImageDraw, frame: int) -> None: + y = dy(WIRE_Y) + draw_rx_wire_label(lambda xy, text: draw_text(draw, xy, text, FONTS["label"], fill=COLORS["text"]), y) + draw_rx_wire(draw, frame, 48, NIC_RECT[0], y, COLORS["wire"], COLORS["header"], pt, s, rgba, box) + + +def draw_device_memory(draw: ImageDraw.ImageDraw, frame: int) -> None: + panels = ( + (HOST_RECT, "Host memory", COLORS["host"]), + (GPU_RECT, "GPU memory", COLORS["gpu"]), + ) + for rect, title, accent in panels: + x1, y1, x2, y2 = rect + draw.rounded_rectangle(box(rect), radius=s(18), fill=COLORS["panel_2"], outline=rgba(accent, 170), width=s(3)) + draw_text(draw, (x1 + 22, y1 + 24), title, FONTS["label"], fill=COLORS["text"], anchor="lt") + + header_slots = {pkt.num for pkt in PACKETS if frame >= pkt.header_end} + payload_slots = {pkt.num for pkt in PACKETS if frame >= pkt.payload_end} + draw_segment_slots(draw, HEADER_SLOT_AREA, COLORS["header"], header_slots, HEADER_SEGMENT_WIDTH, HEADER_SEGMENT_HEIGHT, "hdr", "horizontal") + draw_segment_slots(draw, PAYLOAD_SLOT_AREA, COLORS["payload"], payload_slots, PAYLOAD_SEGMENT_WIDTH, PAYLOAD_SEGMENT_HEIGHT, "payload", "vertical") + + +def draw_static_background(draw: ImageDraw.ImageDraw) -> None: + pass + + +def draw_kernel_bypass(draw: ImageDraw.ImageDraw) -> None: + x1, y1, x2, y2 = KERNEL_BOX + draw.rounded_rectangle(box(KERNEL_BOX), radius=s(12), fill=rgba(COLORS["kernel_fill"], 150), outline=rgba(COLORS["kernel"], 120), width=s(2)) + draw_text(draw, ((x1 + x2) / 2, y1 + 23), "Linux kernel", FONTS["small"], fill=COLORS["muted"], anchor="mm") + draw_text(draw, ((x1 + x2) / 2, y1 + 47), "bypassed", FONTS["small"], fill=COLORS["muted"], anchor="mm") + + +def split_origin() -> tuple[float, float]: + return NIC_RECT[2], dy(WIRE_Y) + + +def segment_pair_gap() -> float: + return (HEADER_SEGMENT_WIDTH + PAYLOAD_SEGMENT_WIDTH) / 2 + + +def payload_wire_start() -> tuple[float, float]: + return 58 + PAYLOAD_SEGMENT_WIDTH / 2, dy(WIRE_Y) + + +def header_wire_start() -> tuple[float, float]: + x, y = payload_wire_start() + return x + segment_pair_gap(), y + + +def payload_split_origin() -> tuple[float, float]: + x, y = split_origin() + return x - segment_pair_gap(), y + + +def header_split_origin() -> tuple[float, float]: + return split_origin() + + +def segment_split_origin(kind: str) -> tuple[float, float]: + return header_split_origin() if kind == "header" else payload_split_origin() + + +def segment_hub(kind: str) -> tuple[float, float]: + if kind == "header": + return HEADER_SLOT_AREA[0] - 20, (HEADER_SLOT_AREA[1] + HEADER_SLOT_AREA[3]) / 2 + return payload_boundary_point() + + +def payload_boundary_point() -> tuple[float, float]: + return PAYLOAD_SLOT_AREA[0], (PAYLOAD_SLOT_AREA[1] + PAYLOAD_SLOT_AREA[3]) / 2 + + +def segment_trunk_path(kind: str) -> list[tuple[float, float]]: + if kind == "payload": + return payload_trunk_path() + start_x, start_y = split_origin() + hub_x, hub_y = segment_hub(kind) + bend_x = start_x + 88 + return rectangular_path( + ( + (start_x, start_y), + (bend_x, start_y), + (bend_x, hub_y), + (hub_x, hub_y), + ) + ) + + +def segment_wire_path(kind: str) -> list[tuple[float, float]]: + start = header_wire_start() if kind == "header" else payload_wire_start() + end = segment_split_origin(kind) + return rectangular_path((start, end)) + + +def segment_slot_center(kind: str, slot_idx: int) -> tuple[float, float]: + return rect_center(header_slot_rect(slot_idx) if kind == "header" else payload_slot_rect(slot_idx)) + + +def header_post_split_path(slot_idx: int) -> list[tuple[float, float]]: + start_x, start_y = header_split_origin() + slot_x, slot_y = segment_slot_center("header", slot_idx) + bend_x = start_x + 88 + return rectangular_path( + ( + (start_x, start_y), + (bend_x, start_y), + (bend_x, slot_y), + (slot_x, slot_y), + ) + ) + + +def payload_trunk_path() -> list[tuple[float, float]]: + start_x, start_y = payload_split_origin() + boundary_x, boundary_y = payload_boundary_point() + bend_x = start_x + segment_pair_gap() + 88 + return rectangular_path( + ( + (start_x, start_y), + (start_x + segment_pair_gap(), start_y), + (bend_x, start_y), + (bend_x, boundary_y), + (boundary_x, boundary_y), + ) + ) + + +def payload_post_split_path(slot_idx: int) -> list[tuple[float, float]]: + slot_x, slot_y = segment_slot_center("payload", slot_idx) + trunk = payload_trunk_path() + diagonal = rectangular_path((trunk[-1], (slot_x, slot_y))) + return trunk + diagonal[1:] + + +def segment_path(kind: str, slot_idx: int) -> list[tuple[float, float]]: + wire = segment_wire_path(kind) + if kind == "header": + post_split = header_post_split_path(slot_idx) + return wire + post_split[1:] + if kind == "payload": + post_split = payload_post_split_path(slot_idx) + return wire + post_split[1:] + trunk = segment_trunk_path(kind) + branch = rectangular_path((trunk[-1], segment_slot_center(kind, slot_idx))) + return wire + trunk[1:] + branch[1:] + + +def segment_post_split_path(kind: str, slot_idx: int) -> list[tuple[float, float]]: + if kind == "header": + return header_post_split_path(slot_idx) + if kind == "payload": + return payload_post_split_path(slot_idx) + trunk = segment_trunk_path(kind) + branch = rectangular_path((trunk[-1], segment_slot_center(kind, slot_idx))) + return trunk + branch[1:] + + +def segment_position(kind: str, pkt: PacketSpec, frame: int) -> tuple[float, float]: + if frame < pkt.header_start: + return point_on_path(segment_wire_path(kind), progress_linear(frame, pkt.start, pkt.header_start)) + + end = pkt.header_end if kind == "header" else pkt.payload_end + return point_on_path(segment_post_split_path(kind, pkt.num), progress_linear(frame, pkt.header_start, end)) + + +def draw_dma_paths(base: Image.Image, draw: ImageDraw.ImageDraw, frame: int) -> None: + header_path = segment_trunk_path("header") + payload_path = segment_trunk_path("payload") + draw_arrow(draw, header_path, rgba(str(COLORS["arrow"]), 130), 3, 12) + draw_arrow(draw, payload_path, rgba(str(COLORS["arrow"]), 130), 3, 12) + bend_x = NIC_RECT[2] + 88 + draw_path_tag(draw, (bend_x + 14, (HEADER_SLOT_AREA[1] + HEADER_SLOT_AREA[3]) / 2 - 30), "header DMA", COLORS["header"]) + draw_path_tag(draw, (bend_x + 14, (PAYLOAD_SLOT_AREA[1] + PAYLOAD_SLOT_AREA[3]) / 2 + 32), "GPUDirect payload DMA", COLORS["payload"]) + + +def draw_flowing_segments( + base: Image.Image, + draw: ImageDraw.ImageDraw, + frame: int, +) -> None: + for pkt in PACKETS: + if pkt.start <= frame < pkt.payload_end: + draw_packet_segment( + draw, + segment_position("payload", pkt, frame), + PAYLOAD_SEGMENT_WIDTH, + PAYLOAD_SEGMENT_HEIGHT, + COLORS["payload"], + f"payload {pkt.num + 1}", + ) + + if pkt.start <= frame < pkt.header_end: + draw_packet_segment( + draw, + segment_position("header", pkt, frame), + HEADER_SEGMENT_WIDTH, + HEADER_SEGMENT_HEIGHT, + COLORS["header"], + f"hdr {pkt.num + 1}", + ) + + +PACKETS = make_packets() +FRAMES = max(190, max(pkt.payload_end for pkt in PACKETS) + 36) + + +def render_frame(frame: int) -> Image.Image: + bg = COLORS.get("bg") + img = Image.new("RGBA", (WIDTH * SCALE, HEIGHT * SCALE), rgba(str(bg), 255) if bg else TRANSPARENT) + draw = ImageDraw.Draw(img) + draw_static_background(draw) + draw_wire(draw, frame) + draw_kernel_bypass(draw) + draw_dma_paths(img, draw, frame) + draw_device_memory(draw, frame) + nic_active = any(pkt.nic_arrive - 3 <= frame < pkt.payload_start for pkt in PACKETS) + draw_nic_chip_base(draw, NIC_RECT, COLORS["nvidia"], pulse=0.8 if nic_active else 0.0) + draw_flowing_segments(img, draw, frame) + draw_nic_engine_box(draw) + return img.resize((WIDTH, HEIGHT), Image.Resampling.LANCZOS) + + +def render_theme(theme: str) -> None: + global COLORS + COLORS = THEMES[theme] + suffix = THEME_OUTPUT_SUFFIX[theme] + animation_path = OUTPUT_DIR / f"header-data-split{suffix}.webp" + poster_path = OUTPUT_DIR / f"header-data-split{suffix}-poster.png" + + frames = [render_frame(i) for i in range(FRAMES)] + save_webp_animation(frames, animation_path, DURATION_MS) + render_frame(min(FRAMES - 1, max(pkt.payload_end for pkt in PACKETS) + 10)).save(poster_path, optimize=True) + print(f"Wrote {animation_path}") + print(f"Wrote {poster_path}") + + +def main() -> None: + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + for theme in THEMES: + render_theme(theme) + + +if __name__ == "__main__": + main() diff --git a/docs/images/packet_diagrams/incoming_wire.py b/docs/images/packet_diagrams/incoming_wire.py new file mode 100644 index 00000000..a90671e8 --- /dev/null +++ b/docs/images/packet_diagrams/incoming_wire.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +import math +from typing import Callable + +from PIL import ImageDraw + +RX_WIRE_LABEL = "RX Packets over Wire" +WIRE_LABEL_OFFSET = 95.0 + +PtFn = Callable[[tuple[float, float]], tuple[int, int]] +ScaleFn = Callable[[float], int] +RgbaFn = Callable[[str, int], tuple[int, int, int, int]] +BoxFn = Callable[[tuple[float, float, float, float]], tuple[int, int, int, int]] + + +def center_wire_y(x: float, x0: float, y_center: float, frame: int, *, wave: bool = False) -> float: + if not wave: + return y_center + i = (x - x0) / 3.55 + return y_center + math.sin((i + frame * 0.95) / 5.4) * 2.4 + + +def draw_rx_wire_label( + draw_text: Callable[..., None], + wire_y: float, + *, + x: float = 48.0, + text: str = RX_WIRE_LABEL, +) -> None: + draw_text((x, wire_y - WIRE_LABEL_OFFSET), text) + + +def draw_rx_wire( + draw: ImageDraw.ImageDraw, + frame: int, + x0: float, + x1: float, + y_center: float, + wire_color: str, + marker_color: str, + pt: PtFn, + s: ScaleFn, + rgba: RgbaFn, + box: BoxFn, + *, + ambient_markers: bool = True, + animate_markers: bool = False, + wave: bool = False, +) -> None: + span = max(x1 - x0, 1.0) + count = max(20, int(span / 3.55) + 1) + for offset in (-13, 0, 13): + pts: list[tuple[float, float]] = [] + for i in range(count): + x = x0 + i * (span / max(1, count - 1)) + wave_offset = math.sin((i + frame * 0.95) / 5.4) * 2.4 if wave else 0.0 + pts.append((x, y_center + offset + wave_offset)) + draw.line([pt(p) for p in pts], fill=wire_color, width=s(2)) + + if not ambient_markers: + return + loop = max(44.0, span - 44.0) + marker_frame = frame if animate_markers else 0 + for i in range(5): + x = x0 + 22 + ((marker_frame * 7 + i * 68) % loop) + cy = center_wire_y(x, x0, y_center, frame, wave=wave) + alpha = 160 - i * 16 + draw.ellipse(box((x - 4, cy - 4, x + 4, cy + 4)), fill=rgba(marker_color, alpha)) + + +def draw_rx_packet_marker( + draw: ImageDraw.ImageDraw, + frame: int, + cx: float, + wire_y: float, + wire_x0: float, + color: str, + rgba: RgbaFn, + box: BoxFn, + *, + radius: float = 5.0, + alpha: int = 255, + snap_to_wire: bool = True, + cy: float | None = None, + wave: bool = False, +) -> None: + marker_y = center_wire_y(cx, wire_x0, wire_y, frame, wave=wave) if snap_to_wire else (cy if cy is not None else wire_y) + draw.ellipse( + box((cx - radius, marker_y - radius, cx + radius, marker_y + radius)), + fill=rgba(color, alpha), + ) + + +def path_lengths(points: list[tuple[float, float]]) -> tuple[list[float], float]: + lengths = [0.0] + for i in range(1, len(points)): + x0, y0 = points[i - 1] + x1, y1 = points[i] + lengths.append(lengths[-1] + math.hypot(x1 - x0, y1 - y0)) + return lengths, lengths[-1] + + +def path_point_and_tangent( + points: list[tuple[float, float]], + t: float, +) -> tuple[float, float, float]: + if len(points) < 2: + p = points[0] if points else (0.0, 0.0) + return p[0], p[1], 0.0 + t = max(0.0, min(1.0, t)) + lengths, total = path_lengths(points) + if total <= 0: + return points[0][0], points[0][1], 0.0 + target = t * total + for i in range(1, len(points)): + if lengths[i] >= target: + seg_len = lengths[i] - lengths[i - 1] + local = 0.0 if seg_len <= 0 else (target - lengths[i - 1]) / seg_len + x0, y0 = points[i - 1] + x1, y1 = points[i] + cx = x0 + (x1 - x0) * local + cy = y0 + (y1 - y0) * local + return cx, cy, math.atan2(y1 - y0, x1 - x0) + x0, y0 = points[-2] + x1, y1 = points[-1] + return x1, y1, math.atan2(y1 - y0, x1 - x0) + + +def wire_t_end(wire: list[tuple[float, float]], full: list[tuple[float, float]]) -> float: + _, wire_total = path_lengths(wire) + _, full_total = path_lengths(full) + if full_total <= 0: + return 1.0 + return min(1.0, wire_total / full_total) diff --git a/docs/images/packet_diagrams/reorder/packet-reorder-light-poster.png b/docs/images/packet_diagrams/reorder/packet-reorder-light-poster.png new file mode 100644 index 00000000..4ea6300c Binary files /dev/null and b/docs/images/packet_diagrams/reorder/packet-reorder-light-poster.png differ diff --git a/docs/images/packet_diagrams/reorder/packet-reorder-light.webp b/docs/images/packet_diagrams/reorder/packet-reorder-light.webp new file mode 100644 index 00000000..9da4b150 Binary files /dev/null and b/docs/images/packet_diagrams/reorder/packet-reorder-light.webp differ diff --git a/docs/images/packet_diagrams/reorder/packet-reorder-poster.png b/docs/images/packet_diagrams/reorder/packet-reorder-poster.png new file mode 100644 index 00000000..aec035a2 Binary files /dev/null and b/docs/images/packet_diagrams/reorder/packet-reorder-poster.png differ diff --git a/docs/images/packet_diagrams/reorder/packet-reorder.webp b/docs/images/packet_diagrams/reorder/packet-reorder.webp new file mode 100644 index 00000000..10405021 Binary files /dev/null and b/docs/images/packet_diagrams/reorder/packet-reorder.webp differ diff --git a/docs/images/packet_diagrams/reorder_animation.py b/docs/images/packet_diagrams/reorder_animation.py new file mode 100644 index 00000000..49dc78b4 --- /dev/null +++ b/docs/images/packet_diagrams/reorder_animation.py @@ -0,0 +1,868 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path + +from PIL import Image, ImageDraw + +from anim_common import font_scheme, save_webp_animation +from incoming_wire import draw_rx_wire, draw_rx_wire_label + + +ROOT = Path(__file__).resolve().parent + +WIDTH = 1180 +HEIGHT = 660 +SCALE = 2 +DURATION_MS = 40 + +PACKET_COUNT = 5 +ARRIVAL_ORDER = (2, 0, 4, 1, 3) +WIRE_FRAMES = 34 +STAGE_TRAVEL_FRAMES = 20 +ARRIVAL_GAP = 16 +KERNEL_DELAY = 22 +WRITE_FRAMES = 42 +WRITE_STAGGER = 2 + +NIC_RECT = (300, 250, 500, 410) +REORDER_RECT = (610, 150, 940, 546) +REORDER_STAGING_COL = (635, 260, 755, 512) +REORDER_QUEUE_COL = (790, 260, 910, 512) +CONVERT_STAGING_COL = (635, 260, 720, 512) +CONVERT_QUEUE_COL = (755, 260, 915, 512) + +WIRE_Y = 330 + +TRANSPARENT = (0, 0, 0, 0) + +PACKET_GRAD_START = "#59d4ff" +PACKET_GRAD_END = "#9b8cff" + +ACCENTS = { + "nvidia": "#76b900", + "reorder": "#76b900", + "ink": "#07111f", +} + +THEMES: dict[str, dict[str, str | None]] = { + "default": { + **ACCENTS, + "bg": None, + "panel": "#0d1b2e", + "panel_2": "#101f34", + "line": "#35516f", + "text": "#f6fbff", + "muted": "#afc0cf", + "wire": "#6d7e92", + "row_bg": "#061220", + "stroke": "#ffffff", + "slot_empty": "#64748b", + "arrow": "#ffffff", + }, + "light": { + **ACCENTS, + "bg": "#ffffff", + "panel": "#f5f5f5", + "panel_2": "#eeeeee", + "line": "#1a1a1a", + "text": "#1a1a1a", + "muted": "#404040", + "wire": "#404040", + "row_bg": "#e8e8e8", + "stroke": "#1a1a1a", + "slot_empty": "#9ca3af", + "arrow": "#1a1a1a", + }, +} + +THEME_OUTPUT_SUFFIX = { + "default": "", + "light": "-light", +} + +COLORS = THEMES["default"] + + +@dataclass(frozen=True) +class VisualConfig: + output_dir: Path + base_name: str + title: str + subtitle: str + output_heading: str + nic_idle: str + nic_tracking_prefix: str + staging_col: tuple[float, float, float, float] + queue_col: tuple[float, float, float, float] + convert_payload: bool = False + + +VISUALIZATIONS = ( + VisualConfig( + output_dir=ROOT / "reorder", + base_name="packet-reorder", + title="GPU reorder", + subtitle="batch staging -> fixed slots", + output_heading="slot = seq % N", + nic_idle="staging", + nic_tracking_prefix="stage seq", + staging_col=REORDER_STAGING_COL, + queue_col=REORDER_QUEUE_COL, + ), + VisualConfig( + output_dir=ROOT / "reorder_quantize", + base_name="packet-reorder-quantize", + title="GPU reorder + convert", + subtitle="batch staging -> fixed slots\nint4 -> fp32 conversion", + output_heading="fp32 slots", + nic_idle="int4 staging", + nic_tracking_prefix="stage int4 seq", + staging_col=CONVERT_STAGING_COL, + queue_col=CONVERT_QUEUE_COL, + convert_payload=True, + ), +) + +CURRENT_VISUAL = VISUALIZATIONS[0] + + +@dataclass(frozen=True) +class PacketSpec: + num: int + start: int + nic_arrive: int + stage_end: int + write_start: int + write_end: int + + +NIC = NIC_RECT +REORDER = REORDER_RECT +NIC_CX = (NIC[0] + NIC[2]) / 2 +NIC_CY = (NIC[1] + NIC[3]) / 2 +REORDER_CY = (REORDER[1] + REORDER[3]) / 2 +REORDER_ENTRY = (REORDER[0], NIC_CY) + + +def rgb(hex_color: str) -> tuple[int, int, int]: + hex_color = hex_color.strip("#") + return tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4)) + + +def rgb_to_hex(r: int, g: int, b: int) -> str: + return f"#{r:02x}{g:02x}{b:02x}" + + +def rgba(hex_color: str, alpha: int = 255) -> tuple[int, int, int, int]: + r, g, b = rgb(hex_color) + return r, g, b, alpha + + +def lerp(a: float, b: float, t: float) -> float: + return a + (b - a) * t + + +def lerp_hex(a: str, b: str, t: float) -> str: + r1, g1, b1 = rgb(a) + r2, g2, b2 = rgb(b) + return rgb_to_hex( + int(lerp(r1, r2, t)), + int(lerp(g1, g2, t)), + int(lerp(b1, b2, t)), + ) + + +def clamp(value: float, lo: float = 0.0, hi: float = 1.0) -> float: + return max(lo, min(hi, value)) + + +def progress_linear(frame: int, start: int, end: int) -> float: + return clamp((frame - start) / max(1, end - start)) + + +FONTS = font_scheme(SCALE) + + +def s(value: float) -> int: + return int(round(value * SCALE)) + + +def pt(point: tuple[float, float]) -> tuple[int, int]: + return s(point[0]), s(point[1]) + + +def box(rect: tuple[float, float, float, float]) -> tuple[int, int, int, int]: + return tuple(s(v) for v in rect) + + +def draw_text( + draw: ImageDraw.ImageDraw, + xy: tuple[float, float], + text: str, + face: ImageFont.FreeTypeFont, + fill: str | tuple[int, int, int, int] = COLORS["text"], + anchor: str = "la", +) -> None: + draw.text(pt(xy), text, font=face, fill=fill, anchor=anchor) + + +def centered_text( + draw: ImageDraw.ImageDraw, + rect: tuple[float, float, float, float], + text: str, + face: ImageFont.FreeTypeFont, + fill: str | tuple[int, int, int, int] = COLORS["text"], +) -> None: + x1, y1, x2, y2 = rect + draw.text((s((x1 + x2) / 2), s((y1 + y2) / 2)), text, font=face, fill=fill, anchor="mm") + + +def centered_multiline_text( + draw: ImageDraw.ImageDraw, + rect: tuple[float, float, float, float], + text: str, + face: ImageFont.FreeTypeFont, + fill: str | tuple[int, int, int, int] = COLORS["text"], +) -> None: + x1, y1, x2, y2 = rect + draw.multiline_text( + (s((x1 + x2) / 2), s((y1 + y2) / 2)), + text, + font=face, + fill=fill, + anchor="mm", + align="center", + spacing=s(4), + ) + + +def rectangular_path(waypoints: tuple[tuple[float, float], ...], steps: int = 28) -> list[tuple[float, float]]: + if len(waypoints) < 2: + return list(waypoints) + path: list[tuple[float, float]] = [] + for i in range(len(waypoints) - 1): + x0, y0 = waypoints[i] + x1, y1 = waypoints[i + 1] + for step in range(steps + 1): + if i > 0 and step == 0: + continue + t = step / steps + path.append((lerp(x0, x1, t), lerp(y0, y1, t))) + return path + + +def path_lengths(points: list[tuple[float, float]]) -> tuple[list[float], float]: + lengths = [0.0] + for i in range(1, len(points)): + x0, y0 = points[i - 1] + x1, y1 = points[i] + lengths.append(lengths[-1] + math.hypot(x1 - x0, y1 - y0)) + return lengths, lengths[-1] + + +def draw_polyline( + draw: ImageDraw.ImageDraw, + points: list[tuple[float, float]], + color: str | tuple[int, int, int, int], + width: int = 4, +) -> None: + draw.line([pt(p) for p in points], fill=color, width=s(width)) + + +def draw_glow_line( + base: Image.Image, + points: list[tuple[float, float]], + color: str, + alpha: int, + width: int, +) -> None: + layer = Image.new("RGBA", base.size, (0, 0, 0, 0)) + d = ImageDraw.Draw(layer) + draw_polyline(d, points, rgba(color, alpha), width) + base.alpha_composite(layer) + + +def draw_arrow_head( + draw: ImageDraw.ImageDraw, + end: tuple[float, float], + prev: tuple[float, float], + color: str | tuple[int, int, int, int], + size: float = 12, +) -> None: + angle = math.atan2(end[1] - prev[1], end[0] - prev[0]) + wing = math.radians(28) + p1 = (end[0] - size * math.cos(angle - wing), end[1] - size * math.sin(angle - wing)) + p2 = (end[0] - size * math.cos(angle + wing), end[1] - size * math.sin(angle + wing)) + draw.polygon([pt(end), pt(p1), pt(p2)], fill=color) + + +def draw_arrow( + draw: ImageDraw.ImageDraw, + points: list[tuple[float, float]], + color: str | tuple[int, int, int, int], + width: int = 4, + arrow_size: float = 13, +) -> None: + draw_polyline(draw, points, color, width) + if len(points) >= 2: + prev_idx = max(0, len(points) - 5) + draw_arrow_head(draw, points[-1], points[prev_idx], color, arrow_size) + + +def point_on_path_distance(points: list[tuple[float, float]], dist: float) -> tuple[float, float]: + if len(points) < 2: + return points[0] if points else (0.0, 0.0) + lengths, total = path_lengths(points) + if total <= 0: + return points[0] + target = clamp(dist, 0.0, total) + for i in range(1, len(points)): + if lengths[i] >= target: + seg_len = lengths[i] - lengths[i - 1] + local = 0.0 if seg_len <= 0 else (target - lengths[i - 1]) / seg_len + x0, y0 = points[i - 1] + x1, y1 = points[i] + return lerp(x0, x1, local), lerp(y0, y1, local) + return points[-1] + + +def draw_path_arrows( + draw: ImageDraw.ImageDraw, + points: list[tuple[float, float]], + color: str | tuple[int, int, int, int], + *, + width: int = 3, + arrow_size: float = 11, + spacing: float = 34, +) -> None: + if len(points) < 2: + return + _, total = path_lengths(points) + draw_polyline(draw, points, color, width) + dist = spacing * 0.6 + while dist < total - arrow_size: + tip = point_on_path_distance(points, dist) + prev = point_on_path_distance(points, max(0.0, dist - 10)) + draw_arrow_head(draw, tip, prev, color, arrow_size) + dist += spacing + if total > arrow_size: + draw_arrow_head(draw, points[-1], points[-2], color, arrow_size + 2) + + +def draw_glow_arrows( + base: Image.Image, + points: list[tuple[float, float]], + color: str, + alpha: int, + width: int, +) -> None: + layer = Image.new("RGBA", base.size, (0, 0, 0, 0)) + d = ImageDraw.Draw(layer) + draw_path_arrows(d, points, rgba(color, alpha), width=width, arrow_size=11, spacing=32) + base.alpha_composite(layer) + + +def column_row_layout(rect: tuple[float, float, float, float]) -> list[tuple[float, float, float, float]]: + x1, y1, x2, y2 = rect + cx = (x1 + x2) / 2 + pad_y = 6 + gap = 6 + inner_h = y2 - y1 - 2 * pad_y + chip_h = (inner_h - gap * (PACKET_COUNT - 1)) / PACKET_COUNT + chip_w = x2 - x1 - 8 + rows: list[tuple[float, float, float, float]] = [] + y = y1 + pad_y + chip_h / 2 + for _ in range(PACKET_COUNT): + rows.append((cx, y, chip_w, chip_h)) + y += chip_h + gap + return rows + + +def staging_col_rect() -> tuple[float, float, float, float]: + return CURRENT_VISUAL.staging_col + + +def queue_col_rect() -> tuple[float, float, float, float]: + return CURRENT_VISUAL.queue_col + + +def staging_row_pos(row: int) -> tuple[float, float, float, float]: + return column_row_layout(staging_col_rect())[row] + + +def queue_row_pos(row: int) -> tuple[float, float, float, float]: + return column_row_layout(queue_col_rect())[row] + + +def slot_index(num: int) -> int: + return num % PACKET_COUNT + + +def packet_order_index(pkt: PacketSpec) -> int: + return PACKETS.index(pkt) + + +def completed_nums(frame: int) -> list[int]: + return [p.num for p in PACKETS if frame >= p.write_end] + + +def staged_packets(frame: int) -> list[PacketSpec]: + return [p for p in PACKETS if p.stage_end <= frame < p.write_start] + + +def writing_packet(frame: int) -> PacketSpec | None: + for pkt in PACKETS: + if pkt.write_start <= frame < pkt.write_end: + return pkt + return None + + +def queue_display(frame: int) -> list[int | None]: + row: list[int | None] = [None] * PACKET_COUNT + for num in completed_nums(frame): + row[slot_index(num)] = num + return row + + +def packet_label(num: int, *, output: bool = False) -> str: + if CURRENT_VISUAL.convert_payload: + return f"fp32 s{num}" if output else f"int4 s{num}" + return f"seq {num}" + + +def before_placement(pkt: PacketSpec) -> list[int]: + return [p.num for p in PACKETS if p.write_end <= pkt.write_start] + + +def after_placement(pkt: PacketSpec) -> list[int]: + return before_placement(pkt) + [pkt.num] + + +def placed_packet_position(frame: int, num: int) -> tuple[float, float, float, float] | None: + if num not in completed_nums(frame): + return None + return queue_row_pos(slot_index(num)) + + +def shuffle_entry_for_row(row: int) -> tuple[float, float]: + _, cy, _, _ = staging_row_pos(row) + return staging_col_rect()[0] + 12, cy + + +def build_paths() -> dict[str, list[tuple[float, float]]]: + wire = rectangular_path(((48, WIRE_Y), (NIC[0], WIRE_Y))) + nic_route = rectangular_path(((NIC[2], NIC_CY), REORDER_ENTRY)) + to_reorder = rectangular_path( + ( + (48, WIRE_Y), + (NIC[0], WIRE_Y), + (NIC[2], NIC_CY), + REORDER_ENTRY, + ) + ) + return { + "wire": wire, + "to_reorder": to_reorder, + "nic_route": nic_route, + } + + +PATHS = build_paths() + + +def make_packets() -> tuple[PacketSpec, ...]: + staged: list[tuple[int, int, int, int]] = [] + cursor = 8 + for num in ARRIVAL_ORDER: + nic_arrive = cursor + WIRE_FRAMES + stage_end = nic_arrive + STAGE_TRAVEL_FRAMES + staged.append((num, cursor, nic_arrive, stage_end)) + cursor += ARRIVAL_GAP + + kernel_start = max(stage_end for _, _, _, stage_end in staged) + KERNEL_DELAY + packets: list[PacketSpec] = [] + for idx, (num, start, nic_arrive, stage_end) in enumerate(staged): + write_start = kernel_start + idx * WRITE_STAGGER + write_end = write_start + WRITE_FRAMES + packets.append(PacketSpec(num, start, nic_arrive, stage_end, write_start, write_end)) + return tuple(packets) + + +PACKETS = make_packets() +KERNEL_START = min(p.write_start for p in PACKETS) +KERNEL_END = max(p.write_end for p in PACKETS) +FRAMES = KERNEL_END + 40 + + +def active_packet(frame: int) -> PacketSpec | None: + for pkt in PACKETS: + if pkt.nic_arrive <= frame < pkt.stage_end: + return pkt + return None + + +def nic_pulse(frame: int, pkt: PacketSpec) -> float: + return max( + progress_linear(frame, pkt.nic_arrive, pkt.nic_arrive + 6), + 1.0 - progress_linear(frame, pkt.nic_arrive + 8, pkt.nic_arrive + 14), + ) + + +def packet_base_color(num: int) -> str: + if PACKET_COUNT <= 1: + return PACKET_GRAD_START + t = clamp(num / (PACKET_COUNT - 1)) + return lerp_hex(PACKET_GRAD_START, PACKET_GRAD_END, t) + + +def packet_gradient_ends(num: int) -> tuple[str, str]: + base = packet_base_color(num) + return lerp_hex(base, "#ffffff", 0.22), lerp_hex(base, PACKET_GRAD_END, 0.3) + + +def point_on_path(points: list[tuple[float, float]], t: float) -> tuple[float, float]: + if len(points) < 2: + return points[0] if points else (0.0, 0.0) + t = clamp(t) + lengths, total = path_lengths(points) + if total <= 0: + return points[0] + target = t * total + for i in range(1, len(points)): + if lengths[i] >= target: + seg_len = lengths[i] - lengths[i - 1] + local = 0.0 if seg_len <= 0 else (target - lengths[i - 1]) / seg_len + x0, y0 = points[i - 1] + x1, y1 = points[i] + return lerp(x0, x1, local), lerp(y0, y1, local) + return points[-1] + + +def stage_path(pkt: PacketSpec) -> list[tuple[float, float]]: + _, sy, _, _ = staging_row_pos(packet_order_index(pkt)) + staging_col = staging_col_rect() + return rectangular_path( + ( + (NIC[2], NIC_CY), + (REORDER[0], NIC_CY), + (staging_col[0] + 10, sy), + ((staging_col[0] + staging_col[2]) / 2, sy), + ), + steps=16, + ) + + +def packet_stage_path(pkt: PacketSpec) -> list[tuple[float, float]]: + _, sy, _, _ = staging_row_pos(packet_order_index(pkt)) + staging_col = staging_col_rect() + return rectangular_path( + ( + (48, WIRE_Y), + (NIC[0], WIRE_Y), + (NIC[2], NIC_CY), + (REORDER[0], NIC_CY), + (staging_col[0] + 10, sy), + ((staging_col[0] + staging_col[2]) / 2, sy), + ), + steps=16, + ) + + +def write_path(pkt: PacketSpec) -> list[tuple[float, float]]: + sx, sy, _, _ = staging_row_pos(packet_order_index(pkt)) + tx, ty, _, _ = queue_row_pos(slot_index(pkt.num)) + return rectangular_path(((sx, sy), (tx, ty)), steps=18) + + +def packet_screen_pos(frame: int, pkt: PacketSpec) -> tuple[float, float, float] | None: + if frame < pkt.start or frame >= pkt.write_end: + return None + width = 92.0 + if frame < pkt.stage_end: + t = progress_linear(frame, pkt.start, pkt.stage_end) + cx, cy = point_on_path(packet_stage_path(pkt), t) + return cx, cy, width + if frame < pkt.write_start: + return None + t = progress_linear(frame, pkt.write_start, pkt.write_end) + sx, sy, sw, _ = staging_row_pos(packet_order_index(pkt)) + tx, ty, tw, _ = queue_row_pos(slot_index(pkt.num)) + return lerp(sx, tx, t), lerp(sy, ty, t), lerp(sw - 6, tw - 6, t) + + +def received_count(frame: int) -> int: + return sum(1 for p in PACKETS if frame >= p.stage_end) + + +def draw_gradient_packet( + base: Image.Image, + draw: ImageDraw.ImageDraw, + cx: float, + cy: float, + width: float, + num: int, + alpha: int = 255, + *, + label: str | None = None, + height: float = 28, +) -> None: + h = height + x1, y1 = cx - width / 2, cy - h / 2 + x2, y2 = cx + width / 2, cy + h / 2 + radius = s(6) + left, right = packet_gradient_ends(num) + accent = packet_base_color(num) + + chip = Image.new("RGBA", (max(1, s(x2 - x1)), max(1, s(y2 - y1))), (0, 0, 0, 0)) + cd = ImageDraw.Draw(chip) + inset = s(1) + cw, ch = chip.size + steps = 16 + for i in range(steps): + t0 = i / steps + t1 = (i + 1) / steps + color = lerp_hex(left, right, (t0 + t1) / 2) + sx1 = int(inset + (cw - 2 * inset) * t0) + sx2 = int(inset + (cw - 2 * inset) * t1) + cd.rectangle((sx1, inset, sx2, ch - inset), fill=rgba(color, alpha)) + mask = Image.new("L", chip.size, 0) + md = ImageDraw.Draw(mask) + md.rounded_rectangle((0, 0, cw - 1, ch - 1), radius=radius, fill=255) + chip.putalpha(mask) + base.paste(chip, (s(x1), s(y1)), chip) + draw.rounded_rectangle( + box((x1, y1, x2, y2)), + radius=radius, + outline=rgba(accent, min(255, 170 + alpha // 3)), + width=s(2), + ) + if label: + if len(label) > 10: + face = FONTS["tiny"] + else: + face = FONTS["packet"] + centered_text(draw, (x1, y1, x2, y2), label, face, fill=rgba(COLORS["ink"], alpha)) + + +def draw_wires(draw: ImageDraw.ImageDraw, frame: int) -> None: + draw_rx_wire_label(lambda xy, text: draw_text(draw, xy, text, FONTS["label"], fill=COLORS["text"]), WIRE_Y) + draw_rx_wire(draw, frame, 48, NIC[0], WIRE_Y, str(COLORS["wire"]), PACKET_GRAD_START, pt, s, rgba, box) + + +def draw_static_route(draw: ImageDraw.ImageDraw) -> None: + draw_arrow(draw, PATHS["nic_route"], rgba(COLORS["reorder"], 255), 3, 12) + + +def draw_sequence_queue(base: Image.Image, draw: ImageDraw.ImageDraw, frame: int) -> None: + row = queue_display(frame) + for idx, value in enumerate(row): + cx, cy, chip_w, chip_h = queue_row_pos(idx) + slot_box = (cx - chip_w / 2, cy - chip_h / 2, cx + chip_w / 2, cy + chip_h / 2) + draw.rounded_rectangle( + box(slot_box), + radius=s(7), + fill=COLORS["row_bg"], + outline=rgba(COLORS["line"], 210), + width=s(1), + ) + if value is None: + centered_text(draw, slot_box, f"slot {idx}", FONTS["slot"], fill=COLORS["slot_empty"]) + continue + draw_gradient_packet( + base, + draw, + cx, + cy, + chip_w - 6, + value, + label=packet_label(value, output=True), + height=chip_h - 6, + ) + + +def draw_staging_list(base: Image.Image, draw: ImageDraw.ImageDraw, frame: int) -> None: + staged_by_row = {packet_order_index(pkt): pkt for pkt in staged_packets(frame)} + for idx in range(PACKET_COUNT): + cx, cy, chip_w, chip_h = staging_row_pos(idx) + row_box = (cx - chip_w / 2, cy - chip_h / 2, cx + chip_w / 2, cy + chip_h / 2) + draw.rounded_rectangle( + box(row_box), + radius=s(7), + fill=COLORS["row_bg"], + outline=rgba(COLORS["line"], 190), + width=s(1), + ) + pkt = staged_by_row.get(idx) + if pkt is None: + centered_text(draw, row_box, "ptr", FONTS["tiny"], fill=COLORS["slot_empty"]) + continue + draw_gradient_packet(base, draw, cx, cy, chip_w - 6, pkt.num, label=packet_label(pkt.num), height=chip_h - 6) + + +def draw_reorder_panel(base: Image.Image, draw: ImageDraw.ImageDraw, frame: int) -> None: + x1, y1, _, _ = REORDER + staging_col = staging_col_rect() + queue_col = queue_col_rect() + kernel_active = KERNEL_START <= frame < KERNEL_END + pulse = 0.55 if kernel_active else 0.0 + glow = int(40 + 55 * pulse) + draw.rounded_rectangle(box(REORDER), radius=s(18), fill=COLORS["panel_2"], outline=rgba(COLORS["reorder"], 170 + glow), width=s(3)) + draw_text(draw, (x1 + 14, y1 + 16), CURRENT_VISUAL.title, FONTS["label"], fill=COLORS["text"], anchor="lt") + for idx, subtitle in enumerate(CURRENT_VISUAL.subtitle.splitlines()): + draw_text(draw, (x1 + 14, y1 + 40 + idx * 17), subtitle, FONTS["small"], fill=COLORS["muted"], anchor="lt") + draw_text(draw, (staging_col[0], staging_col[1] - 20), "staged ptrs", FONTS["tiny"], fill=COLORS["muted"], anchor="lt") + draw_text(draw, (queue_col[0], queue_col[1] - 20), CURRENT_VISUAL.output_heading, FONTS["tiny"], fill=COLORS["muted"], anchor="lt") + if CURRENT_VISUAL.convert_payload and kernel_active: + draw_text( + draw, + ((staging_col[2] + queue_col[0]) / 2, staging_col[1] - 20), + "convert", + FONTS["tiny"], + fill=COLORS["reorder"], + anchor="mm", + ) + draw_staging_list(base, draw, frame) + draw_sequence_queue(base, draw, frame) + + +def nic_status(frame: int) -> tuple[str, str, str, float]: + pkt = active_packet(frame) + tracking = pkt is not None + pulse = nic_pulse(frame, pkt) if tracking and pkt else 0.0 + if tracking and pkt: + if CURRENT_VISUAL.convert_payload: + return f"stage int4\nseq {pkt.num}", packet_base_color(pkt.num), "#ffffff", pulse + return f"stage seq\n{pkt.num}", packet_base_color(pkt.num), "#ffffff", pulse + + count = received_count(frame) + if count: + return f"staged\n{count}/{PACKET_COUNT}", COLORS["nvidia"], COLORS["ink"], pulse + idle = "int4\nstaging" if CURRENT_VISUAL.convert_payload else CURRENT_VISUAL.nic_idle + return idle, COLORS["nvidia"], COLORS["ink"], pulse + + +def draw_nic_status_pill(draw: ImageDraw.ImageDraw, frame: int) -> None: + x1, y1, x2, _ = NIC + pill_text, pill_accent, pill_ink, pulse = nic_status(frame) + pill = (x1 + 12, y1 + 38, x2 - 12, y1 + 122) + draw.rounded_rectangle( + box(pill), + radius=s(8), + fill=rgba(pill_accent, 245), + outline=rgba(pill_accent, 220), + width=s(2), + ) + centered_multiline_text(draw, pill, pill_text, FONTS["nic_center"], fill=pill_ink) + + +def draw_nic_chip_base(draw: ImageDraw.ImageDraw, frame: int) -> None: + x1, y1, x2, y2 = NIC + pkt = active_packet(frame) + tracking = pkt is not None + pulse = nic_pulse(frame, pkt) if tracking and pkt else 0.0 + accent = COLORS["nvidia"] + + glow_alpha = int(55 + 120 * pulse) if tracking else 55 + draw.rounded_rectangle(box((x1 - 10, y1 - 10, x2 + 10, y2 + 10)), radius=s(24), fill=rgba(accent, 16 + glow_alpha // 3)) + draw.rounded_rectangle(box(NIC), radius=s(18), fill=COLORS["panel"], outline=rgba(accent, 190 + int(65 * pulse)), width=s(3 + int(2 * pulse))) + for i in range(7): + y = lerp(y1 + 18, y2 - 18, i / 6) + draw.line((s(x1 - 14), s(y), s(x1), s(y)), fill=rgba(COLORS["line"], 220), width=s(3)) + draw.line((s(x2), s(y), s(x2 + 14), s(y)), fill=rgba(COLORS["line"], 220), width=s(3)) + for i in range(5): + x = lerp(x1 + 25, x2 - 25, i / 4) + draw.line((s(x), s(y1 - 12), s(x), s(y1)), fill=rgba(COLORS["line"], 220), width=s(3)) + draw.line((s(x), s(y2), s(x), s(y2 + 12)), fill=rgba(COLORS["line"], 220), width=s(3)) + centered_text(draw, (x1 + 18, y1 + 4, x2 - 18, y1 + 38), "NVIDIA NIC", FONTS["chip"], fill=COLORS["text"]) + + +def draw_flowing_packets(base: Image.Image, draw: ImageDraw.ImageDraw, frame: int) -> None: + for pkt in PACKETS: + if frame < pkt.start or frame >= pkt.write_end: + continue + pos = packet_screen_pos(frame, pkt) + if pos is None: + continue + cx, cy, w = pos + if frame < pkt.stage_end: + chip_h = 32 + label = None if NIC[0] <= cx <= NIC[2] and NIC[1] <= cy <= NIC[3] else packet_label(pkt.num) + else: + _, _, _, chip_h = queue_row_pos(slot_index(pkt.num)) + chip_h -= 6 + label = packet_label(pkt.num, output=frame >= pkt.write_end - 4) + draw_gradient_packet( + base, + draw, + cx, + cy, + w, + pkt.num, + label=label, + height=chip_h, + ) + + +def draw_route_glow(base: Image.Image, frame: int) -> None: + for pkt in PACKETS: + if pkt.nic_arrive <= frame < pkt.stage_end: + glow = progress_linear(frame, pkt.nic_arrive, pkt.stage_end) + draw_glow_line(base, stage_path(pkt), COLORS["reorder"], int(45 + 45 * glow), 6) + elif pkt.write_start <= frame < pkt.write_end: + glow = progress_linear(frame, pkt.write_start, pkt.write_end) + draw_glow_line(base, write_path(pkt), COLORS["reorder"], int(45 + 45 * glow), 5) + if CURRENT_VISUAL.convert_payload: + sx, sy, _, _ = staging_row_pos(packet_order_index(pkt)) + tx, ty, _, _ = queue_row_pos(slot_index(pkt.num)) + draw = ImageDraw.Draw(base) + draw_text( + draw, + (lerp(sx, tx, glow), lerp(sy, ty, glow) - 24), + "int4 -> fp32", + FONTS["tiny"], + fill=COLORS["reorder"], + anchor="mm", + ) + + +def render_frame(frame: int) -> Image.Image: + bg = COLORS.get("bg") + img = Image.new("RGBA", (WIDTH * SCALE, HEIGHT * SCALE), rgba(str(bg), 255) if bg else TRANSPARENT) + draw = ImageDraw.Draw(img) + draw_wires(draw, frame) + draw_static_route(draw) + draw_reorder_panel(img, draw, frame) + draw_nic_chip_base(draw, frame) + draw_route_glow(img, frame) + draw_flowing_packets(img, draw, frame) + draw_nic_status_pill(draw, frame) + return img.resize((WIDTH, HEIGHT), Image.Resampling.LANCZOS) + + +def render_theme(theme: str, visual: VisualConfig) -> None: + global COLORS, CURRENT_VISUAL + COLORS = THEMES[theme] + CURRENT_VISUAL = visual + suffix = THEME_OUTPUT_SUFFIX[theme] + animation_path = visual.output_dir / f"{visual.base_name}{suffix}.webp" + poster_path = visual.output_dir / f"{visual.base_name}{suffix}-poster.png" + + visual.output_dir.mkdir(parents=True, exist_ok=True) + frames = [render_frame(i) for i in range(FRAMES)] + save_webp_animation(frames, animation_path, DURATION_MS) + poster_frame = min(FRAMES - 1, KERNEL_END + 12) + render_frame(poster_frame).save(poster_path, optimize=True) + print(f"Wrote {animation_path}") + print(f"Wrote {poster_path}") + + +def main() -> None: + for visual in VISUALIZATIONS: + for theme in THEMES: + render_theme(theme, visual) + + +if __name__ == "__main__": + main() diff --git a/docs/images/packet_diagrams/reorder_quantize/packet-reorder-quantize-light-poster.png b/docs/images/packet_diagrams/reorder_quantize/packet-reorder-quantize-light-poster.png new file mode 100644 index 00000000..a8934f10 Binary files /dev/null and b/docs/images/packet_diagrams/reorder_quantize/packet-reorder-quantize-light-poster.png differ diff --git a/docs/images/packet_diagrams/reorder_quantize/packet-reorder-quantize-light.webp b/docs/images/packet_diagrams/reorder_quantize/packet-reorder-quantize-light.webp new file mode 100644 index 00000000..b67c24db Binary files /dev/null and b/docs/images/packet_diagrams/reorder_quantize/packet-reorder-quantize-light.webp differ diff --git a/docs/images/packet_diagrams/reorder_quantize/packet-reorder-quantize-poster.png b/docs/images/packet_diagrams/reorder_quantize/packet-reorder-quantize-poster.png new file mode 100644 index 00000000..b0fc7115 Binary files /dev/null and b/docs/images/packet_diagrams/reorder_quantize/packet-reorder-quantize-poster.png differ diff --git a/docs/images/packet_diagrams/reorder_quantize/packet-reorder-quantize.webp b/docs/images/packet_diagrams/reorder_quantize/packet-reorder-quantize.webp new file mode 100644 index 00000000..0ddf786a Binary files /dev/null and b/docs/images/packet_diagrams/reorder_quantize/packet-reorder-quantize.webp differ diff --git a/mkdocs.yml b/mkdocs.yml index 5867d1de..f2310e86 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -51,6 +51,7 @@ extra_javascript: site_dir: site exclude_docs: | landing/ + images/packet_diagrams/README.md nav: - Home: index.md