Gaussian Splatting WebGPU (Math blog post and DEMO)

June 5, 2024 · View on GitHub

I still remember the first time I saw photogrammetry in action. The artist takes a lot of photos of a real object, plugs them into software, and receives a 3D object.

In 2020, "NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis" introduced neural radiant fields. It uses a neural network to learn a 5D function (position x, y, z, viewing angle θ, ϕ) that returns a color and density/opacity. Watch the "Why THIS is the Future of Imagery (and Nobody Knows it Yet)" by Corridor Crew. Only 2 years later, Nvidia released instant-ngp that speeds up the process. Along the way, we learned that neural networks are not even needed for this task - Plenoxels.

This leads us to "3D Gaussian Splatting for Real-Time Radiance Field Rendering" (August 2023). The model learns a point cloud that represents a scene. Each point is a Gaussian (imagine spindle or fusiform). With enough of them, we can represent the objects with stunning accuracy.

In this repo, we have a WebGPU-based online renderer for a Gaussian splat file. I have written a dedicated blog post "Notes for 3D Gaussian Splatting renderer" explaining the math behind this app.

WebGPU is so new it's not available in Firefox. Use Chrome instead.


gaussian-article-main-img

Nike ZoomX Vaporfly Next% model by Alex Carlier.

Features

Usage

WebGPU does not work on Firefox. On Chrome, it requires HTTPS even during development.

  1. openssl req -newkey rsa:2048 -new -nodes -x509 -days 3650 -keyout key.pem -out cert.pem. Move both files to ./static.
  2. yarn install.
  3. yarn dev. Start a dev server that watches and compiles the code into ./build. It also copies stuff from ./static.
  4. yarn serve in a separate terminal. Starts http-server. It's not included in package.json, but let's be honest - you already have it installed globally.

Or yarn build for prod build.

Camera control

Use the [W, S, A, D] keys to move and [Z, SPACEBAR] to fly up or down.

Running unit tests

Node.js does not support WebGPU. Deno does (since version 1.39, December 14, 2023 - 5 months ago). Internally, it uses Firefox's wgpu. It's not amazing: "WebGPU is still considered unstable in Deno". Compute is OK (with unit tests). Rendering struggles.

Instruction:

  1. Download the .zip file from deno/releases.
  2. "<path-to-unzipped-deno>/deno.exe" task test

FAQ

What are the supported file formats?

Only the .splat from antimatter15/splat. Writing file loaders is boring.

How complicated is the renderer?

  1. Load the .ply or .splat file.
  2. Display it as a point cloud. Trivial using instances=splatCount with 4 vertices and triangle fan primitive. Then billboard it in a vertex shader.
  3. Implement depth sorting (on CPU to make it easier). The Gaussians have transparency, so you need to render closest to the camera first. The simplest option is to generate an index buffer for each frame (no instancing or triangle fan). Add appropriate blend mode. At this point, your model should be recognizable.
  4. Implement Gaussians. You can use my notes as a reference.
  5. Sorting on the GPU. The most popular (but hard to implement) is the radix sort. Instead, I've copied the bitonic sorter example code from Wikipedia. Iterate over k and j on the CPU (or as pre-generated uniform buffers) and over i inside the compute shader.
    1. It could also be possible to kernel-fuse it, but WebGPU gets on my nerves (see below). The WGSL offers limited barriers.
    2. In a non-fused approach, you need to generate 2 more shaders. One to precalculate depths by which you sort. The second one is to generate an index buffer after sorting (turn each splatId into 6 vertices).

What could be some further improvements?

  1. Tiled renderer, like in the original paper.
  2. Spherical harmonics other than l=0;
  3. Blog posts from Aras Pranckevičius:
    1. "Making Gaussian Splats smaller".
    2. "Making Gaussian Splats more smaller".
  4. Other improvements in the field:
    1. awesome-3D-gaussian-splatting.
    2. "Gaussian explosion".

Why does it not work on Firefox?

This app uses WebGPU, a new API for writing 3D applications. Chrome has implemented the specification since May 2023. Firefox is still working on its wgpu. Nothing I can do here.

If you want to render Guassians in the browser, use antimatter15's splat or huggingface's gsplat.js. Both use WebGL, an older but more stable technology.

Is WebGPU good?

The answer will be biased, as I have extensive experience with Vulkan. It's like WebGPU, but 10x more complex. It would take something weird to raise my eyebrows.

The WebGPU API is pleasant to use. It has compute shaders. Let's be honest, this is the killer feature. Unfortunately, this is where the positives end.

  • No Firefox support ATM.
  • The API lacks a lot of things. A few examples:
    • Want to check if the shader compiles correctly synchronously? Nope, GPUCompilationInfo is not supported even in Chrome. Failed compilation operation does not even throw an error. Should I invoke the shader to check if the operation failed?! There exists an asynchronous GPUDevice: pushErrorScope() (see below). But as GPUDevice: createShaderModule() is synchronous, it's a weird choice. Especially as it's common to text replace the content of the shaders before compilation. If I was given an error the moment the compilation failed, I could write a better error message.
    • Funny edge case: you can profile only on a per-render/compute-pass basis. Let's say you have a bitonic sort implementation that invokes shader 90 times. This leads to 180 timestamps (start + end). Oops! There was writeTimestamp() to mark the code regions, but it was removed.
  • Clunky error handling. The errors are returned asynchronously, which is fine. ATM you sandwich the code between GPUDevice: pushErrorScope() and GPUDevice: popErrorScope(). Here is how it works in practice:
    • Imagine you create a buffer on the GPU. You do some calculations on it. Read it back to the CPU. It is filled with 0. Ahh, so obvious You forgot the GPUBufferUsage.MAP_READ flag! But, wait. The GPUBufferUsage.MAP_READ is not valid with most of the other combinations of GPUBufferUsage. This can happen to any buffer. The error message does not even contain the label of the GPUBuffer. BTW. How granular are your error scopes?
    • Each unit test has to call popErrorScope() before EVERY assertion. Otherwise, if the assertion fails, the test stops. No WebGPU logs in that case. In Deno, you could try to circumvent this with 'beforeunload' as there is no Jest's afterEach(). The whole mechanism is clunky.
    • In Vulkan, activate validation layers to get tons of high-quality logs. You can even do this from the external program. No need to recompile anything. Imagine having a browser dev tool tab that can (on the developer's request) hook and synchronously display the messages.
  • No tooling. In JS/other languages, I rarely use debuggers. But WebGPU works on the GPU. First, you have to check if the values from the CPU were received correctly. Then you need to e.g. select which pixel to debug. The fullscreen triangle on the 1024x720px image is 737280 fragment shader calls. Sure, Deno helps with unit tests. But you know what is better? RenderDoc, like I've used when writing Vulkan code.
    • Language servers for the code editor? Forget. The ones available are missing a lot of WGLS features e.g. overrides.
  • WGSL - WebGPU Shading Language a Rust-like language for writing shaders. Imagine if *Web*GPU used syntax similar to.. I don't know.. JS?! function instead of fn, const and let instead of let and var (yes, WGSL has this reversed wrt. to JS), : for return type instead of ->, etc. I love it when someone's pet project spills into formal specifications. It says a lot about the authors and the community. Ofc. it's not like there already was GLSL for WebGL and OpenGL. Have fun rewriting all the shaders when porting the apps!

I would not say that WebGPU is dead on arrival. Compute shaders are too strong of a value proposition. Fortunately, given the massive amount of bugs (Chrome is just slightly better than Firefox), and wgpu velocity, the 'arrival' is still years away. The fact that the dev experience is hostile is an icing on the cake.

References