three-gpu-pathtracer 0.0.26 → 0.0.28

package/src/webgpu/EquirectHdrInfoNode.js

@@ -1,20 +1,9 @@
-import { Matrix3 } from 'three/webgpu';
-import { texture, sampler, uniform, wgslFn } from 'three/tsl';
+import { Matrix3, StorageBufferAttribute } from 'three/webgpu';
+import { texture, sampler, uniform, storage, PI } from 'three/tsl';
 import { EquirectHdrInfoUniform } from '../uniforms/EquirectHdrInfoUniform.js';
 import { wgslTagFn } from 'three-mesh-bvh/webgpu';
-import { constants, environmentSampleStruct } from './nodes/structs.wgsl.js';
-import { equirectDirectionToUvFn, equirectUvToDirectionFn, luminanceFn } from './nodes/sampling.wgsl.js';
-
-// totalWeight includes sin( theta ), which cancels in the solid-angle conversion.
-const equirectLuminancePdfFn = wgslFn( /* wgsl */ `
-
-	fn equirectLuminancePdf( luminance: f32, totalWeight: f32, resolution: vec2u ) -> f32 {
-
-		return f32( resolution.x * resolution.y ) * luminance / ( 2.0 * PI * PI * totalWeight );
-
-	}
-
-`, [ constants ] );
+import { environmentSampleStruct } from './nodes/structs.wgsl.js';
+import { equirectDirectionToUvFn, equirectUvToDirectionFn } from './nodes/sampling.wgsl.js';
 
 export class EquirectHdrInfoNode extends EquirectHdrInfoUniform {
 
@@ -22,13 +11,10 @@ export class EquirectHdrInfoNode extends EquirectHdrInfoUniform {
 
 		super();
 
-		// environment map + importance-sampling CDF textures, each with a sampler
+		// environment map with a sampler, and the importance sampling CDFs for the default 1x1 map
 		this.mapNode = texture( this.map );
 		this.mapSampler = sampler( this.mapNode );
-		this.marginalNode = texture( this.marginalWeights );
-		this.marginalSampler = sampler( this.marginalNode );
-		this.conditionalNode = texture( this.conditionalWeights );
-		this.conditionalSampler = sampler( this.conditionalNode );
+		this.cdfNode = storage( new StorageBufferAttribute( new Float32Array( [ 0, 1, 0, 1 ] ), 1 ), 'float' ).toReadOnly();
 
 		// scalar parameters that assemble into the EnvironmentInfo struct in the shader
 		this.rotationNode = uniform( new Matrix3() );
@@ -53,17 +39,35 @@ export class EquirectHdrInfoNode extends EquirectHdrInfoUniform {
 
 		const {
 			mapNode,
-			marginalNode,
-			conditionalNode,
+			cdfNode,
 			totalSumNode,
 		} = this;
 
 		// refresh values in place on the existing nodes so no rebuild is required
 		mapNode.value = this.map;
-		marginalNode.value = this.marginalWeights;
-		conditionalNode.value = this.conditionalWeights;
 		totalSumNode.value = this.totalSum;
 
+		// pack the row CDF and then each row's pixel CDF into one buffer, each led by a zero so a
+		// pixel's range is two neighboring entries
+		const { width, height } = this.map.image;
+		const { cdfMarginal, cdfConditional } = this;
+		const cdf = new Float32Array( height + 1 + height * ( width + 1 ) );
+		cdf.set( cdfMarginal, 1 );
+		for ( let y = 0; y < height; y ++ ) {
+
+			cdf.set( cdfConditional.subarray( y * width, y * width + width ), height + 1 + y * ( width + 1 ) + 1 );
+
+		}
+
+		cdfNode.value = new StorageBufferAttribute( cdf, 1 );
+
+	}
+
+	dispose() {
+
+		super.dispose();
+		this.cdfNode.value.dispose();
+
 	}
 
 	_initFns() {
@@ -71,15 +75,60 @@ export class EquirectHdrInfoNode extends EquirectHdrInfoUniform {
 		const {
 			mapNode,
 			mapSampler,
-			marginalNode,
-			marginalSampler,
-			conditionalNode,
-			conditionalSampler,
+			cdfNode,
 			totalSumNode,
 			rotationNode,
 			intensityNode,
 		} = this;
 
+		// Binary searches the "count" pixel ranges from "offset" for the one holding "r", writing where
+		// "r" falls within it to "fraction". See:
+		// https://pbr-book.org/3ed-2018/Monte_Carlo_Integration/2D_Sampling_with_Multidimensional_Transformations#Piecewise-Constant2DDistributions
+		// https://github.com/blender/cycles/blob/main/src/kernel/light/background.h
+		const searchCdfFn = wgslTagFn/* wgsl */`
+			fn searchEnvCdf( r: f32, offset: u32, count: u32, fraction: ptr<function, f32> ) -> u32 {
+
+				var lower = 0u;
+				var upper = count - 1u;
+				while ( lower < upper ) {
+
+					let mid = ( lower + upper ) >> 1u;
+					if ( ${ cdfNode }[ offset + mid + 1u ] <= r ) {
+
+						lower = mid + 1u;
+
+					} else {
+
+						upper = mid;
+
+					}
+
+				}
+
+				let start = ${ cdfNode }[ offset + lower ];
+				let end = ${ cdfNode }[ offset + lower + 1u ];
+				*fraction = select( 0.5, saturate( ( r - start ) / ( end - start ) ), end > start );
+				return lower;
+
+			}
+		`;
+
+		// the solid angle pdf of a direction at "v" in the given pixel, which is sampled uniformly over its area
+		const getTexelPdfFn = wgslTagFn/* wgsl */`
+			fn getEnvTexelPdf( texel: vec2u, v: f32 ) -> f32 {
+
+				let resolution = textureDimensions( ${ mapNode } );
+				let rowOffset = resolution.y + 1u + texel.y * ( resolution.x + 1u );
+				let rowPdf = ${ cdfNode }[ texel.y + 1u ] - ${ cdfNode }[ texel.y ];
+				let pixelPdf = ${ cdfNode }[ rowOffset + texel.x + 1u ] - ${ cdfNode }[ rowOffset + texel.x ];
+
+				// the equirect projection stretches each pixel by 1 / sin( theta ) over the sphere
+				let sinTheta = sin( ${ PI } * v );
+				return select( 0.0, rowPdf * pixelPdf * f32( resolution.x * resolution.y ) / ( 2.0 * ${ PI } * ${ PI } * sinTheta ), sinTheta > 0.0 );
+
+			}
+		`;
+
 		this.sampleColor = wgslTagFn/* wgsl */`
 			fn sampleEnv( direction: vec3f ) -> vec4f {
 
@@ -97,31 +146,34 @@ export class EquirectHdrInfoNode extends EquirectHdrInfoUniform {
 
 				var result: ${ environmentSampleStruct };
 
-				// walk the CDF: marginal picks the row (v), conditional picks the column (u)
-				let v = textureSampleLevel( ${ marginalNode }, ${ marginalSampler }, vec2f( r.x, 0.0 ), 0 ).x;
-				let u = textureSampleLevel( ${ conditionalNode }, ${ conditionalSampler }, vec2f( r.y, v ), 0 ).x;
-				let uv = vec2f( u, v );
-				let totalSum = ${ totalSumNode };
+				// search the CDFs: marginal picks the row (v), conditional picks the column (u), and the
+				// fractions place the sample within the pixel
+				let resolution = textureDimensions( ${ mapNode } );
+				var rowFraction = 0.0;
+				var columnFraction = 0.0;
+				let row = ${ searchCdfFn }( r.x, 0u, resolution.y, &rowFraction );
+				let column = ${ searchCdfFn }( r.y, resolution.y + 1u + row * ( resolution.x + 1u ), resolution.x, &columnFraction );
+				let texel = vec2u( column, row );
+				let uv = ( vec2f( texel ) + vec2f( columnFraction, rowFraction ) ) / vec2f( resolution );
 
 				let direction = ${ equirectUvToDirectionFn }( uv );
 				let color = textureSampleLevel( ${ mapNode }, ${ mapSampler }, uv, 0 ).rgb;
 
 				result.direction = transpose( ${ rotationNode } ) * direction;
 				result.color = color * ${ intensityNode };
+				result.pdf = ${ getTexelPdfFn }( texel, uv.y );
 
-				if ( totalSum != 0.0 ) {
-
-					let lum = ${ luminanceFn }( color );
-					let resolution = textureDimensions( ${ mapNode } );
-					result.pdf = ${ equirectLuminancePdfFn }( lum, totalSum, resolution );
-
-				} else {
+				return result;
 
-					result.pdf = 0.0;
+			}
+		`;
 
-				}
+		// the environment's average irradiance, for weighing it against the lights
+		this.getWeight = wgslTagFn/* wgsl */`
+			fn getEnvWeight() -> f32 {
 
-				return result;
+				let resolution = textureDimensions( ${ mapNode } );
+				return ${ intensityNode } * ${ PI } * ${ PI } * ${ totalSumNode } / ( 2.0 * f32( resolution.x * resolution.y ) );
 
 			}
 		`;
@@ -137,11 +189,10 @@ export class EquirectHdrInfoNode extends EquirectHdrInfoUniform {
 
 				let rotatedDir = ${ rotationNode } * direction;
 				let mapUv = ${ equirectDirectionToUvFn }( rotatedDir );
-				let color = textureSampleLevel( ${ mapNode }, ${ mapSampler }, mapUv, 0 ).rgb;
-				let lum = ${ luminanceFn }( color );
 				let resolution = textureDimensions( ${ mapNode } );
+				let texel = min( vec2u( mapUv * vec2f( resolution ) ), resolution - 1u );
 
-				return ${ equirectLuminancePdfFn }( lum, ${ totalSumNode }, resolution );
+				return ${ getTexelPdfFn }( texel, mapUv.y );
 
 			}
 		`;

package/src/webgpu/LightsInfoNode.js

@@ -5,6 +5,7 @@ import { AtlasTexture } from './AtlasTexture.js';
 import { LightsInfoUniformStruct } from '../uniforms/LightsInfoUniformStruct.js';
 import { lightStruct, lightRecordStruct } from './nodes/structs.wgsl.js';
 import { sampleTexelFunc } from './nodes/utils.wgsl.js';
+import { luminanceFn } from './nodes/sampling.wgsl.js';
 import {
 	RECT_AREA_LIGHT_TYPE,
 	CIRC_AREA_LIGHT_TYPE,
@@ -17,6 +18,8 @@ import {
 	randomAreaLightSampleFn,
 	randomSpotLightSampleFn,
 	getSpotAttenuationFn,
+	getDistanceAttenuationFn,
+	getSpotLightNearestFn,
 } from './nodes/lights.wgsl.js';
 
 export class LightsInfoNode extends LightsInfoUniformStruct {
@@ -33,7 +36,7 @@ export class LightsInfoNode extends LightsInfoUniformStruct {
 		// ies profiles packed into an atlas alongside their placement rects
 		this.iesAtlas = new AtlasTexture( { type: HalfFloatType } );
 		this.iesProfilesNode = texture( this.iesAtlas.texture );
-		this.iesInfoNode = uniformArray( this.iesAtlas.textureInfo, 'uvec4' );
+		this.iesInfoNode = uniformArray( this.iesAtlas.textureInfo, 'uvec4' ).setName( 'iesInfo' );
 
 		this._initFns();
 
@@ -65,12 +68,13 @@ export class LightsInfoNode extends LightsInfoUniformStruct {
 		const src = this.tex.image.data;
 		this.buffer.array.set( src.subarray( 0, count * stride ) );
 
-		// rewrite the int fields ( lightType, iesProfile ) as i32 bits
+		// rewrite the int fields ( lightType, visibleToCameraRays, iesProfile ) as i32 bits
 		const intView = new Int32Array( this.buffer.array.buffer );
 		for ( let i = 0; i < count; i ++ ) {
 
 			const base = i * stride;
 			intView[ base + 3 ] = Math.round( src[ base + 3 ] );
+			intView[ base + 11 ] = Math.round( src[ base + 11 ] );
 			intView[ base + 21 ] = Math.round( src[ base + 21 ] );
 
 		}
@@ -84,12 +88,146 @@ export class LightsInfoNode extends LightsInfoUniformStruct {
 
 	_initFns() {
 
-		const { bufferNode, iesProfilesNode, iesInfoNode } = this;
+		const { bufferNode, countNode, iesProfilesNode, iesInfoNode } = this;
 
 		// profiles are sampled out of an atlas so filtering must resolve tile-relative wrapping
 		const sampleIesTexelFn = sampleTexelFunc( iesInfoNode, iesProfilesNode, 'sampleIesTexel' );
 
-		// uniformly pick a light and sample it
+		this.isLightVisibleToCamera = wgslTagFn/* wgsl */`
+			fn isLightVisibleToCamera( index: u32 ) -> bool {
+
+				return ${ bufferNode }[ index ].visibleToCameraRays != 0;
+
+			}
+		`;
+
+		// A cheap estimate of how much light reaches a position from a light, ignoring occlusion and
+		// the surface orientation, used to choose which light to sample.
+		// TODO: weigh lights below the surface's horizon to zero using its normal, skipped for
+		// transmissive surfaces. Forward hit MIS would then need the previous vertex's normal too.
+		this.getLightWeight = wgslTagFn/* wgsl */`
+			fn getLightWeight( index: u32, position: vec3f ) -> f32 {
+
+				let light = ${ bufferNode }[ index ];
+				let power = ${ luminanceFn }( light.color * light.intensity );
+				if ( light.lightType == ${ SPOT_LIGHT_TYPE } ) {
+
+					// the most any point on the spot's disc lights the position. Ies profiles are not
+					// evaluated here
+					let nearest = ${ getSpotLightNearestFn }( light, position );
+					var attenuation = 1.0;
+					if ( light.iesProfile < 0 ) {
+
+						attenuation = ${ getSpotAttenuationFn }( light.coneCos, light.penumbraCos, nearest.x );
+
+					}
+
+					return power * attenuation * ${ getDistanceAttenuationFn }( nearest.y, light.distance, light.decay );
+
+				} else if ( light.lightType == ${ POINT_LIGHT_TYPE } ) {
+
+					// the point light's world position is packed into the u slot
+					let dist = length( light.u - position );
+					return power * ${ getDistanceAttenuationFn }( dist, light.distance, light.decay );
+
+				} else if ( light.lightType == ${ DIR_LIGHT_TYPE } ) {
+
+					return power;
+
+				} else {
+
+					let toLight = light.position - position;
+					let distSq = dot( toLight, toLight );
+					if ( distSq == 0.0 ) {
+
+						return 0.0;
+
+					}
+
+					// area lights only emit from their front face, and the distance is held above the size
+					// of the light so positions on or near it do not dominate
+					let normal = normalize( cross( light.u, light.v ) );
+					let cosTheta = dot( toLight, normal ) * inverseSqrt( distSq );
+					return power * light.area * max( cosTheta, 0.0 ) / max( distSq, light.area );
+
+				}
+
+			}
+		`;
+
+		this.getLightsWeight = wgslTagFn/* wgsl */`
+			fn getLightsWeight( position: vec3f ) -> f32 {
+
+				var total = 0.0;
+				for ( var i = 0u; i < ${ countNode }; i ++ ) {
+
+					total += ${ this.getLightWeight }( i, position );
+
+				}
+
+				return total;
+
+			}
+		`;
+
+		// Chooses a light, or the environment when the returned index equals the light count, in
+		// proportion to the estimated weights. The pdf is zero when nothing can light the position.
+		this.selectLight = wgslTagFn/* wgsl */`
+			fn selectLight( position: vec3f, envWeight: f32, r: f32, selectionPdf: ptr<function, f32> ) -> u32 {
+
+				let count = ${ countNode };
+				let totalWeight = ${ this.getLightsWeight }( position ) + envWeight;
+				if ( totalWeight <= 0.0 ) {
+
+					*selectionPdf = 0.0;
+					return 0u;
+
+				}
+
+				let targetWeight = r * totalWeight;
+				var accumulated = 0.0;
+				var lastWeighted = 0u;
+				var index = count;
+				for ( var i = 0u; i < count; i ++ ) {
+
+					let lightWeight = ${ this.getLightWeight }( i, position );
+					accumulated += lightWeight;
+					if ( lightWeight > 0.0 ) {
+
+						lastWeighted = i;
+
+					}
+
+					if ( targetWeight < accumulated ) {
+
+						index = i;
+						break;
+
+					}
+
+				}
+
+				// rounding can carry the target past the last light when there is no environment
+				if ( index == count && envWeight <= 0.0 ) {
+
+					index = lastWeighted;
+
+				}
+
+				var weight = envWeight;
+				if ( index < count ) {
+
+					weight = ${ this.getLightWeight }( index, position );
+
+				}
+
+				*selectionPdf = weight / totalWeight;
+				return index;
+
+			}
+		`;
+
+		// sample the light at the given index
 		this.randomLightSample = wgslTagFn/* wgsl */`
 			fn randomLightSample( lightIndex: u32, rayOrigin: vec3f, ruv: vec2f ) -> ${ lightRecordStruct } {
 

package/src/webgpu/MegaKernelPathTracer.js

@@ -69,7 +69,7 @@ export class MegaKernelPathTracer extends PathTracerBackend {
 
 	setBVHData( bvhData ) {
 
-		this.kernel.bvhData = bvhData;
+		this.kernel.context.bvhData = bvhData;
 		this.rebuild();
 
 	}
@@ -179,6 +179,7 @@ export class MegaKernelPathTracer extends PathTracerBackend {
 			kernel,
 			maxBounces,
 			maxTransparentBounces,
+			maxSubsurfaceBounces,
 
 			outputTarget,
 			sampleCountTarget,
@@ -193,6 +194,10 @@ export class MegaKernelPathTracer extends PathTracerBackend {
 
 		kernel.maxBounces = maxBounces;
 		kernel.maxTransparentBounces = maxTransparentBounces;
+		kernel.maxSubsurfaceBounces = maxSubsurfaceBounces;
+
+		// compile the kernel before its first dispatch
+		yield* this.compileKernels( [ kernel ] );
 
 		// number of tile cycles that have finished, ie. the sample count every pixel has reached
 		let completedSamples = 0;
@@ -201,6 +206,8 @@ export class MegaKernelPathTracer extends PathTracerBackend {
 
 			// skip the tile loop entirely once every pixel is finished
 			kernel.maxSamples = this.maxSamples;
+			kernel.auxTarget = this.auxTarget ?? this._auxPlaceholder;
+			kernel.accumulateAux = this.auxTarget ? 1 : 0;
 			if ( this.maxSamples !== 0 && completedSamples >= this.maxSamples ) {
 
 				yield;
@@ -210,7 +217,7 @@ export class MegaKernelPathTracer extends PathTracerBackend {
 
 			// split the frame into a near square grid of tiles no larger than the budget
 			this.getSize( _size );
-			const budgetTiles = Math.max( 1, Math.ceil( _size.x * _size.y / this.frameBudget ) );
+			const budgetTiles = lowResMode ? 1 : Math.max( 1, Math.ceil( _size.x * _size.y / this.frameBudget ) );
 			tiles.x = Math.ceil( Math.sqrt( budgetTiles ) );
 			tiles.y = Math.ceil( budgetTiles / tiles.x );
 
@@ -245,7 +252,7 @@ export class MegaKernelPathTracer extends PathTracerBackend {
 				for ( let y = 0; y < tiles.y; y ++ ) {
 
 					kernel.offset.set( x, y ).multiply( tileSize );
-					renderer.compute( kernel.kernel, dispatchSize );
+					this.dispatch( kernel, dispatchSize );
 
 					// in low res mode the tiles past the first fall outside the target and do no work
 					completedTiles = Math.min( completedTiles + 1, tileCount );
@@ -271,6 +278,7 @@ export class MegaKernelPathTracer extends PathTracerBackend {
 		// TODO: dispose of all buffers
 		this.envInfo.dispose();
 		this.lightsInfo.dispose();
+		this.kernel.dispose();
 
 	}