diff --git a/INSTRUCTION.md b/INSTRUCTION.md index 9c54787..bcc4a13 100644 --- a/INSTRUCTION.md +++ b/INSTRUCTION.md @@ -70,7 +70,7 @@ WebGPU errors will appear in your browser's developer console (Ctrl + Shift + J ### Part 1: Understanding 3D Gaussian Point Cloud & Add MVP calculation (10pts) - Read over the [3D Gaussian Splatting Paper](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/) to have a basic understanding. -- Then read over `point_cloud` renderer, add MVP calculation to the vertex shader. After that, you can see yellow point cloud rendered to screen. +- Then read over `point_cloud` renderer, add MVP calculation to the vertex shader. After that, you can see yellow point cloud rendered to screen. ### Part 2: Gaussian Renderer (80pts total)(50pts(preprocess)+30pts(renderer)) diff --git a/README.md b/README.md index edffdaf..a01d2ad 100644 --- a/README.md +++ b/README.md @@ -2,30 +2,103 @@ **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) +* Yunhao Qian +* Tested on: (TODO) Google Chrome 141.0 on + * Windows 11, 24H2 + * 13th Gen Intel(R) Core(TM) i7-13700 (2.10 GHz) + * NVIDIA GeForce RTX 4090 -### Live Demo +## Live Demo -[![](img/thumb.png)](http://TODO.github.io/Project4-WebGPU-Forward-Plus-and-Clustered-Deferred) +[![Live demo](img/thumb.png)](https://yunhao-qian.github.io/showcase/gaussian-splatting/) -### Demo Video/GIF +## Demo Video -[![](img/video.mp4)](TODO) +https://github.com/user-attachments/assets/5a173f00-dfd2-4566-b4ca-4921d0ac81bb -### (TODO: Your README) +# Implementation & Performance Report -*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. +## Implementation -This assignment has a considerable amount of performance analysis compared -to implementation work. Complete the implementation early to leave time! +### Point Cloud Renderer -### Credits +- Renders all points as unit yellow dots (fixed size), regardless of distance to camera. +- Color is a uniform solid yellow. -- [Vite](https://vitejs.dev/) -- [tweakpane](https://tweakpane.github.io/docs//v3/monitor-bindings/) -- [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) +### Gaussian Renderer + +- Renders each point as a splat (disk/billboard) whose: + - Color comes from the source PLY model (viewed from different angles you’ll perceive different color mixes due to overlapping/transparency). + - Size and opacity falloff are driven by the point’s covariance: high opacity at the center; opacity decreases smoothly toward the edges. + +## Performance Analysis + +Test assets: `bonsai.ply` and `bicycle.ply`. + +### Point Cloud vs. Gaussian Renderers + +Visuals: + +- Point cloud: fixed-size pixels/points, always solid yellow, no depth-related size attenuation. +- Gaussian: true surface colors from PLY, splat size respects scale, opacity smoothly decays from each Gaussian’s center. + +Throughput: + +- On my setup, both renderers are similar in frame rate in uncongested views, generally ~56-60 FPS. + +### Workgroup Sizes (Preprocess Step) + +Measured total preprocess time while sweeping a single workgroup-size parameter: + +| Workgroup Size | Time (ms) | Relative Speedup vs. 32× | +|---:|---:|---:| +| 32 | 4,627 | 1.00× | +| 64 | 4,734 | 0.98× | +| 128 | 5,651 | 0.82× | +| 256 | 4,365 | 1.06× | +| 512 | 3,305 | 1.40× | +| 1,024 | 4,846 | 0.95× | +| 2,048 | 5,139 | 0.90× | + +Analysis: + +* 512 is the clear sweet spot here (≈1.40× faster than 32; ~30% faster than 256). +* Very small groups (32) underutilize the GPU; very large ones (≥1,024) likely reduce scheduling flexibility. + +> Sorting Kernel (for Gaussian Splatting): Attempted to change sorting workgroup size; program misbehaved. + +### View-Frustum Culling + +* Outside scene (both `bonsai.ply` and `bicycle.ply`): ~56 FPS with or without culling — the scenes are small vs. GPU capability, so savings are masked by fixed overhead. +* Heavy view on `bicycle.ply` (set Gaussian multiplier to max (1.5), zoomed in): + * ~10 FPS without culling -> ~40 FPS with culling (~4× speedup). +* Conclusion: Culling removes gaussians from sorting and from the graphics pipeline, so it can reduce time roughly proportionally to the number of discarded splats. Gains become obvious in dense, overdraw-heavy views. + +### Number of Gaussians (Scaling) + +Preprocess timings + +* `bicycle.ply`: 1,063,091 pts in 24,684 ms -> ~43.07 pts/ms (≈ 43.1 k pts/s). +* `bonsai.ply`: 272,956 pts in 4,365 ms -> ~62.53 pts/ms (≈ 62.5 k pts/s). + +Average throughput: ~52.8 pts/ms (≈ 52.8 k pts/s). + +Analysis: + +* Preprocess time scales roughly linearly with the number of points, indicating the GPU is well saturated. +* The per-scene variation (43k-62k pts/s) likely comes from: + * Different covariance distributions (branching, math intensity). + * Cache locality and memory access patterns (attribute stride/packing). + * Different proportions of filtered/culled points during preprocess. + +Display stage: + +* In typical views for both assets, display holds near ~56-60 FPS; large improvements appear only when enabling culling in dense views. + +## Credits + +* [Vite](https://vitejs.dev/) +* [tweakpane](https://tweakpane.github.io/docs//v3/monitor-bindings/) +* [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) 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/renderers/gaussian-renderer.ts b/src/renderers/gaussian-renderer.ts index 1684523..ca1254d 100644 --- a/src/renderers/gaussian-renderer.ts +++ b/src/renderers/gaussian-renderer.ts @@ -1,11 +1,11 @@ import { PointCloud } from '../utils/load'; import preprocessWGSL from '../shaders/preprocess.wgsl'; import renderWGSL from '../shaders/gaussian.wgsl'; -import { get_sorter,c_histogram_block_rows,C } from '../sort/sort'; +import { get_sorter, c_histogram_block_rows, C } from '../sort/sort'; import { Renderer } from './renderer'; export interface GaussianRenderer extends Renderer { - + set_gaussian_multiplier: (multiplier: number) => void, } // Utility to create GPU buffers @@ -17,7 +17,13 @@ const createBuffer = ( data?: ArrayBuffer | ArrayBufferView ) => { const buffer = device.createBuffer({ label, size, usage }); - if (data) device.queue.writeBuffer(buffer, 0, data); + if (data !== undefined) { + if (data instanceof ArrayBuffer) { + device.queue.writeBuffer(buffer, 0, data); + } else { + device.queue.writeBuffer(buffer, 0, data.buffer as ArrayBuffer, data.byteOffset, data.byteLength); + } + } return buffer; }; @@ -29,19 +35,127 @@ export default function get_renderer( ): GaussianRenderer { const sorter = get_sorter(pc.num_points, device); - + // =============================================== // Initialize GPU Buffers // =============================================== const nulling_data = new Uint32Array([0]); + const render_settings_buffer = createBuffer( + device, + 'render settings buffer', + 8, + GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, + new Float32Array([1.0, pc.sh_deg]), // gaussian_multiplier, sh_degree + ); + + const splats_buffer = createBuffer( + device, + 'splats buffer', + pc.num_points * 5 * 4, // 1 + 2 + 2 u32's per splat + GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, + ); + + const indirect_draw_buffer = createBuffer( + device, + 'indirect draw buffer', + 4 * 4, + GPUBufferUsage.INDIRECT | GPUBufferUsage.COPY_DST, + new Uint32Array([6, pc.num_points, 0, 0]) + ); + + // Bind group layouts + const camera_bind_group_layout = device.createBindGroupLayout({ + label: 'camera bind group layout', + entries: [ + { // camera + binding: 0, + visibility: GPUShaderStage.COMPUTE | GPUShaderStage.VERTEX | GPUShaderStage.FRAGMENT, + buffer: { type: 'uniform' }, + }, + { // render_settings + binding: 1, + visibility: GPUShaderStage.COMPUTE | GPUShaderStage.VERTEX | GPUShaderStage.FRAGMENT, + buffer: { type: 'uniform' }, + }, + ], + }); + const preprocess_gaussian_bind_group_layout = device.createBindGroupLayout({ + label: 'gaussian bind group layout', + entries: [ + { // gaussians + binding: 0, + visibility: GPUShaderStage.COMPUTE, + buffer: { type: 'read-only-storage' }, + }, + { // splats + binding: 1, + visibility: GPUShaderStage.COMPUTE, + buffer: { type: 'storage' }, + }, + { // sh_coeffs + binding: 2, + visibility: GPUShaderStage.COMPUTE, + buffer: { type: 'read-only-storage' }, + }, + ], + }); + const sort_bind_group_layout = device.createBindGroupLayout({ + label: 'sort bind group layout', + entries: [ + { // sort_infos + binding: 0, + visibility: GPUShaderStage.COMPUTE, + buffer: { type: 'storage' }, + }, + { // sort_depths + binding: 1, + visibility: GPUShaderStage.COMPUTE, + buffer: { type: 'storage' }, + }, + { // sort_indices + binding: 2, + visibility: GPUShaderStage.COMPUTE, + buffer: { type: 'storage' }, + }, + { // sort_dispatch + binding: 3, + visibility: GPUShaderStage.COMPUTE, + buffer: { type: 'storage' }, + }, + ], + }); + const render_gaussian_bind_group_layout = device.createBindGroupLayout({ + label: 'render gaussian bind group layout', + entries: [ + { // splats + binding: 0, + visibility: GPUShaderStage.VERTEX | GPUShaderStage.FRAGMENT, + buffer: { type: 'read-only-storage' }, + }, + { + // sort indices + binding: 1, + visibility: GPUShaderStage.VERTEX, + buffer: { type: 'read-only-storage' }, + }, + ], + }); + // =============================================== // Create Compute Pipeline and Bind Groups // =============================================== const preprocess_pipeline = device.createComputePipeline({ label: 'preprocess', - layout: 'auto', + layout: device.createPipelineLayout({ + label: 'preprocess pipeline layout', + bindGroupLayouts: [ + camera_bind_group_layout, + preprocess_gaussian_bind_group_layout, + sort_bind_group_layout, + ], + }), compute: { module: device.createShaderModule({ code: preprocessWGSL }), entryPoint: 'preprocess', @@ -52,9 +166,28 @@ export default function get_renderer( }, }); + const camera_bind_group = device.createBindGroup({ + label: 'camera bind group', + layout: camera_bind_group_layout, + entries: [ + { binding: 0, resource: { buffer: camera_buffer } }, + { binding: 1, resource: { buffer: render_settings_buffer } }, + ], + }); + + const preprocess_gaussian_bind_group = device.createBindGroup({ + label: 'gaussian bind group', + layout: preprocess_gaussian_bind_group_layout, + entries: [ + { binding: 0, resource: { buffer: pc.gaussian_3d_buffer } }, + { binding: 1, resource: { buffer: splats_buffer } }, + { binding: 2, resource: { buffer: pc.sh_buffer } }, + ], + }); + const sort_bind_group = device.createBindGroup({ label: 'sort', - layout: preprocess_pipeline.getBindGroupLayout(2), + layout: sort_bind_group_layout, entries: [ { binding: 0, resource: { buffer: sorter.sort_info_buffer } }, { binding: 1, resource: { buffer: sorter.ping_pong[0].sort_depths_buffer } }, @@ -63,24 +196,109 @@ export default function get_renderer( ], }); - // =============================================== // Create Render Pipeline and Bind Groups // =============================================== - + const render_pipeline = device.createRenderPipeline({ + label: 'render gaussian pipeline', + layout: device.createPipelineLayout({ + label: 'render gaussian pipeline layout', + bindGroupLayouts: [ + camera_bind_group_layout, + render_gaussian_bind_group_layout, + ], + }), + vertex: { + module: device.createShaderModule({ code: renderWGSL }), + entryPoint: 'vs_main', + buffers: [], + }, + fragment: { + module: device.createShaderModule({ code: renderWGSL }), + entryPoint: 'fs_main', + targets: [{ + format: presentation_format, + blend: { + color: { + srcFactor: 'src-alpha', + dstFactor: 'one-minus-src-alpha', + operation: 'add', + }, + alpha: { + srcFactor: 'one', + dstFactor: 'one-minus-src-alpha', + operation: 'add', + }, + }, + writeMask: GPUColorWrite.ALL, + }], + }, + primitive: { + topology: 'triangle-list', + }, + }); + + const render_gaussian_bind_group = device.createBindGroup({ + label: 'render gaussian bind group', + layout: render_gaussian_bind_group_layout, + entries: [ + { binding: 0, resource: { buffer: splats_buffer } }, + { binding: 1, resource: { buffer: sorter.ping_pong[0].sort_indices_buffer } }, + ], + }); // =============================================== // Command Encoder Functions // =============================================== - + const preprocess = (encoder: GPUCommandEncoder) => { + let pass = encoder.beginComputePass({ label: 'preprocess gaussians' }); + pass.setPipeline(preprocess_pipeline); + pass.setBindGroup(0, camera_bind_group); + pass.setBindGroup(1, preprocess_gaussian_bind_group); + pass.setBindGroup(2, sort_bind_group); + const workgroupCount = Math.ceil(pc.num_points / C.histogram_wg_size); + pass.dispatchWorkgroups(workgroupCount); + pass.end(); + }; + + const render = (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { + const pass = encoder.beginRenderPass({ + label: 'render gaussians', + colorAttachments: [ + { + view: texture_view, + loadOp: 'clear', + storeOp: 'store', + } + ], + }); + pass.setPipeline(render_pipeline); + pass.setBindGroup(0, camera_bind_group); + pass.setBindGroup(1, render_gaussian_bind_group); + + pass.drawIndirect(indirect_draw_buffer, 0); + pass.end(); + }; // =============================================== // Return Render Object // =============================================== return { frame: (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { + // Set sort_info.keys_size to 0 + device.queue.writeBuffer(sorter.sort_info_buffer, 0, nulling_data); + // Set sort_dispatch.dispatch_x to 0 + device.queue.writeBuffer(sorter.sort_dispatch_indirect_buffer, 0, nulling_data); + preprocess(encoder); sorter.sort(encoder); + // Copy the updated instance count (keys_size) into the indirect draw arguments without + // clobbering the fixed vertex count (6 vertices per splat quad). + encoder.copyBufferToBuffer(sorter.sort_info_buffer, 0, indirect_draw_buffer, 4, 4); + render(encoder, texture_view); }, camera_buffer, + set_gaussian_multiplier: (multiplier: number) => { + device.queue.writeBuffer(render_settings_buffer, 0, new Float32Array([multiplier, pc.sh_deg])); + }, }; } diff --git a/src/renderers/renderer.ts b/src/renderers/renderer.ts index ffdf9ba..7af6a4e 100644 --- a/src/renderers/renderer.ts +++ b/src/renderers/renderer.ts @@ -3,7 +3,7 @@ import { Pane } from 'tweakpane'; import * as TweakpaneFileImportPlugin from 'tweakpane-plugin-file-import'; import { default as get_renderer_gaussian, GaussianRenderer } from './gaussian-renderer'; import { default as get_renderer_pointcloud } from './point-cloud-renderer'; -import { Camera, load_camera_presets} from '../camera/camera'; +import { Camera, load_camera_presets } from '../camera/camera'; import { CameraControl } from '../camera/camera-control'; import { time, timeReturn } from '../utils/simple-console'; @@ -17,14 +17,14 @@ export default async function init( context: GPUCanvasContext, device: GPUDevice ) { - let ply_file_loaded = false; - let cam_file_loaded = false; + let ply_file_loaded = false; + let cam_file_loaded = false; let renderers: { pointcloud?: Renderer, gaussian?: Renderer } = {}; - let gaussian_renderer: GaussianRenderer | undefined; - let pointcloud_renderer: Renderer | undefined; - let renderer: Renderer | undefined; + let gaussian_renderer: GaussianRenderer | undefined; + let pointcloud_renderer: Renderer | undefined; + let renderer: Renderer | undefined; let cameras; - + const camera = new Camera(canvas, device); const control = new CameraControl(camera); @@ -35,14 +35,14 @@ export default async function init( camera.on_update_canvas(); }); observer.observe(canvas); - + const presentation_format = navigator.gpu.getPreferredCanvasFormat(); context.configure({ device, format: presentation_format, alphaMode: 'opaque', }); - + // Tweakpane: easily adding tweak control for parameters. const params = { @@ -60,7 +60,7 @@ export default async function init( pane.registerPlugin(TweakpaneFileImportPlugin); { pane.addMonitor(params, 'fps', { - readonly:true + readonly: true }); } { @@ -80,22 +80,22 @@ export default async function init( filetypes: ['.ply'], invalidFiletypeMessage: "We can't accept those filetypes!" }) - .on('change', async (file) => { - const uploadedFile = file.value; - if (uploadedFile) { - 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); - renderers = { - pointcloud: pointcloud_renderer, - gaussian: gaussian_renderer, - }; - renderer = renderers[params.renderer]; - ply_file_loaded = true; - }else{ - ply_file_loaded = false; - } - }); + .on('change', async (file) => { + const uploadedFile = file.value; + if (uploadedFile) { + 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); + renderers = { + pointcloud: pointcloud_renderer, + gaussian: gaussian_renderer, + }; + renderer = renderers[params.renderer]; + ply_file_loaded = true; + } else { + ply_file_loaded = false; + } + }); } { pane.addInput(params, 'cam_file', { @@ -104,29 +104,32 @@ export default async function init( filetypes: ['.json'], invalidFiletypeMessage: "We can't accept those filetypes!" }) - .on('change', async (file) => { - const uploadedFile = file.value; - if (uploadedFile) { - cameras=await load_camera_presets(file.value); - camera.set_preset(cameras[0]); - cam_file_loaded = true; - }else{ - cam_file_loaded = false; - } - }); + .on('change', async (file) => { + const uploadedFile = file.value; + if (uploadedFile) { + cameras = await load_camera_presets(file.value); + camera.set_preset(cameras[0]); + cam_file_loaded = true; + } else { + cam_file_loaded = false; + } + }); } { pane.addInput( params, 'gaussian_multiplier', - {min: 0, max: 1.5} + { min: 0, max: 1.5 } ).on('change', (e) => { - //TODO: Bind constants to the gaussian renderer. + // TODO: Bind constants to the gaussian renderer. + if (gaussian_renderer) { + gaussian_renderer.set_gaussian_multiplier(e.value); + } }); } document.addEventListener('keydown', (event) => { - switch(event.key) { + switch (event.key) { case '0': case '1': case '2': @@ -146,7 +149,7 @@ export default async function init( function frame() { if (ply_file_loaded && cam_file_loaded) { - params.fps=1.0/timeReturn()*1000.0; + params.fps = 1.0 / timeReturn() * 1000.0; time(); const encoder = device.createCommandEncoder(); const texture_view = context.getCurrentTexture().createView(); diff --git a/src/shaders/gaussian.wgsl b/src/shaders/gaussian.wgsl index 759226d..ee2854b 100644 --- a/src/shaders/gaussian.wgsl +++ b/src/shaders/gaussian.wgsl @@ -1,22 +1,87 @@ +struct CameraUniforms { + view: mat4x4, + view_inv: mat4x4, + proj: mat4x4, + proj_inv: mat4x4, + viewport: vec2, + focal: vec2 +}; + +struct RenderSettings { + gaussian_scaling: f32, + sh_deg: f32, +} + struct VertexOutput { @builtin(position) position: vec4, - //TODO: information passed from vertex shader to fragment shader + // TODO: information passed from vertex shader to fragment shader + @location(0) center_pixel: vec2, + @location(1) color: vec4, + @location(2) conic: vec3, }; struct Splat { - //TODO: information defined in preprocess compute shader + // TODO: information defined in preprocess compute shader + center: u32, + color: array, + conic_radius: array, }; +@group(0) @binding(0) +var camera: CameraUniforms; +@group(0) @binding(1) +var render_settings: RenderSettings; + +@group(1) @binding(0) +var splats: array; +@group(1) @binding(1) +var sort_indices : array; + @vertex fn vs_main( + @builtin(vertex_index) vertex_index: u32, + @builtin(instance_index) instance_index: u32 ) -> VertexOutput { - //TODO: reconstruct 2D quad based on information from splat, pass + // TODO: reconstruct 2D quad based on information from splat, pass + let splat = splats[sort_indices[instance_index]]; + + let center = unpack2x16float(splat.center); // in NDC space + let color = vec4( + unpack2x16float(splat.color[0]), + unpack2x16float(splat.color[1]) + ); + let conic_0 = unpack2x16float(splat.conic_radius[0]); + let conic_1_radius = unpack2x16float(splat.conic_radius[1]); + let conic = vec3(conic_0, conic_1_radius[0]); + let radius = conic_1_radius[1]; + let quad_size = 2.0 * radius / camera.viewport; // in NDC space + + let offsets = array, 6>( + vec2(-1, -1), + vec2(1, -1), + vec2(-1, 1), + vec2(-1, 1), + vec2(1, -1), + vec2(1, 1) + ); + let offset = offsets[vertex_index] * quad_size; // in NDC space + var out: VertexOutput; - out.position = vec4(1. ,1. , 0., 1.); + out.position = vec4(center + offset, 0.0, 1.0); + out.center_pixel = (center * vec2(0.5, -0.5) + 0.5) * camera.viewport; // in pixel space + out.color = color; + out.conic = conic; return out; } @fragment fn fs_main(in: VertexOutput) -> @location(0) vec4 { - return vec4(1.); -} \ No newline at end of file + // https://github.com/graphdeco-inria/diff-gaussian-rasterization/blob/main/cuda_rasterizer/forward.cu#L330 + let d = in.position.xy - in.center_pixel; // in pixel space + let power = -0.5 * (in.conic.x * d.x * d.x + in.conic.z * d.y * d.y) - in.conic.y * d.x * d.y; + if (power > 0.0) { + discard; + } + let alpha = min(0.99, in.color.a * exp(power)); + return vec4(in.color.rgb, alpha); +} diff --git a/src/shaders/point_cloud.wgsl b/src/shaders/point_cloud.wgsl index 01dded1..ced9e90 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..74574bd 100644 --- a/src/shaders/preprocess.wgsl +++ b/src/shaders/preprocess.wgsl @@ -56,10 +56,25 @@ struct Gaussian { }; struct Splat { - //TODO: store information for 2D splat rendering + // TODO: store information for 2D splat rendering + center: u32, + color: array, + conic_radius: array, }; -//TODO: bind your data here +// 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_coeffs: array; + @group(2) @binding(0) var sort_infos: SortInfos; @group(2) @binding(1) @@ -69,10 +84,32 @@ var sort_indices : array; @group(2) @binding(3) var sort_dispatch: DispatchIndirect; +const MAX_SH_COEFFS: u32 = 16u; +const SH_CHANNELS: u32 = 3u; +const SH_HALFS_PER_GAUSSIAN: u32 = MAX_SH_COEFFS * SH_CHANNELS; // 48 half floats + +fn read_sh_half(splat_idx: u32, half_offset: u32) -> f32 { + // Each u32 in the buffer stores two f16 values packed together. + let global_half_index = splat_idx * SH_HALFS_PER_GAUSSIAN + half_offset; + let packed_index = global_half_index >> 1u; + let packed = sh_coeffs[packed_index]; + let decoded = unpack2x16float(packed); + let is_high = (global_half_index & 1u) == 1u; + return select(decoded.x, decoded.y, is_high); +} + /// 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); + if (c_idx >= MAX_SH_COEFFS) { + return vec3(0.0); + } + + let base_half = c_idx * SH_CHANNELS; + return vec3( + read_sh_half(splat_idx, base_half + 0u), + read_sh_half(splat_idx, base_half + 1u), + read_sh_half(splat_idx, base_half + 2u) + ); } // spherical harmonics evaluation with Condon–Shortley phase @@ -108,11 +145,129 @@ 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)); +} + +fn float_to_u32_key(f: f32) -> u32 { + let u = bitcast(f); + let neg = (u & 0x80000000u) != 0u; + let mask = select(0x80000000u, 0xffffffffu, neg); + return u ^ mask; +} + @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 + // 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 pos_world = vec3(pos_xy, pos_z_opacity[0]); + let opacity = sigmoid(pos_z_opacity[1]); + + let pos_view = (camera.view * vec4(pos_world, 1.0)).xyz; + var pos_ndc = camera.proj * vec4(pos_view, 1.0); + pos_ndc /= pos_ndc.w; + // View-frustum culling: + if ( + !(abs(pos_ndc.x) <= 1.1) || + !(abs(pos_ndc.y) <= 1.1) || + !(pos_ndc.z >= -0.1 && pos_ndc.z <= 1.1) + ) { + return; + } + + let rot_wx = unpack2x16float(gaussian.rot[0]); + let rot_yz = unpack2x16float(gaussian.rot[1]); + let rot = vec4(rot_wx[1], rot_yz, rot_wx[0]); // xyzw + + let scale_xy = exp(unpack2x16float(gaussian.scale[0])); + let scale_z_padding = exp(unpack2x16float(gaussian.scale[1])); + let scale = vec3(scale_xy, scale_z_padding[0]); + + // https://github.com/kwea123/gaussian_splatting_notes + let R = mat3x3( + 1.0 - 2.0 * (rot.y * rot.y + rot.z * rot.z), + 2.0 * (rot.x * rot.y - rot.w * rot.z), + 2.0 * (rot.x * rot.z + rot.w * rot.y), + 2.0 * (rot.x * rot.y + rot.w * rot.z), + 1.0 - 2.0 * (rot.x * rot.x + rot.z * rot.z), + 2.0 * (rot.y * rot.z - rot.w * rot.x), + 2.0 * (rot.x * rot.z - rot.w * rot.y), + 2.0 * (rot.y * rot.z + rot.w * rot.x), + 1.0 - 2.0 * (rot.x * rot.x + rot.y * rot.y) + ); + let S = render_settings.gaussian_scaling * mat3x3( + scale.x, 0.0, 0.0, + 0.0, scale.y, 0.0, + 0.0, 0.0, scale.z + ); + let cov3d = R * S * S * transpose(R); + let Vrk = mat3x3( // TBD + 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] + ); + + // Viewing transformation + let W = transpose(mat3x3( + camera.proj[0].xyz, + camera.proj[1].xyz, + camera.proj[2].xyz + )); + // Jacobian of the affine approximation of the projective transformation + let J = mat3x3( + camera.focal.x / pos_view.z, 0.0, -(camera.focal.x * pos_view.x) / (pos_view.z * pos_view.z), + 0.0, camera.focal.y / pos_view.z, -(camera.focal.y * pos_view.y) / (pos_view.z * pos_view.z), + 0.0, 0.0, 0.0 + ); + let T = W * J; + + var cov2d = transpose(T) * transpose(Vrk) * T; + // A small trick to ensure the numerical stability of the inverse + cov2d[0][0] += 0.3; + cov2d[1][1] += 0.3; + + let cov2d_flat = vec3(cov2d[0][0], cov2d[0][1], cov2d[1][1]); + let det = cov2d_flat.x * cov2d_flat.z - cov2d_flat.y * cov2d_flat.y; + let det_inv = 1.0 / det; + let conic = vec3( + cov2d_flat.z * det_inv, + -cov2d_flat.y * det_inv, + cov2d_flat.x * det_inv + ); + + let mid = 0.5 * (cov2d_flat.x + cov2d_flat.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))); // Make it at least 1 + + let camera_pos = camera.view_inv[3].xyz; + let view_dir = normalize(pos_world - camera_pos); + let color = computeColorFromSH(view_dir, idx, u32(render_settings.sh_deg)); + + let splat_idx = atomicAdd(&sort_infos.keys_size, 1u); + splats[splat_idx].center = pack2x16float(pos_ndc.xy); + splats[splat_idx].color = array( + pack2x16float(color.rg), + pack2x16float(vec2(color.b, opacity)) + ); + splats[splat_idx].conic_radius = array( + pack2x16float(conic.xy), + pack2x16float(vec2(conic.z, radius)) + ); + + sort_indices[splat_idx] = splat_idx; + sort_depths[splat_idx] = float_to_u32_key(-pos_view.z); let keys_per_dispatch = workgroupSize * sortKeyPerThread; // increment DispatchIndirect.dispatchx each time you reach limit for one dispatch of keys + if (splat_idx % keys_per_dispatch == 0u) { + atomicAdd(&sort_dispatch.dispatch_x, 1u); + } } \ No newline at end of file