Async snapshot API for efficient video recording

[hyiip]

Hi! I've been using pygfx (via fastplotlib) for scientific visualization and needed to record animations to video. I noticed that the current snapshot() method blocks on queue.read_texture(), which creates a serial pipeline where GPU and CPU can't overlap.

Benchmark

I experimented with double-buffered staging buffers using map_async() + sync_wait():

GPU: RTX 4080 SUPER
Frame: 1920x1080, 60 frames

Approach                  Time         FPS        Speedup
--------------------------------------------------
Sync (current)            654.3ms      91.7       1.00x
Double-buffered async     396.4ms     151.4       1.65x

How it works

Sync (current):
Frame 0: [Render][Copy][WAIT████][Read]
Frame 1:                              [Render][Copy][WAIT████][Read]

Double-buffered:
Frame 0: [Render][Copy to Buf0]──────────────────[wait+read Buf0]
Frame 1:        [Render][Copy to Buf1]──────────────────[wait+read Buf1]
                       ↑
         GPU works on N while we read N-1

The key insight: map_async() returns immediately. By the time we call sync_wait() for frame N-1, the GPU has already finished the copy, so it returns almost instantly.

Proof of Concept

Here's a standalone benchmark (no pygfx changes needed):

# See gist: https://gist.github.com/hyiip/079b52198f2bcb8251987e1be295e790

Full benchmark code

"""Async Double-Buffered Snapshot Benchmark for pygfx"""
import wgpu
import numpy as np
import time

def align_to_256(value):
    return (value + 255) // 256 * 256

adapter = wgpu.gpu.request_adapter_sync(power_preference="high-performance")
device = adapter.request_device_sync()

WIDTH, HEIGHT = 1920, 1080
BYTES_PER_PIXEL = 4
NUM_FRAMES = 60

aligned_bytes_per_row = align_to_256(WIDTH * BYTES_PER_PIXEL)
buffer_size = aligned_bytes_per_row * HEIGHT

source_texture = device.create_texture(
    size=(WIDTH, HEIGHT, 1),
    format=wgpu.TextureFormat.rgba8unorm,
    usage=wgpu.TextureUsage.COPY_SRC | wgpu.TextureUsage.RENDER_ATTACHMENT,
)

def create_staging_buffer():
    return device.create_buffer(
        size=buffer_size,
        usage=wgpu.BufferUsage.COPY_DST | wgpu.BufferUsage.MAP_READ,
    )

def copy_texture_to_buffer(buf):
    encoder = device.create_command_encoder()
    encoder.copy_texture_to_buffer(
        {"texture": source_texture, "mip_level": 0, "origin": (0, 0, 0)},
        {"buffer": buf, "offset": 0,
         "bytes_per_row": aligned_bytes_per_row, "rows_per_image": HEIGHT},
        (WIDTH, HEIGHT, 1),
    )
    device.queue.submit([encoder.finish()])

# Sync benchmark
staging = create_staging_buffer()
t0 = time.perf_counter()
for _ in range(NUM_FRAMES):
    copy_texture_to_buffer(staging)
    staging.map_sync(mode=wgpu.MapMode.READ)
    data = staging.read_mapped()
    staging.unmap()
sync_time = time.perf_counter() - t0

# Async double-buffered benchmark
bufs = [create_staging_buffer() for _ in range(2)]
idx, pending = 0, None
t0 = time.perf_counter()
for frame in range(NUM_FRAMES + 1):
    if pending:
        pending[0].sync_wait()
        data = pending[1].read_mapped()
        pending[1].unmap()
    if frame < NUM_FRAMES:
        copy_texture_to_buffer(bufs[idx])
        pending = (bufs[idx].map_async(mode=wgpu.MapMode.READ), bufs[idx])
        idx = 1 - idx
async_time = time.perf_counter() - t0

print(f"Sync: {sync_time*1000:.1f}ms | Async: {async_time*1000:.1f}ms | Speedup: {sync_time/async_time:.2f}x")

Questions

1. Would pygfx be interested in an async snapshot API? Something like:

   handle = renderer.snapshot_async_start()
   # ... render next frame ...
   data = renderer.snapshot_async_get(handle)  # fast - copy already done

2. Does this align with the async direction discussed in Support async use of wgpu? #615?

3. Is WgpuRenderer the right place, or would a separate helper class be preferred?

Happy to contribute a PR if there's interest!

Related

• Record Figure to a video file robustly fastplotlib/fastplotlib#918 - Record Figure to a video file robustly
• expose on-demand render mode fastplotlib/fastplotlib#765 - Expose on-demand render mode
• Support async use of wgpu? #615 - Async wgpu support

Thanks for the great library!

[almarklein]

The double-buffer approach is more or less what's implemented in the WIP PR pygfx/rendercanvas#138. I imagine that once we have that, it would also be possible to some a similar mechanism to record the full window of normal canvas.

Would that cover your particular use-case?

Would pygfx be interested in an async snapshot API? Something like:

Yes, that would still probably make sense. But we should probably wait until wgpu and rendercanvas have a next release that have improved async support. Then we can do

image = await renderer.snapshot_asyc()
# or
render.snapshot_async().then(handle_image_func)

Does this align with the async direction discussed in #615?

Yes, eventually I hope to remove all sync calls in Pygfx.

Is WgpuRenderer the right place, or would a separate helper class be preferred?

The renderer holds the internal textures (color and depth), so it seems like the correct object to ask for a snapshot.