@gnsx/three 0.185.15 → 0.186.0

package/examples/jsm/gpgpu/CountingSort.js

@@ -0,0 +1,271 @@
+import { StorageBufferAttribute, DynamicDrawUsage } from '@gnsx/three/webgpu';
+import { Fn, Loop, atomicAdd, atomicLoad, atomicStore, instanceIndex, storage, uint } from '@gnsx/three/tsl';
+
+/**
+ * A reusable GPU counting sort.
+ *
+ * This computes a stable-ish permutation of the integers `[0, count)` that orders them by an
+ * arbitrary, user supplied `uint` key ("bin") in the range `[0, binCount)`. It is a good fit for
+ * approximate ordering of large element counts (hundreds of thousands to millions) where an exact
+ * comparison sort such as a bitonic sort (see {@link BitonicSort}) would be too slow: a counting
+ * sort only requires a fixed number of passes (reset, histogram, prefix sum, scatter) regardless of
+ * `count`, at the cost of only being accurate to the resolution of `binCount` - elements that land
+ * in the same bin end up in an unspecified relative order.
+ *
+ * This class does not compute the sort key itself. Instead, a TSL function is supplied via
+ * {@link CountingSort#setBinNode} that maps the current `instanceIndex` to a bin, and an equivalent
+ * plain JavaScript function can be supplied to {@link CountingSort#computeCPU} for platforms without
+ * compute shader support (e.g. the WebGL backend of {@link WebGPURenderer}).
+ *
+ * ```js
+ * const sort = new CountingSort( count, { binCount: 4096 } );
+ * sort.setBinNode( () => {
+ *
+ * 	// return a `Node<uint>` bin index for `instanceIndex`, e.g. derived from a depth value.
+ *
+ * } );
+ *
+ * sort.compute( renderer );
+ *
+ * // `sort.orderRead` now holds a storage buffer of `count` indices, ordered by bin.
+ * ```
+ *
+ * @three_import import { CountingSort } from 'three/addons/gpgpu/CountingSort.js';
+ */
+class CountingSort {
+
+	/**
+	 * Constructs a new counting sort.
+	 *
+	 * @param {number} count - The number of elements to sort.
+	 * @param {Object} [options={}] - Options that modify the counting sort.
+	 * @param {number} [options.binCount=4096] - The number of bins/buckets the sort key is quantized into. Larger values improve sort accuracy at the cost of a longer (but still single-pass) prefix sum.
+	 * @param {number} [options.workgroupSize=256] - The workgroup size of the compute shaders executed during the sort.
+	 */
+	constructor( count, { binCount = 4096, workgroupSize = 256 } = {} ) {
+
+		/**
+		 * The number of elements to sort.
+		 *
+		 * @type {number}
+		 */
+		this.count = count;
+
+		/**
+		 * The number of bins/buckets the sort key is quantized into.
+		 *
+		 * @type {number}
+		 */
+		this.binCount = binCount;
+
+		/**
+		 * The workgroup size of the compute shaders executed during the sort.
+		 *
+		 * @type {number}
+		 */
+		this.workgroupSize = workgroupSize;
+
+		const orderData = new Uint32Array( count );
+		for ( let i = 0; i < count; i ++ ) orderData[ i ] = i;
+
+		/**
+		 * The buffer attribute holding the sorted order (a permutation of `[0, count)`). This is
+		 * also the attribute that is kept up to date by {@link CountingSort#computeCPU}.
+		 *
+		 * @type {StorageBufferAttribute}
+		 */
+		this.orderAttribute = new StorageBufferAttribute( orderData, 1, Uint32Array );
+
+		const binAttribute = new StorageBufferAttribute( new Uint32Array( count ), 1, Uint32Array );
+		const histogramAttribute = new StorageBufferAttribute( new Uint32Array( binCount ), 1, Uint32Array );
+		const offsetAttribute = new StorageBufferAttribute( new Uint32Array( binCount ), 1, Uint32Array );
+
+		/**
+		 * A read-only storage node for the sorted order buffer.
+		 *
+		 * @type {StorageBufferNode}
+		 */
+		this.orderRead = storage( this.orderAttribute, 'uint', count ).toReadOnly();
+
+		/**
+		 * A writable storage node for the sorted order buffer.
+		 *
+		 * @type {StorageBufferNode}
+		 */
+		this.orderWrite = storage( this.orderAttribute, 'uint', count );
+
+		/**
+		 * A read-only storage node holding each element's bin, computed during the histogram pass.
+		 *
+		 * @type {StorageBufferNode}
+		 */
+		this.binRead = storage( binAttribute, 'uint', count ).toReadOnly();
+
+		/**
+		 * A writable storage node holding each element's bin.
+		 *
+		 * @type {StorageBufferNode}
+		 */
+		this.binWrite = storage( binAttribute, 'uint', count );
+
+		/**
+		 * An atomic storage node used to accumulate the per-bin histogram.
+		 *
+		 * @type {StorageBufferNode}
+		 */
+		this.histogramAtomic = storage( histogramAttribute, 'uint', binCount ).toAtomic();
+
+		/**
+		 * An atomic storage node used both for the exclusive prefix sum of the histogram and, during
+		 * the scatter pass, as a per-bin write cursor.
+		 *
+		 * @type {StorageBufferNode}
+		 */
+		this.offsetAtomic = storage( offsetAttribute, 'uint', binCount ).toAtomic();
+
+		this._webGLBuffersEnabled = false;
+
+		this._cpuBins = new Uint32Array( count );
+		this._cpuCounts = new Uint32Array( binCount );
+		this._cpuOffsets = new Uint32Array( binCount );
+
+		this._resetNode = null;
+		this._histogramNode = null;
+		this._prefixNode = null;
+		this._scatterNode = null;
+
+	}
+
+	/**
+	 * Sets the TSL function used to compute the bin of the element currently referenced by
+	 * `instanceIndex`, and (re)builds the compute nodes used by {@link CountingSort#compute}.
+	 *
+	 * @param {Function} binNode - A parameterless function returning a `Node<uint>` in `[0, binCount)`.
+	 */
+	setBinNode( binNode ) {
+
+		const { binCount, workgroupSize, count } = this;
+
+		this._resetNode = Fn( () => {
+
+			atomicStore( this.histogramAtomic.element( instanceIndex ), uint( 0 ) );
+			atomicStore( this.offsetAtomic.element( instanceIndex ), uint( 0 ) );
+
+		} )().compute( binCount, [ workgroupSize ] ).setName( 'CountingSortReset' );
+
+		this._histogramNode = Fn( () => {
+
+			const bin = binNode().toVar( 'bin' );
+
+			this.binWrite.element( instanceIndex ).assign( bin );
+			atomicAdd( this.histogramAtomic.element( bin ), uint( 1 ) );
+
+		} )().compute( count, [ workgroupSize ] ).setName( 'CountingSortHistogram' );
+
+		this._prefixNode = Fn( () => {
+
+			const sum = uint( 0 ).toVar( 'sum' );
+
+			Loop( { start: 0, end: binCount, type: 'uint', name: 'bin', condition: '<' }, ( { bin } ) => {
+
+				const binCountValue = atomicLoad( this.histogramAtomic.element( bin ) ).toVar( 'count' );
+				atomicStore( this.offsetAtomic.element( bin ), sum );
+				sum.addAssign( binCountValue );
+
+			} );
+
+		} )().compute( 1 ).setName( 'CountingSortPrefix' );
+
+		this._scatterNode = Fn( () => {
+
+			const bin = this.binRead.element( instanceIndex ).toVar( 'bin' );
+			const targetIndex = atomicAdd( this.offsetAtomic.element( bin ), uint( 1 ) ).toVar( 'targetIndex' );
+
+			this.orderWrite.element( targetIndex ).assign( instanceIndex );
+
+		} )().compute( count, [ workgroupSize ] ).setName( 'CountingSortScatter' );
+
+	}
+
+	/**
+	 * Executes a complete counting sort on the GPU, updating {@link CountingSort#orderRead}.
+	 *
+	 * @param {Renderer} renderer - The current scene's renderer.
+	 */
+	compute( renderer ) {
+
+		renderer.compute( this._resetNode );
+		renderer.compute( this._histogramNode );
+		renderer.compute( this._prefixNode );
+		renderer.compute( this._scatterNode );
+
+	}
+
+	/**
+	 * Executes a complete counting sort on the CPU, updating {@link CountingSort#orderAttribute}.
+	 * Intended as a fallback for backends without compute shader support.
+	 *
+	 * @param {Function} binFn - A function taking an element index and returning its bin (a plain number in `[0, binCount)`).
+	 */
+	computeCPU( binFn ) {
+
+		const { count, binCount } = this;
+		const order = this.orderAttribute.array;
+		const bins = this._cpuBins;
+		const counts = this._cpuCounts;
+		const offsets = this._cpuOffsets;
+
+		counts.fill( 0 );
+
+		for ( let i = 0; i < count; i ++ ) {
+
+			const bin = binFn( i );
+
+			bins[ i ] = bin;
+			counts[ bin ] ++;
+
+		}
+
+		let sum = 0;
+
+		for ( let i = 0; i < binCount; i ++ ) {
+
+			offsets[ i ] = sum;
+			sum += counts[ i ];
+
+		}
+
+		for ( let i = 0; i < count; i ++ ) {
+
+			order[ offsets[ bins[ i ] ] ++ ] = i;
+
+		}
+
+		this.orderAttribute.needsUpdate = true;
+
+		if ( this.orderAttribute.pbo !== undefined ) {
+
+			this.orderAttribute.pbo.needsUpdate = true;
+
+		}
+
+	}
+
+	/**
+	 * Enables the WebGL-specific storage buffer path (PBO + dynamic draw usage) for the order buffer.
+	 * Only needed when {@link CountingSort#computeCPU} is used with the WebGL backend of {@link WebGPURenderer}.
+	 */
+	enableWebGLBuffers() {
+
+		if ( this._webGLBuffersEnabled === true ) return;
+
+		this.orderAttribute.setUsage( DynamicDrawUsage );
+		this.orderRead.setPBO( true );
+
+		this._webGLBuffersEnabled = true;
+
+	}
+
+}
+
+export { CountingSort };

package/examples/jsm/helpers/LightProbeGridHelper.js

@@ -1,173 +1,27 @@
 import {
-	Data3DTexture,
-	FloatType,
 	InstancedBufferAttribute,
 	InstancedMesh,
-	LinearFilter,
 	Matrix4,
-	RGBAFormat,
-	ShaderMaterial,
+	NodeMaterial,
 	SphereGeometry,
 	Vector3
-} from '@gnsx/three';
-// WITH_GENESYS
-import { NodeMaterial } from '@gnsx/three/webgpu';
-import {
-	attribute,
-	Fn,
-	getLightProbeGridIrradianceAtUV,
-	normalWorld,
-	texture3D,
-	uniform,
-	vec4
-} from '@gnsx/three/tsl';
-// !WITH_GENESYS
-
-// WITH_GENESYS
-const _emptyTexture = /*@__PURE__*/ new Data3DTexture( new Float32Array( 1 * 1 * 28 * 4 ), 1, 1, 28 );
-
-_emptyTexture.format = RGBAFormat;
-_emptyTexture.type = FloatType;
-_emptyTexture.minFilter = LinearFilter;
-_emptyTexture.magFilter = LinearFilter;
-_emptyTexture.needsUpdate = true;
-
-function createNodeMaterial( probes ) {
-
-	const probesSH = texture3D( probes.texture || _emptyTexture );
-	const probesResolution = uniform( new Vector3( 1, 1, 1 ) );
-	const instanceUVW = attribute( 'instanceUVW', 'vec3' );
-
-	const fragmentNode = Fn( () => {
-
-		const irradiance = getLightProbeGridIrradianceAtUV( {
-			probesSH,
-			probesResolution,
-			gridUVW: instanceUVW,
-			worldNormal: normalWorld
-		} );
-
-		return vec4( irradiance, 1.0 );
-
-	} )();
-
-	const material = new NodeMaterial();
-	material.fragmentNode = fragmentNode;
-
-	return { material, probesSH, probesResolution };
-
-}
-// !WITH_GENESYS
-
-function createShaderMaterial() {
-
-	return new ShaderMaterial( {
-
-		uniforms: {
-
-			probesSH: { value: null },
-			probesResolution: { value: new Vector3() },
-
-		},
-
-		vertexShader: /* glsl */`
-
-			attribute vec3 instanceUVW;
-
-			varying vec3 vWorldNormal;
-			varying vec3 vUVW;
-
-			void main() {
-
-				vUVW = instanceUVW;
-				vWorldNormal = normalize( mat3( modelMatrix ) * normal );
-				gl_Position = projectionMatrix * modelViewMatrix * instanceMatrix * vec4( position, 1.0 );
-
-			}
-
-		`,
-
-		fragmentShader: /* glsl */`
-
-			precision highp sampler3D;
-
-			uniform sampler3D probesSH;
-			uniform vec3 probesResolution;
-
-			varying vec3 vWorldNormal;
-			varying vec3 vUVW;
-
-			void main() {
-
-				// Atlas UV mapping — must match lightprobes_pars_fragment.glsl.js
-				float nz          = probesResolution.z;
-				float paddedSlices = nz + 2.0;
-				float atlasDepth  = 7.0 * paddedSlices;
-				float uvZBase     = vUVW.z * nz + 1.0;
-
-				vec4 s0 = texture( probesSH, vec3( vUVW.xy, ( uvZBase                       ) / atlasDepth ) );
-				vec4 s1 = texture( probesSH, vec3( vUVW.xy, ( uvZBase +       paddedSlices   ) / atlasDepth ) );
-				vec4 s2 = texture( probesSH, vec3( vUVW.xy, ( uvZBase + 2.0 * paddedSlices   ) / atlasDepth ) );
-				vec4 s3 = texture( probesSH, vec3( vUVW.xy, ( uvZBase + 3.0 * paddedSlices   ) / atlasDepth ) );
-				vec4 s4 = texture( probesSH, vec3( vUVW.xy, ( uvZBase + 4.0 * paddedSlices   ) / atlasDepth ) );
-				vec4 s5 = texture( probesSH, vec3( vUVW.xy, ( uvZBase + 5.0 * paddedSlices   ) / atlasDepth ) );
-				vec4 s6 = texture( probesSH, vec3( vUVW.xy, ( uvZBase + 6.0 * paddedSlices   ) / atlasDepth ) );
-
-				// Unpack 9 vec3 SH L2 coefficients
-
-				vec3 c0 = s0.xyz;
-				vec3 c1 = vec3( s0.w, s1.xy );
-				vec3 c2 = vec3( s1.zw, s2.x );
-				vec3 c3 = s2.yzw;
-				vec3 c4 = s3.xyz;
-				vec3 c5 = vec3( s3.w, s4.xy );
-				vec3 c6 = vec3( s4.zw, s5.x );
-				vec3 c7 = s5.yzw;
-				vec3 c8 = s6.xyz;
-
-				vec3 n = normalize( vWorldNormal );
-
-				float x = n.x, y = n.y, z = n.z;
-
-				// band 0
-				vec3 result = c0 * 0.886227;
-
-				// band 1,
-				result += c1 * 2.0 * 0.511664 * y;
-				result += c2 * 2.0 * 0.511664 * z;
-				result += c3 * 2.0 * 0.511664 * x;
-
-				// band 2,
-				result += c4 * 2.0 * 0.429043 * x * y;
-				result += c5 * 2.0 * 0.429043 * y * z;
-				result += c6 * ( 0.743125 * z * z - 0.247708 );
-				result += c7 * 2.0 * 0.429043 * x * z;
-				result += c8 * 0.429043 * ( x * x - y * y );
-
-				gl_FragColor = vec4( max( result, vec3( 0.0 ) ), 1.0 );
-
-				#include <tonemapping_fragment>
-				#include <colorspace_fragment>
-
-			}
-
-		`
-
-	} );
-
-}
+} from '@gnsx/three/webgpu';
+import { array, attribute, Fn, getShIrradianceAt, normalWorld, texture3D, uniform, vec3, vec4 } from '@gnsx/three/tsl';
 
 /**
- * Visualizes an {@link LightProbeGrid} by rendering a sphere at each
- * probe position, shaded with the probe's L2 spherical harmonics.
+ * Visualizes a {@link LightProbeGrid} by rendering a sphere at each probe
+ * position, shaded with the probe's L2 spherical harmonics. Uses a single
+ * `InstancedMesh` draw call for all probes.
  *
- * Uses a single `InstancedMesh` draw call for all probes.
+ * This helper can only be used with {@link WebGPURenderer}.
+ * When using {@link WebGLRenderer}, import from `LightProbeGridHelperWebGL.js`.
  *
  * ```js
  * const helper = new LightProbeGridHelper( probes );
  * scene.add( helper );
  * ```
  *
+ * @private
  * @augments InstancedMesh
  * @three_import import { LightProbeGridHelper } from 'three/addons/helpers/LightProbeGridHelper.js';
  */
@@ -178,18 +32,16 @@ class LightProbeGridHelper extends InstancedMesh {
 	 *
 	 * @param {LightProbeGrid} probes - The probe grid to visualize.
 	 * @param {number} [sphereSize=0.12] - The radius of each probe sphere.
-	 * @param {Object} [options] - The helper options.
-	 * @param {boolean} [options.webgpu=false] - Whether to use a WebGPU-compatible node material.
 	 */
-	constructor( probes, sphereSize = 0.12, options = {} ) {
+	constructor( probes, sphereSize = 0.12 ) {
 
 		const geometry = new SphereGeometry( sphereSize, 16, 16 );
-		const materialState = options.webgpu === true ? createNodeMaterial( probes ) : { material: createShaderMaterial() };
+		const material = new NodeMaterial();
 
 		const res = probes.resolution;
 		const count = res.x * res.y * res.z;
 
-		super( geometry, materialState.material, count );
+		super( geometry, material, count );
 
 		/**
 		 * The probe grid to visualize.
@@ -200,18 +52,49 @@ class LightProbeGridHelper extends InstancedMesh {
 
 		this.type = 'LightProbeGridHelper';
 
-		// WITH_GENESYS
-		this._probesSH = materialState.probesSH || null;
-		this._probesResolution = materialState.probesResolution || null;
-		// !WITH_GENESYS
+		// Atlas and resolution are swappable uniforms, so the shading node builds once.
+
+		this._atlas = texture3D( probes.texture );
+		this._resolution = uniform( new Vector3() );
+
+		const nz = this._resolution.z;
+		const paddedSlices = nz.add( 2.0 );
+		const atlasDepth = paddedSlices.mul( 7.0 );
+
+		material.fragmentNode = Fn( () => {
+
+			const uvw = attribute( 'instanceUVW', 'vec3' );
+			const uvZBase = uvw.z.mul( nz ).add( 1.0 );
+
+			const slice = ( t ) => this._atlas.sample( vec3( uvw.xy, uvZBase.add( paddedSlices.mul( t ) ).div( atlasDepth ) ) );
+
+			const s0 = slice( 0 ), s1 = slice( 1 ), s2 = slice( 2 ), s3 = slice( 3 );
+			const s4 = slice( 4 ), s5 = slice( 5 ), s6 = slice( 6 );
+
+			const sh = array( [
+				s0.xyz,
+				vec3( s0.w, s1.xy ),
+				vec3( s1.zw, s2.x ),
+				s2.yzw,
+				s3.xyz,
+				vec3( s3.w, s4.xy ),
+				vec3( s4.zw, s5.x ),
+				s5.yzw,
+				s6.xyz
+			] );
+
+			return vec4( getShIrradianceAt( normalWorld, sh ).max( vec3( 0.0 ) ), 1.0 );
+
+		} )();
 
 		this.update();
 
 	}
 
 	/**
-	 * Rebuilds instance matrices and UVW attributes from the current probe grid.
-	 * Call this after changing `probes` or after re-baking.
+	 * Rebuilds instance matrices and UVW attributes from the current probe grid,
+	 * and rebinds the shading node to its atlas. Call this after changing
+	 * `probes` or after re-baking.
 	 */
 	update() {
 
@@ -219,17 +102,17 @@ class LightProbeGridHelper extends InstancedMesh {
 		const res = probes.resolution;
 		const count = res.x * res.y * res.z;
 
+		// Resize instance matrix buffer if needed.
+
 		if ( this.instanceMatrix.count !== count ) {
 
-			// WITH_GENESYS
-			this.instanceMatrix.dispose();
-			// !WITH_GENESYS
 			this.instanceMatrix = new InstancedBufferAttribute( new Float32Array( count * 16 ), 16 );
 
 		}
 
 		this.count = count;
 
+		const uvwArray = new Float32Array( count * 3 );
 		const matrix = new Matrix4();
 		const probePos = new Vector3();
 
@@ -241,6 +124,11 @@ class LightProbeGridHelper extends InstancedMesh {
 
 				for ( let ix = 0; ix < res.x; ix ++ ) {
 
+					// Remap to texel centers (must match LightProbeGridNode).
+					uvwArray[ i * 3 ] = ( ix + 0.5 ) / res.x;
+					uvwArray[ i * 3 + 1 ] = ( iy + 0.5 ) / res.y;
+					uvwArray[ i * 3 + 2 ] = ( iz + 0.5 ) / res.z;
+
 					probes.getProbePosition( ix, iy, iz, probePos );
 					matrix.makeTranslation( probePos.x, probePos.y, probePos.z );
 					this.setMatrixAt( i, matrix );
@@ -255,65 +143,10 @@ class LightProbeGridHelper extends InstancedMesh {
 
 		this.instanceMatrix.needsUpdate = true;
 
-		let instanceUVW = this.geometry.getAttribute( 'instanceUVW' );
-
-		if ( instanceUVW === undefined || instanceUVW.count !== count ) {
-
-			// WITH_GENESYS
-			if ( instanceUVW !== undefined ) {
-
-				instanceUVW.dispose();
-
-			}
-			// !WITH_GENESYS
-
-			const uvwArray = new Float32Array( count * 3 );
-
-			i = 0;
-
-			for ( let iz = 0; iz < res.z; iz ++ ) {
-
-				for ( let iy = 0; iy < res.y; iy ++ ) {
-
-					for ( let ix = 0; ix < res.x; ix ++ ) {
-
-						// Remap to texel centers (must match lightprobes_pars_fragment.glsl.js)
-						uvwArray[ i * 3 ] = ( ix + 0.5 ) / res.x;
-						uvwArray[ i * 3 + 1 ] = ( iy + 0.5 ) / res.y;
-						uvwArray[ i * 3 + 2 ] = ( iz + 0.5 ) / res.z;
+		this.geometry.setAttribute( 'instanceUVW', new InstancedBufferAttribute( uvwArray, 3 ) );
 
-						i ++;
-
-					}
-
-				}
-
-			}
-
-			instanceUVW = new InstancedBufferAttribute( uvwArray, 3 );
-			this.geometry.setAttribute( 'instanceUVW', instanceUVW );
-
-		}
-
-		// Update texture uniforms
-
-		// WITH_GENESYS
-		if ( this._probesSH !== null && this._probesResolution !== null ) {
-
-			this._probesSH.value = probes.texture || _emptyTexture;
-			this._probesResolution.value.copy( probes.resolution );
-
-		} else {
-
-			// !WITH_GENESYS
-
-			this.material.uniforms.probesSH.value = probes.texture;
-			this.material.uniforms.probesResolution.value.copy( probes.resolution );
-
-			// WITH_GENESYS
-
-		}
-		// !WITH_GENESYS
+		this._atlas.value = probes.texture;
+		this._resolution.value.copy( res );
 
 	}
 
@@ -323,18 +156,6 @@ class LightProbeGridHelper extends InstancedMesh {
 	 */
 	dispose() {
 
-		// WITH_GENESYS
-		this.instanceMatrix.dispose();
-
-		const instanceUVW = this.geometry.getAttribute( 'instanceUVW' );
-
-		if ( instanceUVW !== undefined ) {
-
-			instanceUVW.dispose();
-
-		}
-		// !WITH_GENESYS
-
 		this.geometry.dispose();
 		this.material.dispose();