diff --git a/README.md b/README.md index f99cdff..95fff3c 100644 --- a/README.md +++ b/README.md @@ -1,26 +1,90 @@ -# Project5-WebGPU-Gaussian-Splat-Viewer +WebGPU Gaussian Splat Viewer +============================ -**University of Pennsylvania, CIS 565: GPU Programming and Architecture, Project 4** +**University of Pennsylvania, CIS 5650: 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) +* Dominik Kau ([LinkedIn](https://www.linkedin.com/in/dominikkau/)) +* Tested on: **Google Chrome 132.0**, macOS Sequoia 15.1.1, Apple M3 Pro -### Live Demo +## Live Demo -[![](img/thumb.png)](http://TODO.github.io/Project4-WebGPU-Forward-Plus-and-Clustered-Deferred) +[Click here to check out my implementation!](https://dominikkau.github.io/Project5-WebGPU-Gaussian-Splat-Viewer/) +You'll need some trained gaussian splat input data. -### Demo Video/GIF +## Overview -[![](img/video.mp4)](TODO) +Gaussian splatting is a technique for reconstructing a 3D scene from images. +The method consists of a training period during which 3-dimensional Gaussian distributions are placed around the scene and optimized to best capture the given image data set. +In this phase the "nice" mathematical properties of a Gaussian distribution can be used to calculate gradients with respect to the target variables such as position, color and size of the Gaussian. +Thus, gradient based optimization techniques can be used. -### (TODO: Your README) +Afterwards the camera can be moved freely throughout the reconstructed scene. +This allows the synthesis of images from new perspectives or videos that can be recorded in real-time with the methods presented in [this paper](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/). +This project implements a renderer to display a scene consisting of 3D Gaussians after having been trained. -*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. +## Features -This assignment has a considerable amount of performance analysis compared -to implementation work. Complete the implementation early to leave time! +### Point Cloud Renderer + +![Image of a scene containing a bench and bicycle visualized only with dots.](images/points.png) + +The point cloud renderer is a simple visualization of the Gaussians that are present in the dataset. +Each visible Gaussian is displayed as a point placed in its center. + +### Gaussian Renderer + +![Reconstructed color image of a scene containing a bench and bicycle.](images/gaussian.png) + +The main feature of this project. +The Gaussian renderer extracts the information from the given dataset and - next to position - reconstructs color (through spherical harmonics) and size of the Gaussians. +The reconstruction of information takes place in a preprocessing compute shader before the render pipeline is carried out. +A simple view-frustum culling is used to reduce the computations in the render pass. +The final rendering pipeline employs alpha blending to composite the Gaussians into an image. +The alpha blending requires the Gaussians to be rendered from furthest back to the closest. +For this, another compute shader is used that sorts the visible Gaussians based on their depth. + +### Half Precision Floating Point Optimization + +This optimization feature changes the preprocessing compute shader to operate on half precision floating point numbers. +This allows for a smaller memory footprint - instead of 32 bits, only 16 bits are needed. +Moreover, modern GPUs can run operations of half precision floats significantly faster. +Not all computations are carried out using 16-bit floats, as some mathematical operations need the higher accuracy to ensure an accurate result without artifacts. + +As a result there are only very slight visual differences between the 16-bit and 32-bit version. + +## Performance Analysis + +### Comparison Between Point Cloud and Gaussian Renderer + +Obviously, the performance of the point cloud renderer is significantly better than the Gaussian renderer. +The point cloud renderer doesn't have a preprocessing or sorting compute shader that needs to run, so it performs drastically fewer computations. + +### Influence of Workgroup Size + +Decreasing the workgroup size leads to a small performance hit. +As this project does not rely on shared memory, this is probably due to how the GPU dispatching is handled. +Smaller workgroups will lead to more work on the dispatching side. + +### Influence of View-Frustum Culling + +The view frustum culling has a significant effect on the performance. +This is especially true for indoor scenes, because there will be many points present in the dataset that don't have to be rendered to the camera. +Because the view-frustum culling is implemented at the beginning of the preprocessing compute shader, a lot of computations can be avoided by returning early. +However, this is also dependent on the degree of thread divergence, because the effect of an early return will only have a big impact on performance if multiple all "connected" threads (in CUDA: a warp) quit early. + +### Influence of Number of Gaussians + +The number of Gaussians has a significant impact on performance. +Firstly, all buffers are sized according to the total number of points - leading to memory bottlenecks. +Secondly, the more Gaussians are visible at the same time, the more computations the compute shaders and the rendering shaders have to carry out. +However, the second point is scene dependent because of view-frustum culling. +A scene with few total but spatially concentrated Gaussians will run at a similar speed as a scene with more but evenly distributed Gaussians, that are not all visible at the same time. +That assumes that memory does not pose an issue. + +### Influence of Using Half Precision Floats + +The usage of half precision floats increases performance, but not as much as expected. +This might be because I'm still using 32-bit computations in operations that I found critical for accurate visual output. ### Credits diff --git a/images/gaussian.png b/images/gaussian.png new file mode 100644 index 0000000..768ef4d Binary files /dev/null and b/images/gaussian.png differ diff --git a/images/points.png b/images/points.png new file mode 100644 index 0000000..991f734 Binary files /dev/null and b/images/points.png differ diff --git a/src/camera/camera-control.ts b/src/camera/camera-control.ts index f6d6d99..01cf0e9 100644 --- a/src/camera/camera-control.ts +++ b/src/camera/camera-control.ts @@ -2,93 +2,93 @@ import { vec3, mat3, mat4, quat } from 'wgpu-matrix'; import { Camera } from './camera'; export class CameraControl { - element: HTMLCanvasElement; - constructor(private camera: Camera) { - this.register_element(camera.canvas); - } + element: HTMLCanvasElement; + constructor(private camera: Camera) { + this.register_element(camera.canvas); + } + + register_element(value: HTMLCanvasElement) { + if (this.element && this.element != value) { + this.element.removeEventListener('pointerdown', this.downCallback.bind(this)); + this.element.removeEventListener('pointermove', this.moveCallback.bind(this)); + this.element.removeEventListener('pointerup', this.upCallback.bind(this)); + this.element.removeEventListener('wheel', this.wheelCallback.bind(this)); + } - register_element(value: HTMLCanvasElement) { - if (this.element && this.element != value) { - this.element.removeEventListener('pointerdown', this.downCallback.bind(this)); - this.element.removeEventListener('pointermove', this.moveCallback.bind(this)); - this.element.removeEventListener('pointerup', this.upCallback.bind(this)); - this.element.removeEventListener('wheel', this.wheelCallback.bind(this)); + this.element = value; + this.element.addEventListener('pointerdown', this.downCallback.bind(this)); + this.element.addEventListener('pointermove', this.moveCallback.bind(this)); + this.element.addEventListener('pointerup', this.upCallback.bind(this)); + this.element.addEventListener('wheel', this.wheelCallback.bind(this)); + this.element.addEventListener('contextmenu', (e) => { e.preventDefault(); }); } - this.element = value; - this.element.addEventListener('pointerdown', this.downCallback.bind(this)); - this.element.addEventListener('pointermove', this.moveCallback.bind(this)); - this.element.addEventListener('pointerup', this.upCallback.bind(this)); - this.element.addEventListener('wheel', this.wheelCallback.bind(this)); - this.element.addEventListener('contextmenu', (e) => { e.preventDefault(); }); - } + private panning = false; + private rotating = false; + private lastX: number; + private lastY: number; - private panning = false; - private rotating = false; - private lastX: number; - private lastY: number; + downCallback(event: PointerEvent) { + if (!event.isPrimary) { + return; + } - downCallback(event: PointerEvent) { - if (!event.isPrimary) { - return; + if (event.button === 0) { + this.rotating = true; + this.panning = false; + } else { + this.rotating = false; + this.panning = true; + } + this.lastX = event.pageX; + this.lastY = event.pageY; } + moveCallback(event: PointerEvent) { + if (!(this.rotating || this.panning)) { + return; + } - if (event.button === 0) { - this.rotating = true; - this.panning = false; - } else { - this.rotating = false; - this.panning = true; + const xDelta = event.pageX - this.lastX; + const yDelta = event.pageY - this.lastY; + this.lastX = event.pageX; + this.lastY = event.pageY; + + if (this.rotating) { + this.rotate(xDelta, yDelta); + } else if (this.panning) { + this.pan(xDelta, yDelta); + } } - this.lastX = event.pageX; - this.lastY = event.pageY; - } - moveCallback(event: PointerEvent) { - if (!(this.rotating || this.panning)) { - return; + upCallback(event: PointerEvent) { + this.rotating = false; + this.panning = false; + event.preventDefault(); } - - const xDelta = event.pageX - this.lastX; - const yDelta = event.pageY - this.lastY; - this.lastX = event.pageX; - this.lastY = event.pageY; - - if (this.rotating) { - this.rotate(xDelta, yDelta); - } else if (this.panning) { - this.pan(xDelta, yDelta); + wheelCallback(event: WheelEvent) { + event.preventDefault(); + const delta = vec3.mulScalar(this.camera.look, -event.deltaY * 0.001); + vec3.add(delta, this.camera.position, this.camera.position); + this.camera.update_buffer(); } - } - upCallback(event: PointerEvent) { - this.rotating = false; - this.panning = false; - event.preventDefault(); - } - wheelCallback(event: WheelEvent) { - event.preventDefault(); - const delta = vec3.mulScalar(this.camera.look, -event.deltaY * 0.001); - vec3.add(delta, this.camera.position, this.camera.position); - this.camera.update_buffer(); - } - rotate(xDelta: number, yDelta: number) { - // const r = mat4.identity(); - // mat4.rotateY(r, -xDelta, r); - // mat4.rotateX(r, yDelta, r); - const r = mat4.fromQuat(quat.fromEuler(yDelta * 0.01, -xDelta * 0.01, 0, 'xyz')); + rotate(xDelta: number, yDelta: number) { + // const r = mat4.identity(); + // mat4.rotateY(r, -xDelta, r); + // mat4.rotateX(r, yDelta, r); + const r = mat4.fromQuat(quat.fromEuler(yDelta * 0.01, -xDelta * 0.01, 0, 'xyz')); - mat4.mul(r, this.camera.rotation, this.camera.rotation); + mat4.mul(r, this.camera.rotation, this.camera.rotation); - this.camera.update_buffer(); - } + this.camera.update_buffer(); + } - pan(xDelta: number, yDelta: number) { - const d = vec3.copy(this.camera.up); - vec3.mulScalar(d, -yDelta * 0.01, d); - vec3.add(d, this.camera.position, this.camera.position); - vec3.copy(this.camera.right, d); - vec3.mulScalar(d, -xDelta * 0.01, d); - vec3.add(d, this.camera.position, this.camera.position); - this.camera.update_buffer(); - } + pan(xDelta: number, yDelta: number) { + const d = vec3.copy(this.camera.up); + vec3.mulScalar(d, -yDelta * 0.01, d); + vec3.add(d, this.camera.position, this.camera.position); + vec3.copy(this.camera.right, d); + vec3.mulScalar(d, -xDelta * 0.01, d); + vec3.add(d, this.camera.position, this.camera.position); + this.camera.update_buffer(); + } }; \ No newline at end of file diff --git a/src/camera/camera.ts b/src/camera/camera.ts index 47ea1dc..7db249e 100644 --- a/src/camera/camera.ts +++ b/src/camera/camera.ts @@ -2,189 +2,195 @@ import { Mat3, mat3, Mat4, mat4, Vec3, vec3, Vec2, vec2 } from 'wgpu-matrix'; import { log, time, timeLog } from '../utils/simple-console'; interface CameraJson { - id: number - img_name: string - width: number - height: number - position: number[] - rotation: number[][] - fx: number - fy: number + id: number + img_name: string + width: number + height: number + position: number[] + rotation: number[][] + fx: number + fy: number }; function focal2fov(focal: number, pixels: number): number { - return 2 * Math.atan(pixels / (2 * focal)); + return 2 * Math.atan(pixels / (2 * focal)); } function fov2focal(fov: number, pixels: number): number { - return pixels / (2 * Math.tan(fov * 0.5)); + return pixels / (2 * Math.tan(fov * 0.5)); } function get_view_matrix(r: Mat4, t: Vec3): Mat4 { - const minus_t = vec3.mulScalar(t, -1); - return mat4.translate(r, minus_t); + const minus_t = vec3.mulScalar(t, -1); + return mat4.translate(r, minus_t); } function get_projection_matrix(znear: number, zfar: number, fov_x: number, fov_y: number) { - // return mat4.perspective(fov_y, 1, znear, zfar); - - const tan_half_fov_y = Math.tan(fov_y / 2.); - const tan_half_fov_x = Math.tan(fov_x / 2.); - - const top = tan_half_fov_y * znear; - const bottom = -top; - const right = tan_half_fov_x * znear; - const left = -right; - - const p = mat4.create(); - p[0] = 2.0 * znear / (right - left); - // p[5] = 2.0 * znear / (top - bottom); - p[5] = -2.0 * znear / (top - bottom); // flip Y - p[2] = (right + left) / (right - left); - p[6] = (top + bottom) / (top - bottom); - p[14] = 1.; - p[10] = zfar / (zfar - znear); - p[11] = -(zfar * znear) / (zfar - znear); - mat4.transpose(p, p); - - // p[0] = 2.0 * znear / (right - left); - // p[5] = 2.0 * znear / (top - bottom); - // p[8] = (right + left) / (right - left); - // p[9] = (top + bottom) / (top - bottom); - // p[10] = zfar / (zfar - znear); - // p[11] = -(zfar * znear) / (zfar - znear); - // p[14] = 1.; - // mat4.transpose(p, p); - return p; + // return mat4.perspective(fov_y, 1, znear, zfar); + + const tan_half_fov_y = Math.tan(fov_y / 2.); + const tan_half_fov_x = Math.tan(fov_x / 2.); + + const top = tan_half_fov_y * znear; + const bottom = -top; + const right = tan_half_fov_x * znear; + const left = -right; + + const p = mat4.create(); + p[0] = 2.0 * znear / (right - left); + // p[5] = 2.0 * znear / (top - bottom); + p[5] = -2.0 * znear / (top - bottom); // flip Y + p[2] = (right + left) / (right - left); + p[6] = (top + bottom) / (top - bottom); + p[14] = 1.; + p[10] = zfar / (zfar - znear); + p[11] = -(zfar * znear) / (zfar - znear); + mat4.transpose(p, p); + + // p[0] = 2.0 * znear / (right - left); + // p[5] = 2.0 * znear / (top - bottom); + // p[8] = (right + left) / (right - left); + // p[9] = (top + bottom) / (top - bottom); + // p[10] = zfar / (zfar - znear); + // p[11] = -(zfar * znear) / (zfar - znear); + // p[14] = 1.; + // mat4.transpose(p, p); + return p; } interface CameraPreset { - position: Vec3, - rotation: Mat4, + position: Vec3, + rotation: Mat4, } export async function load_camera_presets(file: string): Promise { - const blob = new Blob([file]); - const arrayBuffer = await new Promise((resolve, reject) => { - const reader = new FileReader(); - - reader.onload = function(event) { - resolve(event.target.result); // Resolve the promise with the ArrayBuffer - }; - - reader.onerror = reject; // Reject the promise in case of an error - reader.readAsArrayBuffer(blob); - }); - const text = new TextDecoder().decode(arrayBuffer as ArrayBuffer); - const json = JSON.parse(text); - log(`loaded cameras count: ${json.length}`); - - return json.map((j: CameraJson): CameraPreset => { - const position = vec3.clone(j.position); - // const rotation = mat3.create(...j.rotation.flat()); - const rotation = mat4.fromMat3(mat3.create(...j.rotation.flat())); - - return { - position, - rotation, - }; - }); + const blob = new Blob([file]); + const arrayBuffer = await new Promise((resolve, reject) => { + const reader = new FileReader(); + + reader.onload = function (event) { + resolve(event.target.result); // Resolve the promise with the ArrayBuffer + }; + + reader.onerror = reject; // Reject the promise in case of an error + reader.readAsArrayBuffer(blob); + }); + const text = new TextDecoder().decode(arrayBuffer as ArrayBuffer); + const json = JSON.parse(text); + log(`loaded cameras count: ${json.length}`); + + return json.map((j: CameraJson): CameraPreset => { + const position = vec3.clone(j.position); + // const rotation = mat3.create(...j.rotation.flat()); + const rotation = mat4.fromMat3(mat3.create(...j.rotation.flat())); + + return { + position, + rotation, + }; + }); } const c_size_vec2 = 4 * 2; const c_size_mat4 = 4 * 16; // byte size of mat4 (i.e. Float32Array(16)) -const c_size_camera_uniform = 4 * c_size_mat4 + 2 * c_size_vec2; +const c_size_camera_uniform = 4 * c_size_mat4 + 4 * c_size_vec2; interface CameraUniform { - view_matrix: Mat4, - view_inv_matrix: Mat4, - proj_matrix: Mat4, - proj_inv_matrix: Mat4, - - viewport: Vec2, - focal: Vec2, + view_matrix: Mat4, + view_inv_matrix: Mat4, + proj_matrix: Mat4, + proj_inv_matrix: Mat4, + + viewport: Vec2, + focal: Vec2, + clipping_planes: Vec2, } export function create_camera_uniform_buffer(device: GPUDevice) { - return device.createBuffer({ - label: 'camera uniform', - size: c_size_camera_uniform, - usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, - }); + return device.createBuffer({ + label: 'camera uniform', + size: c_size_camera_uniform, + usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, + }); } const intermediate_float_32_array = new Float32Array(c_size_camera_uniform / Float32Array.BYTES_PER_ELEMENT); export class Camera { - constructor( - public readonly canvas: HTMLCanvasElement, - private readonly device: GPUDevice, - ) { - this.uniform_buffer = create_camera_uniform_buffer(device); - this.on_update_canvas(); - } - - on_update_canvas(): void { - const focal = 0.5 * this.canvas.height / Math.tan(this.fovY * 0.5); - this.focal[0] = focal; - this.focal[1] = focal; - this.fovX = focal2fov(focal, this.canvas.width); - this.viewport[0] = this.canvas.width; - this.viewport[1] = this.canvas.height; - // const viewport_ratio = this.canvas.width / this.canvas.height; - - this.update_buffer(); - } - - readonly uniform_buffer: GPUBuffer; - - position = vec3.create(); - rotation = mat4.create(); - private fovY: number = 45 / 180 * Math.PI; - private fovX: number; - private focal: Vec2 = vec2.create(); - private viewport: Vec2 = vec2.create(); - - private view_matrix: Mat4 = mat4.identity(); - private proj_matrix: Mat4 = mat4.identity(); - - look = vec3.create(0, 0, 1); - up = vec3.create(0, 1, 0); - right = vec3.create(1, 0, 0); - - update_buffer(): void { - let offset = 0; - - this.view_matrix = get_view_matrix(this.rotation, this.position); - this.proj_matrix = get_projection_matrix(0.01, 100, this.fovX, this.fovY); - - const inv_view_matrix = mat4.inverse(this.view_matrix); - vec3.transformMat4Upper3x3(vec3.create(0, 0, 1), inv_view_matrix, this.look); - vec3.normalize(this.look, this.look); - - vec3.cross(this.up, this.look, this.right); - vec3.normalize(this.right, this.right); - - intermediate_float_32_array.set(this.view_matrix, offset); - offset += 16; - intermediate_float_32_array.set(inv_view_matrix, offset); - offset += 16; - intermediate_float_32_array.set(this.proj_matrix, offset); - offset += 16; - intermediate_float_32_array.set(mat4.inverse(this.proj_matrix), offset); - offset += 16; - intermediate_float_32_array.set(this.viewport, offset); - offset += 2; - intermediate_float_32_array.set(this.focal, offset); - offset += 2; - - this.device.queue.writeBuffer(this.uniform_buffer, 0, intermediate_float_32_array); - } - set_preset(preset: CameraPreset): void { - vec3.copy(preset.position, this.position); - mat4.copy(preset.rotation, this.rotation); - this.update_buffer(); - } + constructor( + public readonly canvas: HTMLCanvasElement, + private readonly device: GPUDevice, + ) { + this.uniform_buffer = create_camera_uniform_buffer(device); + this.on_update_canvas(); + } + + on_update_canvas(): void { + const focal = 0.5 * this.canvas.height / Math.tan(this.fovY * 0.5); + this.focal[0] = focal; + this.focal[1] = focal; + this.fovX = focal2fov(focal, this.canvas.width); + this.viewport[0] = this.canvas.width; + this.viewport[1] = this.canvas.height; + // const viewport_ratio = this.canvas.width / this.canvas.height; + + this.update_buffer(); + } + + readonly uniform_buffer: GPUBuffer; + + position = vec3.create(); + rotation = mat4.create(); + private fovY: number = 45 / 180 * Math.PI; + private fovX: number; + private focal: Vec2 = vec2.create(); + private viewport: Vec2 = vec2.create(); + + private view_matrix: Mat4 = mat4.identity(); + private proj_matrix: Mat4 = mat4.identity(); + + look = vec3.create(0, 0, 1); + up = vec3.create(0, 1, 0); + right = vec3.create(1, 0, 0); + + update_buffer(): void { + let offset = 0; + + const clippingPlanes = vec2.create(0.01, 100); + + this.view_matrix = get_view_matrix(this.rotation, this.position); + this.proj_matrix = get_projection_matrix(clippingPlanes[0], clippingPlanes[1], this.fovX, this.fovY); + // this.proj_matrix = get_projection_matrix(0.01, 100, this.fovX, this.fovY); + + const inv_view_matrix = mat4.inverse(this.view_matrix); + vec3.transformMat4Upper3x3(vec3.create(0, 0, 1), inv_view_matrix, this.look); + vec3.normalize(this.look, this.look); + + vec3.cross(this.up, this.look, this.right); + vec3.normalize(this.right, this.right); + + intermediate_float_32_array.set(this.view_matrix, offset); + offset += 16; + intermediate_float_32_array.set(inv_view_matrix, offset); + offset += 16; + intermediate_float_32_array.set(this.proj_matrix, offset); + offset += 16; + intermediate_float_32_array.set(mat4.inverse(this.proj_matrix), offset); + offset += 16; + intermediate_float_32_array.set(this.viewport, offset); + offset += 2; + intermediate_float_32_array.set(this.focal, offset); + offset += 2; + intermediate_float_32_array.set(clippingPlanes, offset); + offset += 2; + + this.device.queue.writeBuffer(this.uniform_buffer, 0, intermediate_float_32_array); + } + set_preset(preset: CameraPreset): void { + vec3.copy(preset.position, this.position); + mat4.copy(preset.rotation, this.rotation); + this.update_buffer(); + } }; diff --git a/src/main.ts b/src/main.ts index 25efcb5..2a408ae 100644 --- a/src/main.ts +++ b/src/main.ts @@ -3,32 +3,33 @@ import init from './renderers/renderer'; import { assert } from './utils/util'; (async () => { - - if (navigator.gpu === undefined) { - const h = document.querySelector('#title') as HTMLElement; - h.innerText = 'WebGPU is not supported in this browser.'; - return; - } - const adapter = await navigator.gpu.requestAdapter({ - powerPreference: 'high-performance', - }); - if (adapter === null) { - const h = document.querySelector('#title') as HTMLElement; - h.innerText = 'No adapter is available for WebGPU.'; - return; - } - - const device = await adapter.requestDevice({ - requiredLimits: { - maxComputeWorkgroupStorageSize: adapter.limits.maxComputeWorkgroupStorageSize, - maxStorageBufferBindingSize: adapter.limits. maxStorageBufferBindingSize - }, - }); - const canvas = document.querySelector('#webgpu-canvas'); - assert(canvas !== null); - const context = canvas.getContext('webgpu') as GPUCanvasContext; - - init(canvas, context, device); - + if (navigator.gpu === undefined) { + const h = document.querySelector('#title') as HTMLElement; + h.innerText = 'WebGPU is not supported in this browser.'; + return; + } + const adapter = await navigator.gpu.requestAdapter({ + powerPreference: 'high-performance', + }); + if (adapter === null) { + const h = document.querySelector('#title') as HTMLElement; + h.innerText = 'No adapter is available for WebGPU.'; + return; + } + + const device = await adapter.requestDevice({ + requiredLimits: { + maxComputeWorkgroupStorageSize: adapter.limits.maxComputeWorkgroupStorageSize, + maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize + }, + requiredFeatures: [ 'shader-f16' ] + }); + + const canvas = document.querySelector('#webgpu-canvas'); + assert(canvas !== null); + const context = canvas.getContext('webgpu') as GPUCanvasContext; + + init(canvas, context, device); + })(); \ No newline at end of file diff --git a/src/renderers/gaussian-renderer.ts b/src/renderers/gaussian-renderer.ts index 1684523..e02244e 100644 --- a/src/renderers/gaussian-renderer.ts +++ b/src/renderers/gaussian-renderer.ts @@ -1,86 +1,241 @@ import { PointCloud } from '../utils/load'; import preprocessWGSL from '../shaders/preprocess.wgsl'; +import preprocessWGSL16 from '../shaders/preprocess-16.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 { - -} + renderSettingsBuffer: GPUBuffer; + preprocessPipeline: GPUComputePipeline; +}; // Utility to create GPU buffers const createBuffer = ( - device: GPUDevice, - label: string, - size: number, - usage: GPUBufferUsageFlags, - data?: ArrayBuffer | ArrayBufferView + device: GPUDevice, + label: string, + size: number, + usage: GPUBufferUsageFlags, + data?: ArrayBuffer | ArrayBufferView ) => { - const buffer = device.createBuffer({ label, size, usage }); - if (data) device.queue.writeBuffer(buffer, 0, data); - return buffer; + const buffer = device.createBuffer({ label, size, usage }); + if (data) device.queue.writeBuffer(buffer, 0, data); + return buffer; }; export default function get_renderer( - pc: PointCloud, - device: GPUDevice, - presentation_format: GPUTextureFormat, - camera_buffer: GPUBuffer, + pc: PointCloud, + device: GPUDevice, + presentation_format: GPUTextureFormat, + camera_buffer: GPUBuffer, + use_f16: boolean = false ): GaussianRenderer { - const sorter = get_sorter(pc.num_points, device); - - // =============================================== - // Initialize GPU Buffers - // =============================================== - - const nulling_data = new Uint32Array([0]); - - // =============================================== - // Create Compute Pipeline and Bind Groups - // =============================================== - const preprocess_pipeline = device.createComputePipeline({ - label: 'preprocess', - layout: 'auto', - compute: { - module: device.createShaderModule({ code: preprocessWGSL }), - entryPoint: 'preprocess', - constants: { - workgroupSize: C.histogram_wg_size, - sortKeyPerThread: c_histogram_block_rows, - }, - }, - }); - - const sort_bind_group = device.createBindGroup({ - label: 'sort', - layout: preprocess_pipeline.getBindGroupLayout(2), - entries: [ - { binding: 0, resource: { buffer: sorter.sort_info_buffer } }, - { binding: 1, resource: { buffer: sorter.ping_pong[0].sort_depths_buffer } }, - { binding: 2, resource: { buffer: sorter.ping_pong[0].sort_indices_buffer } }, - { binding: 3, resource: { buffer: sorter.sort_dispatch_indirect_buffer } }, - ], - }); - - - // =============================================== - // Create Render Pipeline and Bind Groups - // =============================================== - - - // =============================================== - // Command Encoder Functions - // =============================================== - - - // =============================================== - // Return Render Object - // =============================================== - return { - frame: (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { - sorter.sort(encoder); - }, - camera_buffer, - }; -} + const workgroupSize = 64; + + const sorter = get_sorter(pc.num_points, device); + + // =============================================== + // Initialize GPU Buffers + // =============================================== + const nullingDataBuffer = createBuffer( + device, + 'nulling data buffer', + 4, + GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST, + new Uint32Array([0]) + ); + + const renderSettingsBuffer = createBuffer( + device, + 'render settings buffer', + 8, + GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, + new Float32Array([1.0, pc.sh_deg]) + ); + + const splatSize = 4 * 6; + const splatBuffer = createBuffer( + device, + "splat buffer", + splatSize * pc.num_points, + GPUBufferUsage.STORAGE + ); + + const indirectDrawBuffer = createBuffer( + device, + 'indirect draw buffer', + 4 * 4, + GPUBufferUsage.INDIRECT | GPUBufferUsage.COPY_DST, + new Uint32Array([4, 0, 0, 0]) + ); + + + // =============================================== + // Create Compute Pipeline and Bind Groups + // =============================================== + let preprocessCode: string; + let prerocessLabel: string; + if (use_f16) { + preprocessCode = preprocessWGSL16; + prerocessLabel = 'preprocess f16'; + } + else { + preprocessCode = preprocessWGSL; + prerocessLabel = 'preprocess f32'; + } + const preprocessPipeline = device.createComputePipeline({ + label: prerocessLabel, + layout: 'auto', + compute: { + module: device.createShaderModule({ code: preprocessCode }), + entryPoint: 'preprocess', + constants: { + workgroupSize: C.histogram_wg_size, + sortKeyPerThread: c_histogram_block_rows + }, + }, + }); + + const uniformsBindGroup = device.createBindGroup({ + label: 'uniforms', + layout: preprocessPipeline.getBindGroupLayout(0), + entries: [ + { binding: 0, resource: { buffer: camera_buffer } }, + { binding: 1, resource: { buffer: renderSettingsBuffer } }, + ] + }); + + const gaussiansBindGroup = device.createBindGroup({ + label: 'gaussians', + layout: preprocessPipeline.getBindGroupLayout(1), + entries: [ + { binding: 0, resource: { buffer: pc.gaussian_3d_buffer } }, + { binding: 1, resource: { buffer: pc.sh_buffer } }, + { binding: 2, resource: { buffer: splatBuffer } }, + ], + }); + + const sortBindGroup = device.createBindGroup({ + label: 'sort', + layout: preprocessPipeline.getBindGroupLayout(2), + entries: [ + { binding: 0, resource: { buffer: sorter.sort_info_buffer } }, + { binding: 1, resource: { buffer: sorter.ping_pong[0].sort_depths_buffer } }, + { binding: 2, resource: { buffer: sorter.ping_pong[0].sort_indices_buffer } }, + { binding: 3, resource: { buffer: sorter.sort_dispatch_indirect_buffer } }, + ], + }); + + // =============================================== + // Create Render Pipeline and Bind Groups + // =============================================== + const renderShader = device.createShaderModule({ + code: renderWGSL, + label: 'gaussian render shader' + }); + + const renderPipeline = device.createRenderPipeline({ + label: 'gaussian render pipeline', + layout: 'auto', + vertex: { + module: renderShader, + entryPoint: 'vs_main' + }, + fragment: { + module: renderShader, + entryPoint: 'fs_main', + targets: [ + { + format: presentation_format, + blend: { + color: { + operation: 'add', + srcFactor: 'one', + dstFactor: 'one-minus-src-alpha', + }, + alpha: { + operation: 'add', + srcFactor: 'one', + dstFactor: 'one-minus-src-alpha', + }, + } + } + ], + }, + primitive: { + topology: "triangle-strip" + } + }); + + const cameraBindGroup = device.createBindGroup({ + label: 'camera', + layout: renderPipeline.getBindGroupLayout(0), + entries: [{ binding: 0, resource: { buffer: camera_buffer } }], + }); + + const splatsBindGroup = device.createBindGroup({ + label: 'gaussian gaussians', + layout: renderPipeline.getBindGroupLayout(1), + entries: [ + { binding: 0, resource: { buffer: splatBuffer } }, + { binding: 1, resource: { buffer: sorter.ping_pong[0].sort_indices_buffer } }, + ] + }); + + // =============================================== + // Command Encoder Functions + // =============================================== + const preprocess = (encoder: GPUCommandEncoder) => { + // initialize two atomic add variables to zero + encoder.copyBufferToBuffer(nullingDataBuffer, 0, sorter.sort_info_buffer, 0, 4); + encoder.copyBufferToBuffer(nullingDataBuffer, 0, sorter.sort_dispatch_indirect_buffer, 0, 4); + + const computePass = encoder.beginComputePass({ + label: 'preprocess', + }); + computePass.setPipeline(preprocessPipeline); + computePass.setBindGroup(0, uniformsBindGroup); + computePass.setBindGroup(1, gaussiansBindGroup); + computePass.setBindGroup(2, sortBindGroup); + computePass.dispatchWorkgroups( + Math.ceil(pc.num_points / workgroupSize) + ); + computePass.end(); + + // copy keys_size into indirect draw buffer + encoder.copyBufferToBuffer(sorter.sort_info_buffer, 0, indirectDrawBuffer, 4, 4); + } + + const render = (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { + const renderPass = encoder.beginRenderPass({ + label: 'gaussian render', + colorAttachments: [ + { + view: texture_view, + loadOp: 'clear', + storeOp: 'store' + } + ], + }); + renderPass.setPipeline(renderPipeline); + renderPass.setBindGroup(0, cameraBindGroup); + renderPass.setBindGroup(1, splatsBindGroup); + renderPass.drawIndirect(indirectDrawBuffer, 0); + renderPass.end(); + }; + + // =============================================== + // Return Render Object + // =============================================== + return { + frame: (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { + preprocess(encoder); + sorter.sort(encoder); + render(encoder, texture_view); + }, + camera_buffer, + renderSettingsBuffer, + preprocessPipeline + }; +} \ No newline at end of file diff --git a/src/renderers/point-cloud-renderer.ts b/src/renderers/point-cloud-renderer.ts index 36a2e8e..66bdbeb 100644 --- a/src/renderers/point-cloud-renderer.ts +++ b/src/renderers/point-cloud-renderer.ts @@ -3,66 +3,66 @@ import pointcloud_wgsl from '../shaders/point_cloud.wgsl'; import { Renderer } from './renderer'; export default function get_renderer( - pc: PointCloud, - device: GPUDevice, - presentation_format: GPUTextureFormat, - camera_buffer: GPUBuffer): Renderer { - const render_shader = device.createShaderModule({code: pointcloud_wgsl}); - const render_pipeline = device.createRenderPipeline({ - label: 'render', - layout: 'auto', - vertex: { - module: render_shader, - entryPoint: 'vs_main', - }, - fragment: { - module: render_shader, - entryPoint: 'fs_main', - targets: [{ format: presentation_format }], - }, - primitive: { - topology: 'point-list', - }, - }); - - const camera_bind_group = device.createBindGroup({ - label: 'point cloud camera', - layout: render_pipeline.getBindGroupLayout(0), - entries: [{binding: 0, resource: { buffer: camera_buffer }}], - }); + pc: PointCloud, + device: GPUDevice, + presentation_format: GPUTextureFormat, + camera_buffer: GPUBuffer): Renderer { + const render_shader = device.createShaderModule({ code: pointcloud_wgsl }); + const render_pipeline = device.createRenderPipeline({ + label: 'render', + layout: 'auto', + vertex: { + module: render_shader, + entryPoint: 'vs_main', + }, + fragment: { + module: render_shader, + entryPoint: 'fs_main', + targets: [{ format: presentation_format }], + }, + primitive: { + topology: 'point-list', + }, + }); - const gaussian_bind_group = device.createBindGroup({ - label: 'point cloud gaussians', - layout: render_pipeline.getBindGroupLayout(1), - entries: [ - {binding: 0, resource: { buffer: pc.gaussian_3d_buffer }}, - ], - }); + const camera_bind_group = device.createBindGroup({ + label: 'point cloud camera', + layout: render_pipeline.getBindGroupLayout(0), + entries: [{ binding: 0, resource: { buffer: camera_buffer } }], + }); - const render = (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { - const pass = encoder.beginRenderPass({ - label: 'point cloud render', - colorAttachments: [ - { - view: texture_view, - loadOp: 'clear', - storeOp: 'store', - } - ], + const gaussian_bind_group = device.createBindGroup({ + label: 'point cloud gaussians', + layout: render_pipeline.getBindGroupLayout(1), + entries: [ + { binding: 0, resource: { buffer: pc.gaussian_3d_buffer } }, + ], }); - pass.setPipeline(render_pipeline); - pass.setBindGroup(0, camera_bind_group); - pass.setBindGroup(1, gaussian_bind_group); - pass.draw(pc.num_points); - pass.end(); - }; + const render = (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { + const pass = encoder.beginRenderPass({ + label: 'point cloud render', + colorAttachments: [ + { + view: texture_view, + loadOp: 'clear', + storeOp: 'store', + } + ], + }); + pass.setPipeline(render_pipeline); + pass.setBindGroup(0, camera_bind_group); + pass.setBindGroup(1, gaussian_bind_group); + + pass.draw(pc.num_points); + pass.end(); + }; - return { - frame: (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { - render(encoder, texture_view); - }, + return { + frame: (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => { + render(encoder, texture_view); + }, - camera_buffer, - }; + camera_buffer, + }; } \ No newline at end of file diff --git a/src/renderers/renderer.ts b/src/renderers/renderer.ts index ffdf9ba..45df05a 100644 --- a/src/renderers/renderer.ts +++ b/src/renderers/renderer.ts @@ -3,158 +3,178 @@ 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'; export interface Renderer { - frame: (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => void, - camera_buffer: GPUBuffer, + frame: (encoder: GPUCommandEncoder, texture_view: GPUTextureView) => void, + camera_buffer: GPUBuffer, } export default async function init( - canvas: HTMLCanvasElement, - context: GPUCanvasContext, - device: GPUDevice + canvas: HTMLCanvasElement, + context: GPUCanvasContext, + device: GPUDevice ) { - 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 cameras; - - const camera = new Camera(canvas, device); - const control = new CameraControl(camera); + let ply_file_loaded = false; + let cam_file_loaded = false; + let renderers: { pointcloud?: Renderer, gaussian?: Renderer , gaussian_f16?: Renderer } = {}; + let gaussian_renderer: GaussianRenderer | undefined; + let gaussian_renderer_f16: GaussianRenderer | undefined; + let pointcloud_renderer: Renderer | undefined; + let renderer: Renderer | undefined; + let cameras; - const observer = new ResizeObserver(() => { - canvas.width = canvas.clientWidth; - canvas.height = canvas.clientHeight; + const camera = new Camera(canvas, device); + const control = new CameraControl(camera); - camera.on_update_canvas(); - }); - observer.observe(canvas); - - const presentation_format = navigator.gpu.getPreferredCanvasFormat(); - context.configure({ - device, - format: presentation_format, - alphaMode: 'opaque', - }); - + const observer = new ResizeObserver(() => { + canvas.width = canvas.clientWidth; + canvas.height = canvas.clientHeight; - // Tweakpane: easily adding tweak control for parameters. - const params = { - fps: 0.0, - gaussian_multiplier: 1, - renderer: 'pointcloud', - ply_file: '', - cam_file: '', - }; - - const pane = new Pane({ - title: 'Config', - expanded: true, - }); - pane.registerPlugin(TweakpaneFileImportPlugin); - { - pane.addMonitor(params, 'fps', { - readonly:true - }); - } - { - pane.addInput(params, 'renderer', { - options: { - pointcloud: 'pointcloud', - gaussian: 'gaussian', - } - }).on('change', (e) => { - renderer = renderers[e.value]; + camera.on_update_canvas(); }); - } - { - pane.addInput(params, 'ply_file', { - view: 'file-input', - lineCount: 3, - 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; - } + observer.observe(canvas); + + const presentation_format = navigator.gpu.getPreferredCanvasFormat(); + context.configure({ + device, + format: presentation_format, + alphaMode: 'opaque', }); - } - { - pane.addInput(params, 'cam_file', { - view: 'file-input', - lineCount: 3, - 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; - } + + // Tweakpane: easily adding tweak control for parameters. + const params = { + fps: 0.0, + gaussian_multiplier: 1, + renderer: 'pointcloud', + ply_file: '', + cam_file: '', + }; + + const pane = new Pane({ + title: 'Config', + expanded: true, }); - } - { - pane.addInput( - params, - 'gaussian_multiplier', - {min: 0, max: 1.5} - ).on('change', (e) => { - //TODO: Bind constants to the gaussian renderer. + pane.registerPlugin(TweakpaneFileImportPlugin); + { + pane.addMonitor(params, 'fps', { + readonly: true + }); + } + { + pane.addInput(params, 'renderer', { + options: { + pointcloud: 'pointcloud', + gaussian: 'gaussian', + gaussian_f16: 'gaussian_f16', + } + }).on('change', (e) => { + renderer = renderers[e.value]; + }); + } + { + pane.addInput(params, 'ply_file', { + view: 'file-input', + lineCount: 3, + 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); + gaussian_renderer_f16 = get_renderer_gaussian(pc, device, presentation_format, camera.uniform_buffer, true); + renderers = { + pointcloud: pointcloud_renderer, + gaussian: gaussian_renderer, + gaussian_f16: gaussian_renderer_f16, + }; + renderer = renderers[params.renderer]; + ply_file_loaded = true; + } else { + ply_file_loaded = false; + } + }); + } + { + pane.addInput(params, 'cam_file', { + view: 'file-input', + lineCount: 3, + 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; + } + }); + } + { + pane.addInput( + params, + 'gaussian_multiplier', + { min: 0, max: 1.5 } + ).on('change', (e) => { + if (gaussian_renderer) { + device.queue.writeBuffer(gaussian_renderer.renderSettingsBuffer, 0, new Float32Array([e.value])); + } + if (gaussian_renderer_f16) { + device.queue.writeBuffer(gaussian_renderer_f16.renderSettingsBuffer, 0, new Float32Array([e.value])); + } + }); + } + + document.addEventListener('keydown', (event) => { + switch (event.key) { + case '0': + case '1': + case '2': + case '3': + case '4': + case '5': + case '6': + case '7': + case '8': + case '9': + const i = parseInt(event.key); + console.log(`set to camera preset ${i}`); + camera.set_preset(cameras[i]); + break; + } }); - } - document.addEventListener('keydown', (event) => { - switch(event.key) { - case '0': - case '1': - case '2': - case '3': - case '4': - case '5': - case '6': - case '7': - case '8': - case '9': - const i = parseInt(event.key); - console.log(`set to camera preset ${i}`); - camera.set_preset(cameras[i]); - break; - } - }); + let countFrame: number = 0; + let totalTime: number = 0; - function frame() { - if (ply_file_loaded && cam_file_loaded) { - params.fps=1.0/timeReturn()*1000.0; - time(); - const encoder = device.createCommandEncoder(); - const texture_view = context.getCurrentTexture().createView(); - renderer.frame(encoder, texture_view); - device.queue.submit([encoder.finish()]); + function frame() { + + if (ply_file_loaded && cam_file_loaded) { + let frameTime = timeReturn(); + totalTime += frameTime + params.fps = 1.0 / frameTime * 1000.0; + countFrame++; + if (totalTime > 60000) { + console.log(totalTime / countFrame); + totalTime = 0; + countFrame = 0; + } + time(); + const encoder = device.createCommandEncoder(); + const texture_view = context.getCurrentTexture().createView(); + renderer.frame(encoder, texture_view); + device.queue.submit([encoder.finish()]); + } + requestAnimationFrame(frame); } - requestAnimationFrame(frame); - } - requestAnimationFrame(frame); + requestAnimationFrame(frame); } diff --git a/src/shaders/gaussian.wgsl b/src/shaders/gaussian.wgsl index 759226d..05982f9 100644 --- a/src/shaders/gaussian.wgsl +++ b/src/shaders/gaussian.wgsl @@ -1,22 +1,86 @@ +struct CameraUniforms { + view: mat4x4f, + viewInv: mat4x4f, + proj: mat4x4f, + projInv: mat4x4f, + viewport: vec2f, + focal: vec2f, + clippingPlanes: vec2f +}; + struct VertexOutput { - @builtin(position) position: vec4, - //TODO: information passed from vertex shader to fragment shader + @builtin(position) position: vec4f, + // No interpolation + @location(0) @interpolate(flat) center: u32, + @location(1) @interpolate(flat) color_opacity: vec2u, + @location(2) @interpolate(flat) conic: vec2u, }; struct Splat { - //TODO: information defined in preprocess compute shader + pos: u32, + size: u32, + color_opacity: vec2u, + conic: vec2u }; +@group(0) @binding(0) +var camera: CameraUniforms; + +@group(1) @binding(0) +var splats: array; +@group(1) @binding(1) +var sortIndices : array; + @vertex fn vs_main( + @builtin(instance_index) instanceIndex: u32, + @builtin(vertex_index) vertexIndex: u32 ) -> VertexOutput { - //TODO: reconstruct 2D quad based on information from splat, pass + const vertices = array( + vec2f(-1, -1), + vec2f( 1, -1), + vec2f(-1, 1), + vec2f( 1, 1) + ); + + let splat = splats[sortIndices[instanceIndex]]; + let position = unpack2x16float(splat.pos); + let size = unpack2x16float(splat.size); + var out: VertexOutput; - out.position = vec4(1. ,1. , 0., 1.); + out.position = vec4f(position + size * vertices[vertexIndex], 0.0, 1.0); + out.center = pack2x16float((vec2f(0.5, -0.5) * position + 0.5) * camera.viewport); + out.color_opacity = splat.color_opacity; + out.conic = splat.conic; return out; } @fragment -fn fs_main(in: VertexOutput) -> @location(0) vec4 { - return vec4(1.); +fn fs_main( + in: VertexOutput +) -> @location(0) vec4f { + let conic = vec4f( + unpack2x16float(in.conic.x), + unpack2x16float(in.conic.y) + ); + let center = vec2f(unpack2x16float(in.center)); + let color_opacity = vec4f( + unpack2x16float(in.color_opacity.x), + unpack2x16float(in.color_opacity.y) + ); + + let d = center - in.position.xy; + let power = 0.5 * (conic.x * d.x * d.x + conic.z * d.y * d.y) + conic.y * d.x * d.y; + + if (power < 0.0) { discard; } + + + + let alpha = min(0.99, color_opacity.a * exp(-power)); + + // Discard if alpha is less than 1/255 + if (alpha < 0.00392156862) { discard; } + + // Premultiplied alpha + return vec4f(color_opacity.rgb * alpha, alpha); } \ No newline at end of file diff --git a/src/shaders/point_cloud.wgsl b/src/shaders/point_cloud.wgsl index 01dded1..7b13e2a 100644 --- a/src/shaders/point_cloud.wgsl +++ b/src/shaders/point_cloud.wgsl @@ -1,46 +1,45 @@ struct CameraUniforms { - view: mat4x4, - view_inv: mat4x4, - proj: mat4x4, - proj_inv: mat4x4, - viewport: vec2, - focal: vec2 + view: mat4x4f, + viewInv: mat4x4f, + proj: mat4x4f, + projInv: mat4x4f, + viewport: vec2f, + focal: vec2f, + clippingPlanes: vec2f }; struct Gaussian { - pos_opacity: array, - rot: array, - scale: array -} + pos_opacity: vec2u, + rot: vec2u, + scale: vec2u +}; @group(0) @binding(0) var camera: CameraUniforms; @group(1) @binding(0) -var gaussians : array; +var gaussians: array; struct VertexOutput { - @builtin(position) position: vec4, + @builtin(position) position: vec4f, }; @vertex fn vs_main( - @builtin(vertex_index) in_vertex_index: u32, + @builtin(vertex_index) vertex_index: u32 ) -> VertexOutput { var out: VertexOutput; - 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.); - - // TODO: MVP calculations - out.position = pos; + let gaussian = gaussians[vertex_index]; + let a = unpack2x16float(gaussian.pos_opacity.x); + let b = unpack2x16float(gaussian.pos_opacity.y); + let pos = vec4f(a.x, a.y, b.x, 1.); + out.position = camera.proj * camera.view * pos; return out; } @fragment -fn fs_main(in: VertexOutput) -> @location(0) vec4 { - return vec4(1., 1., 0., 1.); +fn fs_main(in: VertexOutput) -> @location(0) vec4f { + return vec4f(1., 1., 0., 1.); } \ No newline at end of file diff --git a/src/shaders/preprocess-16.wgsl b/src/shaders/preprocess-16.wgsl new file mode 100644 index 0000000..c93c37e --- /dev/null +++ b/src/shaders/preprocess-16.wgsl @@ -0,0 +1,262 @@ +enable f16; + +const SH_C0: f16 = 0.28209479177387814; +const SH_C1: f16 = 0.4886025119029199; +const SH_C2 = array( + 1.0925484305920792, + -1.0925484305920792, + 0.31539156525252005, + -1.0925484305920792, + 0.5462742152960396 +); +const SH_C3 = array( + -0.5900435899266435, + 2.890611442640554, + -0.4570457994644658, + 0.3731763325901154, + -0.4570457994644658, + 1.445305721320277, + -0.5900435899266435 +); + +override workgroupSize: u32; +override sortKeyPerThread: u32; + +struct DispatchIndirect { + dispatchX: atomic, + dispatchY: u32, + dispatchZ: u32, +} + +struct SortInfos { + keysSize: atomic, // instance_count in DrawIndirect + // data below is for info inside radix sort + paddedSize: u32, + passes: u32, + evenPass: u32, + oddPass: u32, +} + +struct CameraUniforms { + view: mat4x4f, + viewInv: mat4x4f, + proj: mat4x4f, + projInv: mat4x4f, + viewport: vec2f, + focal: vec2f, + clippingPlanes: vec2f +}; + +struct RenderSettings { + gaussianScaling: f32, + shDegree: f32, +} + +struct Gaussian { + pos_opacity: vec2u, + rot: vec2u, + scale: vec2u +}; + +struct Splat { + pos: u32, + size: u32, + color_opacity: vec2u, + conic: vec2u +}; + +@group(0) @binding(0) +var camera: CameraUniforms; +@group(0) @binding(1) +var renderSettings: RenderSettings; + +@group(1) @binding(0) +var gaussians: array; +@group(1) @binding(1) +var shCoefs: array; +@group(1) @binding(2) +var splats: array; + +@group(2) @binding(0) +var sortInfos: SortInfos; +@group(2) @binding(1) +var sortDepths : array; +@group(2) @binding(2) +var sortIndices : array; +@group(2) @binding(3) +var sortDispatch: DispatchIndirect; + +// reads the ith sh coef from the storage buffer +fn readSHCoef(splatIdx: u32, coefIdx: u32) -> vec3h { + const maxNumCoefs = 16u; + let offset = (splatIdx * maxNumCoefs + coefIdx) * 3; + if (offset % 2 == 0) { + return vec3h( + bitcast(shCoefs[offset / 2]), + bitcast(shCoefs[offset / 2 + 1]).x + ); + } + else { + return vec3h( + bitcast(shCoefs[offset / 2]).y, + bitcast(shCoefs[offset / 2 + 1]) + ); + } +} + +fn sigmoid(x: f16) -> f16 { + return 1.0 / (1.0 + exp(-x)); +} + +// spherical harmonics evaluation with Condon–Shortley phase +fn computeColorFromSH(dir: vec3h, v_idx: u32, sh_deg: u32) -> vec3h { + var result = SH_C0 * readSHCoef(v_idx, 0u); + + if sh_deg > 0u { + + let x = dir.x; + let y = dir.y; + let z = dir.z; + + result += - SH_C1 * y * readSHCoef(v_idx, 1u) + + SH_C1 * z * readSHCoef(v_idx, 2u) + - SH_C1 * x * readSHCoef(v_idx, 3u); + + if sh_deg > 1u { + + let xx = dir.x * dir.x; + let yy = dir.y * dir.y; + let zz = dir.z * dir.z; + let xy = dir.x * dir.y; + let yz = dir.y * dir.z; + let xz = dir.x * dir.z; + + result += SH_C2[0] * xy * readSHCoef(v_idx, 4u) + + SH_C2[1] * yz * readSHCoef(v_idx, 5u) + + SH_C2[2] * (2.0 * zz - xx - yy) * readSHCoef(v_idx, 6u) + + SH_C2[3] * xz * readSHCoef(v_idx, 7u) + + SH_C2[4] * (xx - yy) * readSHCoef(v_idx, 8u); + + if sh_deg > 2u { + result += SH_C3[0] * y * (3.0 * xx - yy) * readSHCoef(v_idx, 9u) + + SH_C3[1] * xy * z * readSHCoef(v_idx, 10u) + + SH_C3[2] * y * (4.0 * zz - xx - yy) * readSHCoef(v_idx, 11u) + + SH_C3[3] * z * (2.0 * zz - 3.0 * xx - 3.0 * yy) * readSHCoef(v_idx, 12u) + + SH_C3[4] * x * (4.0 * zz - xx - yy) * readSHCoef(v_idx, 13u) + + SH_C3[5] * z * (xx - yy) * readSHCoef(v_idx, 14u) + + SH_C3[6] * x * (xx - 3.0 * yy) * readSHCoef(v_idx, 15u); + } + } + } + result += 0.5; + + return max(vec3h(0.0), result); +} + +@compute @workgroup_size(workgroupSize, 1, 1) +fn preprocess(@builtin(global_invocation_id) globalIndex: vec3u) { + let index = globalIndex.x; + if (index >= arrayLength(&gaussians)) { return; } + + let gaussian = gaussians[index]; + + // Get position in screen space + let pos_opac = vec4f(unpack2x16float(gaussian.pos_opacity[0]), unpack2x16float(gaussian.pos_opacity[1])); + let posView = camera.view * vec4f(pos_opac.xyz, 1.0); + // Z-clipping + if (posView.z < camera.clippingPlanes[0] || posView.z > camera.clippingPlanes[1]) { return; } + let posClip = mat4x4h(camera.proj) * vec4h(posView); + let posNDC = posClip.xyz / posClip.w; + let depth = posView.z; + + // Simple view frustum culling + if (any(abs(posNDC.xy) > vec2h(1.2))) { return; } + + // Set some f16 variables to prevent repeated casting + let focal = vec2h(camera.focal); + let viewport = vec2h(camera.viewport); + let posView16 = vec3h(posView.xyz); + + // Compute color from spherical harmonics + // Camera world position is last column of inversed view matrix + let direction = normalize(vec3h(pos_opac.xyz - camera.viewInv[3].xyz)); + let color = computeColorFromSH( + direction, index, u32(renderSettings.shDegree) + ); + let opacity = sigmoid(f16(pos_opac.w)); + + // Calculate 3D covariance matrix from scale and rotation + // rotation matrix + let rot = vec4h(bitcast(gaussian.rot[0]), bitcast(gaussian.rot[1])); + let rotMat = mat3x3h( + 1.0 - 2.0 * (rot.z * rot.z + rot.w * rot.w), 2.0 * (rot.y * rot.z + rot.x * rot.w), 2.0 * (rot.y * rot.w - rot.x * rot.z), + 2.0 * (rot.y * rot.z - rot.x * rot.w), 1.0 - 2.0 * (rot.y * rot.y + rot.w * rot.w), 2.0 * (rot.z * rot.w + rot.x * rot.y), + 2.0 * (rot.y * rot.w + rot.x * rot.z), 2.0 * (rot.z * rot.w - rot.x * rot.y), 1.0 - 2.0 * (rot.y * rot.y + rot.z * rot.z) + ); + // diagonal scale matrix, saved in log space + let scale = f16(renderSettings.gaussianScaling) * exp(vec3h(bitcast(gaussian.scale[0]), bitcast(gaussian.scale[1]).x)); + var tmpMat = mat3x3h( + scale.x * rotMat[0], + scale.y * rotMat[1], + scale.z * rotMat[2], + ); + let cov3d = tmpMat * transpose(tmpMat); + + // Calculate 2D covariance matrix from 3D covariance matrix + let W = mat3x3f(mat3x3f(camera.view[0].xyz, camera.view[1].xyz, camera.view[2].xyz)); + + let lim = 0.65 * viewport * focal; + let t = vec3h( + clamp(posView16.xy / posView16.z, -lim, lim) * posView16.z, + posView16.z + ); + let invTz = 1.0 / t.z; + let J = mat3x2f(mat3x2h( + focal.x * invTz, 0.0, + 0.0, focal.y * invTz, + -focal.x * t.x * invTz * invTz, -focal.y * t.y * invTz * invTz + )); + + let tmpMat2 = J * W; + var cov2d = tmpMat2 * mat3x3f(cov3d) * transpose(tmpMat2); + cov2d[0][0] += 0.3; + cov2d[1][1] += 0.3; + + // Calculate max radius via eigenvalues of 2D covariance matrix + // Use f32 for some intermediate calculations + let det = determinant(cov2d); + if (det == 0) { return; } + + let mid = 0.5 * (cov2d[0][0] + cov2d[1][1]); + let dist = sqrt(max(0.1, mid * mid - det)); + let lambda1 = mid + dist; + let lambda2 = mid - dist; + + // Get size from max radius in NDC space + let size = vec2h(ceil(3.0 * sqrt(max(lambda1, lambda2))) * 2.0 / camera.viewport); + + // Store sorting information + let keyIndex = atomicAdd(&sortInfos.keysSize, 1); + sortDepths[keyIndex] = bitcast(camera.clippingPlanes[1] - depth); + sortIndices[keyIndex] = keyIndex; + + // increment sortDispatch.dispatchX each time you reach limit for one dispatch of keys + if (keyIndex % (workgroupSize * sortKeyPerThread) == 0) { + atomicAdd(&sortDispatch.dispatchX, 1); + } + + // Handle output to render pipeline + var output: Splat; + output.pos = bitcast(posNDC.xy); + output.size = bitcast(size); + output.color_opacity = vec2u( + bitcast(color.rg), + bitcast(vec2h(color.b, opacity)) + ); + let detInv = 1.0 / det; + output.conic = vec2u( + pack2x16float(vec2f(detInv * cov2d[1][1], detInv * -cov2d[0][1])), + pack2x16float(vec2f(detInv * cov2d[0][0], 0)) + ); + splats[keyIndex] = output; +} \ No newline at end of file diff --git a/src/shaders/preprocess.wgsl b/src/shaders/preprocess.wgsl index bbc63f5..695e1b4 100644 --- a/src/shaders/preprocess.wgsl +++ b/src/shaders/preprocess.wgsl @@ -1,13 +1,13 @@ const SH_C0: f32 = 0.28209479177387814; -const SH_C1 = 0.4886025119029199; -const SH_C2 = array( +const SH_C1: f32 = 0.4886025119029199; +const SH_C2 = array( 1.0925484305920792, -1.0925484305920792, 0.31539156525252005, -1.0925484305920792, 0.5462742152960396 ); -const SH_C3 = array( +const SH_C3 = array( -0.5900435899266435, 2.890611442640554, -0.4570457994644658, @@ -21,63 +21,94 @@ override workgroupSize: u32; override sortKeyPerThread: u32; struct DispatchIndirect { - dispatch_x: atomic, - dispatch_y: u32, - dispatch_z: u32, + dispatchX: atomic, + dispatchY: u32, + dispatchZ: u32, } struct SortInfos { - keys_size: atomic, // instance_count in DrawIndirect + keysSize: atomic, // instance_count in DrawIndirect //data below is for info inside radix sort - padded_size: u32, + paddedSize: u32, passes: u32, - even_pass: u32, - odd_pass: u32, + evenPass: u32, + oddPass: u32, } struct CameraUniforms { - view: mat4x4, - view_inv: mat4x4, - proj: mat4x4, - proj_inv: mat4x4, - viewport: vec2, - focal: vec2 + view: mat4x4f, + viewInv: mat4x4f, + proj: mat4x4f, + projInv: mat4x4f, + viewport: vec2f, + focal: vec2f, + clippingPlanes: vec2f }; struct RenderSettings { - gaussian_scaling: f32, - sh_deg: f32, + gaussianScaling: f32, + shDegree: f32, } struct Gaussian { - pos_opacity: array, - rot: array, - scale: array + pos_opacity: vec2u, + rot: vec2u, + scale: vec2u }; struct Splat { - //TODO: store information for 2D splat rendering + pos: u32, + size: u32, + color_opacity: vec2u, + conic: vec2u }; -//TODO: bind your data here +@group(0) @binding(0) +var camera: CameraUniforms; +@group(0) @binding(1) +var renderSettings: RenderSettings; + +@group(1) @binding(0) +var gaussians: array; +@group(1) @binding(1) +var shCoefs: array; +@group(1) @binding(2) +var splats: array; + @group(2) @binding(0) -var sort_infos: SortInfos; +var sortInfos: SortInfos; @group(2) @binding(1) -var sort_depths : array; +var sortDepths : array; @group(2) @binding(2) -var sort_indices : array; +var sortIndices : array; @group(2) @binding(3) -var sort_dispatch: DispatchIndirect; +var sortDispatch: DispatchIndirect; + +// reads the ith sh coef from the storage buffer +fn readSHCoef(splatIdx: u32, coefIdx: u32) -> vec3f { + const maxNumCoefs = 16u; + let offset = (splatIdx * maxNumCoefs + coefIdx) * 3; + if (offset % 2 == 0) { + return vec3f( + unpack2x16float(shCoefs[offset / 2]), + unpack2x16float(shCoefs[offset / 2 + 1]).x + ); + } + else { + return vec3f( + unpack2x16float(shCoefs[offset / 2]).y, + unpack2x16float(shCoefs[offset / 2 + 1]) + ); + } +} -/// 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); +fn sigmoid(x: f32) -> f32 { + return 1.0 / (1.0 + exp(-x)); } // spherical harmonics evaluation with Condon–Shortley phase -fn computeColorFromSH(dir: vec3, v_idx: u32, sh_deg: u32) -> vec3 { - var result = SH_C0 * sh_coef(v_idx, 0u); +fn computeColorFromSH(dir: vec3f, v_idx: u32, sh_deg: u32) -> vec3f { + var result = SH_C0 * readSHCoef(v_idx, 0u); if sh_deg > 0u { @@ -85,7 +116,9 @@ fn computeColorFromSH(dir: vec3, v_idx: u32, sh_deg: u32) -> vec3 { let y = dir.y; let z = dir.z; - result += - SH_C1 * y * sh_coef(v_idx, 1u) + SH_C1 * z * sh_coef(v_idx, 2u) - SH_C1 * x * sh_coef(v_idx, 3u); + result += - SH_C1 * y * readSHCoef(v_idx, 1u) + + SH_C1 * z * readSHCoef(v_idx, 2u) + - SH_C1 * x * readSHCoef(v_idx, 3u); if sh_deg > 1u { @@ -96,23 +129,126 @@ fn computeColorFromSH(dir: vec3, v_idx: u32, sh_deg: u32) -> vec3 { let yz = dir.y * dir.z; let xz = dir.x * dir.z; - result += SH_C2[0] * xy * sh_coef(v_idx, 4u) + SH_C2[1] * yz * sh_coef(v_idx, 5u) + SH_C2[2] * (2.0 * zz - xx - yy) * sh_coef(v_idx, 6u) + SH_C2[3] * xz * sh_coef(v_idx, 7u) + SH_C2[4] * (xx - yy) * sh_coef(v_idx, 8u); + result += SH_C2[0] * xy * readSHCoef(v_idx, 4u) + + SH_C2[1] * yz * readSHCoef(v_idx, 5u) + + SH_C2[2] * (2.0 * zz - xx - yy) * readSHCoef(v_idx, 6u) + + SH_C2[3] * xz * readSHCoef(v_idx, 7u) + + SH_C2[4] * (xx - yy) * readSHCoef(v_idx, 8u); if sh_deg > 2u { - result += SH_C3[0] * y * (3.0 * xx - yy) * sh_coef(v_idx, 9u) + SH_C3[1] * xy * z * sh_coef(v_idx, 10u) + SH_C3[2] * y * (4.0 * zz - xx - yy) * sh_coef(v_idx, 11u) + SH_C3[3] * z * (2.0 * zz - 3.0 * xx - 3.0 * yy) * sh_coef(v_idx, 12u) + SH_C3[4] * x * (4.0 * zz - xx - yy) * sh_coef(v_idx, 13u) + SH_C3[5] * z * (xx - yy) * sh_coef(v_idx, 14u) + SH_C3[6] * x * (xx - 3.0 * yy) * sh_coef(v_idx, 15u); + result += SH_C3[0] * y * (3.0 * xx - yy) * readSHCoef(v_idx, 9u) + + SH_C3[1] * xy * z * readSHCoef(v_idx, 10u) + + SH_C3[2] * y * (4.0 * zz - xx - yy) * readSHCoef(v_idx, 11u) + + SH_C3[3] * z * (2.0 * zz - 3.0 * xx - 3.0 * yy) * readSHCoef(v_idx, 12u) + + SH_C3[4] * x * (4.0 * zz - xx - yy) * readSHCoef(v_idx, 13u) + + SH_C3[5] * z * (xx - yy) * readSHCoef(v_idx, 14u) + + SH_C3[6] * x * (xx - 3.0 * yy) * readSHCoef(v_idx, 15u); } } } result += 0.5; - return max(vec3(0.), result); + return max(vec3f(0.0), result); } -@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 - - let keys_per_dispatch = workgroupSize * sortKeyPerThread; - // increment DispatchIndirect.dispatchx each time you reach limit for one dispatch of keys +@compute @workgroup_size(workgroupSize, 1, 1) +fn preprocess(@builtin(global_invocation_id) globalIndex: vec3u) { + let index = globalIndex.x; + if (index >= arrayLength(&gaussians)) { return; } + + let gaussian = gaussians[index]; + + // Get position in screen space + let pos_opac = vec4f(unpack2x16float(gaussian.pos_opacity[0]), unpack2x16float(gaussian.pos_opacity[1])); + let posView = camera.view * vec4f(pos_opac.xyz, 1.0); + // Z-clipping + if (posView.z < camera.clippingPlanes[0] || posView.z > camera.clippingPlanes[1]) { return; } + let posClip = camera.proj * posView; + let posNDC = posClip.xyz / posClip.w; + let depth = posView.z; + + // Simple view frustum culling + if (any(abs(posNDC.xy) > vec2f(1.2))) { return; } + + // Compute color from spherical harmonics + // Camera world position is last column of inversed view matrix + let direction = normalize(pos_opac.xyz - camera.viewInv[3].xyz); + let color = computeColorFromSH( + direction, index, u32(renderSettings.shDegree) + ); + let opacity = sigmoid(pos_opac.w); + + // Calculate 3D covariance matrix from scale and rotation + // rotation matrix + let rot = vec4f(unpack2x16float(gaussian.rot[0]), unpack2x16float(gaussian.rot[1])); + let rotMat = mat3x3f( + 1.0 - 2.0 * (rot.z * rot.z + rot.w * rot.w), 2.0 * (rot.y * rot.z + rot.x * rot.w), 2.0 * (rot.y * rot.w - rot.x * rot.z), + 2.0 * (rot.y * rot.z - rot.x * rot.w), 1.0 - 2.0 * (rot.y * rot.y + rot.w * rot.w), 2.0 * (rot.z * rot.w + rot.x * rot.y), + 2.0 * (rot.y * rot.w + rot.x * rot.z), 2.0 * (rot.z * rot.w - rot.x * rot.y), 1.0 - 2.0 * (rot.y * rot.y + rot.z * rot.z) + ); + // diagonal scale matrix, saved in log space + let scale = renderSettings.gaussianScaling * exp(vec3f(unpack2x16float(gaussian.scale[0]), unpack2x16float(gaussian.scale[1]).x)); + let scaleRot = mat3x3f( + scale.x * rotMat[0], + scale.y * rotMat[1], + scale.z * rotMat[2], + ); + let cov3d = scaleRot * transpose(scaleRot); + + // Calculate 2D covariance matrix from 3D covariance matrix + let WT = mat3x3f(camera.view[0].xyz, camera.view[1].xyz, camera.view[2].xyz); + + let lim = 0.65 * camera.viewport * camera.focal; + let t = vec3f( + clamp(posView.xy / posView.z, -lim, lim) * posView.z, + posView.z + ); + let invTz = 1.0 / t.z; + let JT = mat3x2f( + camera.focal.x * invTz, 0.0, + 0.0, camera.focal.y * invTz, + -camera.focal.x * t.x * invTz * invTz, -camera.focal.y * t.y * invTz * invTz + ); + + let tmpMat = JT * WT; + var cov2d = tmpMat * cov3d * transpose(tmpMat); + cov2d[0][0] += 0.3; + cov2d[1][1] += 0.3; + + // Calculate max radius via eigenvalues of 2D covariance matrix + let det = determinant(cov2d); + if (det == 0) { return; } + + let mid = 0.5 * (cov2d[0][0] + cov2d[1][1]); + let dist = sqrt(max(0.1, mid * mid - det)); + let lambda1 = mid + dist; + let lambda2 = mid - dist; + + // Get size from max radius in NDC space + let size = ceil(3.0 * sqrt(max(lambda1, lambda2))) * 2.0 / camera.viewport; + + // Store sorting information + let keyIndex = atomicAdd(&sortInfos.keysSize, 1); + sortDepths[keyIndex] = bitcast(camera.clippingPlanes[1] - depth); + sortIndices[keyIndex] = keyIndex; + + // increment sortDispatch.dispatchX each time you reach limit for one dispatch of keys + if (keyIndex % (workgroupSize * sortKeyPerThread) == 0) { + atomicAdd(&sortDispatch.dispatchX, 1); + } + + // Handle output to render pipeline + var output: Splat; + output.pos = pack2x16float(posNDC.xy); + output.size = pack2x16float(size); + output.color_opacity = vec2u( + pack2x16float(color.rg), + pack2x16float(vec2f(color.b, opacity)) + ); + let detInv = 1.0 / det; + output.conic = vec2u( + pack2x16float(vec2f(detInv * cov2d[1][1], detInv * -cov2d[0][1])), + pack2x16float(vec2f(detInv * cov2d[0][0], 0)) + ); + splats[keyIndex] = output; } \ No newline at end of file diff --git a/src/sort/sort.ts b/src/sort/sort.ts index 4bbeb23..6d629d8 100644 --- a/src/sort/sort.ts +++ b/src/sort/sort.ts @@ -12,33 +12,33 @@ import radix_sort_wgsl from './radix_sort.wgsl'; import { align } from '../utils/util'; export interface SortStuff { - sort: (encoder: GPUCommandEncoder) => void, - sort_info_buffer: GPUBuffer, - sort_dispatch_indirect_buffer: GPUBuffer, - - // ping-pong - ping_pong: { - sort_indices_buffer: GPUBuffer, - sort_depths_buffer: GPUBuffer, - }[] + sort: (encoder: GPUCommandEncoder) => void, + sort_info_buffer: GPUBuffer, + sort_dispatch_indirect_buffer: GPUBuffer, + + // ping-pong + ping_pong: { + sort_indices_buffer: GPUBuffer, + sort_depths_buffer: GPUBuffer, + }[] } function create_ping_pong_buffer(adjusted_count: number, keysize: number, device: GPUDevice) { - return { - // payload - sort_indices_buffer: device.createBuffer({ - label: 'ping pong sort indices', - size: keysize * 4, - usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, - }), - // key - sort_depths_buffer: device.createBuffer({ - label: 'ping pong sort depths', - size: adjusted_count * 4, - usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, - }), - }; + return { + // payload + sort_indices_buffer: device.createBuffer({ + label: 'ping pong sort indices', + size: keysize * 4, + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, + }), + // key + sort_depths_buffer: device.createBuffer({ + label: 'ping pong sort depths', + size: adjusted_count * 4, + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, + }), + }; } // 8 bit radices @@ -46,321 +46,321 @@ const c_radix_log2 = 8; export const c_histogram_block_rows = 15; export const C = { - histogram_sg_size: 32, - histogram_wg_size: 256, - rs_radix_log2: 8, - rs_radix_size: 1 << c_radix_log2, - rs_keyval_size: 32 / c_radix_log2, - rs_histogram_block_rows: c_histogram_block_rows, - rs_scatter_block_rows: c_histogram_block_rows, - - prefix_wg_size: 1 << 7, - scatter_wg_size: 1 << 8, - - rs_mem_dwords: 0, + histogram_sg_size: 32, + histogram_wg_size: 256, + rs_radix_log2: 8, + rs_radix_size: 1 << c_radix_log2, + rs_keyval_size: 32 / c_radix_log2, + rs_histogram_block_rows: c_histogram_block_rows, + rs_scatter_block_rows: c_histogram_block_rows, + + prefix_wg_size: 1 << 7, + scatter_wg_size: 1 << 8, + + rs_mem_dwords: 0, }; const c_rs_smem_phase_2 = C.rs_radix_size + C.rs_scatter_block_rows * C.scatter_wg_size; C.rs_mem_dwords = c_rs_smem_phase_2; function create_pipelines(device: GPUDevice) { - // storage array length cannot use override, use string concat const instead - const module = device.createShaderModule({ - label: 'radix sort', - // code: radix_sort_wgsl, - code: `const rs_mem_dwords = ${C.rs_mem_dwords}u; + // storage array length cannot use override, use string concat const instead + const module = device.createShaderModule({ + label: 'radix sort', + // code: radix_sort_wgsl, + code: `const rs_mem_dwords = ${C.rs_mem_dwords}u; ${radix_sort_wgsl} ` - }); - - const bind_group_layout = device.createBindGroupLayout({ - entries: [ - // info - { - binding: 0, - visibility: GPUShaderStage.COMPUTE, - buffer: { - type: 'storage', - } - }, - // histograms - { - binding: 1, - visibility: GPUShaderStage.COMPUTE, - buffer: { - type: 'storage', - } - }, - // keys_a - { - binding: 2, - visibility: GPUShaderStage.COMPUTE, - buffer: { - type: 'storage', - } - }, - // keys_b - { - binding: 3, - visibility: GPUShaderStage.COMPUTE, - buffer: { - type: 'storage', - } - }, - // payload_a - { - binding: 4, - visibility: GPUShaderStage.COMPUTE, - buffer: { - type: 'storage', - } - }, - // payload_b - { - binding: 5, - visibility: GPUShaderStage.COMPUTE, - buffer: { - type: 'storage', - } - }, - ] - }); - - const pipeline_layout = device.createPipelineLayout({ - label: 'radix sort', - bindGroupLayouts: [bind_group_layout], - }); - - return { - bind_group_layout, - zero: device.createComputePipeline({ - // label: 'zero histograms', - // layout: 'auto', - layout: pipeline_layout, - compute: { - module: module, - entryPoint: 'zero_histograms', - constants: { - histogram_wg_size: C.histogram_wg_size, - rs_radix_log2: C.rs_radix_log2, - rs_radix_size: C.rs_radix_size, - rs_keyval_size: C.rs_keyval_size, - } - }, - }), - histogram: device.createComputePipeline({ - layout: pipeline_layout, - compute: { - module: module, - entryPoint: 'calculate_histogram', - } - }), - prefix: device.createComputePipeline({ - layout: pipeline_layout, - compute: { - module: module, - entryPoint: 'prefix_histogram', - constants: { - rs_radix_size: C.histogram_wg_size, - prefix_wg_size: C.prefix_wg_size, - } - } - }), - scatter_odd: device.createComputePipeline({ - layout: pipeline_layout, - compute: { - module: module, - entryPoint: 'scatter_odd', - constants: { - histogram_sg_size: C.histogram_sg_size, - histogram_wg_size: C.histogram_wg_size, - rs_radix_log2: C.rs_radix_log2, - rs_radix_size: C.rs_radix_size, - rs_keyval_size: C.rs_keyval_size, - scatter_wg_size: C.scatter_wg_size, - } - } - }), - scatter_even: device.createComputePipeline({ - layout: pipeline_layout, - compute: { - module: module, - entryPoint: 'scatter_even', - constants: { - histogram_sg_size: C.histogram_sg_size, - histogram_wg_size: C.histogram_wg_size, - rs_radix_log2: C.rs_radix_log2, - rs_radix_size: C.rs_radix_size, - rs_keyval_size: C.rs_keyval_size, - scatter_wg_size: C.scatter_wg_size, - } - } - }), - }; + }); + + const bind_group_layout = device.createBindGroupLayout({ + entries: [ + // info + { + binding: 0, + visibility: GPUShaderStage.COMPUTE, + buffer: { + type: 'storage', + } + }, + // histograms + { + binding: 1, + visibility: GPUShaderStage.COMPUTE, + buffer: { + type: 'storage', + } + }, + // keys_a + { + binding: 2, + visibility: GPUShaderStage.COMPUTE, + buffer: { + type: 'storage', + } + }, + // keys_b + { + binding: 3, + visibility: GPUShaderStage.COMPUTE, + buffer: { + type: 'storage', + } + }, + // payload_a + { + binding: 4, + visibility: GPUShaderStage.COMPUTE, + buffer: { + type: 'storage', + } + }, + // payload_b + { + binding: 5, + visibility: GPUShaderStage.COMPUTE, + buffer: { + type: 'storage', + } + }, + ] + }); + + const pipeline_layout = device.createPipelineLayout({ + label: 'radix sort', + bindGroupLayouts: [bind_group_layout], + }); + + return { + bind_group_layout, + zero: device.createComputePipeline({ + // label: 'zero histograms', + // layout: 'auto', + layout: pipeline_layout, + compute: { + module: module, + entryPoint: 'zero_histograms', + constants: { + histogram_wg_size: C.histogram_wg_size, + rs_radix_log2: C.rs_radix_log2, + rs_radix_size: C.rs_radix_size, + rs_keyval_size: C.rs_keyval_size, + } + }, + }), + histogram: device.createComputePipeline({ + layout: pipeline_layout, + compute: { + module: module, + entryPoint: 'calculate_histogram', + } + }), + prefix: device.createComputePipeline({ + layout: pipeline_layout, + compute: { + module: module, + entryPoint: 'prefix_histogram', + constants: { + rs_radix_size: C.histogram_wg_size, + prefix_wg_size: C.prefix_wg_size, + } + } + }), + scatter_odd: device.createComputePipeline({ + layout: pipeline_layout, + compute: { + module: module, + entryPoint: 'scatter_odd', + constants: { + histogram_sg_size: C.histogram_sg_size, + histogram_wg_size: C.histogram_wg_size, + rs_radix_log2: C.rs_radix_log2, + rs_radix_size: C.rs_radix_size, + rs_keyval_size: C.rs_keyval_size, + scatter_wg_size: C.scatter_wg_size, + } + } + }), + scatter_even: device.createComputePipeline({ + layout: pipeline_layout, + compute: { + module: module, + entryPoint: 'scatter_even', + constants: { + histogram_sg_size: C.histogram_sg_size, + histogram_wg_size: C.histogram_wg_size, + rs_radix_log2: C.rs_radix_log2, + rs_radix_size: C.rs_radix_size, + rs_keyval_size: C.rs_keyval_size, + scatter_wg_size: C.scatter_wg_size, + } + } + }), + }; }; function get_scatter_histogram_sizes(keysize: number) { - // as a general rule of thumb, scater_blocks_ru is equal to histo_blocks_ru, except the amount of elements in these two stages is different - - const scatter_block_kvs = C.histogram_wg_size * C.rs_scatter_block_rows; - const scatter_blocks_ru = Math.floor((keysize + scatter_block_kvs - 1) / scatter_block_kvs); - const count_ru_scatter = scatter_blocks_ru * scatter_block_kvs; - - const histo_block_kvs = C.histogram_wg_size * C.rs_histogram_block_rows; - const histo_blocks_ru = Math.floor((count_ru_scatter + histo_block_kvs - 1) / histo_block_kvs); - const count_ru_histo = histo_blocks_ru * histo_block_kvs; - - return { - scatter_block_kvs, - scatter_blocks_ru, - count_ru_scatter, - histo_block_kvs, - histo_blocks_ru, - count_ru_histo, - }; + // as a general rule of thumb, scater_blocks_ru is equal to histo_blocks_ru, except the amount of elements in these two stages is different + + const scatter_block_kvs = C.histogram_wg_size * C.rs_scatter_block_rows; + const scatter_blocks_ru = Math.floor((keysize + scatter_block_kvs - 1) / scatter_block_kvs); + const count_ru_scatter = scatter_blocks_ru * scatter_block_kvs; + + const histo_block_kvs = C.histogram_wg_size * C.rs_histogram_block_rows; + const histo_blocks_ru = Math.floor((count_ru_scatter + histo_block_kvs - 1) / histo_block_kvs); + const count_ru_histo = histo_blocks_ru * histo_block_kvs; + + return { + scatter_block_kvs, + scatter_blocks_ru, + count_ru_scatter, + histo_block_kvs, + histo_blocks_ru, + count_ru_histo, + }; } // caclulates and allocates a buffer that is sufficient for holding all needed information for // sorting. This includes the histograms and the temporary scatter buffer function create_histogram_buffer(keysize: number, device: GPUDevice) { - // currently only a few different key bits are supported, maybe has to be extended - // assert!(key_bits == 32 || key_bits == 64 || key_bits == 16); - - // subgroup and workgroup sizes - const histo_sg_size = C.histogram_sg_size; - const _histo_wg_size = C.histogram_wg_size; - const _prefix_sg_size = histo_sg_size; - const _internal_sg_size = histo_sg_size; - - // The "internal" memory map looks like this: - // +---------------------------------+ <-- 0 - // | histograms[keyval_size] | - // +---------------------------------+ <-- keyval_size * histo_size - // | partitions[scatter_blocks_ru-1] | - // +---------------------------------+ <-- (keyval_size + scatter_blocks_ru - 1) * histo_size - // | workgroup_ids[keyval_size] | - // +---------------------------------+ <-- (keyval_size + scatter_blocks_ru - 1) * histo_size + workgroup_ids_size - - const { scatter_blocks_ru } = get_scatter_histogram_sizes(keysize); - - const histo_size = C.rs_radix_size * Uint32Array.BYTES_PER_ELEMENT; - - // const internal_size = (C.keyval_size + scatter_blocks_ru - 1 + 1) * histo_size; // +1 safety - const internal_size = align((C.rs_keyval_size + scatter_blocks_ru - 1 + 1) * histo_size, 4); // +1 safety - return device.createBuffer({ - label: 'histogram', - size: internal_size, - usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, - }); + // currently only a few different key bits are supported, maybe has to be extended + // assert!(key_bits == 32 || key_bits == 64 || key_bits == 16); + + // subgroup and workgroup sizes + const histo_sg_size = C.histogram_sg_size; + const _histo_wg_size = C.histogram_wg_size; + const _prefix_sg_size = histo_sg_size; + const _internal_sg_size = histo_sg_size; + + // The "internal" memory map looks like this: + // +---------------------------------+ <-- 0 + // | histograms[keyval_size] | + // +---------------------------------+ <-- keyval_size * histo_size + // | partitions[scatter_blocks_ru-1] | + // +---------------------------------+ <-- (keyval_size + scatter_blocks_ru - 1) * histo_size + // | workgroup_ids[keyval_size] | + // +---------------------------------+ <-- (keyval_size + scatter_blocks_ru - 1) * histo_size + workgroup_ids_size + + const { scatter_blocks_ru } = get_scatter_histogram_sizes(keysize); + + const histo_size = C.rs_radix_size * Uint32Array.BYTES_PER_ELEMENT; + + // const internal_size = (C.keyval_size + scatter_blocks_ru - 1 + 1) * histo_size; // +1 safety + const internal_size = align((C.rs_keyval_size + scatter_blocks_ru - 1 + 1) * histo_size, 4); // +1 safety + return device.createBuffer({ + label: 'histogram', + size: internal_size, + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, + }); } const num_pass = 4; export function get_sorter(keysize: number, device: GPUDevice): SortStuff { - const keys_per_workgroup = C.histogram_wg_size * C.rs_histogram_block_rows; - const keys_count_adjusted = (Math.floor((keysize + keys_per_workgroup - 1) / keys_per_workgroup) + 1) * keys_per_workgroup; - - console.log(`keys count adjusted: ${keys_count_adjusted}`); // histogram count - console.log(`key size: ${keysize}`); - - const sort_info_buffer = device.createBuffer({ - label: 'sort info', - size: 5 * 4, - usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, - }); - - const sort_dispatch_indirect_buffer = device.createBuffer({ - label: 'sort dispatch indirect', - size: 3 * 4, - usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.INDIRECT, - }); - - const pipelines = create_pipelines(device); - - const ping_pong = [ - create_ping_pong_buffer(keys_count_adjusted, keysize, device), - create_ping_pong_buffer(keys_count_adjusted, keysize, device), - ]; - - const histogram_buffer = create_histogram_buffer(keysize, device); - - const { scatter_blocks_ru, count_ru_histo } = get_scatter_histogram_sizes(keysize); - device.queue.writeBuffer(sort_info_buffer, 0, new Uint32Array([keysize, count_ru_histo, num_pass, 0, 0])); - device.queue.writeBuffer(sort_dispatch_indirect_buffer, 0, new Uint32Array([scatter_blocks_ru, 1, 1])); - - const bind_group = device.createBindGroup({ - label: 'sort', - layout: pipelines.bind_group_layout, - entries: [ - { binding: 0, resource: { buffer: sort_info_buffer } }, - { binding: 1, resource: { buffer: histogram_buffer } }, - { binding: 2, resource: { buffer: ping_pong[0].sort_depths_buffer } }, - { binding: 3, resource: { buffer: ping_pong[1].sort_depths_buffer } }, - { binding: 4, resource: { buffer: ping_pong[0].sort_indices_buffer } }, - { binding: 5, resource: { buffer: ping_pong[1].sort_indices_buffer } }, - ] - }); - - function record_calculate_histogram_indirect(encoder: GPUCommandEncoder) { - { - const pass = encoder.beginComputePass({ - label: 'zeroing histogram', - }); - pass.setPipeline(pipelines.zero); - pass.setBindGroup(0, bind_group); - pass.dispatchWorkgroupsIndirect(sort_dispatch_indirect_buffer, 0); - pass.end(); - } - { - const pass = encoder.beginComputePass({ - label: 'calculate histogram', - }); - pass.setPipeline(pipelines.histogram); - pass.setBindGroup(0, bind_group); - pass.dispatchWorkgroupsIndirect(sort_dispatch_indirect_buffer, 0); - pass.end(); - } - } + const keys_per_workgroup = C.histogram_wg_size * C.rs_histogram_block_rows; + const keys_count_adjusted = (Math.floor((keysize + keys_per_workgroup - 1) / keys_per_workgroup) + 1) * keys_per_workgroup; + + console.log(`keys count adjusted: ${keys_count_adjusted}`); // histogram count + console.log(`key size: ${keysize}`); - function record_prefix_histogram(encoder: GPUCommandEncoder) { - const pass = encoder.beginComputePass({ - label: 'prefix histogram', + const sort_info_buffer = device.createBuffer({ + label: 'sort info', + size: 5 * 4, + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, }); - pass.setPipeline(pipelines.prefix); - pass.setBindGroup(0, bind_group); - pass.dispatchWorkgroups(4); // passes - pass.end(); - } - - function record_scatter_keys_indirect(encoder: GPUCommandEncoder) { - const pass = encoder.beginComputePass({ - label: 'scatter keyvals', + + const sort_dispatch_indirect_buffer = device.createBuffer({ + label: 'sort dispatch indirect', + size: 3 * 4, + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.INDIRECT, }); - pass.setBindGroup(0, bind_group); - - // assert: passes == 4 - - pass.setPipeline(pipelines.scatter_even); - pass.dispatchWorkgroupsIndirect(sort_dispatch_indirect_buffer, 0); - pass.setPipeline(pipelines.scatter_odd); - pass.dispatchWorkgroupsIndirect(sort_dispatch_indirect_buffer, 0); - pass.setPipeline(pipelines.scatter_even); - pass.dispatchWorkgroupsIndirect(sort_dispatch_indirect_buffer, 0); - pass.setPipeline(pipelines.scatter_odd); - pass.dispatchWorkgroupsIndirect(sort_dispatch_indirect_buffer, 0); - pass.end(); - } - - function sort(encoder: GPUCommandEncoder) { - record_calculate_histogram_indirect(encoder); - record_prefix_histogram(encoder); - record_scatter_keys_indirect(encoder); - }; - - return { - sort_info_buffer, - sort_dispatch_indirect_buffer, - ping_pong, - - sort, - }; + + const pipelines = create_pipelines(device); + + const ping_pong = [ + create_ping_pong_buffer(keys_count_adjusted, keysize, device), + create_ping_pong_buffer(keys_count_adjusted, keysize, device), + ]; + + const histogram_buffer = create_histogram_buffer(keysize, device); + + const { scatter_blocks_ru, count_ru_histo } = get_scatter_histogram_sizes(keysize); + device.queue.writeBuffer(sort_info_buffer, 0, new Uint32Array([keysize, count_ru_histo, num_pass, 0, 0])); + device.queue.writeBuffer(sort_dispatch_indirect_buffer, 0, new Uint32Array([scatter_blocks_ru, 1, 1])); + + const bind_group = device.createBindGroup({ + label: 'sort', + layout: pipelines.bind_group_layout, + entries: [ + { binding: 0, resource: { buffer: sort_info_buffer } }, + { binding: 1, resource: { buffer: histogram_buffer } }, + { binding: 2, resource: { buffer: ping_pong[0].sort_depths_buffer } }, + { binding: 3, resource: { buffer: ping_pong[1].sort_depths_buffer } }, + { binding: 4, resource: { buffer: ping_pong[0].sort_indices_buffer } }, + { binding: 5, resource: { buffer: ping_pong[1].sort_indices_buffer } }, + ] + }); + + function record_calculate_histogram_indirect(encoder: GPUCommandEncoder) { + { + const pass = encoder.beginComputePass({ + label: 'zeroing histogram', + }); + pass.setPipeline(pipelines.zero); + pass.setBindGroup(0, bind_group); + pass.dispatchWorkgroupsIndirect(sort_dispatch_indirect_buffer, 0); + pass.end(); + } + { + const pass = encoder.beginComputePass({ + label: 'calculate histogram', + }); + pass.setPipeline(pipelines.histogram); + pass.setBindGroup(0, bind_group); + pass.dispatchWorkgroupsIndirect(sort_dispatch_indirect_buffer, 0); + pass.end(); + } + } + + function record_prefix_histogram(encoder: GPUCommandEncoder) { + const pass = encoder.beginComputePass({ + label: 'prefix histogram', + }); + pass.setPipeline(pipelines.prefix); + pass.setBindGroup(0, bind_group); + pass.dispatchWorkgroups(4); // passes + pass.end(); + } + + function record_scatter_keys_indirect(encoder: GPUCommandEncoder) { + const pass = encoder.beginComputePass({ + label: 'scatter keyvals', + }); + pass.setBindGroup(0, bind_group); + + // assert: passes == 4 + + pass.setPipeline(pipelines.scatter_even); + pass.dispatchWorkgroupsIndirect(sort_dispatch_indirect_buffer, 0); + pass.setPipeline(pipelines.scatter_odd); + pass.dispatchWorkgroupsIndirect(sort_dispatch_indirect_buffer, 0); + pass.setPipeline(pipelines.scatter_even); + pass.dispatchWorkgroupsIndirect(sort_dispatch_indirect_buffer, 0); + pass.setPipeline(pipelines.scatter_odd); + pass.dispatchWorkgroupsIndirect(sort_dispatch_indirect_buffer, 0); + pass.end(); + } + + function sort(encoder: GPUCommandEncoder) { + record_calculate_histogram_indirect(encoder); + record_prefix_histogram(encoder); + record_scatter_keys_indirect(encoder); + }; + + return { + sort_info_buffer, + sort_dispatch_indirect_buffer, + ping_pong, + + sort, + }; } diff --git a/src/style.css b/src/style.css index f2cdc96..82cedde 100644 --- a/src/style.css +++ b/src/style.css @@ -1,32 +1,35 @@ html, body { - height: 100%; - margin: 0; + height: 100%; + margin: 0; } canvas { - width: 100%; - height: 100%; - display: block; + width: 100%; + height: 100%; + display: block; } .console { - position: absolute; - margin: 20px 20px; + position: absolute; + margin: 20px 20px; - & h1,h2,h3,h4 { - color:#000; - text-shadow: -1px -1px 0 #ccc, 1px -1px 0 #ccc, -1px 1px 0 #ccc, 1px 1px 0 #ccc; - margin: auto; - } + & h1, + h2, + h3, + h4 { + color: #000; + text-shadow: -1px -1px 0 #ccc, 1px -1px 0 #ccc, -1px 1px 0 #ccc, 1px 1px 0 #ccc; + margin: auto; + } - & #log { - color:#fff; - text-shadow: -1px -1px 0 #555, 1px -1px 0 #555, -1px 1px 0 #555, 1px 1px 0 #555; + & #log { + color: #fff; + text-shadow: -1px -1px 0 #555, 1px -1px 0 #555, -1px 1px 0 #555, 1px 1px 0 #555; - & p { - font-size: 12px; - margin: 0.2em; + & p { + font-size: 12px; + margin: 0.2em; + } } - } } \ No newline at end of file diff --git a/src/types.d.ts b/src/types.d.ts index e6cb88c..949dcf0 100644 --- a/src/types.d.ts +++ b/src/types.d.ts @@ -1,4 +1,4 @@ declare module '*.wgsl' { - const shader: string; - export default shader; + const shader: string; + export default shader; } \ No newline at end of file diff --git a/src/utils/load.ts b/src/utils/load.ts index 030b857..2db75ea 100644 --- a/src/utils/load.ts +++ b/src/utils/load.ts @@ -1,121 +1,121 @@ import { Float16Array } from '@petamoriken/float16'; import { log, time, timeLog } from './simple-console'; -import { decodeHeader, readRawVertex ,nShCoeffs} from './plyreader'; +import { decodeHeader, readRawVertex, nShCoeffs } from './plyreader'; const c_size_float = 2; // byte size of f16 const c_size_3d_gaussian = - 3 * c_size_float // x y z (position) - + c_size_float // opacity - + 4 * c_size_float // rotation - + 4 * c_size_float //scale -; + 3 * c_size_float // x y z (position) + + c_size_float // opacity + + 4 * c_size_float // rotation + + 4 * c_size_float //scale + ; export type PointCloud = Awaited>; export async function load(file: string, device: GPUDevice) { - const blob = new Blob([file]); - const arrayBuffer = await new Promise((resolve, reject) => { - const reader = new FileReader(); - - reader.onload = function(event) { - resolve(event.target.result); // Resolve the promise with the ArrayBuffer - }; + const blob = new Blob([file]); + const arrayBuffer = await new Promise((resolve, reject) => { + const reader = new FileReader(); + + reader.onload = function (event) { + resolve(event.target.result); // Resolve the promise with the ArrayBuffer + }; + + reader.onerror = reject; // Reject the promise in case of an error + reader.readAsArrayBuffer(blob); + }); + + const [vertexCount, propertyTypes, vertexData] = decodeHeader(arrayBuffer as ArrayBuffer); + // figure out the SH degree from the number of coefficients + var nRestCoeffs = 0; + for (const propertyName in propertyTypes) { + if (propertyName.startsWith('f_rest_')) { + nRestCoeffs += 1; + } + } + const nCoeffsPerColor = nRestCoeffs / 3; + const sh_deg = Math.sqrt(nCoeffsPerColor + 1) - 1; + const num_coefs = nShCoeffs(sh_deg); + const max_num_coefs = 16; + + const c_size_sh_coef = + 3 * max_num_coefs * c_size_float // 3 channels (RGB) x 16 coefs + ; + + // figure out the order in which spherical harmonics should be read + const shFeatureOrder = []; + for (let rgb = 0; rgb < 3; ++rgb) { + shFeatureOrder.push(`f_dc_${rgb}`); + } + for (let i = 0; i < nCoeffsPerColor; ++i) { + for (let rgb = 0; rgb < 3; ++rgb) { + shFeatureOrder.push(`f_rest_${rgb * nCoeffsPerColor + i}`); + } + } - reader.onerror = reject; // Reject the promise in case of an error - reader.readAsArrayBuffer(blob); - }); - - const [vertexCount, propertyTypes, vertexData] = decodeHeader(arrayBuffer as ArrayBuffer); - // figure out the SH degree from the number of coefficients - var nRestCoeffs = 0; - for (const propertyName in propertyTypes) { - if (propertyName.startsWith('f_rest_')) { - nRestCoeffs += 1; - } - } - const nCoeffsPerColor = nRestCoeffs / 3; - const sh_deg = Math.sqrt(nCoeffsPerColor + 1) - 1; - const num_coefs = nShCoeffs(sh_deg); - const max_num_coefs = 16; - - const c_size_sh_coef = - 3 * max_num_coefs * c_size_float // 3 channels (RGB) x 16 coefs - ; - - // figure out the order in which spherical harmonics should be read - const shFeatureOrder = []; - for (let rgb = 0; rgb < 3; ++rgb) { - shFeatureOrder.push(`f_dc_${rgb}`); - } - for (let i = 0; i < nCoeffsPerColor; ++i) { - for (let rgb = 0; rgb < 3; ++rgb) { - shFeatureOrder.push(`f_rest_${rgb * nCoeffsPerColor + i}`); - } - } - - const num_points = vertexCount; - - log(`num points: ${num_points}`); - log(`processing loaded attributes...`); - time(); - - // xyz (position), opacity, cov (from rot and scale) - const gaussian_3d_buffer = device.createBuffer({ - label: 'ply input 3d gaussians data buffer', - size: num_points * c_size_3d_gaussian, // buffer size multiple of 4? - usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.STORAGE, - mappedAtCreation: true, - }); - const gaussian = new Float16Array(gaussian_3d_buffer.getMappedRange()); - - // Spherical harmonic function coeffs - const sh_buffer = device.createBuffer({ - label: 'ply input 3d gaussians data buffer', - size: num_points * c_size_sh_coef, // buffer size multiple of 4? - usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.STORAGE, - mappedAtCreation: true, - }); - const sh = new Float16Array(sh_buffer.getMappedRange()); - - var readOffset = 0; - for (let i = 0; i < num_points; i++) { - const [newReadOffset, rawVertex] = readRawVertex(readOffset, vertexData, propertyTypes); - readOffset = newReadOffset; - - const o = i * (c_size_3d_gaussian / c_size_float); - const output_offset = i * max_num_coefs * 3; - - for (let order = 0; order < num_coefs; ++order) { - const order_offset = order * 3; - for (let j = 0; j < 3; ++j) { - const coeffName = shFeatureOrder[order * 3 + j]; - sh[output_offset +order_offset+j]=rawVertex[coeffName]; + const num_points = vertexCount; + + log(`num points: ${num_points}`); + log(`processing loaded attributes...`); + time(); + + // xyz (position), opacity, cov (from rot and scale) + const gaussian_3d_buffer = device.createBuffer({ + label: 'ply input 3d gaussians data buffer', + size: num_points * c_size_3d_gaussian, // buffer size multiple of 4? + usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.STORAGE, + mappedAtCreation: true, + }); + const gaussian = new Float16Array(gaussian_3d_buffer.getMappedRange()); + + // Spherical harmonic function coeffs + const sh_buffer = device.createBuffer({ + label: 'ply input 3d gaussians data buffer', + size: num_points * c_size_sh_coef, // buffer size multiple of 4? + usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.STORAGE, + mappedAtCreation: true, + }); + const sh = new Float16Array(sh_buffer.getMappedRange()); + + var readOffset = 0; + for (let i = 0; i < num_points; i++) { + const [newReadOffset, rawVertex] = readRawVertex(readOffset, vertexData, propertyTypes); + readOffset = newReadOffset; + + const o = i * (c_size_3d_gaussian / c_size_float); + const output_offset = i * max_num_coefs * 3; + + for (let order = 0; order < num_coefs; ++order) { + const order_offset = order * 3; + for (let j = 0; j < 3; ++j) { + const coeffName = shFeatureOrder[order * 3 + j]; + sh[output_offset + order_offset + j] = rawVertex[coeffName]; + } } + + gaussian[o + 0] = rawVertex.x; + gaussian[o + 1] = rawVertex.y; + gaussian[o + 2] = rawVertex.z; + gaussian[o + 3] = rawVertex.opacity; + gaussian[o + 4] = rawVertex.rot_0; + gaussian[o + 5] = rawVertex.rot_1; + gaussian[o + 6] = rawVertex.rot_2; + gaussian[o + 7] = rawVertex.rot_3; + gaussian[o + 8] = rawVertex.scale_0; + gaussian[o + 9] = rawVertex.scale_1; + gaussian[o + 10] = rawVertex.scale_2; } - gaussian[o + 0] = rawVertex.x; - gaussian[o + 1] = rawVertex.y; - gaussian[o + 2] = rawVertex.z; - gaussian[o + 3] = rawVertex.opacity; - gaussian[o + 4] = rawVertex.rot_0; - gaussian[o + 5] = rawVertex.rot_1; - gaussian[o + 6] = rawVertex.rot_2; - gaussian[o + 7] = rawVertex.rot_3; - gaussian[o + 8] = rawVertex.scale_0; - gaussian[o + 9] = rawVertex.scale_1; - gaussian[o + 10] = rawVertex.scale_2; - } - - gaussian_3d_buffer.unmap(); - sh_buffer.unmap(); - - timeLog(); - console.log("return result!"); - return { - num_points: num_points, - sh_deg: sh_deg, - gaussian_3d_buffer, - sh_buffer, - }; + gaussian_3d_buffer.unmap(); + sh_buffer.unmap(); + + timeLog(); + console.log("return result!"); + return { + num_points: num_points, + sh_deg: sh_deg, + gaussian_3d_buffer, + sh_buffer, + }; } diff --git a/src/utils/simple-console.ts b/src/utils/simple-console.ts index 213ece1..2966926 100644 --- a/src/utils/simple-console.ts +++ b/src/utils/simple-console.ts @@ -1,24 +1,24 @@ const html_log = document.querySelector('#log') as HTMLDivElement; export async function log(msg: string) { - console.log(msg); - const p = document.createElement('p'); - p.innerText = msg; - html_log.appendChild(p); + console.log(msg); + const p = document.createElement('p'); + p.innerText = msg; + html_log.appendChild(p); } let t: number; export function time() { - t = performance.now(); + t = performance.now(); } export function timeLog() { - const d = performance.now() - t; - log(`${d.toFixed(0)} ms`); + const d = performance.now() - t; + log(`${d.toFixed(0)} ms`); } -export function timeReturn(){ - const d = performance.now() - t; - return d; +export function timeReturn() { + const d = performance.now() - t; + return d; } \ No newline at end of file diff --git a/src/utils/util.ts b/src/utils/util.ts index 7ad96a8..47530fe 100644 --- a/src/utils/util.ts +++ b/src/utils/util.ts @@ -4,30 +4,30 @@ export function assert( condition: boolean, msg?: string | (() => string) - ): asserts condition { +): asserts condition { if (!condition) { - throw new Error(msg && (typeof msg === 'string' ? msg : msg())); + throw new Error(msg && (typeof msg === 'string' ? msg : msg())); } - } +} /** If the argument is an Error, throw it. Otherwise, pass it back. */ export function assertOK(value: Error | T): T { -if (value instanceof Error) { - throw value; -} -return value; + if (value instanceof Error) { + throw value; + } + return value; } /** * Assert this code is unreachable. Unconditionally throws an `Error`. */ export function unreachable(msg?: string): never { -throw new Error(msg); + throw new Error(msg); } /** Round `n` up to the next multiple of `alignment` (inclusive). */ export function align(n: number, alignment: number): number { - assert(Number.isInteger(n) && n >= 0, 'n must be a non-negative integer'); - assert(Number.isInteger(alignment) && alignment > 0, 'alignment must be a positive integer'); - return Math.ceil(n / alignment) * alignment; + assert(Number.isInteger(n) && n >= 0, 'n must be a non-negative integer'); + assert(Number.isInteger(alignment) && alignment > 0, 'alignment must be a positive integer'); + return Math.ceil(n / alignment) * alignment; } \ No newline at end of file