Examples - Update Compute Fluid to use IndirectStorageBuffer (#32176)

Author: cmhhelgeson

examples/webgpu_compute_particles_fluid.html

@@ -33,7 +33,7 @@
 
 			import * as THREE from 'three/webgpu';
 
-			import { Fn, If, Return, instancedArray, instanceIndex, uniform, attribute, uint, float, clamp, struct, atomicStore, int, ivec3, array, vec3, atomicAdd, Loop, atomicLoad, max, pow, mat3, vec4, cross, step } from 'three/tsl';
+			import { Fn, If, Return, instancedArray, instanceIndex, uniform, attribute, uint, float, clamp, struct, atomicStore, int, ivec3, array, vec3, atomicAdd, Loop, atomicLoad, max, pow, mat3, vec4, cross, step, storage } from 'three/tsl';
 
 			import { Inspector } from 'three/addons/inspector/Inspector.js';
 
@@ -48,12 +48,14 @@
 
 			const maxParticles = 8192 * 16;
 			const gridSize1d = 64;
+			const workgroupSize = 64;
 			const gridSize = new THREE.Vector3( gridSize1d, gridSize1d, gridSize1d );
 			const fixedPointMultiplier = 1e7;
 
 			let particleCountUniform, stiffnessUniform, restDensityUniform, dynamicViscosityUniform, dtUniform, gravityUniform, gridSizeUniform;
 			let particleBuffer, cellBuffer, cellBufferFloat;
-			let clearGridKernel, p2g1Kernel, p2g2Kernel, updateGridKernel, g2pKernel;
+			let clearGridKernel, p2g1Kernel, p2g2Kernel, updateGridKernel, g2pKernel, workgroupKernel;
+			let p2g1KernelWorkgroupBuffer, p2g2KernelWorkgroupBuffer, g2pKernelWorkgroupBuffer;
 			let particleMesh;
 			const mouseCoord = new THREE.Vector3();
 			const prevMouseCoord = new THREE.Vector3();
@@ -105,12 +107,29 @@
 				setupParticles();
 
 				const gui = renderer.inspector.createParameters( 'Settings' );
+			
+				const numWorkgroups = Math.ceil( params.particleCount / workgroupSize );
+
+				p2g1KernelWorkgroupBuffer = new THREE.IndirectStorageBufferAttribute( new Uint32Array( [ numWorkgroups, 1, 1 ] ), 1 );
+				p2g2KernelWorkgroupBuffer = new THREE.IndirectStorageBufferAttribute( new Uint32Array( [ numWorkgroups, 1, 1 ] ), 1 );
+				g2pKernelWorkgroupBuffer = new THREE.IndirectStorageBufferAttribute( new Uint32Array( [ numWorkgroups, 1, 1 ] ), 1 );
+
+				const p2g1WorkgroupStorage = storage( p2g1KernelWorkgroupBuffer, 'uint', 3 );
+				const p2g2WorkgroupStorage = storage( p2g2KernelWorkgroupBuffer, 'uint', 3 );
+				const g2pWorkgroupStorage = storage( g2pKernelWorkgroupBuffer, 'uint', 3 );
+
+				workgroupKernel = Fn( () => {
+
+					const workgroupsToDispatch = ( particleCountUniform.sub( 1 ) ).div( workgroupSize ).add( 1 );
+
+					p2g1WorkgroupStorage.element( 0 ).assign( workgroupsToDispatch );
+					p2g2WorkgroupStorage.element( 0 ).assign( workgroupsToDispatch );
+					g2pWorkgroupStorage.element( 0 ).assign( workgroupsToDispatch );
+
+				} )().compute( 1 );
 
 				gui.add( params, 'particleCount', 4096, maxParticles, 4096 ).onChange( value => {
 
-					p2g1Kernel.count = value;
-					p2g2Kernel.count = value;
-					g2pKernel.count = value;
 					particleMesh.count = value;
 					particleCountUniform.value = value;
 
@@ -258,7 +277,7 @@
 
 					} );
 
-				} )().compute( params.particleCount ).setName( 'p2g1Kernel' );
+				} )().compute( params.particleCount, [ workgroupSize, 1, 1 ] ).setName( 'p2g1Kernel' );
 
 				p2g2Kernel = Fn( () => {
 
@@ -329,7 +348,7 @@
 
 					} );
 
-				} )().compute( params.particleCount ).setName( 'p2g2Kernel' );
+				} )().compute( params.particleCount, [ workgroupSize, 1, 1 ] ).setName( 'p2g2Kernel' );
 
 				updateGridKernel = Fn( () => {
 
@@ -477,7 +496,7 @@
 					particleBuffer.element( instanceIndex ).get( 'position' ).assign( particlePosition );
 					particleBuffer.element( instanceIndex ).get( 'velocity' ).assign( particleVelocity );
 
-				} )().compute( params.particleCount ).setName( 'g2pKernel' );
+				} )().compute( params.particleCount, [ workgroupSize, 1, 1 ] ).setName( 'g2pKernel' );
 
 			}
 
@@ -561,12 +580,14 @@
 
 				prevMouseCoord.copy( mouseCoord );
 
+				renderer.compute( workgroupKernel );
+
 				//renderer.compute( [ clearGridKernel, p2g1Kernel, p2g2Kernel, updateGridKernel, g2pKernel ] );
 				renderer.compute( clearGridKernel );
-				renderer.compute( p2g1Kernel );
-				renderer.compute( p2g2Kernel );
+				renderer.compute( p2g1Kernel, p2g1KernelWorkgroupBuffer );
+				renderer.compute( p2g2Kernel, p2g2KernelWorkgroupBuffer );
 				renderer.compute( updateGridKernel );
-				renderer.compute( g2pKernel );
+				renderer.compute( g2pKernel, g2pKernelWorkgroupBuffer );
 
 				renderer.render( scene, camera );