GTAONode: Clip-space ray marching and loop micro-optimizations (#33903)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Mugen87 <michael.herzog@human-interactive.org>

examples/jsm/tsl/display/GTAONode.js

@@ -1,5 +1,5 @@
 import { DataTexture, RenderTarget, RepeatWrapping, Vector2, Vector3, TempNode, QuadMesh, NodeMaterial, RendererUtils, RedFormat } from 'three/webgpu';
-import { reference, logarithmicDepthToViewZ, viewZToPerspectiveDepth, getNormalFromDepth, getScreenPosition, getViewPosition, nodeObject, Fn, float, NodeUpdateType, uv, uniform, Loop, vec2, vec3, int, dot, max, pow, abs, If, textureSize, sin, cos, PI, texture, passTexture, mat3, add, normalize, cross, mix, acos, clamp, interleavedGradientNoise, screenCoordinate, fract, rand } from 'three/tsl';
+import { reference, logarithmicDepthToViewZ, viewZToPerspectiveDepth, getNormalFromDepth, getViewPosition, getScreenPositionFromClip, nodeObject, Fn, float, NodeUpdateType, uv, uniform, Loop, vec2, vec3, vec4, int, dot, max, min, pow, abs, If, textureSize, sin, cos, PI, texture, passTexture, mat3, add, normalize, cross, mix, acos, clamp, interleavedGradientNoise, screenCoordinate, rand } from 'three/tsl';
 
 const _quadMesh = /*@__PURE__*/ new QuadMesh();
 const _size = /*@__PURE__*/ new Vector2();
@@ -369,7 +369,9 @@ class GTAONode extends TempNode {
 			const viewPosition = getViewPosition( uvNode, depth, this._cameraProjectionMatrixInverse ).toVar();
 			const viewNormal = sampleNormal( uvNode ).toVar();
 
-			const radiusToUse = this.radius;
+			const radius = this.radius;
+			const viewDir = normalize( viewPosition.xyz.negate() ).toVar();
+			const clipPosition = this._cameraProjectionMatrix.mul( vec4( viewPosition, 1.0 ) ).toVar();
 
 			const noiseResolution = textureSize( this._noiseNode, 0 );
 			let noiseUv = vec2( uvNode.x, uvNode.y.oneMinus() );
@@ -390,15 +392,14 @@ class GTAONode extends TempNode {
 
 			// Per-step phase jitter for spatio-temporal decorrelation.
 			const noiseJitterIdx = this._temporalDirection.mul( 0.02 );
-			const stepJitter = fract( interleavedGradientNoise( screenCoordinate.add( this._temporalOffset ) ) ).add( rand( uvNode.add( noiseJitterIdx ).mul( 2 ).sub( 1 ) ) );
+			const stepJitter = interleavedGradientNoise( screenCoordinate.add( this._temporalOffset ) ).add( rand( uvNode.add( noiseJitterIdx ).mul( 2 ).sub( 1 ) ) );
 
 			Loop( { start: int( 0 ), end: DIRECTIONS, type: 'int', condition: '<' }, ( { i } ) => {
 
 				const angle = float( i ).div( float( DIRECTIONS ) ).mul( PI ).add( this._temporalDirection ).toVar();
-				const sampleDir = vec3( cos( angle ), sin( angle ), 0 ).toVar();
-				sampleDir.assign( normalize( kernelMatrix.mul( sampleDir ) ) );
+				const sampleDir = kernelMatrix.mul( vec3( cos( angle ), sin( angle ), 0 ) ).toVar();
+				const clipDirRadius = this._cameraProjectionMatrix.mul( vec4( sampleDir, 0.0 ) ).mul( radius ).toVar();
 
-				const viewDir = normalize( viewPosition.xyz.negate() ).toVar();
 				const sliceBitangent = normalize( cross( sampleDir, viewDir ) ).toVar();
 				const sliceTangent = cross( sliceBitangent, viewDir ).toVar();
 
@@ -416,7 +417,8 @@ class GTAONode extends TempNode {
 				const angleN = signNSin.mul( acos( nCos ) ).toVar();
 
 				const tangentToNormalInSlice = cross( projN, sliceBitangent ).toVar();
-				const cosHorizons = vec2( dot( viewDir, tangentToNormalInSlice ), dot( viewDir, tangentToNormalInSlice.negate() ) ).toVar();
+				const cosHorizon = dot( viewDir, tangentToNormalInSlice ).toVar();
+				const cosHorizons = vec2( cosHorizon, cosHorizon.negate() ).toVar();
 
 				// For each slice, the inner loop performs ray marching to find the horizons.
 
@@ -426,13 +428,13 @@ class GTAONode extends TempNode {
 					// near-field. (Blender's Eevee adaptation)
 					const t = float( j ).add( 1.0 ).add( stepJitter ).div( STEPS ).toVar();
 					const sampleDist = t.mul( t );
-					const sampleViewOffset = sampleDir.mul( radiusToUse ).mul( sampleDist );
+					const clipOffset = clipDirRadius.mul( sampleDist ).toVar();
 
 					// The loop marches in two opposite directions (x and y) along the slice's line to find the horizon on both sides.
 
 					// x
 
-					const sampleScreenPositionX = getScreenPosition( viewPosition.add( sampleViewOffset ), this._cameraProjectionMatrix ).toVar();
+					const sampleScreenPositionX = getScreenPositionFromClip( clipPosition.add( clipOffset ) ).toVar();
 					const sampleDepthX = sampleDepth( sampleScreenPositionX ).toVar();
 					const sampleSceneViewPositionX = getViewPosition( sampleScreenPositionX, sampleDepthX, this._cameraProjectionMatrixInverse ).toVar();
 					const viewDeltaX = sampleSceneViewPositionX.sub( viewPosition ).toVar();
@@ -445,7 +447,7 @@ class GTAONode extends TempNode {
 					// back toward the prior horizon as it approaches the radius boundary.
 					// (squared variant of the paper's near-field attenuation;
 					// Activision GTAO paper, Section 4.3 "Bounding the sampling area")
-					const distFacX = clamp( lenX.div( radiusToUse ), 0, 1 );
+					const distFacX = min( lenX.div( radius ), 1 );
 					const distFacSqX = distFacX.mul( distFacX );
 
 					If( abs( viewDeltaX.z ).lessThan( this.thickness ), () => {
@@ -456,7 +458,7 @@ class GTAONode extends TempNode {
 
 					// y
 
-					const sampleScreenPositionY = getScreenPosition( viewPosition.sub( sampleViewOffset ), this._cameraProjectionMatrix ).toVar();
+					const sampleScreenPositionY = getScreenPositionFromClip( clipPosition.sub( clipOffset ) ).toVar();
 					const sampleDepthY = sampleDepth( sampleScreenPositionY ).toVar();
 					const sampleSceneViewPositionY = getViewPosition( sampleScreenPositionY, sampleDepthY, this._cameraProjectionMatrixInverse ).toVar();
 					const viewDeltaY = sampleSceneViewPositionY.sub( viewPosition ).toVar();
@@ -464,7 +466,7 @@ class GTAONode extends TempNode {
 
 					const sHY = dot( viewDir, viewDeltaY.div( max( lenY, float( 0.0001 ) ) ) );
 
-					const distFacY = clamp( lenY.div( radiusToUse ), 0, 1 );
+					const distFacY = min( lenY.div( radius ), 1 );
 					const distFacSqY = distFacY.mul( distFacY );
 
 					If( abs( viewDeltaY.z ).lessThan( this.thickness ), () => {

src/Three.TSL.js

@@ -209,6 +209,7 @@ export const vogelDiskSample = TSL.vogelDiskSample;
 export const getParallaxCorrectNormal = TSL.getParallaxCorrectNormal;
 export const getRoughness = TSL.getRoughness;
 export const getScreenPosition = TSL.getScreenPosition;
+export const getScreenPositionFromClip = TSL.getScreenPositionFromClip;
 export const getShIrradianceAt = TSL.getShIrradianceAt;
 export const getShadowMaterial = TSL.getShadowMaterial;
 export const getShadowRenderObjectFunction = TSL.getShadowRenderObjectFunction;

src/nodes/utils/PostProcessingUtils.js

@@ -54,6 +54,27 @@ export const getScreenPosition = /*@__PURE__*/ Fn( ( [ viewPosition, projectionM
 
 } );
 
+/**
+ * Converts a clip-space position into a screen position expressed as uv coordinates.
+ *
+ * @tsl
+ * @function
+ * @param {Node<vec4>} clipPosition - The position in clip space.
+ * @return {Node<vec2>} The screen position expressed as uv coordinates.
+ */
+export const getScreenPositionFromClip = /*@__PURE__*/ Fn( ( [ clipPosition ] ) => {
+
+	const screen = clipPosition.xy.div( clipPosition.w ).mul( 0.5 ).add( 0.5 ).toVar();
+	return vec2( screen.x, screen.y.oneMinus() );
+
+} ).setLayout( {
+	name: 'getScreenPositionFromClip',
+	type: 'vec2',
+	inputs: [
+		{ name: 'clipPosition', type: 'vec4' }
+	]
+} );
+
 /**
  * Computes a normal vector based on depth data. Can be used as a fallback when no normal render
  * target is available or if flat surface normals are required.