Addons: Add `TRAAPassNode`. (#29636)

* TRAAPass: Initial setup.

* TRAAPassNode: Implement clamping, fix velocity.

* TRAAPassNode: Finalize initial code.

* TRAAPassNode: Clean up.

* Examples: Clean up.

* TRAAPassNode: Clean up.

* TRAAPassNode: Fix dispose().

* TRAANodePass: Refactor MRT setup.

* TRAAPassNode: Clean up.

* TRAAPassNode: Make MRT handling more robust.

* TRAAPassNode: Implement prepass.

* Revert "TRAAPassNode: Implement prepass."

This reverts commit 172c2eeabc446facc53fb3f7f6dd134ba18e7009.

* TRAAPassNode: Clean up.

* WebGLBackend: Support render target textures in copyTextureToTexture().

* TRAAPassNode: Update comment.

* TRAAPassNode: Remove redundant clear().

examples/files.json

@@ -400,6 +400,7 @@
 		"webgpu_postprocessing_sobel",
 		"webgpu_postprocessing_ssaa",
 		"webgpu_postprocessing_ssr",
+		"webgpu_postprocessing_traa",
 		"webgpu_postprocessing_transition",
 		"webgpu_postprocessing",
 		"webgpu_procedural_texture",

examples/jsm/tsl/display/TRAAPassNode.js

@@ -0,0 +1,347 @@
+import { Color, Vector2, PostProcessingUtils, NearestFilter, Matrix4 } from 'three';
+import { add, float, If, Loop, int, Fn, min, max, clamp, nodeObject, PassNode, QuadMesh, texture, NodeMaterial, uniform, uv, vec2, vec4, luminance } from 'three/tsl';
+
+const _quadMesh = /*@__PURE__*/ new QuadMesh();
+const _size = /*@__PURE__*/ new Vector2();
+
+let _rendererState;
+
+/**
+* Temporal Reprojection Anti-Aliasing (TRAA).
+*
+* References:
+* https://alextardif.com/TAA.html
+* https://www.elopezr.com/temporal-aa-and-the-quest-for-the-holy-trail/
+*
+*/
+class TRAAPassNode extends PassNode {
+
+	static get type() {
+
+		return 'TRAAPassNode';
+
+	}
+
+	constructor( scene, camera ) {
+
+		super( PassNode.COLOR, scene, camera );
+
+		this.isTRAAPassNode = true;
+
+		this.clearColor = new Color( 0x000000 );
+		this.clearAlpha = 0;
+
+		this._jitterIndex = 0;
+		this._originalProjectionMatrix = new Matrix4();
+
+		// uniforms
+
+		this._invSize = uniform( new Vector2() );
+
+		// render targets
+
+		this._sampleRenderTarget = null;
+		this._historyRenderTarget = null;
+
+		// materials
+
+		this._resolveMaterial = new NodeMaterial();
+		this._resolveMaterial.name = 'TRAA.Resolve';
+
+	}
+
+	setSize( width, height ) {
+
+		super.setSize( width, height );
+
+		let needsRestart = false;
+
+		if ( this.renderTarget.width !== this._sampleRenderTarget.width || this.renderTarget.height !== this._sampleRenderTarget.height ) {
+
+			this._sampleRenderTarget.setSize( this.renderTarget.width, this.renderTarget.height );
+			this._historyRenderTarget.setSize( this.renderTarget.width, this.renderTarget.height );
+
+			this._invSize.value.set( 1 / this.renderTarget.width, 1 / this.renderTarget.height );
+
+			needsRestart = true;
+
+		}
+
+		return needsRestart;
+
+	}
+
+	updateBefore( frame ) {
+
+		const { renderer } = frame;
+		const { scene, camera } = this;
+
+		_rendererState = PostProcessingUtils.resetRendererAndSceneState( renderer, scene, _rendererState );
+
+		//
+
+		this._pixelRatio = renderer.getPixelRatio();
+		const size = renderer.getSize( _size );
+
+		const needsRestart = this.setSize( size.width, size.height, renderer );
+
+		//
+
+		this._cameraNear.value = camera.near;
+		this._cameraFar.value = camera.far;
+
+		const viewOffset = {
+
+			fullWidth: this.renderTarget.width,
+			fullHeight: this.renderTarget.height,
+			offsetX: 0,
+			offsetY: 0,
+			width: this.renderTarget.width,
+			height: this.renderTarget.height
+
+		};
+
+		const originalViewOffset = Object.assign( {}, camera.view );
+
+		if ( originalViewOffset.enabled ) Object.assign( viewOffset, originalViewOffset );
+
+		const jitterOffset = _JitterVectors[ this._jitterIndex ];
+
+		camera.updateProjectionMatrix();
+		this._originalProjectionMatrix.copy( camera.projectionMatrix );
+
+		camera.setViewOffset(
+
+			viewOffset.fullWidth, viewOffset.fullHeight,
+
+			viewOffset.offsetX + jitterOffset[ 0 ] * 0.0625, viewOffset.offsetY + jitterOffset[ 1 ] * 0.0625, // 0.0625 = 1 / 16
+
+			viewOffset.width, viewOffset.height
+
+		);
+
+		const mrt = this.getMRT();
+		const velocityOutput = mrt.get( 'velocity' );
+
+		if ( velocityOutput !== undefined ) {
+
+			velocityOutput.setProjectionMatrix( this._originalProjectionMatrix );
+
+		} else {
+
+			throw new Error( 'THREE:TRAAPassNode: Missing velocity output in MRT configuration.' );
+
+		}
+
+		renderer.setMRT( mrt );
+
+		renderer.setClearColor( this.clearColor, this.clearAlpha );
+		renderer.setRenderTarget( this._sampleRenderTarget );
+		renderer.render( scene, camera );
+
+		renderer.setRenderTarget( null );
+		renderer.setMRT( null );
+
+		// every time when the dimensions change we need fresh history data. Copy the sample
+		// into the history and final render target (no AA happens at that point).
+
+		if ( needsRestart === true ) {
+
+			// bind and clear render target to make sure they are initialized after the resize which triggers a dispose()
+
+			renderer.setRenderTarget( this._historyRenderTarget );
+			renderer.clear();
+
+			renderer.setRenderTarget( this.renderTarget );
+			renderer.clear();
+
+			renderer.setRenderTarget( null );
+
+			renderer.copyTextureToTexture( this._sampleRenderTarget.texture, this._historyRenderTarget.texture );
+			renderer.copyTextureToTexture( this._sampleRenderTarget.texture, this.renderTarget.texture );
+
+		} else {
+
+			// resolve
+
+			renderer.setRenderTarget( this.renderTarget );
+			_quadMesh.material = this._resolveMaterial;
+			_quadMesh.render( renderer );
+			renderer.setRenderTarget( null );
+
+			// update history
+
+			renderer.copyTextureToTexture( this.renderTarget.texture, this._historyRenderTarget.texture );
+
+		}
+
+		// copy depth
+
+		renderer.copyTextureToTexture( this._sampleRenderTarget.depthTexture, this.renderTarget.depthTexture );
+
+		// update jitter index
+
+		this._jitterIndex ++;
+		this._jitterIndex = this._jitterIndex % ( _JitterVectors.length - 1 );
+
+		// restore
+
+		if ( originalViewOffset.enabled ) {
+
+			camera.setViewOffset(
+
+				originalViewOffset.fullWidth, originalViewOffset.fullHeight,
+
+				originalViewOffset.offsetX, originalViewOffset.offsetY,
+
+				originalViewOffset.width, originalViewOffset.height
+
+			);
+
+		} else {
+
+			camera.clearViewOffset();
+
+		}
+
+		velocityOutput.setProjectionMatrix( null );
+
+		PostProcessingUtils.restoreRendererAndSceneState( renderer, scene, _rendererState );
+
+	}
+
+	setup( builder ) {
+
+		if ( this._sampleRenderTarget === null ) {
+
+			this._sampleRenderTarget = this.renderTarget.clone();
+			this._historyRenderTarget = this.renderTarget.clone();
+
+			this._sampleRenderTarget.texture.minFiler = NearestFilter;
+			this._sampleRenderTarget.texture.magFilter = NearestFilter;
+
+			const velocityTarget = this._sampleRenderTarget.texture.clone();
+			velocityTarget.isRenderTargetTexture = true;
+			velocityTarget.name = 'velocity';
+
+			this._sampleRenderTarget.textures.push( velocityTarget );
+
+		}
+
+		// textures
+
+		const historyTexture = texture( this._historyRenderTarget.texture );
+		const sampleTexture = texture( this._sampleRenderTarget.textures[ 0 ] );
+		const velocityTexture = texture( this._sampleRenderTarget.textures[ 1 ] );
+		const depthTexture = texture( this._sampleRenderTarget.depthTexture );
+
+		const resolve = Fn( () => {
+
+			const uvNode = uv();
+
+			const minColor = vec4( 10000 ).toVar();
+			const maxColor = vec4( - 10000 ).toVar();
+			const closestDepth = float( 1 ).toVar();
+			const closestDepthPixelPosition = vec2( 0 ).toVar();
+
+			// sample a 3x3 neighborhood to create a box in color space
+			// clamping the history color with the resulting min/max colors mitigates ghosting
+
+			Loop( { start: int( - 1 ), end: int( 1 ), type: 'int', condition: '<=', name: 'x' }, ( { x } ) => {
+
+				Loop( { start: int( - 1 ), end: int( 1 ), type: 'int', condition: '<=', name: 'y' }, ( { y } ) => {
+
+					const uvNeighbor = uvNode.add( vec2( float( x ), float( y ) ).mul( this._invSize ) ).toVar();
+					const colorNeighbor = max( vec4( 0 ), sampleTexture.uv( uvNeighbor ) ).toVar(); // use max() to avoid propagate garbage values
+
+					minColor.assign( min( minColor, colorNeighbor ) );
+					maxColor.assign( max( maxColor, colorNeighbor ) );
+
+					const currentDepth = depthTexture.uv( uvNeighbor ).r.toVar();
+
+					// find the sample position of the closest depth in the neighborhood (used for velocity)
+
+					If( currentDepth.lessThan( closestDepth ), () => {
+
+						closestDepth.assign( currentDepth );
+						closestDepthPixelPosition.assign( uvNeighbor );
+
+					} );
+
+				} );
+
+			} );
+
+			// sampling/reprojection
+
+			const offset = velocityTexture.uv( closestDepthPixelPosition ).xy.mul( vec2( 0.5, - 0.5 ) ); // NDC to uv offset
+
+			const currentColor = sampleTexture.uv( uvNode );
+			const historyColor = historyTexture.uv( uvNode.sub( offset ) );
+
+			// clamping
+
+			const clampedHistoryColor = clamp( historyColor, minColor, maxColor );
+
+			// flicker reduction based on luminance weighing
+
+			const currentWeight = float( 0.05 ).toVar();
+			const historyWeight = currentWeight.oneMinus().toVar();
+
+			const compressedCurrent = currentColor.mul( float( 1 ).div( ( max( max( currentColor.r, currentColor.g ), currentColor.b ).add( 1.0 ) ) ) );
+			const compressedHistory = clampedHistoryColor.mul( float( 1 ).div( ( max( max( clampedHistoryColor.r, clampedHistoryColor.g ), clampedHistoryColor.b ).add( 1.0 ) ) ) );
+
+			const luminanceCurrent = luminance( compressedCurrent.rgb );
+			const luminanceHistory = luminance( compressedHistory.rgb );
+
+			currentWeight.mulAssign( float( 1.0 ).div( luminanceCurrent.add( 1 ) ) );
+			historyWeight.mulAssign( float( 1.0 ).div( luminanceHistory.add( 1 ) ) );
+
+			return add( currentColor.mul( currentWeight ), clampedHistoryColor.mul( historyWeight ) ).div( max( currentWeight.add( historyWeight ), 0.00001 ) );
+
+		} );
+
+		// materials
+
+		this._resolveMaterial.fragmentNode = resolve();
+
+		return super.setup( builder );
+
+	}
+
+	dispose() {
+
+		super.dispose();
+
+		if ( this._sampleRenderTarget !== null ) {
+
+			this._sampleRenderTarget.dispose();
+			this._historyRenderTarget.dispose();
+
+		}
+
+		this._resolveMaterial.dispose();
+
+	}
+
+}
+
+export default TRAAPassNode;
+
+// These jitter vectors are specified in integers because it is easier.
+// I am assuming a [-8,8) integer grid, but it needs to be mapped onto [-0.5,0.5)
+// before being used, thus these integers need to be scaled by 1/16.
+//
+// Sample patterns reference: https://msdn.microsoft.com/en-us/library/windows/desktop/ff476218%28v=vs.85%29.aspx?f=255&MSPPError=-2147217396
+const _JitterVectors = [
+	[ - 4, - 7 ], [ - 7, - 5 ], [ - 3, - 5 ], [ - 5, - 4 ],
+	[ - 1, - 4 ], [ - 2, - 2 ], [ - 6, - 1 ], [ - 4, 0 ],
+	[ - 7, 1 ], [ - 1, 2 ], [ - 6, 3 ], [ - 3, 3 ],
+	[ - 7, 6 ], [ - 3, 6 ], [ - 5, 7 ], [ - 1, 7 ],
+	[ 5, - 7 ], [ 1, - 6 ], [ 6, - 5 ], [ 4, - 4 ],
+	[ 2, - 3 ], [ 7, - 2 ], [ 1, - 1 ], [ 4, - 1 ],
+	[ 2, 1 ], [ 6, 2 ], [ 0, 4 ], [ 4, 4 ],
+	[ 2, 5 ], [ 7, 5 ], [ 5, 6 ], [ 3, 7 ]
+];
+
+export const traaPass = ( scene, camera ) => nodeObject( new TRAAPassNode( scene, camera ) );

examples/webgpu_postprocessing_traa.html

@@ -0,0 +1,127 @@
+<!DOCTYPE html>
+<html lang="en">
+	<head>
+		<title>three.js webgpu - postprocessing traa</title>
+		<meta charset="utf-8">
+		<meta name="viewport" content="width=device-width, user-scalable=no, minimum-scale=1.0, maximum-scale=1.0">
+		<link type="text/css" rel="stylesheet" href="main.css">
+	</head>
+	<body>
+		<div id="info">
+			<a href="https://threejs.org" target="_blank" rel="noopener">three.js</a> - Temporal Reprojection Anti-Aliasing (TRAA) 
+		</div>
+
+		<script type="importmap">
+			{
+				"imports": {
+					"three": "../build/three.webgpu.js",
+					"three/tsl": "../build/three.webgpu.js",
+					"three/addons/": "./jsm/"
+				}
+			}
+		</script>
+
+		<script type="module">
+
+			import * as THREE from 'three';
+			import { mrt, output, velocity } from 'three/tsl';
+			import { traaPass } from 'three/addons/tsl/display/TRAAPassNode.js';
+
+			import Stats from 'three/addons/libs/stats.module.js';
+
+			let camera, scene, renderer, postProcessing;
+			let stats;
+			let index = 0;
+
+			init();
+
+			function init() {
+
+				renderer = new THREE.WebGPURenderer( { forceWebGL: true } );
+				renderer.setPixelRatio( window.devicePixelRatio );
+				renderer.setSize( window.innerWidth, window.innerHeight );
+				renderer.setAnimationLoop( animate );
+				document.body.appendChild( renderer.domElement );
+
+				stats = new Stats();
+				document.body.appendChild( stats.dom );
+
+				camera = new THREE.PerspectiveCamera( 70, window.innerWidth / window.innerHeight, 0.1, 10 );
+				camera.position.z = 2.5;
+
+				scene = new THREE.Scene();
+
+				const geometry = new THREE.BoxGeometry();
+				const material1 = new THREE.MeshBasicMaterial( { color: 0xffffff, wireframe: true } );
+
+				const mesh1 = new THREE.Mesh( geometry, material1 );
+				mesh1.position.x = - 1;
+				scene.add( mesh1 );
+
+				const texture = new THREE.TextureLoader().load( 'textures/brick_diffuse.jpg' );
+				texture.minFilter = THREE.NearestFilter;
+				texture.magFilter = THREE.NearestFilter;
+				texture.generateMipmaps = false;
+				texture.colorSpace = THREE.SRGBColorSpace;
+
+				const material2 = new THREE.MeshBasicMaterial( { map: texture } );
+
+				const mesh2 = new THREE.Mesh( geometry, material2 );
+				mesh2.position.x = 1;
+				scene.add( mesh2 );
+
+				// postprocessing
+
+				postProcessing = new THREE.PostProcessing( renderer );
+				const scenePass = traaPass( scene, camera );
+				scenePass.setMRT( mrt( {
+					output: output,
+					velocity: velocity
+				} ) );
+
+				postProcessing.outputNode = scenePass;
+
+				//
+
+				window.addEventListener( 'resize', onWindowResize );
+
+			}
+
+			function onWindowResize() {
+
+				const width = window.innerWidth;
+				const height = window.innerHeight;
+
+				camera.aspect = width / height;
+				camera.updateProjectionMatrix();
+
+				renderer.setSize( width, height );
+
+			}
+
+			function animate() {
+
+				index ++;
+
+				if ( Math.round( index / 200 ) % 2 === 0 ) {
+
+					for ( let i = 0; i < scene.children.length; i ++ ) {
+
+						const child = scene.children[ i ];
+
+						child.rotation.x += 0.005;
+						child.rotation.y += 0.01;
+
+					}
+
+				}
+
+				postProcessing.render();
+
+				stats.update();
+
+			}
+
+		</script>
+	</body>
+</html>

src/nodes/accessors/VelocityNode.js

@@ -23,6 +23,8 @@ class VelocityNode extends TempNode {
 
 		super( 'vec2' );
 
+		this.projectionMatrix = null;
+
 		this.updateType = NodeUpdateType.OBJECT;
 		this.updateAfterType = NodeUpdateType.OBJECT;
 
@@ -32,6 +34,12 @@ class VelocityNode extends TempNode {
 
 	}
 
+	setProjectionMatrix( projectionMatrix ) {
+
+		this.projectionMatrix = projectionMatrix;
+
+	}
+
 	update( { frameId, camera, object } ) {
 
 		const previousModelMatrix = getPreviousMatrix( object );
@@ -54,7 +62,7 @@ class VelocityNode extends TempNode {
 				cameraData.currentProjectionMatrix = new Matrix4();
 				cameraData.currentCameraViewMatrix = new Matrix4();
 
-				cameraData.previousProjectionMatrix.copy( camera.projectionMatrix );
+				cameraData.previousProjectionMatrix.copy( this.projectionMatrix || camera.projectionMatrix );
 				cameraData.previousCameraViewMatrix.copy( camera.matrixWorldInverse );
 
 			} else {
@@ -64,7 +72,7 @@ class VelocityNode extends TempNode {
 
 			}
 
-			cameraData.currentProjectionMatrix.copy( camera.projectionMatrix );
+			cameraData.currentProjectionMatrix.copy( this.projectionMatrix || camera.projectionMatrix );
 			cameraData.currentCameraViewMatrix.copy( camera.matrixWorldInverse );
 
 			this.previousProjectionMatrix.value.copy( cameraData.previousProjectionMatrix );
@@ -82,9 +90,11 @@ class VelocityNode extends TempNode {
 
 	setup( /*builder*/ ) {
 
+		const projectionMatrix = ( this.projectionMatrix === null ) ? cameraProjectionMatrix : uniform( this.projectionMatrix );
+
 		const previousModelViewMatrix = this.previousCameraViewMatrix.mul( this.previousModelWorldMatrix );
 
-		const clipPositionCurrent = cameraProjectionMatrix.mul( modelViewMatrix ).mul( positionLocal );
+		const clipPositionCurrent = projectionMatrix.mul( modelViewMatrix ).mul( positionLocal );
 		const clipPositionPrevious = this.previousProjectionMatrix.mul( previousModelViewMatrix ).mul( positionPrevious );
 
 		const ndcPositionCurrent = clipPositionCurrent.xy.div( clipPositionCurrent.w );

src/renderers/webgl-fallback/WebGLBackend.js

@@ -1281,9 +1281,9 @@ class WebGLBackend extends Backend {
 
 	}
 
-	copyTextureToTexture( position, srcTexture, dstTexture, level ) {
+	copyTextureToTexture( srcTexture, dstTexture, srcRegion, dstPosition, level ) {
 
-		this.textureUtils.copyTextureToTexture( position, srcTexture, dstTexture, level );
+		this.textureUtils.copyTextureToTexture( srcTexture, dstTexture, srcRegion, dstPosition, level );
 
 	}
 
@@ -1355,6 +1355,7 @@ class WebGLBackend extends Backend {
 						const texture = textures[ i ];
 						const textureData = this.get( texture );
 						textureData.renderTarget = descriptor.renderTarget;
+						textureData.cacheKey = cacheKey; // required for copyTextureToTexture()
 
 						const attachment = gl.COLOR_ATTACHMENT0 + i;
 
@@ -1370,6 +1371,8 @@ class WebGLBackend extends Backend {
 
 					const textureData = this.get( descriptor.depthTexture );
 					const depthStyle = stencilBuffer ? gl.DEPTH_STENCIL_ATTACHMENT : gl.DEPTH_ATTACHMENT;
+					textureData.renderTarget = descriptor.renderTarget;
+					textureData.cacheKey = cacheKey; // required for copyTextureToTexture()
 
 					gl.framebufferTexture2D( gl.FRAMEBUFFER, depthStyle, gl.TEXTURE_2D, textureData.textureGPU, 0 );
 

src/renderers/webgl-fallback/utils/WebGLTextureUtils.js

@@ -659,20 +659,46 @@ class WebGLTextureUtils {
 		gl.pixelStorei( gl.UNPACK_SKIP_PIXELS, minX );
 		gl.pixelStorei( gl.UNPACK_SKIP_ROWS, minY );
 
+		if ( srcTexture.isRenderTargetTexture || srcTexture.isDepthTexture ) {
 
-		if ( srcTexture.isDataTexture ) {
+			const srcTextureData = backend.get( srcTexture );
+			const dstTextureData = backend.get( dstTexture );
 
-			gl.texSubImage2D( gl.TEXTURE_2D, level, dstX, dstY, width, height, glFormat, glType, image.data );
+			const srcRenderContextData = backend.get( srcTextureData.renderTarget );
+			const dstRenderContextData = backend.get( dstTextureData.renderTarget );
+
+			const srcFramebuffer = srcRenderContextData.framebuffers[ srcTextureData.cacheKey ];
+			const dstFramebuffer = dstRenderContextData.framebuffers[ dstTextureData.cacheKey ];
+
+			state.bindFramebuffer( gl.READ_FRAMEBUFFER, srcFramebuffer );
+			state.bindFramebuffer( gl.DRAW_FRAMEBUFFER, dstFramebuffer );
+
+			let mask = gl.COLOR_BUFFER_BIT;
+
+			if ( srcTexture.isDepthTexture ) mask = gl.DEPTH_BUFFER_BIT;
+
+			gl.blitFramebuffer( minX, minY, width, height, dstX, dstY, width, height, mask, gl.NEAREST );
+
+			state.bindFramebuffer( gl.READ_FRAMEBUFFER, null );
+			state.bindFramebuffer( gl.DRAW_FRAMEBUFFER, null );
 
 		} else {
 
-			if ( srcTexture.isCompressedTexture ) {
+			if ( srcTexture.isDataTexture ) {
 
-				gl.compressedTexSubImage2D( gl.TEXTURE_2D, level, dstX, dstY, image.width, image.height, glFormat, image.data );
+				gl.texSubImage2D( gl.TEXTURE_2D, level, dstX, dstY, width, height, glFormat, glType, image.data );
 
 			} else {
 
-				gl.texSubImage2D( gl.TEXTURE_2D, level, dstX, dstY, width, height, glFormat, glType, image );
+				if ( srcTexture.isCompressedTexture ) {
+
+					gl.compressedTexSubImage2D( gl.TEXTURE_2D, level, dstX, dstY, image.width, image.height, glFormat, image.data );
+
+				} else {
+
+					gl.texSubImage2D( gl.TEXTURE_2D, level, dstX, dstY, width, height, glFormat, glType, image );
+
+				}
 
 			}
 

test/e2e/puppeteer.js

@@ -122,6 +122,7 @@ const exceptionList = [
 	'webgpu_video_panorama',
 	'webgpu_postprocessing_bloom_emissive',
 	'webgpu_lights_tiled',
+	'webgpu_postprocessing_traa',
 
 	// Awaiting for WebGPU Backend support in Puppeteer
 	'webgpu_storage_buffer',