WebGPURenderer: Add WebGPUProgrammableStage. (#21750)

* WebGPURenderer: Add WebGPUProgrammableStage.

* WebGPUProgrammableStage: Clean up.

* WebGPURenderPipelines: Clean up.

examples/jsm/renderers/webgpu/WebGPUProgrammableStage.js

@@ -0,0 +1,24 @@
+let _id = 0;
+
+class WebGPUProgrammableStage {
+
+	constructor( device, glslang, code, type ) {
+
+		this.id = _id ++;
+
+		this.code = code;
+		this.type = type;
+		this.usedTimes = 0;
+
+		const byteCode = glslang.compileGLSL( code, type );
+
+		this.stage = {
+			module: device.createShaderModule( { code: byteCode } ),
+			entryPoint: 'main'
+		};
+
+	}
+
+}
+
+export default WebGPUProgrammableStage;

examples/jsm/renderers/webgpu/WebGPURenderPipeline.js

@@ -11,10 +11,12 @@ import {
 
 class WebGPURenderPipeline {
 
-	constructor( cacheKey, device, renderer, sampleCount ) {
+	constructor( device, renderer, sampleCount ) {
 
-		this.cacheKey = cacheKey;
+		this.cacheKey = null;
 		this.shaderAttributes = null;
+		this.stageVertex = null;
+		this.stageFragment = null;
 		this.usedTimes = 0;
 
 		this._device = device;
@@ -23,7 +25,7 @@ class WebGPURenderPipeline {
 
 	}
 
-	init( moduleVertex, moduleFragment, object, nodeBuilder ) {
+	init( cacheKey, stageVertex, stageFragment, object, nodeBuilder ) {
 
 		const material = object.material;
 		const geometry = object.geometry;
@@ -50,7 +52,10 @@ class WebGPURenderPipeline {
 
 		}
 
+		this.cacheKey = cacheKey;
 		this.shaderAttributes = shaderAttributes;
+		this.stageVertex = stageVertex;
+		this.stageFragment = stageFragment;
 
 		// blending
 
@@ -88,8 +93,8 @@ class WebGPURenderPipeline {
 		const depthStencilFormat = this._renderer.getCurrentDepthStencilFormat();
 
 		this.pipeline = this._device.createRenderPipeline( {
-			vertex: Object.assign( {}, moduleVertex, { buffers: vertexBuffers } ),
-			fragment: Object.assign( {}, moduleFragment, { targets: [ {
+			vertex: Object.assign( {}, stageVertex.stage, { buffers: vertexBuffers } ),
+			fragment: Object.assign( {}, stageFragment.stage, { targets: [ {
 				format: colorFormat,
 				blend: {
 					alpha: alphaBlend,

examples/jsm/renderers/webgpu/WebGPURenderPipelines.js

@@ -1,6 +1,5 @@
 import WebGPURenderPipeline from './WebGPURenderPipeline.js';
-
-let _id = 0;
+import WebGPUProgrammableStage from './WebGPUProgrammableStage.js';
 
 class WebGPURenderPipelines {
 
@@ -16,7 +15,7 @@ class WebGPURenderPipelines {
 		this.pipelines = [];
 		this.objectCache = new WeakMap();
 
-		this.shaderModules = {
+		this.stages = {
 			vertex: new Map(),
 			fragment: new Map()
 		};
@@ -26,6 +25,7 @@ class WebGPURenderPipelines {
 	get( object ) {
 
 		const device = this.device;
+		const glslang = this.glslang;
 		const properties = this.properties;
 
 		const material = object.material;
@@ -41,45 +41,32 @@ class WebGPURenderPipelines {
 
 			const nodeBuilder = this.nodes.get( object );
 
-			// shader modules
-
-			const glslang = this.glslang;
-
-			let moduleVertex = this.shaderModules.vertex.get( nodeBuilder.vertexShader );
-
-			if ( moduleVertex === undefined ) {
+			// programmable stages
 
-				const byteCodeVertex = glslang.compileGLSL( nodeBuilder.vertexShader, 'vertex' );
+			let stageVertex = this.stages.vertex.get( nodeBuilder.vertexShader );
 
-				moduleVertex = {
-					module: device.createShaderModule( { code: byteCodeVertex } ),
-					entryPoint: 'main',
-					id: _id ++
-				};
+			if ( stageVertex === undefined ) {
 
-				this.shaderModules.vertex.set( nodeBuilder.vertexShader, moduleVertex );
+				stageVertex = new WebGPUProgrammableStage( device, glslang, nodeBuilder.vertexShader, 'vertex' );
+				this.stages.vertex.set( nodeBuilder.vertexShader, stageVertex );
 
 			}
 
-			let moduleFragment = this.shaderModules.fragment.get( nodeBuilder.fragmentShader );
+			let stageFragment = this.stages.fragment.get( nodeBuilder.fragmentShader );
 
-			if ( moduleFragment === undefined ) {
+			if ( stageFragment === undefined ) {
 
-				const byteCodeFragment = glslang.compileGLSL( nodeBuilder.fragmentShader, 'fragment' );
-
-				moduleFragment = {
-					module: device.createShaderModule( { code: byteCodeFragment } ),
-					entryPoint: 'main',
-					id: _id ++
-				};
-
-				this.shaderModules.fragment.set( nodeBuilder.fragmentShader, moduleFragment );
+				stageFragment = new WebGPUProgrammableStage( device, glslang, nodeBuilder.fragmentShader, 'fragment' );
+				this.stages.fragment.set( nodeBuilder.fragmentShader, stageFragment );
 
 			}
 
+			stageVertex.usedTimes ++;
+			stageFragment.usedTimes ++;
+
 			// determine render pipeline
 
-			currentPipeline = this._acquirePipeline( moduleVertex, moduleFragment, object, nodeBuilder );
+			currentPipeline = this._acquirePipeline( stageVertex, stageFragment, object, nodeBuilder );
 			objectProperties.currentPipeline = currentPipeline;
 
 			// keep track of all pipelines which are used by a material
@@ -128,14 +115,14 @@ class WebGPURenderPipelines {
 
 	}
 
-	_acquirePipeline( moduleVertex, moduleFragment, object, nodeBuilder ) {
+	_acquirePipeline( stageVertex, stageFragment, object, nodeBuilder ) {
 
 		let pipeline;
 		const pipelines = this.pipelines;
 
 		// check for existing pipeline
 
-		const cacheKey = this._computeCacheKey( moduleVertex, moduleFragment, object );
+		const cacheKey = this._computeCacheKey( stageVertex, stageFragment, object );
 
 		for ( let i = 0, il = pipelines.length; i < il; i ++ ) {
 
@@ -152,8 +139,9 @@ class WebGPURenderPipelines {
 
 		if ( pipeline === undefined ) {
 
-			pipeline = new WebGPURenderPipeline( cacheKey, this.device, this.renderer, this.sampleCount );
-			pipeline.init( moduleVertex, moduleFragment, object, nodeBuilder );
+			pipeline = new WebGPURenderPipeline( this.device, this.renderer, this.sampleCount );
+			pipeline.init( cacheKey, stageVertex, stageFragment, object, nodeBuilder );
+
 			pipelines.push( pipeline );
 
 		}
@@ -162,13 +150,13 @@ class WebGPURenderPipelines {
 
 	}
 
-	_computeCacheKey( moduleVertex, moduleFragment, object ) {
+	_computeCacheKey( stageVertex, stageFragment, object ) {
 
 		const material = object.material;
 		const renderer = this.renderer;
 
 		const parameters = [
-			moduleVertex.id, moduleFragment.id,
+			stageVertex.id, stageFragment.id,
 			material.transparent, material.blending, material.premultipliedAlpha,
 			material.blendSrc, material.blendDst, material.blendEquation,
 			material.blendSrcAlpha, material.blendDstAlpha, material.blendEquationAlpha,
@@ -196,6 +184,22 @@ class WebGPURenderPipelines {
 			pipelines[ i ] = pipelines[ pipelines.length - 1 ];
 			pipelines.pop();
 
+			this._relaseStage( pipeline.stageVertex );
+			this._relaseStage( pipeline.stageFragment );
+
+		}
+
+	}
+
+	_relaseStage( stage ) {
+
+		if ( -- stage.usedTimes === 0 ) {
+
+			const code = stage.code;
+			const type = stage.type;
+
+			this.stages[ type ].delete( code );
+
 		}
 
 	}
@@ -283,20 +287,18 @@ function onMaterialDispose( event ) {
 
 	// remove references to pipelines
 
-	const pipelines = properties.get( material ).pipelines;
+	const pipelines = materialProperties.pipelines;
 
 	if ( pipelines !== undefined ) {
 
-		pipelines.forEach( function ( pipeline ) {
+		for ( const pipeline of pipelines ) {
 
 			this._releasePipeline( pipeline );
 
-		} );
+		}
 
 	}
 
-	// @TODO: still need remove nodes and bindings
-
 }
 
 export default WebGPURenderPipelines;