blender-bridge / test / test_bench_tile_baker.py

#!/usr/bin/env python3
"""Performance benchmarks for the tile baker.

Measures:
  1. Tile extraction 128×128 (target: < 0.5ms per tile)
  2. Mip generation 2x2→1 (target: < 1ms)
  3. Full image bake 1024×1024 (target: < 50ms)
  4. Ring buffer write throughput (target: > 1GB/s)

Run: python3 test/test_bench_tile_baker.py
"""

import importlib.util
import os
import sys
import time

import numpy as np

# ── Module loading (same as test_tile_baker.py) ──────────────────────────────
_addon_dir = os.path.join(os.path.dirname(__file__), "..", "blender_addon")

_shm_spec = importlib.util.spec_from_file_location(
    "blender_addon.shm_protocol",
    os.path.join(_addon_dir, "shm_protocol.py"),
    submodule_search_locations=[],
)
shm_protocol = importlib.util.module_from_spec(_shm_spec)
sys.modules["blender_addon"] = type(sys)("blender_addon")
sys.modules["blender_addon"].__path__ = [_addon_dir]
sys.modules["blender_addon.shm_protocol"] = shm_protocol
_shm_spec.loader.exec_module(shm_protocol)

_tb_spec = importlib.util.spec_from_file_location(
    "blender_addon.tile_baker",
    os.path.join(_addon_dir, "tile_baker.py"),
    submodule_search_locations=[],
)
tile_baker = importlib.util.module_from_spec(_tb_spec)
tile_baker.__package__ = "blender_addon"
sys.modules["blender_addon.tile_baker"] = tile_baker
_tb_spec.loader.exec_module(tile_baker)

ShmRegion = shm_protocol.ShmRegion
TileBaker = tile_baker.TileBaker

SHM_NAME = "/bench_tile_baker"
SHM_SIZE = 16 * 1024 * 1024  # 16 MB
TILE_SIZE = 128


def report(name, ms, target_ms):
    status = "OK" if ms <= target_ms else "SLOW"
    print(f"  {name:<45s} {ms:8.3f} ms  (target < {target_ms:.1f})  [{status}]")


# ── Bench 1: Single tile extraction ──────────────────────────────────────────

print("=== Tile Baker Benchmarks ===\n")

shm = ShmRegion()
assert shm.create(SHM_NAME, SHM_SIZE)

try:
    baker = TileBaker(shm, tile_size=TILE_SIZE, max_mips=8)
    baker.init_ring_buffer(SHM_SIZE // 2, SHM_SIZE // 2)

    img = np.random.rand(1024, 1024, 4).astype(np.float32)

    # Warm up
    baker._extract_tile(img, 1024, 1024, 0, 0, TILE_SIZE)

    ITERS = 1000
    start = time.perf_counter()
    for i in range(ITERS):
        baker._extract_tile(img, 1024, 1024, i % 8, i % 8, TILE_SIZE)
    elapsed = (time.perf_counter() - start) * 1000 / ITERS
    report("Tile extract 128x128 (per tile)", elapsed, 0.5)

finally:
    shm.close()


# ── Bench 2: Mip generation ──────────────────────────────────────────────────

shm = ShmRegion()
assert shm.create(SHM_NAME, SHM_SIZE)

try:
    baker = TileBaker(shm, tile_size=TILE_SIZE, max_mips=4)
    baker.init_ring_buffer(SHM_SIZE // 2, SHM_SIZE // 2)

    # Create base tiles (2x2 = 4 tiles at 128x128)
    tile_bytes = TILE_SIZE * TILE_SIZE * 4
    base_tiles = [
        (0, 0, np.random.bytes(tile_bytes)),
        (1, 0, np.random.bytes(tile_bytes)),
        (0, 1, np.random.bytes(tile_bytes)),
        (1, 1, np.random.bytes(tile_bytes)),
    ]

    ITERS = 100
    start = time.perf_counter()
    for i in range(ITERS):
        baker._pending_messages = []
        baker._ring_offset = 0
        baker._generate_mips(base_tiles, 1, 256, 256)
    elapsed = (time.perf_counter() - start) * 1000 / ITERS
    report("Mip gen from 2x2 base tiles (per gen)", elapsed, 1.0)

finally:
    shm.close()


# ── Bench 3: Full image bake 1024×1024 ───────────────────────────────────────

shm = ShmRegion()
assert shm.create(SHM_NAME, SHM_SIZE)

try:
    baker = TileBaker(shm, tile_size=TILE_SIZE, max_mips=4)
    baker.init_ring_buffer(SHM_SIZE // 2, SHM_SIZE // 2)

    img = np.random.rand(1024, 1024, 4).astype(np.float32)

    # Warm up
    baker._pending_messages = []
    baker._ring_offset = 0
    baker.bake_from_array(img, 1024, 1024, 1, 0, 0, 1024, 1024)

    ITERS = 10
    start = time.perf_counter()
    for i in range(ITERS):
        baker._pending_messages = []
        baker._ring_offset = 0
        baker.bake_from_array(img, 1024, 1024, 1, 0, 0, 1024, 1024)
    elapsed = (time.perf_counter() - start) * 1000 / ITERS
    report("Full bake 1024x1024 (per bake)", elapsed, 50.0)

    # Count tiles produced
    msgs = baker.get_pending_messages()
    mip_counts = {}
    for m in msgs:
        parts = m.split(":")[1].split(",")
        mip = int(parts[3])
        mip_counts[mip] = mip_counts.get(mip, 0) + 1

    total_tiles = sum(mip_counts.values())
    print(f"\n  Tiles produced per bake: {total_tiles} (mips: {dict(sorted(mip_counts.items()))})")

finally:
    shm.close()


# ── Bench 4: Ring buffer write throughput ─────────────────────────────────────

shm = ShmRegion()
assert shm.create(SHM_NAME, SHM_SIZE)

try:
    baker = TileBaker(shm, tile_size=TILE_SIZE, max_mips=1)
    ring_size = SHM_SIZE // 2
    baker.init_ring_buffer(SHM_SIZE // 2, ring_size)

    tile_data = np.random.bytes(TILE_SIZE * TILE_SIZE * 4)  # 64KB per tile

    ITERS = 1000
    start = time.perf_counter()
    for i in range(ITERS):
        baker._ring_offset = 0
        baker._write_to_ring(tile_data)
    elapsed_s = time.perf_counter() - start
    throughput_gbs = (len(tile_data) * ITERS) / elapsed_s / (1024 * 1024 * 1024)
    print(f"\n  Ring buffer throughput:   {throughput_gbs:.2f} GB/s")
    status = "OK" if throughput_gbs > 1.0 else "SLOW"
    print(f"  (target > 1.0 GB/s)      [{status}]")

finally:
    shm.close()


print("\nDone.")