WebGPURenderer: Introduce compileComputeAsync()

examples/webgpu_compute_particles_fluid.html

@@ -135,6 +135,10 @@
 
 				} );
 
+				await renderer.init();
+				await renderer.compileComputeAsync([ clearGridKernel, p2g1Kernel, p2g2Kernel, updateGridKernel, g2pKernel ]);
+				await renderer.compile(scene, camera);
+
 				window.addEventListener( 'resize', onWindowResize );
 				controls.update();
 				renderer.setAnimationLoop( render );

src/renderers/common/Pipelines.js

@@ -73,9 +73,10 @@ class Pipelines extends DataMap {
 	 *
 	 * @param {Node} computeNode - The compute node.
 	 * @param {Array<BindGroup>} bindings - The bindings.
+	 * @param {Promise[]} promises - An array of compilation promises which is only relevant in context of `Renderer.compileAsync()`.
 	 * @return {ComputePipeline} The compute pipeline.
 	 */
-	getForCompute( computeNode, bindings ) {
+	getForCompute( computeNode, bindings, promises = null ) {
 
 		const { backend } = this;
 
@@ -121,7 +122,7 @@ class Pipelines extends DataMap {
 
 				if ( previousPipeline && previousPipeline.usedTimes === 0 ) this._releasePipeline( previousPipeline );
 
-				pipeline = this._getComputePipeline( computeNode, stageCompute, cacheKey, bindings );
+				pipeline = this._getComputePipeline( computeNode, stageCompute, cacheKey, bindings, promises );
 
 			}
 
@@ -311,9 +312,10 @@ class Pipelines extends DataMap {
 	 * @param {ProgrammableStage} stageCompute - The programmable stage representing the compute shader.
 	 * @param {string} cacheKey - The cache key.
 	 * @param {Array<BindGroup>} bindings - The bindings.
+	 * @param {Promise[]} promises - An array of compilation promises which is only relevant in context of `Renderer.compileAsync()`.
 	 * @return {ComputePipeline} The compute pipeline.
 	 */
-	_getComputePipeline( computeNode, stageCompute, cacheKey, bindings ) {
+	_getComputePipeline( computeNode, stageCompute, cacheKey, bindings, promises ) {
 
 		// check for existing pipeline
 
@@ -327,7 +329,7 @@ class Pipelines extends DataMap {
 
 			this.caches.set( cacheKey, pipeline );
 
-			this.backend.createComputePipeline( pipeline, bindings );
+			this.backend.createComputePipeline( pipeline, bindings, promises );
 
 		}
 

src/renderers/common/Renderer.js

@@ -831,9 +831,10 @@ class Renderer {
 	 * @param {Object3D} scene - The scene or 3D object to precompile.
 	 * @param {Camera} camera - The camera that is used to render the scene.
 	 * @param {?Scene} targetScene - If the first argument is a 3D object, this parameter must represent the scene the 3D object is going to be added.
+	 * @param {?Function} onProgress - A callback that is called when the compile has been finished.
 	 * @return {Promise} A Promise that resolves when the compile has been finished.
 	 */
-	async compileAsync( scene, camera, targetScene = null ) {
+	async compileAsync( scene, camera, targetScene = null, onProgress = null ) {
 
 		if ( this._isDeviceLost === true ) return;
 
@@ -861,6 +862,61 @@ class Renderer {
 
 		const compilationPromises = [];
 
+		let compilationTotal = 0;
+		let compilationDone = 0;
+
+		function reportProgress() {
+
+			if ( onProgress !== null ) onProgress( compilationDone / compilationTotal * 100 );
+
+		}
+
+		const originalPush = compilationPromises.push;
+
+		compilationPromises.push = function () {
+
+			for ( let i = 0; i < arguments.length; i ++ ) {
+
+				let promise = arguments[ i ];
+
+				if ( promise && promise.then ) {
+
+					// ok
+
+				} else {
+
+					promise = Promise.resolve( promise );
+
+				}
+
+				compilationTotal ++;
+
+				if ( onProgress !== null ) {
+
+					promise.then( function () {
+
+						compilationDone ++;
+						reportProgress();
+
+					}, function () {
+
+						compilationDone ++;
+						reportProgress();
+
+					} );
+
+				}
+
+				originalPush.call( compilationPromises, promise );
+
+			}
+
+			return compilationPromises.length;
+
+		};
+
+		//
+
 		this._currentRenderContext = renderContext;
 		this._currentRenderObjectFunction = this.renderObject;
 
@@ -965,6 +1021,140 @@ class Renderer {
 
 	}
 
+	/**
+	 * Compile compute programs. This can be useful to avoid a
+	 * phenomenon which is called "shader compilation stutter", which occurs when
+	 * rendering an object with a new shader for the first time.
+	 *
+	 * @async
+	 * @param {Node|Array<Node>} computeNodes - The compute node(s).
+	 * @param {Function} onProgress - A callback that is called when the compile has been finished.
+	 * @return {Promise} A Promise that resolves when the compile has been finished.
+	 */
+	async compileComputeAsync( computeNodes, onProgress = null ) {
+
+		if ( this._isDeviceLost === true ) return;
+
+		if ( this._initialized === false ) await this.init();
+
+		const previousCompilationPromises = this._compilationPromises;
+
+		const compilationPromises = [];
+
+		let compilationTotal = 0;
+		let compilationDone = 0;
+
+		function reportProgress() {
+
+			if ( onProgress === null ) return;
+
+			const ratio = ( compilationTotal === 0 ) ? 0 : ( compilationDone / compilationTotal );
+
+			onProgress( ratio * 100 );
+
+		}
+
+		const originalPush = compilationPromises.push;
+
+		compilationPromises.push = function () {
+
+			for ( let i = 0; i < arguments.length; i ++ ) {
+
+				let promise = arguments[ i ];
+
+				if ( promise && promise.then ) {
+
+					// ok
+
+				} else {
+
+					promise = Promise.resolve( promise );
+
+				}
+
+				compilationTotal ++;
+
+				if ( onProgress !== null ) {
+
+					promise.then( function () {
+
+						compilationDone ++;
+						reportProgress();
+
+					}, function () {
+
+						compilationDone ++;
+						reportProgress();
+
+					} );
+
+				}
+
+				originalPush.call( compilationPromises, promise );
+
+			}
+
+			return compilationPromises.length;
+
+		};
+
+		this._compilationPromises = compilationPromises;
+
+		//
+
+		const pipelines = this._pipelines;
+		const bindings = this._bindings;
+		const nodes = this._nodes;
+
+		const computeList = Array.isArray( computeNodes ) ? computeNodes : [ computeNodes ];
+
+		if ( computeList[ 0 ] === undefined || computeList[ 0 ].isComputeNode !== true ) {
+
+			throw new Error( 'THREE.Renderer: .compileCompute() expects a ComputeNode.' );
+
+		}
+
+		for ( const computeNode of computeList ) {
+
+			if ( pipelines.has( computeNode ) === false ) {
+
+				const dispose = () => {
+
+					computeNode.removeEventListener( 'dispose', dispose );
+
+					pipelines.delete( computeNode );
+					bindings.deleteForCompute( computeNode );
+					nodes.delete( computeNode );
+
+				};
+
+				computeNode.addEventListener( 'dispose', dispose );
+
+				const onInitFn = computeNode.onInitFunction;
+
+				if ( onInitFn !== null ) {
+
+					onInitFn.call( computeNode, { renderer: this } );
+
+				}
+
+			}
+
+			nodes.updateForCompute( computeNode );
+			bindings.updateForCompute( computeNode );
+
+			const computeBindings = bindings.getForCompute( computeNode );
+
+			pipelines.getForCompute( computeNode, computeBindings, compilationPromises );
+
+		}
+
+		await Promise.all( compilationPromises );
+
+		this._compilationPromises = previousCompilationPromises;
+
+	}
+
 	/**
 	 * Renders the scene in an async fashion.
 	 *
@@ -3297,6 +3487,7 @@ class Renderer {
 	 * @param {Object3D} scene - The scene or 3D object to precompile.
 	 * @param {Camera} camera - The camera that is used to render the scene.
 	 * @param {Scene} targetScene - If the first argument is a 3D object, this parameter must represent the scene the 3D object is going to be added.
+	 * @param {?Function} onProgress - A callback that is called when the compile has been finished.
 	 * @return {function(Object3D, Camera, ?Scene): Promise|undefined} A Promise that resolves when the compile has been finished.
 	 */
 	get compile() {

src/renderers/webgl-fallback/WebGLBackend.js

@@ -1593,10 +1593,11 @@ class WebGLBackend extends Backend {
 	 *
 	 * @param {ComputePipeline} computePipeline - The compute pipeline.
 	 * @param {Array<BindGroup>} bindings - The bindings.
+	 * @param {?Array<Promise>} [promises=null] - Optional compilation promises.
 	 */
-	createComputePipeline( computePipeline, bindings ) {
+	createComputePipeline( computePipeline, bindings, promises = null ) {
 
-		const { state, gl } = this;
+		const { gl } = this;
 
 		// Program
 
@@ -1639,19 +1640,6 @@ class WebGLBackend extends Backend {
 
 		gl.linkProgram( programGPU );
 
-		if ( gl.getProgramParameter( programGPU, gl.LINK_STATUS ) === false ) {
-
-			this._logProgramError( programGPU, fragmentShader, vertexShader );
-
-
-		}
-
-		state.useProgram( programGPU );
-
-		// Bindings
-
-		this._setupBindings( bindings, programGPU );
-
 		const attributeNodes = computeProgram.attributes;
 		const attributes = [];
 		const transformBuffers = [];
@@ -1678,14 +1666,73 @@ class WebGLBackend extends Backend {
 
 		}
 
-		//
+		// Store pipeline data
 
 		this.set( computePipeline, {
 			programGPU,
+			fragmentShader,
+			vertexShader,
 			transformBuffers,
 			attributes
 		} );
 
+		if ( promises !== null && this.parallel ) {
+
+			const parallel = this.parallel;
+
+			const p = new Promise( ( resolve ) => {
+
+				const checkStatus = () => {
+
+					if ( gl.getProgramParameter( programGPU, parallel.COMPLETION_STATUS_KHR ) ) {
+
+						this._completeComputeCompile( computePipeline, bindings );
+						resolve();
+
+					} else {
+
+						requestAnimationFrame( checkStatus );
+
+					}
+
+				};
+
+				checkStatus();
+
+			} );
+
+			promises.push( p );
+			return;
+
+		}
+
+		// Sync fallback
+		this._completeComputeCompile( computePipeline, bindings );
+
+	}
+
+	/**
+	 * Completes the compute pipeline setup for the given compute pipeline.
+	 *
+	 * @param {ComputePipeline} pipeline - The compute pipeline.
+	 * @param {Array<BindGroup>} bindings - Array of bind groups.
+	 */
+	_completeComputeCompile( computePipeline, bindings ) {
+
+		const { state, gl } = this;
+		const { programGPU, fragmentShader, vertexShader } = this.get( computePipeline );
+
+		if ( gl.getProgramParameter( programGPU, gl.LINK_STATUS ) === false ) {
+
+			this._logProgramError( programGPU, fragmentShader, vertexShader );
+
+		}
+
+		state.useProgram( programGPU );
+
+		// Bindings (must be after link completion)
+		this._setupBindings( bindings, programGPU );
+
 	}
 
 	/**

src/renderers/webgpu/WebGPUBackend.js

@@ -2060,10 +2060,11 @@ class WebGPUBackend extends Backend {
 	 *
 	 * @param {ComputePipeline} computePipeline - The compute pipeline.
 	 * @param {Array<BindGroup>} bindings - The bindings.
+	 * @param {?Array<Promise>} [promises=null] - Optional compilation promises.
 	 */
-	createComputePipeline( computePipeline, bindings ) {
+	createComputePipeline( computePipeline, bindings, promises = null ) {
 
-		this.pipelineUtils.createComputePipeline( computePipeline, bindings );
+		this.pipelineUtils.createComputePipeline( computePipeline, bindings, promises );
 
 	}