Backend: Browser (remote)

[almarklein]

Would be interesting to have a generic browser backend. A bit like jupyter_rfb but without Jupyter. It may also make it easier to test / benchmark new techniques.

Overview in short

The simplified and naive approach:

• Context must present as image.
• Canvas backend sends image over network.
• Browser presents the image.

More detailed steps and notes on performance

                                                                                     
  ┌───────────────┐                                                                  
  │Compress/reduce│      download                                                    
  │   on GPU      │─────┐                                       ┌──────────────────┐ 
  └───────────────┘     │                                       │Present via       │ 
                        │                                       │ <img> or <canvas>│ 
                        ▼                                       │                  │ 
                     ┌───────────────┐              ┌──────┐    └──────────────────┘ 
                     │Compress/encode│              │Decode│            ▲            
                     │   on CPU      │ ───────────► │      │ ───────────┘            
                     └───────────────┘              └──────┘                         

See also pygfx/wgpu-py#378

We can also use a custom present-method, so that the image can be encoded. That broadens the options quite a bot. Assuming rendering with wgpu here:

• Encoding the image on the GPU (optional):
  • The purpose is to reduce the image size (in bytes), so the next steps are faster.
  • To make this step fast, we have to stick to encodings that parallelize well.
  • There are compressed texture formats in wgpu.
  • Can use YUV (yuv420 results in 37.5% the size compared to rgba)
  • JPEG encoding with variable quality
  • Diff images
  • mp4 encoding
• Downloading from GPU into RAM:
  • The bottleneck is the amount of bytes.
  • For small images, as typical in a notebook, this step does not cost a lot of time.
  • For larger images, going towards full screen, the impact becomes very significant (when downloading raw rgba data)
  • A lot of time is spent on waiting for the copy-buffer to become mapped. An async approach would matter a lot.
  • Also see the little benchmark in Refactor canvas context to allow presenting as image wgpu-py#586
• Encoding the image on the CPU:
  • Can be done by the context, or in canvas._rc_present_xx(), or a bit in both.
  • Let's assume we have a binary websocket, but if we don't than base64 encoding is part of this step too.
  • jupyter_rfb does png and jpg encoding.
  • Making the image size smaller will make the next step faster.
  • But the encoding also greatly affects the presentation in the browser.
• Sending over the network, io:
  • The speed of this step is pretty good if you're on localhost.
  • Otherwise, the speed can vary a lot, depending on where the server and browser are.
  • Speed can vary over time too.
  • It would be nice to measure this speed and use it to control the encoding steps above.
• Presenting in the browser:
  • This step includes any necessary decoding. Browsers have builtin support for some decoding mechanisms.
  • jupyter_rfb sets the src of an <img> object to the base64 encoded png/jpg. This is surprisingly fast.
  • Using a canvas. IIRC drawing an image to a canvas is pretty slow.
  • Using WebGL / WebGPU.

There are multiple paths to go through these steps. What's tricky is that a choice in one step affects the other steps, so
we cannot simply pick "the best" from each step and stack them together.

MP4 has nice properties in being pretty good in a situations with low/variable bandwidth, and it very naturally improves the resul over time for static images.

Something like WebCRT can possibly be used to implement some of these steps.

Some unsorted thoughts and links

• Partial jpg encoding/decoding on the GPU: https://github.com/negge/jpeg_gpu
• Can also run deflate on the GPU (GDeflate), so basically encode to png.
• Simply down-sampling the image at the GPU will reduce the size by a factor 4, which affects all steps. For moving images this is barely noticeably (because by default our images are 2x larger for aa). We could even go for a factor of 9 or 16 in low-latency cases.
• Jpeg encoding is a lot faster and smaller than png.
• Formats like webp and avif have good compression, which help with the file size, but they are rather slow.
• Advanced solutions like webCRT are cool, but maybe heavy on the dependencies.

Open questions

• Can we jpeg-compress on the GPU?
• How good are compressed texture formats?
• How do game streaming platforms do this?

[almarklein]

I've been doing some research to image file formats in almarklein/ppng#2. Some conclusions relevant to this topic:

• Compressing with jpeg is a lot faster than png, and produces smaller blobs (and easy to tune size). But we already know this.
• With png compression, the size of the blob depends heavily on how well the image can be compressed (i.e. the content of the image).
• The above actually holds for lossless formats (like jpg) as well, but to a lesser degree.
• Compression with (lossy) WebP can produce smaller blobs than jpeg, but its much slower.
• When sending png, dropping the filter and lowering the compression help speed a lot, with relatively small effect on the file size.

One viable idea might be to zlib-compress the image on the GPU (people have done this) and then sending as png. The downside is that if the image has high entropy, it does not compress well. So if possible, I'd rather move towards jpeg/mpeg compression.

[almarklein]

I got a tip from @hmaarrfk to align the buffer to 4096 when downloading from the GPU, which ought to result in a faster download. I have yet to try this.

[almarklein]

I got a tip from @hmaarrfk to align the buffer to 4096 when downloading from the GPU, which ought to result in a faster download. I have yet to try this.

I had a look at this, but we don't allocate the data for the buffer to copy to. The wgpu API does that:

        src_ptr = libf.wgpuBufferGetMappedRange(self._internal, offset, size)
        src_address = int(ffi.cast("intptr_t", src_ptr))

Interestingly, the src_address always seems aligned to 4096 😉

I also tried padding the buffer size to make it a multiple of 4096, but I did not observe a difference in performance.

[almarklein]

I spent some time investigating possible ways to improve the speed by which we can download data from the GPU.

From what I read, the steps that we do are the correct ones:

• Allocate (or re-use) a buffer large enough to contain the pixel data.
• Use command_encoder.copy_texture_to_buffer() to copy the data over.
• Map the buffer (async op).
• Wait for it to be mapped ⏳
• Read the mapped data, and copy to a new array (the mapped data is released when the buffer is unmapped).
• Unmap the buffer.

Re-using the target texture and the buffer to copy to, helps significantly. I also found that one data-copy can be avoided in some cases (when bytes-per-row is not a multiple of 256). And Using numpy to copy the data to avoid a for-loop and enabling allocating an empty array (instead of zeros).

Together these bring the FPS for full screen apps from 24-ish to 30-ish (on a 4k screen).

I also found a much bigger way to improve performance. The part where we wait for the buffer to be mapped is (apparently) a significant portion of the total spent time. As an experiment, I mapped the buffer, but then wait for it to be mapped in the next frame, and then read data from it. FPS goes up to 50+ 🚀 Sweet! However, this by itself is not a viable approach, because in cases like the notebook, where the network can be the bottleneck, being one frame behind may become noticeable. What it means is that we need some form of async / callbacks here, if we want to have this kind of fps for rendering via a bitmap. And we need some more work on this front. Nevertheless, this is promising!

[almarklein]

I also tried doing the copy_texture_to_buffer in the Pygfx flusher, using the command_encoder that's already doing the flushing, to avoid creating a new encoder and doing an extra queue.submit().

[kushalkolar]

@apasarkar and I have been making some progress on JPEG encoding, I think we'll have a webgpu jpeg encoder soon!

Some notes and things to think about:

• By jpeg encoding on the GPU we also significantly reduce the amount of data that has to be downloaded from the GPU.
  • Caveat, the number of bytes that represent a jpeg image are not consistent between frames. AFAIK, in wgpu you can't download an arbitrary sized buffer from the GPU to the CPU?
  • Naive idea to get around this: for a given quality level, set the max number of bytes allowed based on the most complex expected frame based on testing the size for a wide variety of plots?
  • Another idea: chunk it into multiple large buffers, download the buffers and see if a stop byte is at the end of any buffer, if so don't download the later chunks.
• On the client end, I'm using anywidget which is amazing. Started prototyping webgpu canvases with anywidget. This will give us the best flexibility for use in the browser across as many notebook platforms as possible, including marimo which has native anywidget support.
  • Questions: What is the fastest way to send binary data from the server, get that on the js end, and upload it to the GPU for jpeg decoding?
  • Some general ideas here maybe: https://toji.dev/webgpu-best-practices/buffer-uploads.html
  • there are some jpeg decoders people have written on github if we need ideas: https://github.com/kushalkolar/Compeg, https://github.com/tamani-coding/webgpu-computing-jpeg

[almarklein]

FAIK, in wgpu you can't download an arbitrary sized buffer from the GPU to the CPU?

Yeah, you can download any size, except for some alignment constraints (nbytes must be a multiple of 8 or 16). But you have to specify it on the CPU, which (I assume) does not know the actual size.

You could try encoding white noise and see how large that becomes as a max. That could provide a sensible upper bound, but would be nice if images that encode well are also fast.

Downloading in multiple steps does result in extra wait time, but this may not be so bad, since you can start sending the first chunk over the network while the GPU is downloading the next chunk. Steaming! Definitely need benchmarks to test what works well and what not.

Questions: What is the fastest way to send binary data from the server, get that on the js end, and upload it to the GPU for jpeg decoding?

Binary websocket. Then probably queue.write_buffer. Or device.createBuffer(..., mapped_at_creation=True).

Mmm, if we stream the encoded data in chunks, the client does not know the size until the last chunk has arrived. In that case it would be good if there is a big-enough buffer on the GPU, and just send chunks to it with queue.write_buffer.

On the GPU, after encoding, what does the data look like? Is it stored in a buffer or texture? Does it already naturally come in chunks of some sort? I'm half wondering if the client can start processing chunks once the first arrives 😉

[kushalkolar]

Thanks for the thoughts! Will keep updating.

On the GPU, after encoding, what does the data look like? Is it stored in a buffer or texture?

JPEG encoded output is just raw bytes, so I imagine the file output would just be from a storage buffer.

Meanwhile, almost have DCT working in wgsl 🥳 . Need to debug why the basis weights are not in the right place, keeping track of [row, col] -> [x, y] in wgsl is tedious 😅 . Thank goodness we have numpy, signal processing with a low level language is always tedious, last time I did this was in C and ARM assembly; at least wgsl is easier to debug because sometimes ARM assembly will just spew out random unicode errors 🤦

EDIT: apprently I can't write a basic for loop, I've been spoiled by writing loops with python's enumerate, and itertools. And textureLoad() does not raise if you try to load a texture from an index that is beyond the size of that dimension 🤦

and we have DCT in wgsl 🥳 🚀

prototyping here: https://github.com/kushalkolar/webgpu-rfb-prototype

[kushalkolar]

Made some more progress on this and finally started templating. Wow texture fetches are expensive, 35% faster if we just dump the DCT basis as a const.