diff --git a/.gitignore b/.gitignore index 054e565..f9156ce 100644 --- a/.gitignore +++ b/.gitignore @@ -24,4 +24,5 @@ dist-ssr *.sw? /.vite -*/scenes \ No newline at end of file +*/scenes +scenes/* \ No newline at end of file diff --git a/README.md b/README.md index f99cdff..0aa4ce1 100644 --- a/README.md +++ b/README.md @@ -1,26 +1,120 @@ # Project5-WebGPU-Gaussian-Splat-Viewer -**University of Pennsylvania, CIS 565: GPU Programming and Architecture, Project 4** +**University of Pennsylvania, CIS 565: GPU Programming and Architecture, Project 5** -* (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) +* CARLOS LOPEZ GARCES + * [LinkedIn](https://www.linkedin.com/in/clopezgarces/) + * [Personal website](https://carlos-lopez-garces.github.io/) +* Tested on: Windows 11, 13th Gen Intel(R) Core(TM) i9-13900HX @ 2.20 GHz, RAM 32GB, NVIDIA GeForce RTX 4060, personal laptop. -### Live Demo -[![](img/thumb.png)](http://TODO.github.io/Project4-WebGPU-Forward-Plus-and-Clustered-Deferred) +### [Live Demo](http://carlos-lopez-garces.github.io/Penn-CIS-5650-Project5-WebGPU-Gaussian-Splat-Viewer) + +[![](images/bicycle_1.png)](http://carlos-lopez-garces.github.io/Penn-CIS-5650-Project5-WebGPU-Gaussian-Splat-Viewer) ### Demo Video/GIF -[![](img/video.mp4)](TODO) +| Bonsai | Bicycle | +|----------|----------| +| ![](images/bonsai_1.gif) | ![](images/bicycle_1.gif) | +| **Truck** | **Train** | +| ![](images/truck_1.gif) | ![](images/train_1.gif) | + +## WebGPU Gaussian Splat and Point Cloud Viewer + +This project implements a viewer and renderer for point clouds and Gaussian splats. The renderer implements the rasterizer of the paper "3D Gaussian Splatting +for Real-Time Radiance Field Rendering" ([paper website](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/)). + +The viewer handles well many of the common scenes typically used by academic papers to present their results. + +| Bonsai | Bicycle | +|----------|----------| +| ![](images/bonsai_1.png) | ![](images/bicycle_1.png) | +| **Truck** | **Train** | +| ![](images/truck_1.png) | ![](images/train_1.png) | + +### Features + +- Real-Time Point Cloud Rendering: Visualize dense point cloud data from .ply files. + +- Gaussian Splatting: Reconstruct the radiance field of a scene from its point cloud using Gaussian Splatting. + +- Camera Controls: Pan, zoom, and rotate for detailed examination of Gaussian splats and point clouds. + +![](images/bonsai_3.gif) + +A point cloud is the main input to the Gaussian splatting rendering process. + +| Bonsai | Bicycle | +|----------|----------| +| ![](images/pc_bonsai_1.gif) | ![](images/pc_bicycle_1.gif) | +| **Truck** | **Train** | +| ![](images/pc_truck_1.gif) | ![](images/pc_train_1.gif) | + +### Gaussian Splat Preprocesing + +The preprocessing step transforms point data into splat representations for rendering by calculating 2D covariances, projecting spherical-harmonics-based color, and organizing/sorting splats for rendering based on depth. + +- Frustum culling checks if the splat falls within the visible area. Culling early reduces unnecessary computations for off-screen points. + +- The Gaussian’s scale and rotation are transformed into a 3D covariance matrix. This 3D covariance is projected to 2D to represent screen-space influence, helping to control the splat’s shape based on viewpoint. + +- Conic parameters are derived from the 2D covariance, resulting in a conic equation that approximates Gaussian distribution across x and y directions. Eigenvectors of the covariance provide major and minor axes, which are scaled to determine the Gaussian’s radius on the screen. + +- Spherical harmonics coefficients approximate lighting effects for each Gaussian based on its orientation. + +![](images/bicycle_2.gif) + +### Gaussian Splat Rendering + +For each splat, six vertices are generated to define a bounding quadrilateral in screen space. These vertices form a rectangle around each splat, with edges aligned based on the splat’s screen-space radius. + +In the fragment shader, the screen-space offset from the splat’s center to the fragment position is computed. A linear combination with the coefficients of a conic form is used to obtain the exponent by which the splat's opacity decays. Its color is blended with the splat's opacity, which finalizes the reconstruction of the radiance field. + +![](images/squares.png) + +### Performance Analysis + +The following table and chart show a few statistics about each of the scenes that were tested and the frames per second range achieved by the viewer for both the point cloud and the Gaussian splats. Because the process to render Gaussian splats is much more involved (preprocessing, data transfer, more complex computations per vertex and fragment), they consistently reach a much lower FPS. Also, FPS decreases with the number of points in the point cloud. FPS ranges are very wide and depend significantly on the camera angle; here they are shown from the default camera angle of the respective scene. + +Observe also that the size of the point cloud has an effect on FPS, but it's not the only factor to consider (note that the truck scene is twice the size of the train scene and they have comparable FPS ranges in the Gaussian splat case). In general, as the number of Gaussians increases, so does the FPS; the preprocessing step processes each of them (culled excluded), so FPS increases proportional. But as I said before, it's not the only factor; I investigate this phenomenon further below. + +| Scene | Points | Time (ms) | Point Cloud FPS | Gaussian Splat FPS | +|--------|---------|-----------|-----------------|--------------| +| Bonsai | 1,244,819 | 19,597 | 160 | 84 | +| Bicycle| 1,063,091 | 25,271 | 150-181 | 54-85 | +| Train | 1,026,508 | 15,285 | 121-255 | 32-45 | +| Truck | 2,541,226 | 40,027 | 166-175 | 32-41 | + +![](images/Point%20Cloud%20vs%20Gaussian%20Splats%20FPS%20Across%20Scenes.png) + +Frustum culling is gives a very modest improvement in FPS, cutting down on the number of Gaussians processed per frame. In this example, where I've zoomed into Bicycle scene, leaving some points out of view, 15 fps was registered without frustum culling compared to 21 with frustum culling. Since the preprocessing compute shader runs simultaneously on the GPU, perhaps the work group size in this case (256) is such that Gaussians that are otherwise culled are processed simultaneously with the visible ones, so that both with and without culling the workgroups finish roughly in the same time. + +![](images/bicycle_closeup.png) + +![](images/bonsai_2.gif) + +Work group size has a two-fold effect: (1) render quality is optimal for a group size of 256 (sizes 64, 128, and 512 result in artifacts); (2) the FPS range reached with each of the sizes varies, with size 256 resulting in the highest FPS. A work group size of 256 allows for a high degree of parallelism, which is crucial for achieving high FPS. Smaller sizes might not leverage the full parallel processing power of the GPU, while larger sizes might introduce overhead that reduces performance. + +| Work Group Size | FPS Range | +|-----------------|-----------| +| 256 | 51 - 84 | +| 128 | 11 - 71 | +| 64 | 14 - 32 | + +### Artifacts During Development + +![](images/blooper_4.png) + +From not indexing the spherical harmonics coefficients correctly. -### (TODO: Your README) +![](images/blooper_1.png) +![](images/blooper_2.png) +![](images/blooper_3.png) -*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. +Flickering. -This assignment has a considerable amount of performance analysis compared -to implementation work. Complete the implementation early to leave time! +![](images/bloopers_5.png) ### Credits diff --git a/images/Point Cloud vs Gaussian Splats FPS Across Scenes.png b/images/Point Cloud vs Gaussian Splats FPS Across Scenes.png new file mode 100644 index 0000000..6be8ac4 Binary files /dev/null and b/images/Point Cloud vs Gaussian Splats FPS Across Scenes.png differ diff --git a/images/bicycle_1.gif b/images/bicycle_1.gif new file mode 100644 index 0000000..ef431c2 Binary files /dev/null and b/images/bicycle_1.gif differ diff --git a/images/bicycle_1.png b/images/bicycle_1.png new file mode 100644 index 0000000..5a1ff56 Binary files /dev/null and b/images/bicycle_1.png differ diff --git a/images/bicycle_2.gif b/images/bicycle_2.gif new file mode 100644 index 0000000..bdaf0d6 Binary files /dev/null and b/images/bicycle_2.gif differ diff --git a/images/bicycle_closeup.png b/images/bicycle_closeup.png new file mode 100644 index 0000000..fcee136 Binary files /dev/null and b/images/bicycle_closeup.png differ diff --git a/images/blooper_1.png b/images/blooper_1.png new file mode 100644 index 0000000..8497d69 Binary files /dev/null and b/images/blooper_1.png differ diff --git a/images/blooper_2.png b/images/blooper_2.png new file mode 100644 index 0000000..63b06f4 Binary files /dev/null and b/images/blooper_2.png differ diff --git a/images/blooper_3.png b/images/blooper_3.png new file mode 100644 index 0000000..0ffa6aa Binary files /dev/null and b/images/blooper_3.png differ diff --git a/images/blooper_4.png b/images/blooper_4.png new file mode 100644 index 0000000..82b6c26 Binary files /dev/null and b/images/blooper_4.png differ diff --git a/images/bloopers_5.png b/images/bloopers_5.png new file mode 100644 index 0000000..53e52ee Binary files /dev/null and b/images/bloopers_5.png differ diff --git a/images/bonsai_1.gif b/images/bonsai_1.gif new file mode 100644 index 0000000..a5dbe62 Binary files /dev/null and b/images/bonsai_1.gif differ diff --git a/images/bonsai_1.png b/images/bonsai_1.png new file mode 100644 index 0000000..8c65b4e Binary files /dev/null and b/images/bonsai_1.png differ diff --git a/images/bonsai_2.gif b/images/bonsai_2.gif new file mode 100644 index 0000000..a88c548 Binary files /dev/null and b/images/bonsai_2.gif differ diff --git a/images/bonsai_3.gif b/images/bonsai_3.gif new file mode 100644 index 0000000..e50ed62 Binary files /dev/null and b/images/bonsai_3.gif differ diff --git a/images/pc_bicycle_1.gif b/images/pc_bicycle_1.gif new file mode 100644 index 0000000..77cea8e Binary files /dev/null and b/images/pc_bicycle_1.gif differ diff --git a/images/pc_bonsai_1.gif b/images/pc_bonsai_1.gif new file mode 100644 index 0000000..c9b2d43 Binary files /dev/null and b/images/pc_bonsai_1.gif differ diff --git a/images/pc_train_1.gif b/images/pc_train_1.gif new file mode 100644 index 0000000..3e179f2 Binary files /dev/null and b/images/pc_train_1.gif differ diff --git a/images/pc_truck_1.gif b/images/pc_truck_1.gif new file mode 100644 index 0000000..f104fad Binary files /dev/null and b/images/pc_truck_1.gif differ diff --git a/images/squares.png b/images/squares.png new file mode 100644 index 0000000..96c4d94 Binary files /dev/null and b/images/squares.png differ diff --git a/images/train_1.gif b/images/train_1.gif new file mode 100644 index 0000000..9c36043 Binary files /dev/null and b/images/train_1.gif differ diff --git a/images/train_1.png b/images/train_1.png new file mode 100644 index 0000000..6f77dcc Binary files /dev/null and b/images/train_1.png differ diff --git a/images/truck_1.gif b/images/truck_1.gif new file mode 100644 index 0000000..0466144 Binary files /dev/null and b/images/truck_1.gif differ diff --git a/images/truck_1.png b/images/truck_1.png new file mode 100644 index 0000000..7319119 Binary files /dev/null and b/images/truck_1.png differ diff --git a/src/renderers/gaussian-renderer.ts b/src/renderers/gaussian-renderer.ts index 1684523..790fc86 100644 --- a/src/renderers/gaussian-renderer.ts +++ b/src/renderers/gaussian-renderer.ts @@ -4,6 +4,8 @@ import renderWGSL from '../shaders/gaussian.wgsl'; import { get_sorter,c_histogram_block_rows,C } from '../sort/sort'; import { Renderer } from './renderer'; +const c_size_splat = 4 * 2 * 3; + export interface GaussianRenderer extends Renderer { } @@ -35,6 +37,23 @@ export default function get_renderer( // =============================================== const nulling_data = new Uint32Array([0]); + const reset_buffer = createBuffer( + device, 'reset buffer', 4, GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST, nulling_data + ); + + const splat_buffer = createBuffer( + device, 'splat buffer', pc.num_points*c_size_splat, GPUBufferUsage.STORAGE + ); + + const indirect_render_buffer = createBuffer( + device, 'indirect buffer', 16, GPUBufferUsage.INDIRECT | GPUBufferUsage.COPY_DST, + new Uint32Array([6, 0, 0, 0]) + ); + + const settings_buffer = createBuffer( + device, 'settings buffer', 8, GPUBufferUsage.COPY_DST | GPUBufferUsage.UNIFORM, + new Float32Array([1.0, pc.sh_deg]) + ); // =============================================== // Create Compute Pipeline and Bind Groups @@ -52,9 +71,21 @@ export default function get_renderer( }, }); + const preprocess_bind_group = device.createBindGroup({ + label: 'preprocess bind group', + layout: preprocess_pipeline.getBindGroupLayout(0), + entries: [ + { binding: 0, resource: { buffer: camera_buffer } }, + { binding: 1, resource: { buffer: settings_buffer } }, + { binding: 2, resource: { buffer: pc.sh_buffer } }, + { binding: 3, resource: { buffer: pc.gaussian_3d_buffer } }, + { binding: 4, resource: { buffer: splat_buffer } } + ], + }); + const sort_bind_group = device.createBindGroup({ - label: 'sort', - layout: preprocess_pipeline.getBindGroupLayout(2), + label: 'sort bind group', + layout: preprocess_pipeline.getBindGroupLayout(1), entries: [ { binding: 0, resource: { buffer: sorter.sort_info_buffer } }, { binding: 1, resource: { buffer: sorter.ping_pong[0].sort_depths_buffer } }, @@ -63,23 +94,83 @@ export default function get_renderer( ], }); - // =============================================== // Create Render Pipeline and Bind Groups // =============================================== - + + const render_pipeline = device.createRenderPipeline({ + label: 'render pipeline', + layout: 'auto', + vertex: { + entryPoint: 'vs_main', + module: device.createShaderModule({code: renderWGSL}) + }, + fragment: { + entryPoint: 'fs_main', + module: device.createShaderModule({code: renderWGSL}), + targets: [{ + format: presentation_format, + blend: { + color: { srcFactor: 'one', dstFactor: 'one-minus-src-alpha' }, + alpha: { srcFactor: 'one', dstFactor: 'one-minus-src-alpha' }, + }, + }] + }, +}); + + const render_bind_group = device.createBindGroup({ + label: 'render bind group', + layout: render_pipeline.getBindGroupLayout(0), + entries: [ + { binding: 0, resource: { buffer: camera_buffer } }, + { binding: 1, resource: { buffer: sorter.ping_pong[1].sort_indices_buffer } }, + { binding: 2, resource: { buffer: splat_buffer } } + ] + }); // =============================================== // Command Encoder Functions // =============================================== - + + const preprocess = (encoder: GPUCommandEncoder) => { + encoder.copyBufferToBuffer(reset_buffer, 0, sorter.sort_info_buffer, 0, 4); + encoder.copyBufferToBuffer(reset_buffer, 0, sorter.sort_dispatch_indirect_buffer, 0, 4); + + const pass = encoder.beginComputePass({ label: 'preprocess pass' }); + pass.setPipeline(preprocess_pipeline); + pass.setBindGroup(0, preprocess_bind_group); + pass.setBindGroup(1, sort_bind_group); + pass.dispatchWorkgroups(Math.ceil(pc.num_points / C.histogram_wg_size)); + pass.end(); + }; + + const render = (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { + encoder.copyBufferToBuffer(sorter.sort_info_buffer, 0, indirect_render_buffer, 4, 4); + + const pass = encoder.beginRenderPass({ + label: 'render pass', + colorAttachments: [{ + view: texture_view, + loadOp: 'clear', + clearValue: { r: 0, g: 0, b: 0, a: 1 }, + storeOp: 'store' + }], + }); + + pass.setPipeline(render_pipeline); + pass.setBindGroup(0, render_bind_group); + pass.drawIndirect(indirect_render_buffer, 0); + pass.end(); + }; // =============================================== // Return Render Object // =============================================== return { frame: (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { + preprocess(encoder); sorter.sort(encoder); + render(encoder, texture_view); }, camera_buffer, }; diff --git a/src/shaders/gaussian.wgsl b/src/shaders/gaussian.wgsl index 759226d..f7a2a32 100644 --- a/src/shaders/gaussian.wgsl +++ b/src/shaders/gaussian.wgsl @@ -1,22 +1,76 @@ -struct VertexOutput { - @builtin(position) position: vec4, - //TODO: information passed from vertex shader to fragment shader +struct CameraUniforms { + view: mat4x4, + view_inv: mat4x4, + proj: mat4x4, + proj_inv: mat4x4, + viewport: vec2, + focal: vec2 }; struct Splat { - //TODO: information defined in preprocess compute shader + mean_xy: u32, + radii: u32, + conic_xy: u32, + conic_z: u32, + rgb_rg: u32, + rgb_b_opacity: u32 +}; + +@group(0) @binding(0) +var camera: CameraUniforms; +@group(0) @binding(1) +var sort_indices : array; +@group(0) @binding(2) +var splats: array; + +const signs = array( + vec2f(-1.0, 1.0), vec2f(-1.0, -1.0), + vec2f(1.0, -1.0), vec2f(1.0, -1.0), + vec2f(1.0, 1.0), vec2f(-1.0, 1.0) +); + +struct VertexOutput { + @builtin(position) position: vec4, + @location(0) mean: vec2f, + @location(1) rgb_opacity: vec4f, + @location(2) conic: vec3f }; @vertex -fn vs_main( -) -> VertexOutput { - //TODO: reconstruct 2D quad based on information from splat, pass - var out: VertexOutput; - out.position = vec4(1. ,1. , 0., 1.); - return out; +fn vs_main(@builtin(vertex_index) vertex_index: u32, @builtin(instance_index) instance_index: u32) -> VertexOutput { + let splat = splats[sort_indices[instance_index]]; + + let mean_xy = unpack2x16float(splat.mean_xy); + let diameter = 2.0 * unpack2x16float(splat.radii); + let rgb_rg = unpack2x16float(splat.rgb_rg); + let rgb_b_opacity = unpack2x16float(splat.rgb_b_opacity);; + let conic_xy = unpack2x16float(splat.conic_xy); + let conic_z = unpack2x16float(splat.conic_z); + + let sign = signs[vertex_index]; + + return VertexOutput( + vec4(mean_xy.x + sign.x*diameter.x, mean_xy.y + sign.y*diameter.y, 0.0, 1.0), + vec2f(mean_xy.x, mean_xy.y), + vec4f(rgb_rg.x, rgb_rg.y, rgb_b_opacity.x, 1.0/(1.0+exp(-rgb_b_opacity.y))), + vec3f(conic_xy.x, conic_xy.y, conic_z.x), + ); } @fragment -fn fs_main(in: VertexOutput) -> @location(0) vec4 { - return vec4(1.); +fn fs_main(input: VertexOutput) -> @location(0) vec4 { + var pos_ndc = 2.0 * (input.position.xy / camera.viewport) - vec2(1.0, 1.0); + pos_ndc.y = -pos_ndc.y; + + // Screen-space offset from the fragment position, with x-coordinate reversed. + var offset_screen = pos_ndc - input.mean; + offset_screen.x = -offset_screen.x; + offset_screen *= camera.viewport * 0.5; + + var exponent = + input.conic.x * offset_screen.x * offset_screen.x + + input.conic.z * offset_screen.y * offset_screen.y + + input.conic.y * offset_screen.x * offset_screen.y; + + return vec4f(input.rgb_opacity.xyz, 1.0) * input.rgb_opacity.w * exp(-exponent/2.0); } \ No newline at end of file diff --git a/src/shaders/point_cloud.wgsl b/src/shaders/point_cloud.wgsl index 01dded1..c61faf3 100644 --- a/src/shaders/point_cloud.wgsl +++ b/src/shaders/point_cloud.wgsl @@ -7,9 +7,13 @@ struct CameraUniforms { focal: vec2 }; +// On the CPU side, a Float16Array of 11 entries. Interpreted as follows: struct Gaussian { + // 4 16-bit floats for position (x,y,z) and opacity packed as 2 u32. pos_opacity: array, + // 4 16-bit floats for rotation (x,y,z,w) packed as 2 u32. rot: array, + // 3 16-bit floats for scale (x,y,z) packed as 2 u32. scale: array } @@ -32,9 +36,14 @@ fn vs_main( let vertex = gaussians[in_vertex_index]; let a = unpack2x16float(vertex.pos_opacity[0]); let b = unpack2x16float(vertex.pos_opacity[1]); - let pos = vec4(a.x, a.y, b.x, 1.); + var pos = vec4(a.x, a.y, b.x, 1.); // TODO: MVP calculations + + // Scale and rotation applies to the Gaussian splats at this point, + // but not to the point itself. + + pos = camera.proj * camera.view * pos; out.position = pos; return out; diff --git a/src/shaders/preprocess.wgsl b/src/shaders/preprocess.wgsl index bbc63f5..807969e 100644 --- a/src/shaders/preprocess.wgsl +++ b/src/shaders/preprocess.wgsl @@ -50,29 +50,51 @@ struct RenderSettings { } struct Gaussian { + // 4 16-bit floats for position (x,y,z) and opacity packed as 2 u32. pos_opacity: array, + // 4 16-bit floats for rotation (x,y,z,w) packed as 2 u32. rot: array, + // 3 16-bit floats for scale (x,y,z) packed as 2 u32. scale: array }; struct Splat { - //TODO: store information for 2D splat rendering + mean_xy: u32, + radii: u32, + conic_xy: u32, + conic_z: u32, + rgb_rg: u32, + rgb_b_opacity: u32, }; -//TODO: bind your data here -@group(2) @binding(0) +@group(0) @binding(0) +var camera: CameraUniforms; +@group(0) @binding(1) +var settings: RenderSettings; +@group(0) @binding(2) +var sh_coefficients: array>; +@group(0) @binding(3) +var gaussians: array; +@group(0) @binding(4) +var splats: array; + +@group(1) @binding(0) var sort_infos: SortInfos; -@group(2) @binding(1) +@group(1) @binding(1) var sort_depths : array; -@group(2) @binding(2) +@group(1) @binding(2) var sort_indices : array; -@group(2) @binding(3) +@group(1) @binding(3) var sort_dispatch: DispatchIndirect; /// reads the ith sh coef from the storage buffer 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 c_idx3 = 3u*c_idx; + return vec3( + unpack2x16float(sh_coefficients[splat_idx][c_idx3 / 2u])[c_idx3 % 2u], + unpack2x16float(sh_coefficients[splat_idx][(c_idx3 + 1u) / 2u])[(c_idx3 + 1u) % 2u], + unpack2x16float(sh_coefficients[splat_idx][(c_idx3 + 2u) / 2u])[(c_idx3 + 2u) % 2u] + ); } // spherical harmonics evaluation with Condon–Shortley phase @@ -108,11 +130,218 @@ fn computeColorFromSH(dir: vec3, v_idx: u32, sh_deg: u32) -> vec3 { return max(vec3(0.), result); } +fn normalizePlane(plane: vec4) -> vec4 { + // Normalize so that the normal is a unit vector. + let length = length(plane.xyz); + return plane / length; +} + +// Gribb-Hartmann. www.gamedevs.org/uploads/fast-extraction-viewing-frustum-planes-from-world-view-projection-matrix.pdf. +fn extractFrustumPlanes(view_proj: mat4x4) -> array, 6> { + var planes: array, 6>; + + // Left. + planes[0] = normalizePlane(view_proj[3] + view_proj[0]); + // Right. + planes[1] = normalizePlane(view_proj[3] - view_proj[0]); + // Bottom. + planes[2] = normalizePlane(view_proj[3] + view_proj[1]); + // Top. + planes[3] = normalizePlane(view_proj[3] - view_proj[1]); + // Near. + planes[4] = normalizePlane(view_proj[3] + view_proj[2]); + // Far. + planes[5] = normalizePlane(view_proj[3] - view_proj[2]); + + // In clip space. + return planes; +} + +fn cull(pos: vec4, scale: vec3, view_proj: mat4x4) -> bool { + var enlarged_scale = abs(scale) * 1.1; + let pos_in_view_space = (camera.view * pos).xyz; + let min_bounds = pos_in_view_space.xyz - enlarged_scale; + let max_bounds = pos_in_view_space.xyz + enlarged_scale; + + let min_clip = view_proj * vec4(min_bounds, 1.0); + let max_clip = view_proj * vec4(max_bounds, 1.0); + + let frustum_planes = extractFrustumPlanes(view_proj); + + for (var i = 0u; i < 6u; i = i + 1u) { + let plane = frustum_planes[i]; + + let p = vec3( + select(min_clip.x, max_clip.x, plane.x > 0.0), + select(min_clip.y, max_clip.y, plane.y > 0.0), + select(min_clip.z, max_clip.z, plane.z > 0.0) + ); + + if (dot(plane.xyz, p) + plane.w < 0.0) { + // Cull. + return false; + } + } + + return true; +} + +fn frustumCull(mean_clip: vec4) -> bool { + let clip_sz = 1.3 * mean_clip.w; + let z = mean_clip.z / mean_clip.w; + return + mean_clip.x < -clip_sz || mean_clip.x > clip_sz + || mean_clip.y < -clip_sz || mean_clip.y > clip_sz + || z <= 0.0 || z >= 1.0; +} + +fn quaternionToMatrix(q: vec4) -> mat3x3 { + let x2 = q.x * q.x; + let y2 = q.y * q.y; + let z2 = q.z * q.z; + let xy = q.x * q.y; + let xz = q.x * q.z; + let yz = q.y * q.z; + let wx = q.w * q.x; + let wy = q.w * q.y; + let wz = q.w * q.z; + + return mat3x3( + 1.0 - 2.0 * (y2 + z2), 2.0 * (xy - wz), 2.0 * (xz + wy), + 2.0 * (xy + wz), 1.0 - 2.0 * (x2 + z2), 2.0 * (yz - wx), + 2.0 * (xz - wy), 2.0 * (yz + wx), 1.0 - 2.0 * (x2 + y2) + ); +} + +// Based on original author's CUDA implementation. +// github.com/graphdeco-inria/diff-gaussian-rasterization/blob/59f5f77e3ddbac3ed9db93ec2cfe99ed6c5d121d/cuda_rasterizer/forward.cu#L118 +fn computeCov3D(gaussian: Gaussian) -> array { + var scale = vec3(unpack2x16float(gaussian.scale[0]).xy, unpack2x16float(gaussian.scale[1]).x); + scale = exp(scale); + let S = mat3x3f( + scale.x, 0.0f, 0.0f, + 0.0f, scale.y, 0.0f, + 0.0f, 0.0f, scale.z + ) * settings.gaussian_scaling; + + let rot_ax = unpack2x16float(gaussian.rot[0]); + let rot_yz = unpack2x16float(gaussian.rot[1]); + let quaternion = vec4(rot_ax.y, rot_yz.x, rot_yz.y, rot_ax.x); + let R = quaternionToMatrix(quaternion); + + let M = S * R; + + let Sigma = transpose(M) * M; + + let cov3D = array( + Sigma[0][0], + Sigma[0][1], + Sigma[0][2], + Sigma[1][1], + Sigma[1][2], + Sigma[2][2], + ); + + return cov3D; +} + +// Based on original author's CUDA implementation. +// github.com/graphdeco-inria/diff-gaussian-rasterization/blob/59f5f77e3ddbac3ed9db93ec2cfe99ed6c5d121d/cuda_rasterizer/forward.cu#L74 +fn computeCov2D(cov3D: array, mean: vec4f) -> vec3f { + var t = (camera.view * mean).xyz; + + let limx = 0.65f * camera.viewport.x / camera.focal.x; + let limy = 0.65f * camera.viewport.y / camera.focal.y; + let txtz = t.x / t.z; + let tytz = t.y / t.z; + t.x = min(limx, max(-limx, txtz)) * t.z; + t.y = min(limy, max(-limy, tytz)) * t.z; + + let J = mat3x3f( + camera.focal.x / t.z, 0.0f, -(camera.focal.x * t.x) / (t.z * t.z), + 0.0f, camera.focal.y / t.z, -(camera.focal.y * t.y) / (t.z * t.z), + 0.0f, 0.0f, 0.0f + ); + + let W = mat3x3f( + 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 = W * J; + + let Vrk = mat3x3f( + cov3D[0], cov3D[1], cov3D[2], + cov3D[1], cov3D[3], cov3D[4], + cov3D[2], cov3D[4], cov3D[5] + ); + + var cov2d = transpose(T) * transpose(Vrk) * T; + cov2d[0][0] += 0.3f; + cov2d[1][1] += 0.3f; + + return vec3f(cov2d[0][0], cov2d[0][1], cov2d[1][1]); +} + @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 gaussian = gaussians[idx]; + + let pos_xy = unpack2x16float(gaussian.pos_opacity[0]); + let pos_z_opacity = unpack2x16float(gaussian.pos_opacity[1]); + let mean = vec4f(pos_xy.x, pos_xy.y, pos_z_opacity.x, 1.0f); + + let mean_clip = camera.proj * camera.view * mean; + let mean_screen = mean_clip.xy / mean_clip.w; + if (frustumCull(mean_clip)) { + return; + } + + // Covariance, radii, and conic computation is based on the paper author's + // CUDA implementation. + // github.com/graphdeco-inria/diff-gaussian-rasterization/blob/main/cuda_rasterizer/forward.cu. + + let cov = computeCov2D(computeCov3D(gaussian), mean); + + var det = cov.x * cov.z - cov.y * cov.y; + if (det == 0.0) { + return; + } + let det_inv = 1.0 / det; + + let conic = vec3f(cov.z*det_inv, -cov.y*det_inv, cov.x*det_inv); + + 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)); + let radius = ceil(3.0 * sqrt(max(lambda1, lambda2))); + + let cam_pos_world = camera.view_inv[3].xyz; + let rgb = computeColorFromSH( + normalize(mean.xyz - cam_pos_world), + idx, + u32(settings.sh_deg) + ); + + let key = atomicAdd(&sort_infos.keys_size, 1u); + let z_far = camera.proj[3][2] / (1.0 - camera.proj[2][2]); + sort_depths[key] = bitcast(z_far - (camera.view * mean).z); + sort_indices[key] = key; + + splats[key].mean_xy = pack2x16float(mean_screen); + splats[key].radii = pack2x16float(vec2f(radius, radius) / camera.viewport); + splats[key].conic_xy = pack2x16float(conic.xy); + splats[key].conic_z = pack2x16float(vec2f(conic.z, 0.0)); + splats[key].rgb_rg = pack2x16float(rgb.rg); + splats[key].rgb_b_opacity = pack2x16float(vec2f(rgb.b, pos_z_opacity.y)); let keys_per_dispatch = workgroupSize * sortKeyPerThread; // increment DispatchIndirect.dispatchx each time you reach limit for one dispatch of keys + if (key % keys_per_dispatch == 0) { + atomicAdd(&sort_dispatch.dispatch_x, 1u); + } } \ No newline at end of file