WebGPURenderer: Add shared `BindGroup` via hash (#32936)

src/nodes/core/NodeBuilder.js

@@ -8,7 +8,7 @@ import ParameterNode from './ParameterNode.js';
 import StructType from './StructType.js';
 import FunctionNode from '../code/FunctionNode.js';
 import NodeMaterial from '../../materials/nodes/NodeMaterial.js';
-import { getDataFromObject, getTypeFromLength } from './NodeUtils.js';
+import { getDataFromObject, getTypeFromLength, hashString } from './NodeUtils.js';
 import { NodeUpdateType, defaultBuildStages, shaderStages } from './constants.js';
 
 import {
@@ -20,7 +20,6 @@ import { stack } from './StackNode.js';
 import { getCurrentStack, setCurrentStack } from '../tsl/TSLBase.js';
 
 import CubeRenderTarget from '../../renderers/common/CubeRenderTarget.js';
-import ChainMap from '../../renderers/common/ChainMap.js';
 
 import BindGroup from '../../renderers/common/BindGroup.js';
 
@@ -35,9 +34,9 @@ import { warn, error } from '../../utils.js';
 
 let _id = 0;
 
-const sharedNodeData = new WeakMap();
+const _bindingGroupsCache = new WeakMap();
 
-const rendererCache = new WeakMap();
+const sharedNodeData = new WeakMap();
 
 const typeFromArray = new Map( [
 	[ Int8Array, 'int' ],
@@ -491,27 +490,6 @@ class NodeBuilder {
 
 	}
 
-	/**
-	 * Returns the bind groups of the current renderer.
-	 *
-	 * @return {ChainMap} The cache.
-	 */
-	getBindGroupsCache() {
-
-		let bindGroupsCache = rendererCache.get( this.renderer );
-
-		if ( bindGroupsCache === undefined ) {
-
-			bindGroupsCache = new ChainMap();
-
-			rendererCache.set( this.renderer, bindGroupsCache );
-
-		}
-
-		return bindGroupsCache;
-
-	}
-
 	/**
 	 * Factory method for creating an instance of {@link RenderTarget} with the given
 	 * dimensions and options.
@@ -572,19 +550,21 @@ class NodeBuilder {
 	 */
 	_getBindGroup( groupName, bindings ) {
 
-		const bindGroupsCache = this.getBindGroupsCache();
+		const groupNode = bindings[ 0 ].groupNode;
 
-		//
+		let sharedGroup = groupNode.shared;
 
-		const bindingsArray = [];
+		if ( sharedGroup ) {
 
-		let sharedGroup = true;
+			for ( let i = 1; i < bindings.length; i ++ ) {
 
-		for ( const binding of bindings ) {
+				if ( groupNode !== bindings[ i ].groupNode ) {
 
-			bindingsArray.push( binding );
+					sharedGroup = false;
 
-			sharedGroup = sharedGroup && binding.groupNode.shared;
+				}
+
+			}
 
 		}
 
@@ -594,19 +574,67 @@ class NodeBuilder {
 
 		if ( sharedGroup ) {
 
-			bindGroup = bindGroupsCache.get( bindingsArray );
+			let cacheKeyString = '';
+
+			for ( const binding of bindings ) {
+
+				if ( binding.isNodeUniformsGroup ) {
+
+					binding.uniforms.sort( ( a, b ) => a.nodeUniform.node.id - b.nodeUniform.node.id );
+
+					for ( const uniform of binding.uniforms ) {
+
+						cacheKeyString += uniform.nodeUniform.node.id;
+
+					}
+
+				} else {
+
+					cacheKeyString += binding.nodeUniform.id;
+
+				}
+
+			}
+
+			// TODO: Remove this hack ._currentRenderContext
+
+			const currentContext = this.renderer._currentRenderContext || this.renderer; // use renderer as fallback until we have a compute context
+
+			let bindingGroupsCache = _bindingGroupsCache.get( currentContext );
+
+			if ( bindingGroupsCache === undefined ) {
+
+				bindingGroupsCache = new Map();
+
+				_bindingGroupsCache.set( currentContext, bindingGroupsCache );
+
+			}
+
+			//
+
+			const cacheKey = hashString( cacheKeyString );
+
+			bindGroup = bindingGroupsCache.get( cacheKey );
 
 			if ( bindGroup === undefined ) {
 
-				bindGroup = new BindGroup( groupName, bindingsArray, this.bindingsIndexes[ groupName ].group );
+				bindGroup = new BindGroup( groupName, bindings, this.bindingsIndexes[ groupName ].group );
 
-				bindGroupsCache.set( bindingsArray, bindGroup );
+				bindingGroupsCache.set( cacheKey, bindGroup );
+
+			} else {
+
+				for ( let i = 0; i < bindings.length; i ++ ) {
+
+					bindGroup.bindings[ i ].visibility |= bindings[ i ].visibility;
+
+				}
 
 			}
 
 		} else {
 
-			bindGroup = new BindGroup( groupName, bindingsArray, this.bindingsIndexes[ groupName ].group );
+			bindGroup = new BindGroup( groupName, bindings, this.bindingsIndexes[ groupName ].group );
 
 		}
 

src/renderers/webgpu/nodes/WGSLNodeBuilder.js

@@ -1146,23 +1146,18 @@ class WGSLNodeBuilder extends NodeBuilder {
 				if ( uniformsGroup === undefined ) {
 
 					uniformsGroup = new NodeUniformsGroup( groupName, group );
-					uniformsGroup.setVisibility( gpuShaderStageLib[ shaderStage ] );
 
 					this.uniformGroups[ groupName ] = uniformsGroup;
 
-					bindings.push( uniformsGroup );
-
-				} else {
-
-					// Update visibility to include this shader stage (bitwise OR)
-					uniformsGroup.setVisibility( uniformsGroup.getVisibility() | gpuShaderStageLib[ shaderStage ] );
+				}
 
-					// Add to bindings for this stage if not already present
-					if ( bindings.indexOf( uniformsGroup ) === - 1 ) {
+				// Update visibility to include this shader stage (bitwise OR)
+				uniformsGroup.setVisibility( uniformsGroup.getVisibility() | gpuShaderStageLib[ shaderStage ] );
 
-						bindings.push( uniformsGroup );
+				// Add to bindings for this stage if not already present
+				if ( bindings.indexOf( uniformsGroup ) === - 1 ) {
 
-					}
+					bindings.push( uniformsGroup );
 
 				}
 

src/renderers/webgpu/utils/WebGPUBindingUtils.js

@@ -6,6 +6,7 @@ import {
 import { FloatType, IntType, UnsignedIntType, Compatibility } from '../../../constants.js';
 import { NodeAccess } from '../../../nodes/core/constants.js';
 import { isTypedArray, error } from '../../../utils.js';
+import { hashString } from '../../../nodes/core/NodeUtils.js';
 
 /**
  * Class representing a WebGPU bind group layout.
@@ -88,36 +89,63 @@ class WebGPUBindingUtils {
 
 		const bindingsData = backend.get( bindGroup );
 
-		// check if the the bind group already has a layout
+		const entries = this._createLayoutEntries( bindGroup );
+		const bindGroupLayoutHash = hashString( JSON.stringify( entries ) );
+
+		let layoutChanged = false;
+
+		// check if the bind group already has a layout and if it's still valid
 
 		if ( bindingsData.layout ) {
 
-			return bindingsData.layout.layoutGPU;
+			// if the layout hash changed (e.g. visibility was updated), invalidate the old layout
 
-		}
+			if ( bindingsData.layoutHash !== bindGroupLayoutHash ) {
 
-		// if not, assing one
+				bindingsData.layout.usedTimes --;
 
-		const entries = this._createLayoutEntries( bindGroup );
-		const bindGroupLayoutKey = JSON.stringify( entries );
+				if ( bindingsData.layout.usedTimes === 0 ) {
 
-		// try to find an existing layout in the cache
+					this._bindGroupLayoutCache.delete( bindingsData.layoutHash );
 
-		let bindGroupLayout = this._bindGroupLayoutCache.get( bindGroupLayoutKey );
+				}
 
-		// if not create a new one
+				bindingsData.layout = undefined;
+				bindingsData.layoutHash = undefined;
+
+				layoutChanged = true;
+
+			} else {
+
+				return bindingsData.layout.layoutGPU;
+
+			}
+
+		}
+
+		// create or reuse a bind group layout from the cache
+
+		let bindGroupLayout = this._bindGroupLayoutCache.get( bindGroupLayoutHash );
 
 		if ( bindGroupLayout === undefined ) {
 
 			bindGroupLayout = new BindGroupLayout( device.createBindGroupLayout( { entries } ) );
-			this._bindGroupLayoutCache.set( bindGroupLayoutKey, bindGroupLayout );
+			this._bindGroupLayoutCache.set( bindGroupLayoutHash, bindGroupLayout );
 
 		}
 
 		bindGroupLayout.usedTimes ++;
 
 		bindingsData.layout = bindGroupLayout;
-		bindingsData.layoutKey = bindGroupLayoutKey;
+		bindingsData.layoutHash = bindGroupLayoutHash;
+
+		// if layout changed, recreate the GPU bind group with the new layout
+
+		if ( layoutChanged ) {
+
+			bindingsData.group = this.createBindGroup( bindGroup, bindGroupLayout.layoutGPU );
+
+		}
 
 		return bindGroupLayout.layoutGPU;
 
@@ -616,12 +644,12 @@ class WebGPUBindingUtils {
 
 			if ( bindingsData.layout.usedTimes === 0 ) {
 
-				this._bindGroupLayoutCache.delete( bindingsData.layoutKey );
+				this._bindGroupLayoutCache.delete( bindingsData.layoutHash );
 
 			}
 
 			bindingsData.layout = undefined;
-			bindingsData.layoutKey = undefined;
+			bindingsData.layoutHash = undefined;
 
 		}
 

src/renderers/webgpu/utils/WebGPUPipelineUtils.js

@@ -105,8 +105,8 @@ class WebGPUPipelineUtils {
 
 		for ( const bindGroup of renderObject.getBindings() ) {
 
-			const bindingsData = backend.get( bindGroup );
-			const { layoutGPU } = bindingsData.layout;
+			// ensure layout is up to date (visibility may have changed due to shared bind groups)
+			const layoutGPU = backend.bindingUtils.createBindingsLayout( bindGroup );
 
 			bindGroupLayouts.push( layoutGPU );
 
@@ -366,8 +366,8 @@ class WebGPUPipelineUtils {
 
 		for ( const bindingsGroup of bindings ) {
 
-			const bindingsData = backend.get( bindingsGroup );
-			const { layoutGPU } = bindingsData.layout;
+			// ensure layout is up to date (visibility may have changed due to shared bind groups)
+			const layoutGPU = backend.bindingUtils.createBindingsLayout( bindingsGroup );
 
 			bindGroupLayouts.push( layoutGPU );