three-gpu-pathtracer 0.0.24 → 0.0.26

package/src/webgpu/EquirectHdrInfoNode.js

@@ -0,0 +1,151 @@
+import { Matrix3 } from 'three/webgpu';
+import { texture, sampler, uniform, wgslFn } 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 ] );
+
+export class EquirectHdrInfoNode extends EquirectHdrInfoUniform {
+
+	constructor() {
+
+		super();
+
+		// environment map + importance-sampling CDF textures, each with a sampler
+		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 );
+
+		// scalar parameters that assemble into the EnvironmentInfo struct in the shader
+		this.rotationNode = uniform( new Matrix3() );
+		this.intensityNode = uniform( 1 );
+		this.totalSumNode = uniform( this.totalSum );
+
+		this._initFns();
+
+	}
+
+	getPixelWeight( r, g, b, row, height ) {
+
+		// weight the pixel contribution by its spherical solid angle.
+		const theta = Math.PI * ( row + 0.5 ) / height;
+		return super.getPixelWeight( r, g, b ) * Math.sin( theta );
+
+	}
+
+	updateFrom( envMap ) {
+
+		super.updateFrom( envMap );
+
+		const {
+			mapNode,
+			marginalNode,
+			conditionalNode,
+			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;
+
+	}
+
+	_initFns() {
+
+		const {
+			mapNode,
+			mapSampler,
+			marginalNode,
+			marginalSampler,
+			conditionalNode,
+			conditionalSampler,
+			totalSumNode,
+			rotationNode,
+			intensityNode,
+		} = this;
+
+		this.sampleColor = wgslTagFn/* wgsl */`
+			fn sampleEnv( direction: vec3f ) -> vec4f {
+
+				let sampleDir = ${ rotationNode } * direction;
+				let mapUv = ${ equirectDirectionToUvFn }( sampleDir );
+				let col = textureSampleLevel( ${ mapNode }, ${ mapSampler }, mapUv, 0 );
+
+				return vec4f( ${ intensityNode } * col.rgb, col.a );
+
+			}
+		`;
+
+		this.sampleDir = wgslTagFn/* wgsl */`
+			fn sampleEnvDir( r: vec2f ) -> ${ environmentSampleStruct } {
+
+				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 };
+
+				let direction = ${ equirectUvToDirectionFn }( uv );
+				let color = textureSampleLevel( ${ mapNode }, ${ mapSampler }, uv, 0 ).rgb;
+
+				result.direction = transpose( ${ rotationNode } ) * direction;
+				result.color = color * ${ intensityNode };
+
+				if ( totalSum != 0.0 ) {
+
+					let lum = ${ luminanceFn }( color );
+					let resolution = textureDimensions( ${ mapNode } );
+					result.pdf = ${ equirectLuminancePdfFn }( lum, totalSum, resolution );
+
+				} else {
+
+					result.pdf = 0.0;
+
+				}
+
+				return result;
+
+			}
+		`;
+
+		this.getDirPdf = wgslTagFn/* wgsl */`
+			fn getEnvDirPdf( direction: vec3f ) -> f32 {
+
+				if ( ${ totalSumNode } == 0.0 ) {
+
+					return 0.0;
+
+				}
+
+				let rotatedDir = ${ rotationNode } * direction;
+				let mapUv = ${ equirectDirectionToUvFn }( rotatedDir );
+				let color = textureSampleLevel( ${ mapNode }, ${ mapSampler }, mapUv, 0 ).rgb;
+				let lum = ${ luminanceFn }( color );
+				let resolution = textureDimensions( ${ mapNode } );
+
+				return ${ equirectLuminancePdfFn }( lum, ${ totalSumNode }, resolution );
+
+			}
+		`;
+
+	}
+
+}

package/src/webgpu/LightsInfoNode.js

@@ -0,0 +1,224 @@
+import { storage, uniform, uniformArray, texture } from 'three/tsl';
+import { StorageBufferAttribute, HalfFloatType } from 'three/webgpu';
+import { wgslTagFn } from 'three-mesh-bvh/webgpu';
+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 {
+	RECT_AREA_LIGHT_TYPE,
+	CIRC_AREA_LIGHT_TYPE,
+	SPOT_LIGHT_TYPE,
+	DIR_LIGHT_TYPE,
+	POINT_LIGHT_TYPE,
+	LIGHT_FAR_DISTANCE,
+	intersectsRectangleFn,
+	intersectsCircleFn,
+	randomAreaLightSampleFn,
+	randomSpotLightSampleFn,
+	getSpotAttenuationFn,
+} from './nodes/lights.wgsl.js';
+
+export class LightsInfoNode extends LightsInfoUniformStruct {
+
+	constructor() {
+
+		super();
+
+		// lights packed into a storage buffer of Light structs
+		this.countNode = uniform( this.count, 'uint' );
+		this.buffer = new StorageBufferAttribute( new Float32Array( 2 * lightStruct.getLength() ), lightStruct.getLength() );
+		this.bufferNode = storage( this.buffer, lightStruct ).toReadOnly().setName( 'lights' );
+
+		// 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._initFns();
+
+	}
+
+	updateFrom( renderer, lights ) {
+
+		// the unique set of ies textures referenced by the lights' "iesMap" fields
+		const iesTextures = Array.from( new Set( lights.map( l => l.iesMap ).filter( t => t ) ) );
+		const changed = super.updateFrom( lights, iesTextures );
+
+		this.iesAtlas.setTextures( renderer, iesTextures );
+		this.iesProfilesNode.value = this.iesAtlas.texture;
+
+		const stride = lightStruct.getLength();
+		const count = this.count;
+		const capacity = Math.max( count, 2 );
+
+		// resize the buffer to the exact light count, keeping the same binding node
+		if ( this.buffer.array.length !== capacity * stride ) {
+
+			this.buffer.dispose();
+			this.buffer = new StorageBufferAttribute( new Float32Array( capacity * stride ), stride );
+			this.bufferNode.value = this.buffer;
+
+		}
+
+		// the texture's packed float layout already matches lightStruct's std layout, so copy it in
+		const src = this.tex.image.data;
+		this.buffer.array.set( src.subarray( 0, count * stride ) );
+
+		// rewrite the int fields ( lightType, 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 + 21 ] = Math.round( src[ base + 21 ] );
+
+		}
+
+		this.buffer.needsUpdate = true;
+		this.countNode.value = this.count;
+
+		return changed;
+
+	}
+
+	_initFns() {
+
+		const { bufferNode, 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.randomLightSample = wgslTagFn/* wgsl */`
+			fn randomLightSample( lightIndex: u32, rayOrigin: vec3f, ruv: vec2f ) -> ${ lightRecordStruct } {
+
+				let light = ${ bufferNode }[ lightIndex ];
+
+				var result: ${ lightRecordStruct };
+				if ( light.lightType == ${ SPOT_LIGHT_TYPE } ) {
+
+					result = ${ randomSpotLightSampleFn }( light, rayOrigin, ruv );
+
+					let spotNormal = normalize( cross( light.u, light.v ) );
+					let cosTheta = dot( result.direction, spotNormal );
+					var spotAttenuation: f32;
+					if ( light.iesProfile >= 0 ) {
+
+						// half a texel of the profile's tile in the atlas
+						let profileInfo = ${ iesInfoNode }[ u32( light.iesProfile ) ];
+						let texelInset = 0.5 / vec2f( vec2u( profileInfo.y & 0xFFFFu, profileInfo.y >> 16u ) );
+
+						// tilt spans [ 0, 180 ]deg inset to the texel centers, twist covers [ 0, 359 ]deg and wraps
+						let tiltAngle = mix( texelInset.x, 1.0 - texelInset.x, acos( cosTheta ) / PI );
+						let twistAngle = ( atan2( dot( result.direction, light.v ), dot( result.direction, light.u ) ) + PI ) / ( 2.0 * PI ) + texelInset.y;
+						let packedProfile = ( 1 << 26 ) | light.iesProfile;
+						spotAttenuation = ${ sampleIesTexelFn }( vec2f( tiltAngle, twistAngle ), packedProfile, 0.0 ).r;
+
+					} else {
+
+						spotAttenuation = ${ getSpotAttenuationFn }( light.coneCos, light.penumbraCos, cosTheta );
+
+					}
+
+					result.emission *= spotAttenuation;
+
+				} else if ( light.lightType == ${ POINT_LIGHT_TYPE } ) {
+
+					// the point light's world position is packed into the u slot
+					let lightRay = light.u - rayOrigin;
+					let lightDist = length( lightRay );
+					let cutoffDistance = light.distance;
+					var distanceFalloff = 1.0 / max( pow( lightDist, light.decay ), 0.01 );
+					if ( cutoffDistance > 0.0 ) {
+
+						let window = clamp( 1.0 - pow( lightDist / cutoffDistance, 4.0 ), 0.0, 1.0 );
+						distanceFalloff *= window * window;
+
+					}
+
+					result.direction = normalize( lightRay );
+					result.dist = lightDist;
+					result.pdf = 1.0;
+					result.emission = light.color * light.intensity * distanceFalloff;
+					result.lightType = light.lightType;
+
+				} else if ( light.lightType == ${ DIR_LIGHT_TYPE } ) {
+
+					// the directional light's direction is packed into the u slot
+					result.dist = ${ LIGHT_FAR_DISTANCE };
+					result.direction = light.u;
+					result.pdf = 1.0;
+					result.emission = light.color * light.intensity;
+					result.lightType = light.lightType;
+
+				} else {
+
+					result = ${ randomAreaLightSampleFn }( light, rayOrigin, ruv );
+
+				}
+
+				return result;
+
+			}
+		`;
+
+		// forward intersection of a ray with a single area light ( rect / circ only ), used for MIS
+		// TODO: support hitting the spot light disk here and move spot lights into the
+		// MIS-weighted set so they appear in sharp reflections
+		this.intersectLightAtIndex = wgslTagFn/* wgsl */`
+			fn intersectLightAtIndex( rayOrigin: vec3f, rayDirection: vec3f, index: u32, lightRec: ptr<function, ${ lightRecordStruct }> ) -> bool {
+
+				let light = ${ bufferNode }[ index ];
+
+				var u = light.u;
+				var v = light.v;
+				let normal = normalize( cross( u, v ) );
+
+				// only front-facing area lights can be hit
+				if ( dot( normal, rayDirection ) > 0.0 ) {
+
+					u *= 1.0 / dot( u, u );
+					v *= 1.0 / dot( v, v );
+
+					var dist = - 1.0;
+					if ( light.lightType == ${ RECT_AREA_LIGHT_TYPE } ) {
+
+						dist = ${ intersectsRectangleFn }( light.position, normal, u, v, rayOrigin, rayDirection );
+
+					} else if ( light.lightType == ${ CIRC_AREA_LIGHT_TYPE } ) {
+
+						dist = ${ intersectsCircleFn }( light.position, normal, u, v, rayOrigin, rayDirection );
+
+					}
+
+					if ( dist > 0.0 ) {
+
+						let cosTheta = dot( rayDirection, normal );
+						lightRec.dist = dist;
+						lightRec.pdf = ( dist * dist ) / ( light.area * cosTheta );
+						lightRec.emission = light.color * light.intensity;
+						lightRec.direction = rayDirection;
+						lightRec.lightType = light.lightType;
+						return true;
+
+					}
+
+				}
+
+				return false;
+
+			}
+		`;
+
+	}
+
+	dispose() {
+
+		this.tex.dispose();
+		this.buffer.dispose();
+		this.iesAtlas.dispose();
+
+	}
+
+}

package/src/webgpu/MegaKernelPathTracer.js

@@ -0,0 +1,277 @@
+import { Matrix4, Vector2 } from 'three/webgpu';
+import { PathTracerMegaKernel } from './compute/PathTracerMegaKernel.js';
+import { EquirectHdrInfoNode } from './EquirectHdrInfoNode.js';
+import { EquirectBackgroundInfo } from './EquirectBackgroundInfo.js';
+import { LightsInfoNode } from './LightsInfoNode.js';
+import { FILTER_GLOSSY_DISABLED } from './nodes/material.wgsl.js';
+import { PathTracerBackend } from './PathTracerBackend.js';
+
+const _size = new Vector2();
+
+export class MegaKernelPathTracer extends PathTracerBackend {
+
+	constructor( renderer ) {
+
+		super( renderer );
+
+		// options
+		this.envInfo = new EquirectHdrInfoNode();
+
+		// every pixel in a tile finishes its sample in one dispatch, so the counts follow from how
+		// far through the tile cycle the render is and need no GPU work to measure
+		this._completedSamples = 0;
+		this._tileProgress = 0;
+
+		// kernels
+		this.kernel = new PathTracerMegaKernel( ).setWorkgroupSize( 8, 8, 1 );
+		this.kernel.envInfo = this.envInfo;
+
+		this.backgroundInfo = new EquirectBackgroundInfo();
+		this.kernel.backgroundInfo = this.backgroundInfo;
+
+		this.lightsInfo = new LightsInfoNode();
+		this.kernel.lightsInfo = this.lightsInfo;
+
+	}
+
+	resetSeed() {
+
+		this.kernel.seed = 0;
+
+	}
+
+	reset() {
+
+		super.reset();
+		this._completedSamples = 0;
+		this._tileProgress = 0;
+
+	}
+
+	async getSampleCountsAsync() {
+
+		if ( this.lowResMode ) {
+
+			return { min: 0, max: 0, avg: 0 };
+
+		}
+
+		const completed = this._completedSamples;
+		const progress = this._tileProgress;
+
+		return {
+			min: progress === 1 ? completed + 1 : completed,
+			max: completed + 1,
+			avg: completed + progress,
+		};
+
+	}
+
+	setBVHData( bvhData ) {
+
+		this.kernel.bvhData = bvhData;
+		this.rebuild();
+
+	}
+
+	rebuild() {
+
+		super.rebuild();
+		this.kernel.needsUpdate = true;
+		this.reset();
+
+	}
+
+	setRandom( random ) {
+
+		this.kernel.context.random = random;
+		this.kernel.needsUpdate = true;
+		this.reset();
+
+	}
+
+	setMaterial( material ) {
+
+		this.kernel.material = material.getData();
+		this.kernel.needsUpdate = true;
+		this.reset();
+
+	}
+
+	setFilterGlossy( value ) {
+
+		// the kernel takes the inverted threshold, mirroring Cycles
+		this.kernel.filterGlossy = value === 0 ? FILTER_GLOSSY_DISABLED : 1 / value;
+		this.reset();
+
+	}
+
+	setClamping( direct, indirect ) {
+
+		this.kernel.clampDirect = direct;
+		this.kernel.clampIndirect = indirect;
+		this.reset();
+
+	}
+
+	setEnvironment( envMap ) {
+
+		this.envInfo.updateFrom( envMap );
+
+	}
+
+	setLights( lights ) {
+
+		this.lightsInfo.updateFrom( this.renderer, lights );
+		this.reset();
+
+	}
+
+	setMultipleImportanceSampling( enabled ) {
+
+		this.kernel.misEnabled = enabled ? 1 : 0;
+		this.reset();
+
+	}
+
+	setEnvironmentParams( envMapIntensity, envMapRotation ) {
+
+		const { envInfo } = this;
+		const rotationMatrix = new Matrix4().makeRotationFromEuler( envMapRotation ).invert();
+		envInfo.rotationNode.value.setFromMatrix4( rotationMatrix );
+		envInfo.intensityNode.value = envMapIntensity;
+
+	}
+
+	setBackground( background ) {
+
+		const { backgroundInfo } = this;
+		backgroundInfo.dispose();
+		backgroundInfo.map = background;
+
+	}
+
+	setBackgroundParams(
+		backgroundIntensity,
+		backgroundRotation,
+		backgroundBlurriness,
+	) {
+
+		const { backgroundInfo } = this;
+		const rotationMatrix = new Matrix4().makeRotationFromEuler( backgroundRotation ).invert();
+		backgroundInfo.rotationNode.value.copy( rotationMatrix );
+		backgroundInfo.intensity = backgroundIntensity;
+		backgroundInfo.blur = backgroundBlurriness;
+
+	}
+
+	setTransmissiveBackground( value ) {
+
+		this.kernel.transmissiveBackground = value;
+		this.reset();
+
+	}
+
+	*createRenderTask() {
+
+		const {
+			renderer,
+			kernel,
+			maxBounces,
+			maxTransparentBounces,
+
+			outputTarget,
+			sampleCountTarget,
+			lowResMode,
+		} = this;
+
+		const tiles = new Vector2();
+
+		// init parameters
+		kernel.outputTarget = outputTarget;
+		kernel.sampleCountTarget = sampleCountTarget;
+
+		kernel.maxBounces = maxBounces;
+		kernel.maxTransparentBounces = maxTransparentBounces;
+
+		// number of tile cycles that have finished, ie. the sample count every pixel has reached
+		let completedSamples = 0;
+
+		while ( true ) {
+
+			// skip the tile loop entirely once every pixel is finished
+			kernel.maxSamples = this.maxSamples;
+			if ( this.maxSamples !== 0 && completedSamples >= this.maxSamples ) {
+
+				yield;
+				continue;
+
+			}
+
+			// 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 ) );
+			tiles.x = Math.ceil( Math.sqrt( budgetTiles ) );
+			tiles.y = Math.ceil( budgetTiles / tiles.x );
+
+			// every pixel in a tile finishes its sample in a single dispatch, so the sample counts
+			// can be derived from how far through the tile cycle we are without reading anything back
+			const tileCount = this.lowResMode ? 1 : tiles.x * tiles.y;
+			let completedTiles = 0;
+
+			const tileSize = kernel.tileSize;
+			if ( lowResMode ) {
+
+				this.getSize( tileSize );
+
+			} else {
+
+				this.getSize( tileSize ).divide( tiles ).ceil();
+
+			}
+
+			const dispatchSize = kernel.getDispatchSize( tileSize.x, tileSize.y );
+			kernel.seed += 1;
+
+			// Swap targets to support devices without <rgba32float, read_write> textures
+			// Copy latest data to a new outputTarget to keep the appearance
+			renderer.copyTextureToTexture( this.outputTarget, this.prevOutputTarget );
+			[ this.outputTarget, this.prevOutputTarget ] = [ this.prevOutputTarget, this.outputTarget ];
+			kernel.prevOutputTarget = this.prevOutputTarget;
+			kernel.outputTarget = this.outputTarget;
+
+			for ( let x = 0; x < tiles.x; x ++ ) {
+
+				for ( let y = 0; y < tiles.y; y ++ ) {
+
+					kernel.offset.set( x, y ).multiply( tileSize );
+					renderer.compute( kernel.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 );
+					this._completedSamples = completedSamples;
+					this._tileProgress = completedTiles / tileCount;
+
+					yield;
+
+				}
+
+			}
+
+			completedSamples ++;
+
+		}
+
+	}
+
+	dispose() {
+
+		super.dispose();
+
+		// TODO: dispose of all buffers
+		this.envInfo.dispose();
+		this.lightsInfo.dispose();
+
+	}
+
+}