TSL: Pre-pass using global context (#32276)

examples/jsm/tsl/display/GTAONode.js

@@ -568,4 +568,4 @@ function generateMagicSquare( size ) {
  * @param {Camera} camera - The camera the scene is rendered with.
  * @returns {GTAONode}
  */
-export const ao = ( depthNode, normalNode, camera ) => nodeObject( new GTAONode( nodeObject( depthNode ), nodeObject( normalNode ), camera ) );
+export const ao = ( depthNode, normalNode, camera ) => new GTAONode( nodeObject( depthNode ), nodeObject( normalNode ), camera );

examples/webgpu_postprocessing_ao.html

@@ -36,7 +36,7 @@
 		<script type="module">
 
 			import * as THREE from 'three/webgpu';
-			import { pass, mrt, output, normalView, velocity, vec3, vec4, directionToColor, colorSpaceToWorking } from 'three/tsl';
+			import { sample, pass, mrt, context, screenUV, normalView, velocity, vec3, vec4, directionToColor, colorToDirection, colorSpaceToWorking } from 'three/tsl';
 			import { ao } from 'three/addons/tsl/display/GTAONode.js';
 			import { traa } from 'three/addons/tsl/display/TRAANode.js';
 
@@ -49,7 +49,7 @@
 
 			let camera, scene, renderer, postProcessing, controls;
 
-			let aoPass, traaPass, blendPassAO, scenePassColor;
+			let aoPass, traaPass;
 
 			const params = {
 				samples: 16,
@@ -79,15 +79,7 @@
 
 				await renderer.init();
 
-				const environment = new RoomEnvironment();
-				const pmremGenerator = new THREE.PMREMGenerator( renderer );
-
-				scene.background = new THREE.Color( 0x666666 );
-				scene.environment = pmremGenerator.fromScene( environment ).texture;
-				environment.dispose();
-				pmremGenerator.dispose();
-
-				//
+				// controls
 
 				controls = new OrbitControls( camera, renderer.domElement );
 				controls.target.set( 0, 0.5, - 1 );
@@ -97,43 +89,70 @@
 				controls.minDistance = 2;
 				controls.maxDistance = 8;
 
-				//
+				// environment
+
+				const environment = new RoomEnvironment();
+				const pmremGenerator = new THREE.PMREMGenerator( renderer );
+
+				scene.background = new THREE.Color( 0x666666 );
+				scene.environment = pmremGenerator.fromScene( environment ).texture;
+				environment.dispose();
+				pmremGenerator.dispose();
+
+				// post-processing
 
 				postProcessing = new THREE.PostProcessing( renderer );
 
-				const scenePass = pass( scene, camera );
-				scenePass.setMRT( mrt( {
-					output: output,
-					normal: normalView,
+				// pre-pass
+
+				const prePass = pass( scene, camera ).toInspector( 'Normal', ( inspectNode ) => colorSpaceToWorking( inspectNode, THREE.SRGBColorSpace ) );
+				prePass.name = 'Pre-Pass';
+				prePass.transparent = false;
+
+				prePass.setMRT( mrt( {
+					output: directionToColor( normalView ),
 					velocity: velocity
 				} ) );
 
-				scenePassColor = scenePass.getTextureNode( 'output' ).toInspector( 'Color' );
-				const scenePassDepth = scenePass.getTextureNode( 'depth' ).toInspector( 'Depth', () => {
+				const prePassNormal = sample( ( uv ) => {
 
-					return scenePass.getLinearDepthNode();
+					return colorToDirection( prePass.getTextureNode().sample( uv ) );
 
 				} );
-				const scenePassNormal = scenePass.getTextureNode( 'normal' ).toInspector( 'Normal', () => {
 
-					return colorSpaceToWorking( directionToColor( scenePassNormal ), THREE.SRGBColorSpace );
+				const prePassDepth = prePass.getTextureNode( 'depth' ).toInspector( 'Depth', () => prePass.getLinearDepthNode() );
+				const prePassVelocity = prePass.getTextureNode( 'velocity' ).toInspector( 'Velocity' );
 
-				} );
-				const scenePassVelocity = scenePass.getTextureNode( 'velocity' ).toInspector( 'Velocity' );
+				// pre-pass - bandwidth optimization
+
+				const normalTexture = prePass.getTexture( 'output' );
+				normalTexture.type = THREE.UnsignedByteType;
+
+				// scene pass
+
+				const scenePass = pass( scene, camera ).toInspector( 'Color' );
 
 				// ao
 
-				aoPass = ao( scenePassDepth, scenePassNormal, camera );
+				aoPass = ao( prePassDepth, prePassNormal, camera ).toInspector( 'GTAO', ( inspectNode ) => inspectNode.r );
 				aoPass.resolutionScale = 0.5; // running AO in half resolution is often sufficient
 				aoPass.useTemporalFiltering = true;
-				blendPassAO = vec4( scenePassColor.rgb.mul( aoPass.r.toInspector( 'AO' ) ), scenePassColor.a ); // the AO is stored only in the red channel
 
-				// traa
+				const aoPassOutput = aoPass.getTextureNode( 'output' );
+
+				// scene context
+
+				scenePass.contextNode = context( {
+					ao: aoPassOutput.sample( screenUV ).r
+				} );
+
+				// final output + traa
+
+				traaPass = traa( scenePass, prePassDepth, prePassVelocity, camera );
 
-				traaPass = traa( blendPassAO, scenePassDepth, scenePassVelocity, camera );
 				postProcessing.outputNode = traaPass;
 
-				//
+				// models
 
 				const dracoLoader = new DRACOLoader();
 				dracoLoader.setDecoderPath( 'jsm/libs/draco/' );
@@ -148,18 +167,16 @@
 				model.position.set( 0, 1, 0 );
 				scene.add( model );
 
-				model.traverse( ( o ) => {
-
-					// Transparent objects (e.g. loaded via GLTFLoader) might have "depthWrite" set to "false".
-					// This is wanted when rendering the beauty pass however it produces wrong results when computing
-					// AO since depth and normal data are out of sync. Computing normals from depth by not using MRT
-					// can mitigate the issue although the depth information (and thus the normals) are not correct in
-					// first place. Besides, normal estimation is computationally more expensive than just sampling a
-					// normal texture. So depending on your scene, consider to enable "depthWrite" for all transparent objects.
+				//
 
-					if ( o.material ) o.material.depthWrite = true;
+				const transparentMesh = new THREE.Mesh( new THREE.PlaneGeometry( 1.8, 2 ), new THREE.MeshStandardNodeMaterial( { transparent: true, opacity: .1 } ) );
+				transparentMesh.material.transparent = true;
+				transparentMesh.position.z = 0;
+				transparentMesh.position.y = 0.5;
+				transparentMesh.visible = false;
+				scene.add( transparentMesh );
 
-				} );
+				// events
 
 				window.addEventListener( 'resize', onWindowResize );
 
@@ -173,6 +190,7 @@
 				gui.add( params, 'scale', 0.01, 2 ).onChange( updateParameters );
 				gui.add( params, 'thickness', 0.01, 2 ).onChange( updateParameters );
 				gui.add( aoPass, 'useTemporalFiltering' ).name( 'temporal filtering' );
+				gui.add( transparentMesh, 'visible' ).name( 'show transparent mesh' );
 				gui.add( params, 'aoOnly' ).onChange( ( value ) => {
 
 					if ( value === true ) {

src/materials/nodes/NodeMaterial.js

@@ -1060,9 +1060,17 @@ class NodeMaterial extends Material {
 
 		}
 
+		let aoNode = builder.context.ao || null;
+
 		if ( this.aoNode !== null || builder.material.aoMap ) {
 
-			const aoNode = this.aoNode !== null ? this.aoNode : materialAO;
+			const mtlAO = this.aoNode !== null ? this.aoNode : materialAO;
+
+			aoNode = aoNode !== null ? aoNode.mul( mtlAO ) : mtlAO;
+
+		}
+
+		if ( aoNode !== null ) {
 
 			materialLightsNode.push( new AONode( aoNode ) );
 

src/materials/nodes/manager/NodeMaterialObserver.js

@@ -259,7 +259,7 @@ class NodeMaterialObserver {
 
 		}
 
-		if ( builder.context.modelViewMatrix || builder.context.modelNormalViewMatrix )
+		if ( builder.context.modelViewMatrix || builder.context.modelNormalViewMatrix || builder.context.ao )
 			return true;
 
 		return false;

src/nodes/core/NodeFrame.js

@@ -155,9 +155,13 @@ class NodeFrame {
 
 			if ( nodeUpdateBeforeMap.frameId !== this.frameId ) {
 
-				if ( node.updateBefore( this ) !== false ) {
+				const previousFrameId = nodeUpdateBeforeMap.frameId;
 
-					nodeUpdateBeforeMap.frameId = this.frameId;
+				nodeUpdateBeforeMap.frameId = this.frameId;
+
+				if ( node.updateBefore( this ) === false ) {
+
+					nodeUpdateBeforeMap.frameId = previousFrameId;
 
 				}
 
@@ -169,9 +173,13 @@ class NodeFrame {
 
 			if ( nodeUpdateBeforeMap.renderId !== this.renderId ) {
 
-				if ( node.updateBefore( this ) !== false ) {
+				const previousRenderId = nodeUpdateBeforeMap.renderId;
+
+				nodeUpdateBeforeMap.renderId = this.renderId;
+
+				if ( node.updateBefore( this ) === false ) {
 
-					nodeUpdateBeforeMap.renderId = this.renderId;
+					nodeUpdateBeforeMap.renderId = previousRenderId;
 
 				}
 

src/nodes/display/PassNode.js

@@ -248,6 +248,29 @@ class PassNode extends TempNode {
 		 */
 		this.renderTarget = renderTarget;
 
+		/**
+		 * An optional override material for the pass.
+		 *
+		 * @type {Material|null}
+		 */
+		this.overrideMaterial = null;
+
+		/**
+		 * Whether the pass is transparent.
+		 *
+		 * @type {boolean}
+		 * @default false
+		 */
+		this.transparent = true;
+
+		/**
+		 * Whether the pass is opaque.
+		 *
+		 * @type {boolean}
+		 * @default true
+		 */
+		this.opaque = true;
+
 		/**
 		 * An optional global context for the pass.
 		 *
@@ -754,8 +777,11 @@ class PassNode extends TempNode {
 		const currentRenderTarget = renderer.getRenderTarget();
 		const currentMRT = renderer.getMRT();
 		const currentAutoClear = renderer.autoClear;
+		const currentTransparent = renderer.transparent;
+		const currentOpaque = renderer.opaque;
 		const currentMask = camera.layers.mask;
 		const currentContextNode = renderer.contextNode;
+		const currentOverrideMaterial = scene.overrideMaterial;
 
 		this._cameraNear.value = camera.near;
 		this._cameraFar.value = camera.far;
@@ -772,9 +798,17 @@ class PassNode extends TempNode {
 
 		}
 
+		if ( this.overrideMaterial !== null ) {
+
+			scene.overrideMaterial = this.overrideMaterial;
+
+		}
+
 		renderer.setRenderTarget( this.renderTarget );
 		renderer.setMRT( this._mrt );
 		renderer.autoClear = true;
+		renderer.transparent = this.transparent;
+		renderer.opaque = this.opaque;
 
 		if ( this.contextNode !== null ) {
 
@@ -782,7 +816,7 @@ class PassNode extends TempNode {
 
 				this._contextNodeCache = {
 					version: this.version,
-					context: context( { ...renderer.contextNode.value, ...this.contextNode.getFlowContextData() } )
+					context: context( { ...renderer.contextNode.getFlowContextData(), ...this.contextNode.getFlowContextData() } )
 				};
 
 			}
@@ -798,10 +832,13 @@ class PassNode extends TempNode {
 		renderer.render( scene, camera );
 
 		scene.name = currentSceneName;
+		scene.overrideMaterial = currentOverrideMaterial;
 
 		renderer.setRenderTarget( currentRenderTarget );
 		renderer.setMRT( currentMRT );
 		renderer.autoClear = currentAutoClear;
+		renderer.transparent = currentTransparent;
+		renderer.opaque = currentOpaque;
 		renderer.contextNode = currentContextNode;
 
 		camera.layers.mask = currentMask;