WebGPUPipelines: Improve reuse pipeline and ComputeNode dispose (#26361)

* WebGPUPipelines: Reuse cache

* ComputeNode: Add .needsUpdate

* fix definition used in tests

* ComputeNode: Added .dispose()

examples/jsm/nodes/core/Node.js

@@ -1,3 +1,4 @@
+import { EventDispatcher } from 'three';
 import { NodeUpdateType } from './constants.js';
 import { getNodeChildren, getCacheKey } from './NodeUtils.js';
 import { MathUtils } from 'three';
@@ -6,11 +7,11 @@ const NodeClasses = new Map();
 
 let _nodeId = 0;
 
-class Node {
+class Node extends EventDispatcher {
 
 	constructor( nodeType = null ) {
 
-		this.isNode = true;
+		super();
 
 		this.nodeType = nodeType;
 
@@ -19,6 +20,8 @@ class Node {
 
 		this.uuid = MathUtils.generateUUID();
 
+		this.isNode = true;
+
 		Object.defineProperty( this, 'id', { value: _nodeId ++ } );
 
 	}
@@ -52,6 +55,12 @@ class Node {
 
 	}
 
+	dispose() {
+
+		this.dispatchEvent( { type: 'dispose' } );
+
+	}
+
 	traverse( callback, replaceNode = null ) {
 
 		callback( this, replaceNode );

examples/jsm/nodes/gpgpu/ComputeNode.js

@@ -16,12 +16,25 @@ class ComputeNode extends Node {
 		this.workgroupSize = workgroupSize;
 		this.dispatchCount = 0;
 
+		this.version = 1;
 		this.updateType = NodeUpdateType.OBJECT;
 
 		this.updateDispatchCount();
 
 	}
 
+	dispose() {
+
+		this.dispatchEvent( { type: 'dispose' } );
+
+	}
+
+	set needsUpdate( value ) {
+
+		if ( value === true ) this.version ++;
+
+	}
+
 	updateDispatchCount() {
 
 		const { count, workgroupSize } = this;

examples/jsm/renderers/common/Pipelines.js

@@ -29,11 +29,16 @@ class Pipelines extends DataMap {
 
 		const data = this.get( computeNode );
 
-		if ( data.pipeline === undefined ) {
+		if ( this._needsComputeUpdate( computeNode ) ) {
 
-			// release previous cache
+			const previousPipeline = data.pipeline;
 
-			const previousPipeline = this._releasePipeline( computeNode );
+			if ( previousPipeline ) {
+
+				previousPipeline.usedTimes --;
+				previousPipeline.computeProgram.usedTimes --;
+
+			}
 
 			// get shader
 
@@ -45,7 +50,7 @@ class Pipelines extends DataMap {
 
 			if ( stageCompute === undefined ) {
 
-				if ( previousPipeline ) this._releaseProgram( previousPipeline.computeShader );
+				if ( previousPipeline && previousPipeline.computeProgram.usedTimes === 0 ) this._releaseProgram( previousPipeline.computeProgram );
 
 				stageCompute = new ProgrammableStage( nodeBuilder.computeShader, 'compute' );
 				this.programs.compute.set( nodeBuilder.computeShader, stageCompute );
@@ -56,7 +61,17 @@ class Pipelines extends DataMap {
 
 			// determine compute pipeline
 
-			const pipeline = this._getComputePipeline( stageCompute );
+			const cacheKey = this._getComputeCacheKey( computeNode, stageCompute );
+
+			let pipeline = this.caches.get( cacheKey );
+
+			if ( pipeline === undefined ) {
+
+				if ( previousPipeline && previousPipeline.usedTimes === 0 ) this._releasePipeline( computeNode );
+
+				pipeline = this._getComputePipeline( computeNode, stageCompute, cacheKey );
+
+			}
 
 			// keep track of all used times
 
@@ -65,6 +80,7 @@ class Pipelines extends DataMap {
 
 			//
 
+			data.version = computeNode.version;
 			data.pipeline = pipeline;
 
 		}
@@ -79,11 +95,17 @@ class Pipelines extends DataMap {
 
 		const data = this.get( renderObject );
 
-		if ( this._needsUpdate( renderObject ) ) {
+		if ( this._needsRenderUpdate( renderObject ) ) {
 
-			// release previous cache
+			const previousPipeline = data.pipeline;
 
-			const previousPipeline = this._releasePipeline( renderObject );
+			if ( previousPipeline ) {
+
+				previousPipeline.usedTimes --;
+				previousPipeline.vertexProgram.usedTimes --;
+				previousPipeline.fragmentProgram.usedTimes --;
+
+			}
 
 			// get shader
 
@@ -95,7 +117,7 @@ class Pipelines extends DataMap {
 
 			if ( stageVertex === undefined ) {
 
-				if ( previousPipeline ) this._releaseProgram( previousPipeline.vertexProgram );
+				if ( previousPipeline && previousPipeline.vertexProgram.usedTimes === 0 ) this._releaseProgram( previousPipeline.vertexProgram );
 
 				stageVertex = new ProgrammableStage( nodeBuilder.vertexShader, 'vertex' );
 				this.programs.vertex.set( nodeBuilder.vertexShader, stageVertex );
@@ -108,7 +130,7 @@ class Pipelines extends DataMap {
 
 			if ( stageFragment === undefined ) {
 
-				if ( previousPipeline ) this._releaseProgram( previousPipeline.fragmentShader );
+				if ( previousPipeline && previousPipeline.fragmentProgram.usedTimes === 0 ) this._releaseProgram( previousPipeline.fragmentProgram );
 
 				stageFragment = new ProgrammableStage( nodeBuilder.fragmentShader, 'fragment' );
 				this.programs.fragment.set( nodeBuilder.fragmentShader, stageFragment );
@@ -119,7 +141,21 @@ class Pipelines extends DataMap {
 
 			// determine render pipeline
 
-			const pipeline = this._getRenderPipeline( renderObject, stageVertex, stageFragment );
+			const cacheKey = this._getRenderCacheKey( renderObject, stageVertex, stageFragment );
+
+			let pipeline = this.caches.get( cacheKey );
+
+			if ( pipeline === undefined ) {
+
+				if ( previousPipeline && previousPipeline.usedTimes === 0 ) this._releasePipeline( previousPipeline );
+
+				pipeline = this._getRenderPipeline( renderObject, stageVertex, stageFragment, cacheKey );
+
+			} else {
+
+				renderObject.pipeline = pipeline;
+
+			}
 
 			// keep track of all used times
 
@@ -139,18 +175,31 @@ class Pipelines extends DataMap {
 
 	delete( object ) {
 
-		const pipeline = this._releasePipeline( object );
+		const pipeline = this.get( object ).pipeline;
+
+		if ( pipeline ) {
 
-		if ( pipeline && pipeline.usedTimes === 0 ) {
+			// pipeline
+
+			pipeline.usedTimes --;
+
+			if ( pipeline.usedTimes === 0 ) this._releasePipeline( pipeline );
+
+			// programs
 
 			if ( pipeline.isComputePipeline ) {
 
-				this._releaseProgram( pipeline.computeProgram );
+				pipeline.computeProgram.usedTimes --;
+
+				if ( pipeline.computeProgram.usedTimes === 0 ) this._releaseProgram( pipeline.computeProgram );
 
 			} else {
 
-				this._releaseProgram( pipeline.vertexProgram );
-				this._releaseProgram( pipeline.fragmentProgram );
+				pipeline.fragmentProgram.usedTimes --;
+				pipeline.vertexProgram.usedTimes --;
+
+				if ( pipeline.vertexProgram.usedTimes === 0 ) this._releaseProgram( pipeline.vertexProgram );
+				if ( pipeline.fragmentProgram.usedTimes === 0 ) this._releaseProgram( pipeline.fragmentProgram );
 
 			}
 
@@ -173,11 +222,11 @@ class Pipelines extends DataMap {
 
 	}
 
-	_getComputePipeline( stageCompute ) {
+	_getComputePipeline( computeNode, stageCompute, cacheKey ) {
 
 		// check for existing pipeline
 
-		const cacheKey = 'compute:' + stageCompute.id;
+		cacheKey = cacheKey || this._getComputeCacheKey( computeNode, stageCompute );
 
 		let pipeline = this.caches.get( cacheKey );
 
@@ -195,11 +244,11 @@ class Pipelines extends DataMap {
 
 	}
 
-	_getRenderPipeline( renderObject, stageVertex, stageFragment ) {
+	_getRenderPipeline( renderObject, stageVertex, stageFragment, cacheKey ) {
 
 		// check for existing pipeline
 
-		const cacheKey = this._getRenderCacheKey( renderObject, stageVertex, stageFragment );
+		cacheKey = cacheKey || this._getRenderCacheKey( renderObject, stageVertex, stageFragment );
 
 		let pipeline = this.caches.get( cacheKey );
 
@@ -213,15 +262,15 @@ class Pipelines extends DataMap {
 
 			this.backend.createRenderPipeline( renderObject );
 
-		} else {
+		}
 
-			// assign a shared pipeline to renderObject
+		return pipeline;
 
-			renderObject.pipeline = pipeline;
+	}
 
-		}
+	_getComputeCacheKey( computeNode, stageCompute ) {
 
-		return pipeline;
+		return 'compute' + computeNode.id + stageCompute.id;
 
 	}
 
@@ -247,36 +296,30 @@ class Pipelines extends DataMap {
 
 	}
 
-	_releasePipeline( object ) {
+	_releasePipeline( pipeline ) {
 
-		const pipeline = this.get( object ).pipeline;
-
-		//this.bindings.delete( object );
+		this.caches.delete( pipeline.cacheKey );
 
-		if ( pipeline && -- pipeline.usedTimes === 0 ) {
+	}
 
-			this.caches.delete( pipeline.cacheKey );
+	_releaseProgram( program ) {
 
-		}
+		const code = program.code;
+		const stage = program.stage;
 
-		return pipeline;
+		this.programs[ stage ].delete( code );
 
 	}
 
-	_releaseProgram( program ) {
-
-		if ( -- program.usedTimes === 0 ) {
-
-			const code = program.code;
-			const stage = program.stage;
+	_needsComputeUpdate( computeNode ) {
 
-			this.programs[ stage ].delete( code );
+		const data = this.get( computeNode );
 
-		}
+		return data.pipeline === undefined || data.version !== computeNode.version;
 
 	}
 
-	_needsUpdate( renderObject ) {
+	_needsRenderUpdate( renderObject ) {
 
 		const data = this.get( renderObject );
 		const material = renderObject.material;
@@ -312,7 +355,7 @@ class Pipelines extends DataMap {
 
 		}
 
-		return needsUpdate || data.pipeline !== undefined;
+		return needsUpdate || data.pipeline === undefined;
 
 	}
 

examples/jsm/renderers/common/RenderObjects.js

@@ -4,12 +4,13 @@ import RenderObject from './RenderObject.js';
 
 class RenderObjects {
 
-	constructor( renderer, nodes, geometries, pipelines, info ) {
+	constructor( renderer, nodes, geometries, pipelines, bindings, info ) {
 
 		this.renderer = renderer;
 		this.nodes = nodes;
 		this.geometries = geometries;
 		this.pipelines = pipelines;
+		this.bindings = bindings;
 		this.info = info;
 
 		this.chainMaps = {};
@@ -77,6 +78,7 @@ class RenderObjects {
 			dataMap.delete( renderObject );
 
 			this.pipelines.delete( renderObject );
+			this.bindings.delete( renderObject );
 			this.nodes.delete( renderObject );
 
 			chainMap.delete( renderObject.getChainArray() );

examples/jsm/renderers/common/Renderer.js

@@ -137,7 +137,7 @@ class Renderer {
 			this._textures = new Textures( backend, this._info );
 			this._pipelines = new Pipelines( backend, this._nodes );
 			this._bindings = new Bindings( backend, this._nodes, this._textures, this._attributes, this._pipelines, this._info );
-			this._objects = new RenderObjects( this, this._nodes, this._geometries, this._pipelines, this._info );
+			this._objects = new RenderObjects( this, this._nodes, this._geometries, this._pipelines, this._bindings, this._info );
 			this._renderLists = new RenderLists();
 			this._renderContexts = new RenderContexts();
 
@@ -595,31 +595,47 @@ class Renderer {
 
 		const backend = this.backend;
 		const pipelines = this._pipelines;
-		const computeGroup = Array.isArray( computeNodes ) ? computeNodes : [ computeNodes ];
+		const bindings = this._bindings;
+		const nodes = this._nodes;
+		const computeList = Array.isArray( computeNodes ) ? computeNodes : [ computeNodes ];
 
-		backend.beginCompute( computeGroup );
+		backend.beginCompute( computeNodes );
 
-		for ( const computeNode of computeGroup ) {
+		for ( const computeNode of computeList ) {
 
 			// onInit
 
 			if ( pipelines.has( computeNode ) === false ) {
 
+				const dispose = () => {
+
+					computeNode.removeEventListener( 'dispose', dispose );
+
+					pipelines.delete( computeNode );
+					bindings.delete( computeNode );
+					nodes.delete( computeNode );
+
+				};
+
+				computeNode.addEventListener( 'dispose', dispose );
+
+				//
+
 				computeNode.onInit( { renderer: this } );
 
 			}
 
-			this._nodes.updateForCompute( computeNode );
-			this._bindings.updateForCompute( computeNode );
+			nodes.updateForCompute( computeNode );
+			bindings.updateForCompute( computeNode );
 
 			const computePipeline = pipelines.getForCompute( computeNode );
-			const computeBindings = this._bindings.getForCompute( computeNode );
+			const computeBindings = bindings.getForCompute( computeNode );
 
-			backend.compute( computeGroup, computeNode, computeBindings, computePipeline );
+			backend.compute( computeNodes, computeNode, computeBindings, computePipeline );
 
 		}
 
-		backend.finishCompute( computeGroup );
+		backend.finishCompute( computeNodes );
 
 	}