diff --git a/.gitignore b/.gitignore index 054e565..8e5c882 100644 --- a/.gitignore +++ b/.gitignore @@ -24,4 +24,6 @@ dist-ssr *.sw? /.vite -*/scenes \ No newline at end of file +*/scenes +scenes +videos \ No newline at end of file diff --git a/README.md b/README.md index edffdaf..8cf52a1 100644 --- a/README.md +++ b/README.md @@ -1,31 +1,193 @@ -# Project5-WebGPU-Gaussian-Splat-Viewer +**Vismay Churiwala** -**University of Pennsylvania, CIS 565: GPU Programming and Architecture, Project 5** +- [LinkedIn](https://www.linkedin.com/in/vismay-churiwala-8b0073190/) | [Website](https://vismaychuriwala.com/) +- Tested on: + - Microsoft Edge 141.0.3537.99 + - Windows 11, AMD Ryzen 7 5800H @ 3.2GHz (8C/16T) + - 32GB DDR4 RAM + - NVIDIA GeForce RTX 3060 Laptop GPU (6GB GDDR6) -* (TODO) YOUR NAME HERE -* Tested on: (TODO) **Google Chrome 222.2** on - Windows 22, i7-2222 @ 2.22GHz 22GB, GTX 222 222MB (Moore 2222 Lab) +# WebGPU Gaussian Splat Viewer -### Live Demo +| [![](images/img/bicycle_final.png)](https://vismaychuriwala.github.io/WebGPU-Gaussian-Splat-Viewer/) | +| :----------------------------------------------------------------------------------------------: | +| Bicycle scene, 1920×1080 | -[![](img/thumb.png)](http://TODO.github.io/Project4-WebGPU-Forward-Plus-and-Clustered-Deferred) +## Table of Contents -### Demo Video/GIF +- [Live Demo](#live-demo) +- [Project Overview](#project-overview) +- [How It Works](#how-it-works) + - [Stage 1: Preprocessing Gaussians](#stage-1-preprocessing-gaussians) + - [Stage 2: Depth Sorting](#stage-2-depth-sorting) + - [Stage 3: Drawing Splats](#stage-3-drawing-splats) +- [Performance Analysis](#performance-analysis) + - [Point Cloud vs Gaussian Renderer](#point-cloud-vs-gaussian-renderer) + - [Workgroup Size Impact](#workgroup-size-impact) + - [View Frustum Culling](#view-frustum-culling) + - [Number of Gaussians](#number-of-gaussians) +- [Bloopers](#bloopers) +- [Credits](#credits) -[![](img/video.mp4)](TODO) +## Live Demo -### (TODO: Your README) +Try it yourself at [vismaychuriwala.github.io/WebGPU-Gaussian-Splat-Viewer](https://vismaychuriwala.github.io/WebGPU-Gaussian-Splat-Viewer/). You'll need to provide your own PLY scene file and camera configuration. -*DO NOT* leave the README to the last minute! It is a crucial part of the -project, and we will not be able to grade you without a good README. +### Demo Video -This assignment has a considerable amount of performance analysis compared -to implementation work. Complete the implementation early to leave time! +https://github.com/user-attachments/assets/193cc355-4029-461c-aeb1-e8b5ea554344 -### Credits +## Project Overview -- [Vite](https://vitejs.dev/) -- [tweakpane](https://tweakpane.github.io/docs//v3/monitor-bindings/) +For this project, I implemented a browser-based gaussian splatting renderer using WebGPU. The goal was to recreate the rasterization pipeline from the [3D Gaussian Splatting paper](https://github.com/graphdeco-inria/gaussian-splatting)—taking pre-trained gaussian data and rendering it interactively. The training and reconstruction aspects weren't part of this assignment. + +Rather than rendering traditional triangle meshes, gaussian splatting works with clouds of 3D Gaussians. Each gaussian is basically an ellipsoid with color information stored as spherical harmonics. During rendering, I transform these 3D shapes into 2D splats on the screen, then blend them together to create the final image. When done right, you get photorealistic results even though there's no explicit geometry. + +I worked from the [original CUDA/C++ implementation](https://github.com/graphdeco-inria/diff-gaussian-rasterization) as a reference, but had to restructure everything for WebGPU. That codebase doesn't use graphics APIs at all—it's written for direct GPU control. Adapting it meant rethinking how data flows through compute shaders, sort passes, and the render pipeline while dealing with WebGPU's binding groups and buffer management. + +Here's an example with the banana dataset from [OpenSplat](https://github.com/pierotofy/opensplat): + +| ![](images/img/banana.png) | +| :--------------------: | +| Banana, 1920×1080 | + +## How It Works + +Every frame, I run three main stages: + +1. **Compute shader preprocessing** - transform gaussians to 2D splats +2. **GPU radix sort** - order splats by depth +3. **Render pass** - draw and blend the splats + +### Stage 1: Preprocessing Gaussians + +The preprocessing compute shader is where most of the heavy lifting happens. Each thread handles one gaussian and decides whether it's visible, then computes all the data needed to render it as a 2D splat. + +**Building the covariance matrices**: I start by reconstructing the 3D covariance matrix from the gaussian's rotation quaternion and scale vector. The rotation gives me an orthonormal basis (via `getR()`), the scale gives me the axis lengths (via `getS()`), and multiplying them produces the world-space covariance. Then I project this into screen space using the Jacobian of the projection transformation—this is how I figure out the 2D shape each gaussian will have on screen. + +**Computing splat parameters**: For each visible gaussian, I calculate: +- The splat's center position in normalized device coordinates +- Its radius (I use 3 standard deviations, covering ~99.7% of the distribution) +- The conic matrix (inverse covariance) for evaluating the gaussian function per-pixel +- Color by evaluating spherical harmonics based on the view direction +- Opacity after sigmoid activation + +**Culling invisible splats**: I implemented frustum culling to skip splats that won't appear on screen. I check if the depth is valid (between 0 and 1) and whether the center is roughly within NDC bounds. I actually use ±1.2 instead of ±1.0 because splats near the screen edge can still contribute pixels even when their centers are slightly outside. Culled splats never make it to the sort or render stages, which saves a ton of work. + +I use atomic counters to track how many splats pass culling—this count becomes the instance count for the render pass. + +#### Optimization: Half-Precision Packing + +One thing I'm proud of is the memory optimization. The original gaussian data is stored as `f16` (half-precision floats), so I realized I could keep most of the splat data at that precision too without noticeable quality loss. Using `pack2x16float`, I stuff two halves into each `u32`. + +The splat structure stores: +- Color as full `f32` (quality matters here) +- Conic matrix + opacity packed into 2 `u32` values +- Center position + radius packed into 2 `u32` values + +This cuts memory usage dramatically and improves bandwidth, which makes both preprocessing and rendering faster: + +| ![](images/img/packvsunpack.png) | +| :--------------------------: | +| Memory layout comparison | + +### Stage 2: Depth Sorting + +Transparency requires back-to-front rendering order, so I sort all visible splats by their view-space depth each frame. During preprocessing, I write each splat's depth and index into separate buffers, which then feed into a GPU radix sort. + +The radix sort runs entirely on the GPU and reorders the splat indices based on their depth keys. I have to do this every frame because camera movement constantly changes the depth ordering. + +### Stage 3: Drawing Splats + +The render pass uses instanced indirect drawing—one quad per visible splat. The indirect buffer gets its instance count from the preprocessing stage's atomic counter. + +**In the vertex shader**: I use the instance index to fetch the sorted splat data. Each splat gets expanded into a quad by unpacking its center and radius, then offsetting four corners appropriately. I pass the unpacked attributes (color, conic, opacity, center) to the fragment shader. + +**In the fragment shader**: This is where I evaluate the actual gaussian. For each pixel, I compute the distance from the splat center, then use the conic matrix to calculate the gaussian's value at that point: + +```wgsl +let power = -0.5 * ( + in.conic.x * d.x * d.x + + in.conic.z * d.y * d.y + - 2.0 * in.conic.y * d.x * d.y +); +``` + +The result is an exponential falloff that makes the splat fade smoothly toward its edges. I multiply by opacity and cap the contribution to avoid numerical issues, then let alpha blending do its job. + +Since splats are already sorted, the blending naturally produces correct transparency and occlusion: + +| ![](images/img/quad_color_mine.png) | ![](images/img/quad_pixeldiff_mine.png) | +| :------------------------: | :----------------------------: | +| Quad boundaries visible | Gaussian falloff applied | + +## Performance Analysis + +All tests run at 1920×1080 on my system (specs above). Test scenes: +- **Bicycle** - 1,063,091 gaussians +- **Banana** - 464,017 gaussians +- **Bonsai** - 272,956 gaussians + +### Point Cloud vs Gaussian Renderer + +| ![](images/img/bonsai_pc.png) | ![](images/img/bonsai.png) | +| :-----------------------------: | :---------------------------: | +| Point cloud | Gaussian splats | + +| Configuration | FPS | +| :------------------------ | :--: | +| Point cloud | 165+ | +| Gaussian | 51 | + +> **Note**: FPS capped at 165 Hz (monitor refresh rate). + +Point cloud uses native point primitives. Gaussian renderer draws instanced quads, evaluates per-pixel falloff, and requires full preprocessing pipeline (covariance computation, projection, sorting). + +### Workgroup Size Impact + +| Workgroup Size | FPS | +| :------------: | :-: | +| 64 | 40 | +| 128 | 73 | +| 256 | 51 | + +128 threads performs best. Likely due to balance between memory coalescing and GPU occupancy. 256 may have reduced occupancy, 64 underutilizes memory bandwidth. + +### View Frustum Culling + +![](images/img/frustum.png) + +| Configuration | FPS | +| :------------------- | :-: | +| No frustum culling | 53 | +| With frustum culling | 97 | + +1.8× speedup. Skips ~43% of gaussians outside view frustum. Reduces sorting overhead and preprocessing work (expensive covariance projection). + +### Number of Gaussians + +| Scene | Gaussian Count | FPS | +| :------ | :------------: | :-: | +| Banana | 464,017 | 117 | +| Bonsai | 272,956 | 153 | +| Bicycle | 1,063,091 | 51 | + +Performance scales roughly with gaussian count. Bicycle has 4× more gaussians than banana but <0.5× FPS due to increased memory bandwidth, larger buffer allocations, and radix sort overhead. + + +## Bloopers + +| ![](images/bloopers/wrong_color.png) | ![](images/bloopers/no_flip_cov_x.png) | ![](images/bloopers/bicycle_wrong_sigmoid.png) | +| :----------------------------------: | :------------------------------------: | :--------------------------------------------: | +| Multiple bugs at once | After fixing color computation | After fixing covariance sign | + +Left: Wrong SH indexing + covariance sign error + missing sigmoid on opacity. Middle: Fixed SH colors, still have covariance orientation bug causing spiky artifacts. Right: Fixed covariance, but opacity still in log-space (missing sigmoid). + +## Credits + +- [Vite](https://vite.dev/) +- [Tweakpane](https://tweakpane.github.io/docs/v3/) - [stats.js](https://github.com/mrdoob/stats.js) - [wgpu-matrix](https://github.com/greggman/wgpu-matrix) -- Special Thanks to: Shrek Shao (Google WebGPU team) & [Differential Guassian Renderer](https://github.com/graphdeco-inria/diff-gaussian-rasterization) +- Shrek Shao from the Google WebGPU team +- The [differential gaussian rasterization](https://github.com/graphdeco-inria/diff-gaussian-rasterization) reference implementation diff --git a/images/bloopers/bicycle_wrong_sigmoid.png b/images/bloopers/bicycle_wrong_sigmoid.png new file mode 100644 index 0000000..0776094 Binary files /dev/null and b/images/bloopers/bicycle_wrong_sigmoid.png differ diff --git a/images/bloopers/no_flip_cov_x.png b/images/bloopers/no_flip_cov_x.png new file mode 100644 index 0000000..84ed8b6 Binary files /dev/null and b/images/bloopers/no_flip_cov_x.png differ diff --git a/images/bloopers/wrong_color.png b/images/bloopers/wrong_color.png new file mode 100644 index 0000000..9dad67c Binary files /dev/null and b/images/bloopers/wrong_color.png differ diff --git a/images/img/banana.png b/images/img/banana.png new file mode 100644 index 0000000..494e253 Binary files /dev/null and b/images/img/banana.png differ diff --git a/images/img/bicycle_final.png b/images/img/bicycle_final.png new file mode 100644 index 0000000..5d17b3b Binary files /dev/null and b/images/img/bicycle_final.png differ diff --git a/images/img/bonsai.png b/images/img/bonsai.png new file mode 100644 index 0000000..42de0fa Binary files /dev/null and b/images/img/bonsai.png differ diff --git a/images/img/bonsai_pc.png b/images/img/bonsai_pc.png new file mode 100644 index 0000000..5e34384 Binary files /dev/null and b/images/img/bonsai_pc.png differ diff --git a/images/img/frustum.png b/images/img/frustum.png new file mode 100644 index 0000000..b981af9 Binary files /dev/null and b/images/img/frustum.png differ diff --git a/images/img/packvsunpack.png b/images/img/packvsunpack.png new file mode 100644 index 0000000..fe14919 Binary files /dev/null and b/images/img/packvsunpack.png differ diff --git a/images/img/point_cloud.png b/images/img/point_cloud.png new file mode 100644 index 0000000..6523b76 Binary files /dev/null and b/images/img/point_cloud.png differ diff --git a/images/img/quad_color_mine.png b/images/img/quad_color_mine.png new file mode 100644 index 0000000..b5bcc5e Binary files /dev/null and b/images/img/quad_color_mine.png differ diff --git a/images/img/quad_pixeldiff_mine.png b/images/img/quad_pixeldiff_mine.png new file mode 100644 index 0000000..80d492d Binary files /dev/null and b/images/img/quad_pixeldiff_mine.png differ diff --git a/package-lock.json b/package-lock.json index 04843bd..06ff36e 100644 --- a/package-lock.json +++ b/package-lock.json @@ -12,6 +12,7 @@ "@loaders.gl/ply": "^4.2.2", "@petamoriken/float16": "^3.8.7", "tweakpane": "^3.1.8", + "tweakpane-plugin-file-import": "^0.2.0", "wgpu-matrix": "^3.2.0" }, "devDependencies": { @@ -379,6 +380,7 @@ "version": "4.2.2", "resolved": "https://registry.npmjs.org/@loaders.gl/core/-/core-4.2.2.tgz", "integrity": "sha512-d3YElSsqL29MaiOwzGB97v994SPotbTvJnooCqoQsYGoYYrECdIetv1/zlfDsh5UB2Wl/NaUMJrzyOqlLmDz5Q==", + "peer": true, "dependencies": { "@loaders.gl/loader-utils": "4.2.2", "@loaders.gl/schema": "4.2.2", @@ -602,6 +604,7 @@ "version": "3.1.10", "resolved": "https://registry.npmjs.org/tweakpane/-/tweakpane-3.1.10.tgz", "integrity": "sha512-rqwnl/pUa7+inhI2E9ayGTqqP0EPOOn/wVvSWjZsRbZUItzNShny7pzwL3hVlaN4m9t/aZhsP0aFQ9U5VVR2VQ==", + "peer": true, "funding": { "url": "https://github.com/sponsors/cocopon" } diff --git a/src/main.ts b/src/main.ts index 25efcb5..4c57c0a 100644 --- a/src/main.ts +++ b/src/main.ts @@ -19,7 +19,7 @@ import { assert } from './utils/util'; } const device = await adapter.requestDevice({ - requiredLimits: { + requiredLimits: { maxComputeWorkgroupStorageSize: adapter.limits.maxComputeWorkgroupStorageSize, maxStorageBufferBindingSize: adapter.limits. maxStorageBufferBindingSize }, diff --git a/src/renderers/gaussian-renderer.ts b/src/renderers/gaussian-renderer.ts index 1684523..87d0eb7 100644 --- a/src/renderers/gaussian-renderer.ts +++ b/src/renderers/gaussian-renderer.ts @@ -5,7 +5,7 @@ import { get_sorter,c_histogram_block_rows,C } from '../sort/sort'; import { Renderer } from './renderer'; export interface GaussianRenderer extends Renderer { - + setGaussianMultiplier: (value: number) => void; } // Utility to create GPU buffers @@ -36,6 +36,50 @@ export default function get_renderer( const nulling_data = new Uint32Array([0]); + const quadVerts = new Float32Array([ + -0.5, -0.5, + 0.5, -0.5, + -0.5, 0.5, + 0.5, 0.5, + ]); + const quadVBO = createBuffer( + device, + 'quad vertices', + quadVerts.byteLength, + GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST, + quadVerts + ); + + const drawArgsData = new Uint32Array([4, pc.num_points, 0, 0]); + const indirect_buffer = createBuffer( + device, + 'indirect buffer', + drawArgsData.byteLength, + GPUBufferUsage.INDIRECT | GPUBufferUsage.COPY_DST, + drawArgsData + ); + + const splatBufferSize = pc.num_points * 8 * Float32Array.BYTES_PER_ELEMENT; + const splat_buffer = createBuffer( + device, + 'gaussian splats', + splatBufferSize, + GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST + ); + + const renderSettingsBlock = new ArrayBuffer(16); + const renderSettingsF32 = new Float32Array(renderSettingsBlock); + const renderSettingsU32 = new Uint32Array(renderSettingsBlock); + renderSettingsF32[0] = 1.0; + renderSettingsU32[1] = pc.sh_deg; + const render_settings_buffer = createBuffer( + device, + 'render settings', + renderSettingsBlock.byteLength, + GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, + renderSettingsBlock + ); + // =============================================== // Create Compute Pipeline and Bind Groups // =============================================== @@ -52,6 +96,8 @@ export default function get_renderer( }, }); + const preprocess_workgroup_count = Math.ceil(pc.num_points / C.histogram_wg_size); + const sort_bind_group = device.createBindGroup({ label: 'sort', layout: preprocess_pipeline.getBindGroupLayout(2), @@ -63,24 +109,132 @@ export default function get_renderer( ], }); + const preprocess_camera_bind_group = device.createBindGroup({ + label: 'preprocess camera', + layout: preprocess_pipeline.getBindGroupLayout(0), + entries: [ + { binding: 0, resource: { buffer: camera_buffer } }, + { binding: 1, resource: { buffer: render_settings_buffer } }, + ], + }); + + const preprocess_data_bind_group = device.createBindGroup({ + label: 'preprocess data', + layout: preprocess_pipeline.getBindGroupLayout(1), + entries: [ + { binding: 0, resource: { buffer: pc.gaussian_3d_buffer } }, + { binding: 1, resource: { buffer: splat_buffer } }, + { binding: 2, resource: { buffer: pc.sh_buffer } }, + ], + }); + // =============================================== // Create Render Pipeline and Bind Groups // =============================================== + const render_shader = device.createShaderModule({code: renderWGSL}); + const render_pipeline = device.createRenderPipeline({ + label: 'gaussian render', + layout: 'auto', + vertex: { + module: render_shader, + entryPoint: 'vs_main', + }, + fragment: { + module: render_shader, + entryPoint: 'fs_main', + targets: [{ + format: presentation_format, + blend: { + color: { + srcFactor: 'one', + dstFactor: 'one-minus-src-alpha', + operation: 'add', + }, + alpha: { + srcFactor: 'one', + dstFactor: 'one-minus-src-alpha', + operation: 'add', + }, + }, + }], + }, + primitive: { + topology: 'triangle-strip', + cullMode: 'none' + }, + }); + + const splat_bind_group = device.createBindGroup({ + label: 'gaussian splats', + layout: render_pipeline.getBindGroupLayout(0), + entries: [ + { binding: 0, resource: { buffer: splat_buffer } }, + { binding: 1, resource: { buffer: sorter.ping_pong[0].sort_indices_buffer } }, + { binding: 2, resource: { buffer: camera_buffer } }, + ], + }); // =============================================== // Command Encoder Functions // =============================================== - + const render = (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { + const pass = encoder.beginRenderPass({ + label: 'gaussian render', + colorAttachments: [ + { + view: texture_view, + loadOp: 'clear', + storeOp: 'store', + } + ], + }); + pass.setPipeline(render_pipeline); + pass.setBindGroup(0, splat_bind_group); + + pass.setVertexBuffer(0, quadVBO); + + pass.drawIndirect(indirect_buffer, 0); + + // pass.draw(pc.num_points); + pass.end(); + }; + + const preprocess = (encoder: GPUCommandEncoder) => { + const pass = encoder.beginComputePass({ + label: 'gaussian preprocess', + }); + pass.setPipeline(preprocess_pipeline); + pass.setBindGroup(0, preprocess_camera_bind_group); + pass.setBindGroup(1, preprocess_data_bind_group); + pass.setBindGroup(2, sort_bind_group); + pass.dispatchWorkgroups(preprocess_workgroup_count); + pass.end(); + }; // =============================================== // Return Render Object // =============================================== return { frame: (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { + device.queue.writeBuffer(sorter.sort_info_buffer, 0, new Uint32Array([0])); + device.queue.writeBuffer(sorter.sort_dispatch_indirect_buffer, 0, new Uint32Array([0, 1, 1])) + preprocess(encoder); + encoder.copyBufferToBuffer( + sorter.sort_info_buffer, // source: keys_size + 0, // src offset (bytes) + indirect_buffer, // dest: indirect draw args + Uint32Array.BYTES_PER_ELEMENT, // dst offset (skip vertexCount) + Uint32Array.BYTES_PER_ELEMENT // size (just the instanceCount) + ); sorter.sort(encoder); + render(encoder, texture_view); }, camera_buffer, + setGaussianMultiplier: (value: number) => { + renderSettingsF32[0] = value; + device.queue.writeBuffer(render_settings_buffer, 0, renderSettingsBlock); + }, }; } diff --git a/src/renderers/renderer.ts b/src/renderers/renderer.ts index ffdf9ba..8771eeb 100644 --- a/src/renderers/renderer.ts +++ b/src/renderers/renderer.ts @@ -86,6 +86,9 @@ export default async function init( const pc = await load(uploadedFile, device); pointcloud_renderer = get_renderer_pointcloud(pc, device, presentation_format, camera.uniform_buffer); gaussian_renderer = get_renderer_gaussian(pc, device, presentation_format, camera.uniform_buffer); + if (gaussian_renderer) { + gaussian_renderer.setGaussianMultiplier(params.gaussian_multiplier); + } renderers = { pointcloud: pointcloud_renderer, gaussian: gaussian_renderer, @@ -121,7 +124,9 @@ export default async function init( 'gaussian_multiplier', {min: 0, max: 1.5} ).on('change', (e) => { - //TODO: Bind constants to the gaussian renderer. + if (gaussian_renderer) { + gaussian_renderer.setGaussianMultiplier(e.value); + } }); } diff --git a/src/shaders/gaussian.wgsl b/src/shaders/gaussian.wgsl index 759226d..c9f5ed8 100644 --- a/src/shaders/gaussian.wgsl +++ b/src/shaders/gaussian.wgsl @@ -1,22 +1,96 @@ struct VertexOutput { - @builtin(position) position: vec4, + @builtin(position) position: vec4, // clip-space position written to rasterizer //TODO: information passed from vertex shader to fragment shader + @location(0) color: vec4, + @location(1) conic: vec3, + @location(2) opacity: f32, + @location(3) center_ndc: vec2, }; struct Splat { - //TODO: information defined in preprocess compute shader + // store information for 2D splat rendering + color: vec4, + conic_opacity: array, + centre_radius_ndc: array, }; +struct CameraUniforms { + view: mat4x4, + view_inv: mat4x4, + proj: mat4x4, + proj_inv: mat4x4, + viewport: vec2, + focal: vec2 +}; + +fn sigmoid(x: f32) -> f32 { + return 1.0 / (1.0 + exp(-x)); +} + +@group(0) @binding(0) +var splats : array; +@group(0) @binding(1) +var sorted_indices : array; +@group(0) @binding(2) +var camera : CameraUniforms; + @vertex fn vs_main( + @builtin(vertex_index) in_vertex_index: u32, + @builtin(instance_index) in_instance_index: u32 ) -> VertexOutput { - //TODO: reconstruct 2D quad based on information from splat, pass + // Reconstruct the quad for each splat in clip space. var out: VertexOutput; - out.position = vec4(1. ,1. , 0., 1.); + let splat_index = sorted_indices[in_instance_index]; + let splat = splats[splat_index]; + + let centre_radius_ndc = splat.centre_radius_ndc; // clip space + let center_ndc = unpack2x16float(centre_radius_ndc[0]); + let radius_ndc = unpack2x16float(centre_radius_ndc[1]); + + var corners = array, 4>( + vec2(-1.0, -1.0), + vec2( 1.0, -1.0), + vec2(-1.0, 1.0), + vec2( 1.0, 1.0) + ); + + let offset_ndc = corners[in_vertex_index] * radius_ndc; // NDC delta + + let a = unpack2x16float(splat.conic_opacity[0]); + let b = unpack2x16float(splat.conic_opacity[1]); + + let conic = vec3(a.x, a.y, b.x); + let opacity = b.y; + + out.position = vec4(center_ndc + offset_ndc, 0.0, 1.0); + out.color = splat.color; + out.conic = conic; + out.opacity = opacity; + out.center_ndc = center_ndc; + return out; } @fragment -fn fs_main(in: VertexOutput) -> @location(0) vec4 { - return vec4(1.); -} \ No newline at end of file +fn fs_main( + in: VertexOutput, +) -> @location(0) vec4 { + + var frag_pos_ndc = (in.position.xy / camera.viewport) * 2.0 - 1.0; + frag_pos_ndc.y *= -1.0; + var d = frag_pos_ndc - in.center_ndc; + d *= camera.viewport * 0.5; + + let power = -0.5 * ( + in.conic.x * d.x * d.x + + in.conic.z * d.y * d.y + - 2.0 * in.conic.y * d.x * d.y + ); + + if (power > 0.0) { + return vec4(0.0, 0.0, 0.0, 0.0); + } + let weight = exp(power) * in.opacity; + return in.color * min(weight, 0.99); +} diff --git a/src/shaders/point_cloud.wgsl b/src/shaders/point_cloud.wgsl index 01dded1..617171e 100644 --- a/src/shaders/point_cloud.wgsl +++ b/src/shaders/point_cloud.wgsl @@ -35,7 +35,7 @@ fn vs_main( let pos = vec4(a.x, a.y, b.x, 1.); // TODO: MVP calculations - out.position = pos; + out.position = camera.proj * camera.view * pos; return out; } diff --git a/src/shaders/preprocess.wgsl b/src/shaders/preprocess.wgsl index bbc63f5..b439b53 100644 --- a/src/shaders/preprocess.wgsl +++ b/src/shaders/preprocess.wgsl @@ -46,7 +46,7 @@ struct CameraUniforms { struct RenderSettings { gaussian_scaling: f32, - sh_deg: f32, + sh_deg: u32, } struct Gaussian { @@ -56,10 +56,27 @@ struct Gaussian { }; struct Splat { - //TODO: store information for 2D splat rendering + // store information for 2D splat rendering + color: vec4, + conic_opacity: array, + centre_radius_ndc: array, }; //TODO: bind your data here + +@group(0) @binding(0) +var camera: CameraUniforms; +@group(0) @binding(1) +var render_settings: RenderSettings; + +@group(1) @binding(0) +var gaussians : array; +@group(1) @binding(1) +var splats : array; +@group(1) @binding(2) +var sh_buffer : array; + + @group(2) @binding(0) var sort_infos: SortInfos; @group(2) @binding(1) @@ -69,10 +86,112 @@ var sort_indices : array; @group(2) @binding(3) var sort_dispatch: DispatchIndirect; -/// reads the ith sh coef from the storage buffer + +fn getR(rot: array) -> mat3x3 { + var r = unpack2x16float(rot[0]).x; + var x = unpack2x16float(rot[0]).y; + var y = unpack2x16float(rot[1]).x; + var z = unpack2x16float(rot[1]).y; + + // Normalize + let len = sqrt(x*x + y*y + z*z + r*r); + x = x / len; + y = y / len; + z = z / len; + r = r / len; + + let R = mat3x3( + 1. - 2. * (y * y + z * z), 2. * (x * y - r * z), 2. * (x * z + r * y), + 2. * (x * y + r * z), 1. - 2. * (x * x + z * z), 2. * (y * z - r * x), + 2. * (x * z - r * y), 2. * (y * z + r * x), 1. - 2. * (x * x + y * y) + ); + return R; +} + +fn getS(scale: array, gaussian_scaling: f32) -> mat3x3 { + var x = unpack2x16float(scale[0]).x; + var y = unpack2x16float(scale[0]).y; + var z = unpack2x16float(scale[1]).x; + + x = exp(x); + y = exp(y); + z = exp(z); + + let S = mat3x3( + x * gaussian_scaling, 0., 0., + 0., y * gaussian_scaling, 0., + 0., 0., z * gaussian_scaling + ); + return S; +} + +fn computeCov3d(rot: array, scale: array, gaussian_scaling: f32) -> mat3x3 { + let R = getR(rot); + let S = getS(scale, gaussian_scaling); + let A = S * R; + return transpose(A) * A; +} + +fn transformPoint4x3(p: vec3, m: mat4x4) -> vec3 { + return vec3( + m[0].x * p.x + m[1].x * p.y + m[2].x * p.z + m[3].x, + m[0].y * p.x + m[1].y * p.y + m[2].y * p.z + m[3].y, + m[0].z * p.x + m[1].z * p.y + m[2].z * p.z + m[3].z + ); +} + +fn computeCov2D(mean: vec3, cov3D: mat3x3) -> vec3 { + let W = mat3x3( + camera.view[0].x, camera.view[1].x, camera.view[2].x, + camera.view[0].y, camera.view[1].y, camera.view[2].y, + camera.view[0].z, camera.view[1].z, camera.view[2].z + ); + let t = transformPoint4x3(mean, camera.view); + + let J = mat3x3( + camera.focal.x / t.z, 0., -(camera.focal.x * t.x) / (t.z * t.z), + 0., camera.focal.y / t.z, -(camera.focal.y * t.y) / (t.z * t.z), + 0., 0., 0. + ); + let T = W * J; + let Vrk = mat3x3( + cov3D[0][0], cov3D[0][1], cov3D[0][2], + cov3D[0][1], cov3D[1][1], cov3D[1][2], + cov3D[0][2], cov3D[1][2], cov3D[2][2], + ); + var cov = transpose(T) * Vrk * T; + cov[0][0] += 0.3; + cov[1][1] += 0.3; + return vec3(cov[0][0], cov[0][1], cov[1][1]); +} + +fn computeConic(cov: vec3) -> vec3 { + let det = (cov.x * cov.z - cov.y * cov.y); + let det_inv = 1.0/ det; + return vec3 (cov.z * det_inv, -cov.y * det_inv, cov.x * det_inv); +} + +fn computeRadius(cov: vec3) -> f32 { + let det = (cov.x * cov.z - cov.y * cov.y); + let mid = 0.5 * (cov.x + cov.z); + let lambda1 = mid + sqrt(max(0.1, mid * mid - det)); + let lambda2 = mid - sqrt(max(0.1, mid * mid - det)); + return ceil(3.0 * sqrt(max(lambda1, lambda2))); +} + fn sh_coef(splat_idx: u32, c_idx: u32) -> vec3 { - //TODO: access your binded sh_coeff, see load.ts for how it is stored - return vec3(0.0); + let half_idx = c_idx >> 1u; // c_idx / 2 + let odd = c_idx & 1u; // c_idx % 2 + + let base_index = splat_idx * 24u + half_idx * 3u + odd; + + let color01 = unpack2x16float(sh_buffer[base_index + 0u]); + let color23 = unpack2x16float(sh_buffer[base_index + 1u]); + + let even_vec = vec3(color01.x, color01.y, color23.x); + let odd_vec = vec3(color01.y, color23.x, color23.y); + + return mix(even_vec, odd_vec, f32(odd)); } // spherical harmonics evaluation with Condon–Shortley phase @@ -108,11 +227,81 @@ fn computeColorFromSH(dir: vec3, v_idx: u32, sh_deg: u32) -> vec3 { return max(vec3(0.), result); } +fn sigmoid(x: f32) -> f32 { + return 1.0 / (1.0 + exp(-x)); +} + @compute @workgroup_size(workgroupSize,1,1) fn preprocess(@builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) wgs: vec3) { let idx = gid.x; //TODO: set up pipeline as described in instruction + if (idx >= arrayLength(&gaussians)) { + return; + } - let keys_per_dispatch = workgroupSize * sortKeyPerThread; + let keys_per_dispatch = workgroupSize * sortKeyPerThread; // increment DispatchIndirect.dispatchx each time you reach limit for one dispatch of keys -} \ No newline at end of file + + var splat_out: Splat; + let vertex = gaussians[idx]; + + let a = unpack2x16float(vertex.pos_opacity[0]); + let b = unpack2x16float(vertex.pos_opacity[1]); + + let pos = vec4(a.x, a.y, b.x, 1.); + let opacity = sigmoid(b.y); + + // // MVP calculations + let clip_pos = camera.proj * camera.view * pos; // clip space + let ndc = clip_pos.xy / clip_pos.w; // NDC space + let depth = clip_pos.z / clip_pos.w; + if (depth < 0.0 || depth > 1.0) { + return; + } + if (abs(ndc.x) > 1.2 || abs(ndc.y) > 1.2) { + return; + } + + let gaussian_scaling = render_settings.gaussian_scaling; + + let out_idx = atomicAdd(&sort_infos.keys_size, 1u); + + if (out_idx >= arrayLength(&splats)) { return; } + + let processed = out_idx + 1u; + if (keys_per_dispatch != 0u) { + let required_dispatch = (processed + keys_per_dispatch - 1u) / keys_per_dispatch; + atomicMax(&sort_dispatch.dispatch_x, required_dispatch); + } + + let cov3D = computeCov3d(vertex.rot, vertex.scale, gaussian_scaling); // world space covariance + let cov2D = computeCov2D(pos.xyz, cov3D); + let conic = computeConic(cov2D); + let radius = computeRadius(cov2D); + let conic_opacity: array = array( + pack2x16float(conic.xy), + pack2x16float(vec2(conic.z, opacity)) + ); + splat_out.conic_opacity = conic_opacity; + + let quad_size = vec2( // NDC half-extent per axis + radius * 2.0 / camera.viewport.x, + radius * 2.0 / camera.viewport.y + ); + + let centre_radius_ndc: array = array( + pack2x16float(ndc), + pack2x16float(quad_size) + ); + + splat_out.centre_radius_ndc = centre_radius_ndc; + + let cam_pos = camera.view_inv[3].xyz; + let dir = normalize(pos.xyz - cam_pos); + splat_out.color = vec4(computeColorFromSH(dir, idx, render_settings.sh_deg), 1.0); + + let pos_depth = 1.0 - depth; + splats[out_idx] = splat_out; + sort_indices[out_idx] = out_idx; + sort_depths[out_idx] = bitcast(pos_depth); +}