Rendering 29,000 videos in real-time with Rust & WGPU

Revisiting the problem

"An image, but every pixel is actually a video."

A while ago, I wrote about my initial version of a videomosaic renderer, written in Python / PyGame, which is a framework that primarily focuses on software rendering (rendering on the CPU). Back then, I managed to achieve interactive framerates by reinventing a lot of what a GPU is already designed to do from first principles, so naturally this begs the question - why not just use a GPU instead?

While that was more of an experimental project rather than a by-the-book implementation, I've recently been learning Rust and the wgpu crate - and I thought it would be a great learning experience to try and actually explore what an answer to that question would look like.

Let's refresh our memory on how much data we're going to need to deal with:

• The above image uses 29,000 distinct clips from 121 sitcom episodes.
• Even though the average clip is only 1.7 seconds long, that's still around 14 hours of video!
• Even with chroma subsampling, we'd need 1,750 gigabytes of VRAM to store it all!

The solution? We need to stop thinking like a video player, and start thinking like an operating system.

Virtual Reality

On a modern OS, the address space for RAM is 2^(48) bytes, far more than any modern computer has in actual RAM (especially with these prices!) - but the OS, for a variety of reasons related to security and performance, frees the process from having to worry about where its memory is actually physically located on the RAM chip.

Instead, when the process asks for memory, the OS hands it an address to a location in virtual memory address space, which can be sparsely distributed (large distances between one allocation and the next). When the program tries to read or write to this address, the OS uses one or more page tables to map a virtual address to a physical address in RAM.

You can think of it similar to how a hashmap lets you store keys 0 and 10,000 together without needing to worry about all the other ones in between. Indeed, a hashmap was the data structure I used last time to solve this problem - but GPUs don't play well with hashmaps.

So let's zoom out a bit.

Test image - a promo image for Hotel Transylvania 2 matched against the entirety of Parks & Recreation
(note the banding in the grass and clown's hair)

In this image, we're using 29197 distinct colours. Let that be our virtual address space. A key realisation we made last time is that we don't necessarily need to store the full resolution of each tile just because it's in the image. When we're zoomed out, we can actually render the entire image just using a 171x171 palette texture, along with a texture that maps each pixel from the original image to a palette colour.

On the left, a black and white image where each pixel value corresponds to a colour from the palette texture
(note the 29196 distinct levels)

Why does it look like that?

For the sake of visualisation, I sort the tiles by colour (r << 16 | g << 8 | b) beforehand.
Otherwise, they look like this...

As we zoom in, the number of tiles on screen naturally get smaller, but the tiles themselves get bigger. This inverse correlation means there's only ever going to be so many pixels on screen at once, regardless of tile size or how many tiles there are on screen.

We can therefore express how many pixels we could possibly need to store on the GPU for all zoom cases mathematically: for a given maximum tile size (here it's 1024x), and any given video frame, we're going to need to support storing (ceil(monitor_res / 1024) + 1) * 1024 pixels. Divide monitor resolution (not just viewport, because it can be resized) by maximum tile resolution, take the ceiling (we can't store half a tile), and then add 1 to account for panning in-between tiles.

For a 1920x1080 monitor, we'll need 3072x3072 pixels of data. Using chroma subsampling, that means we only need to actually allocate 13.5 megabytes of VRAM to store our tiles for any given rendered frame. Let this be our physical memory - we'll store this on the GPU in an atlas-like texture.

Now we have a problem. Our address space is 3,244x larger than the no. of supported tiles at full resolution. We can't simply give each tile a fixed position in the atlas according to its index - we'd run out of space starting at 18x18 tiles. We need to position each tile dynamically, according to which ones are currently on screen, and then create a separate reference texture our GPU will consult to locate the tile for that frame. That reference texture will be our page table.

Let's explore what an implementation of this architecture would actually look like - starting with figuring out what's on screen to begin with.

are you seeing what i'm seeing?

Because of our limited atlas space, we need to figure out which unique tile IDs are on screen and stream just those frames into the atlas. The simplest solution (and indeed, the one I used last time) would be to just slice the visible section of the mosaic every time the camera moves, and deduplicate. This isn't very fast, though - and it gets quadratically slower as resolution increases!

Instead, we can exploit the fact that, in a given frame, the user is most likely not going to pan the entire screen's width worth of pixels. The actual difference between frames is usually a fraction of the visible rectangle, and we can diff and process just the new areas instead for substantially less work.

While processing the new regions, we can check against a set of visible tiles to track which tiles are actually on screen, and to identify exactly when a new tile becomes visible, vs. just instances of ones that are already on screen (for example, panning through colours that all map to the same tile). Instead of re-uploading the atlas every frame just in case, now we have an accurate signal as to when we should refresh the atlas.

Before we can update the atlas, though, we first need to stream in the content that will be uploaded to it - which means somehow fitting all of these videos into RAM...

pyramid scheming

We can think of our streaming problem in terms of levels of detail. When we're zoomed out, there are many tiles on screen - but they're not very big! Needing to store a full resolution video for every tile would quickly become untenable.

Instead, we store our tiles as a pyramid of mip-maps. When we're zoomed out, we only load the videos in at a small resolution, like 64x64. When we zoom in, since we know which specific tiles are actually visible, we can load the high-res version of just those tiles, on demand.

What about the ones that go off-screen though? Ideally, we don't want to waste the work by immediately deleting them, especially if the tile is about to reappear on screen later (e.g. a repeating colour pattern). Instead, we keep a cache of videos in CPU-side RAM, and only "evict" tiles when we're truly constrained for space.

To keep the renderer responsive, we delegate the work of actually streaming in the videos to a pool of worker threads (made significantly easier by the fantastic Rayon crate). A background thread receives messages from the main thread, figures out which tiles need to be streamed in & which need to be evicted to make room, and then dispatches the necessary jobs for the tiles to stream in.

We can then remain responsive by rendering a fallback (in this case, the colour of the tile from the palette) instead, and "popping" the tile into view once we get a message that it's loaded in.

An example of panning around, showing some tiles popping in and some tiles hitting cache and appearing seamlessly

Staying up-to-date

Now that our tiles are present and accounted for, we can finally take the last step and upload them to the GPU. Luckily, this part is refreshingly easy - but with a little twist.

Every time the atlas is "dirty" (either because camera state changed, new tiles have been streamed in, or video playback advances a frame), we simply iterate over the visible set of IDs, copy them over from RAM, write them into the atlas as we go and record their position in our page table.

The catch? Recall our textures are 3072x3072x1 + 1536x1536x2 (we use chroma subsampling to save RAM and bandwidth). That's a 13.5 megabyte copy on the CPU. Doing a copy of this size is possible but also wasteful, reduces framerates, and also gets quadratically worse as the resolution increases.

We can again exploit a property of the mosaic to get a speed boost. Because we're reusing so many colours (29k colours over 2 million pixels, each one is reused an average of 71 times), the actual amount of atlas usage is usually a fraction of its total maximum theoretical size. Therefore, we don't need to copy and upload the whole thing. We can, instead, simply write a slice of the texture and upload just that slice to the GPU instead. This is what it looks like:

Atlas luma + chroma plane

Note how that laziness manifests in these debug textures. The lower resolution mips overwrite the higher resolution ones but cut off where the visible set iteration ends. That memory isn't touched by either the GPU or the CPU until necessary, greatly reducing the number of pixels we send over the PCIE bus in some cases.

What does the page table look like?

The page table and colour palette textures are both the same size, and indeed refer to the same tiles. For a pixel in the page table, that position is the same position that tile's colour will be in the palette.

This means a lookup can fall back on the same coordinates. Every "world space" mosaic pixel, is looked up in the page table first. If it's 0, fall back to the same position in the palette. Anything greater than 0 is the position in the atlas, modulo current tile mip resolution.

It's also written to lazily - I don't go back and write a 0 for a tile that goes off screen. By nature, it won't be queried, so we can avoid zeroing it out or extra loop iterations. It does mean that this debug texture is technically inaccurate, but it does look pretty.

Making a match

At this point, we've gone over the entire process a tile takes from RAM to screen. For this final section I want to go into detail as to how the videos are processed and colour-matched to begin with - from disk to RAM.

We start with a standard definition rip of the sitcom. Using FFMPEG Rust bindings, the entire show is processed in a thread pool and decoded with the scaler set to 64x64 (faster than full resolution). Each frame gets its average colour calculated and then saved in a database keyed to the video hash itself, so that starting the program again doesn't reprocess the entire show (which takes around 6 minutes on my machine).

Then, the video is sliced, with the goal of trying to detect cuts or things that would otherwise be deviations. The algorithm is actually pretty simple:

• Pin the colour of the first frame,
• Iterate through the colours until we find a colour such that the mean squared error between it and our pinned colour is greater than some arbitrary threshold (I settled with 300 based on purely eyeballing it),
• Push the index of the pinned colour and the index of the frame before this colour to a list,
• Pin this new colour instead,
• Repeat until the video is over.

A tile is then created with the start and end frame for each entry in that list, and we match against the first frame with the assumption that each frame stays roughly the same colour as the start of the sequence. Increase the threshold, it starts becoming noisy. Decrease the threshold, and the clips get shorter.

From there, the matching is done simply and performantly using a k-d tree from the excellent Kiddo crate. It may not be entirely accurate to how human eyes perceive colour, but there are so many tiles (around 300,000 if threshold is 0, goes down to 100,000 at 300) that the images don't look too dissimilar.

Conclusion

The performance gains against the Python implementation are, of course, substantial. Moving from software -> hardware rendering nets us anywhere from 30-80x faster rendering depending on screen resolution. More manual control over memory means we get less stuttering, less allocator exhaustion, and deterministic control over when memory is freed and allocated. True "fearless" parallelism allows us to write much simpler, direct logic without needing to worry about trying to route it all through layers that release the GIL.

Is it a fair comparison? I wouldn't say so. The goal of both projects was never to be the fastest, most idiomatic implementations of their respective techniques, but rather just to experiment with both software & hardware paradigms and to explore what goes into virtualised rendering & memory. Not to mention, arriving at the right answer the wrong way worked wonders for understanding the foundations of parallelism, GPU programming, memory management and image + video processing.

As such, I look forward to reinventing more wheels in the future.

Here are some other examples:

The Swanson Lisa

With debug textures enabled - note the tiles in the atlas stop moving as the camera zooms out, as that part of the atlas is no longer being written to or read from.

and, of course,

An episode of Parks & Recreation, but each pixel is replaced with a frame from an episode of Parks & Recreation.

Read more about the first time I tried to solve this problem.