Allows the specification of workgroups for compute shaders and the specification of the dispatchSize

Author: Attila Schroeder

src/nodes/gpgpu/ComputeNode.js

@@ -20,9 +20,9 @@ class ComputeNode extends Node {
 	 *
 	 * @param {Node} computeNode - TODO
 	 * @param {number} count - TODO.
-	 * @param {Array<number>} [workgroupSize=[64]] - TODO.
+	 * @param {Array<number>} [workgroupSize = [ 64, 1, 1 ]]
 	 */
-	constructor( computeNode, count, workgroupSize = [ 64 ] ) {
+	constructor( computeNode, count, workgroupSize = [ 64, 1, 1 ] ) {
 
 		super( 'void' );
 
@@ -53,7 +53,7 @@ class ComputeNode extends Node {
 		 * TODO
 		 *
 		 * @type {Array<number>}
-		 * @default [64]
+		 * @default [ 64, 1, 1 ]
 		 */
 		this.workgroupSize = workgroupSize;
 
@@ -220,9 +220,22 @@ export default ComputeNode;
  * @function
  * @param {Node} node - TODO
  * @param {number} count - TODO.
- * @param {Array<number>} [workgroupSize=[64]] - TODO.
+ * @param {Array<number>} [workgroupSize=[ 64, 1, 1 ]]
  * @returns {AtomicFunctionNode}
  */
-export const compute = ( node, count, workgroupSize ) => nodeObject( new ComputeNode( nodeObject( node ), count, workgroupSize ) );
+export const compute = ( node, countOrWorkgroupSize, workgroupSize ) => {
+
+	let count = countOrWorkgroupSize;
+
+	if ( Array.isArray( countOrWorkgroupSize ) ) {
+
+		workgroupSize = countOrWorkgroupSize;
+		count = null;
+
+	}
+
+	return nodeObject( new ComputeNode( nodeObject( node ), count, workgroupSize ) );
+
+};
 
 addMethodChaining( 'compute', compute );

src/renderers/common/Renderer.js

@@ -1243,24 +1243,13 @@ class Renderer {
 
 		frameBufferTarget.depthBuffer = depth;
 		frameBufferTarget.stencilBuffer = stencil;
-		if ( outputRenderTarget !== null ) {
-
-			frameBufferTarget.setSize( outputRenderTarget.width, outputRenderTarget.height, outputRenderTarget.depth );
-
-		} else {
-
-			frameBufferTarget.setSize( width, height, 1 );
-
-		}
-
+		frameBufferTarget.setSize( width, height, outputRenderTarget !== null ? outputRenderTarget.depth : 1 );
 		frameBufferTarget.viewport.copy( this._viewport );
 		frameBufferTarget.scissor.copy( this._scissor );
 		frameBufferTarget.viewport.multiplyScalar( this._pixelRatio );
 		frameBufferTarget.scissor.multiplyScalar( this._pixelRatio );
 		frameBufferTarget.scissorTest = this._scissorTest;
 		frameBufferTarget.multiview = outputRenderTarget !== null ? outputRenderTarget.multiview : false;
-		frameBufferTarget.resolveDepthBuffer = outputRenderTarget !== null ? outputRenderTarget.resolveDepthBuffer : true;
-		frameBufferTarget._autoAllocateDepthBuffer = outputRenderTarget !== null ? outputRenderTarget._autoAllocateDepthBuffer : false;
 
 		return frameBufferTarget;
 
@@ -1516,15 +1505,6 @@ class Renderer {
 
 	}
 
-	_setXRLayerSize( width, height ) {
-
-		this._width = width;
-		this._height = height;
-
-		this.setViewport( 0, 0, width, height );
-
-	}
-
 	/**
 	 * The output pass performs tone mapping and color space conversion.
 	 *
@@ -2310,7 +2290,7 @@ class Renderer {
 	 * @param {Node|Array<Node>} computeNodes - The compute node(s).
 	 * @return {Promise|undefined} A Promise that resolve when the compute has finished. Only returned when the renderer has not been initialized.
 	 */
-	compute( computeNodes ) {
+	compute( computeNodes, dispatchSize = [ 0, 0, 0 ] ) {
 
 		if ( this._isDeviceLost === true ) return;
 
@@ -2389,7 +2369,7 @@ class Renderer {
 			const computeBindings = bindings.getForCompute( computeNode );
 			const computePipeline = pipelines.getForCompute( computeNode, computeBindings );
 
-			backend.compute( computeNodes, computeNode, computeBindings, computePipeline );
+			backend.compute( computeNodes, computeNode, computeBindings, computePipeline, dispatchSize );
 
 		}
 

src/renderers/webgpu/WebGPUBackend.js

@@ -1319,9 +1319,13 @@ class WebGPUBackend extends Backend {
 	 * @param {Array<BindGroup>} bindings - The bindings.
 	 * @param {ComputePipeline} pipeline - The compute pipeline.
 	 */
-	compute( computeGroup, computeNode, bindings, pipeline ) {
+	compute( computeGroup, computeNode, bindings, pipeline, dispatchSize ) {
 
+		const computeNodeData = this.get( computeNode );
 		const { passEncoderGPU } = this.get( computeGroup );
+		const isValid = dispatchSize[ 0 ] > 0 && dispatchSize[ 1 ] > 0 && dispatchSize[ 2 ] > 0;
+
+		dispatchSize = isValid ? dispatchSize : computeNodeData;
 
 		// pipeline
 
@@ -1340,30 +1344,38 @@ class WebGPUBackend extends Backend {
 
 		}
 
-		const maxComputeWorkgroupsPerDimension = this.device.limits.maxComputeWorkgroupsPerDimension;
+		if ( isValid ) {
 
-		const computeNodeData = this.get( computeNode );
+			passEncoderGPU.dispatchWorkgroups(
+				dispatchSize[ 0 ],
+				dispatchSize[ 1 ],
+				dispatchSize[ 2 ]
+			);
 
-		if ( computeNodeData.dispatchSize === undefined ) computeNodeData.dispatchSize = { x: 0, y: 1, z: 1 };
+		} else {
 
-		const { dispatchSize } = computeNodeData;
+			const maxComputeWorkgroupsPerDimension = this.device.limits.maxComputeWorkgroupsPerDimension;
 
-		if ( computeNode.dispatchCount > maxComputeWorkgroupsPerDimension ) {
+			if ( computeNodeData.dispatchSize === undefined ) computeNodeData.dispatchSize = { x: 0, y: 1, z: 1 };
 
-			dispatchSize.x = Math.min( computeNode.dispatchCount, maxComputeWorkgroupsPerDimension );
-			dispatchSize.y = Math.ceil( computeNode.dispatchCount / maxComputeWorkgroupsPerDimension );
+			if ( computeNode.dispatchCount > maxComputeWorkgroupsPerDimension ) {
 
-		} else {
+				dispatchSize.x = Math.min( computeNode.dispatchCount, maxComputeWorkgroupsPerDimension );
+				dispatchSize.y = Math.ceil( computeNode.dispatchCount / maxComputeWorkgroupsPerDimension );
+
+			} else {
 
-			dispatchSize.x = computeNode.dispatchCount;
+				dispatchSize.x = computeNode.dispatchCount;
 
-		}
+			}
 
-		passEncoderGPU.dispatchWorkgroups(
-			dispatchSize.x,
-			dispatchSize.y,
-			dispatchSize.z
-		);
+			passEncoderGPU.dispatchWorkgroups(
+				dispatchSize.x,
+				dispatchSize.y,
+				dispatchSize.z
+			);
+
+		}
 
 	}
 

src/renderers/webgpu/nodes/WGSLNodeBuilder.js

@@ -1898,7 +1898,13 @@ ${ flowData.code }
 
 		} else {
 
-			this.computeShader = this._getWGSLComputeCode( shadersData.compute, ( this.object.workgroupSize || [ 64 ] ).join( ', ' ) );
+			//this.computeShader = this._getWGSLComputeCode( shadersData.compute, ( this.object.workgroupSize || [ 64 ] ).join( ', ' ) );
+
+			const workgroupSize = this.object.workgroupSize || [ 8, 8, 1 ];
+
+			if ( workgroupSize.length !== 3 ) throw new Error( "workgroupSize must have 3 elements" );
+
+			this.computeShader = this._getWGSLComputeCode( shadersData.compute, workgroupSize );
 
 		}
 
@@ -2103,6 +2109,8 @@ fn main( ${shaderData.varyings} ) -> ${shaderData.returnType} {
 	 */
 	_getWGSLComputeCode( shaderData, workgroupSize ) {
 
+		const [ workgroupSizeX, workgroupSizeY, workgroupSizeZ ] = workgroupSize;
+
 		return `${ this.getSignature() }
 // directives
 ${shaderData.directives}
@@ -2122,11 +2130,13 @@ ${shaderData.uniforms}
 // codes
 ${shaderData.codes}
 
-@compute @workgroup_size( ${workgroupSize} )
+@compute @workgroup_size( ${workgroupSizeX}, ${workgroupSizeY}, ${workgroupSizeZ} )
 fn main( ${shaderData.attributes} ) {
 
 	// system
-	instanceIndex = globalId.x + globalId.y * numWorkgroups.x * u32(${workgroupSize}) + globalId.z * numWorkgroups.x * numWorkgroups.y * u32(${workgroupSize});
+	instanceIndex = globalId.x
+    + globalId.y * (${workgroupSizeX} * numWorkgroups.x)
+    + globalId.z * (${workgroupSizeX} * numWorkgroups.x) * (${workgroupSizeY} * numWorkgroups.y);
 
 	// vars
 	${shaderData.vars}