three-gpu-pathtracer 0.0.23 → 0.0.25

package/src/webgpu/compute/TallySampleCountsKernel.js

@@ -0,0 +1,111 @@
+import { StorageBufferAttribute, StorageTexture, StructTypeNode } from 'three/webgpu';
+import { storage, textureStore, globalId } from 'three/tsl';
+import { wgslTagFn } from 'three-mesh-bvh/webgpu';
+import { ComputeKernel } from './ComputeKernel.js';
+import { SAMPLE_COUNT_MASK, SAMPLE_DISPATCHED_FLAG } from '../constants.js';
+
+// field order of the sample counter buffer, matching the struct below
+export const SAMPLE_COUNTER_MIN = 0;
+export const SAMPLE_COUNTER_MAX = 1;
+export const SAMPLE_COUNTER_TOTAL_LO = 2;
+export const SAMPLE_COUNTER_TOTAL_HI = 3;
+export const SAMPLE_COUNTER_PIXEL_COUNT = 4;
+export const SAMPLE_COUNTER_LENGTH = 5;
+
+export const U32_RANGE = 4294967296;
+
+// Reduction target for the sample counts. The total outgrows 32 bits, so it is split across
+// "totalLo" and "totalHi" and reassembled on read.
+export const sampleCountersStruct = new StructTypeNode( {
+	minSamples: { type: 'uint', atomic: true },
+	maxSamples: { type: 'uint', atomic: true },
+	totalLo: { type: 'uint', atomic: true },
+	totalHi: { type: 'uint', atomic: true },
+	pixelCount: { type: 'uint', atomic: true },
+}, 'SampleCounters' );
+
+// Clears the counters ahead of the tally. "minSamples" starts at the u32 max so the first
+// "atomicMin" replaces it.
+export class PrimeSampleCountersKernel extends ComputeKernel {
+
+	constructor() {
+
+		const params = {
+			counters: storage( new StorageBufferAttribute( 1, 1 ), sampleCountersStruct ),
+		};
+
+		const fn = wgslTagFn/* wgsl */`
+			fn compute() -> void {
+
+				atomicStore( &${ params.counters }.minSamples, 0xFFFFFFFFu );
+				atomicStore( &${ params.counters }.maxSamples, 0u );
+				atomicStore( &${ params.counters }.totalLo, 0u );
+				atomicStore( &${ params.counters }.totalHi, 0u );
+				atomicStore( &${ params.counters }.pixelCount, 0u );
+
+			}
+		`;
+
+		super( fn( params ) );
+
+		this.defineUniformAccessors( params );
+
+	}
+
+}
+
+// Reduces the sample count target to a min, max, and total so the CPU reads back a few values
+// instead of the whole texture.
+export class TallySampleCountsKernel extends ComputeKernel {
+
+	constructor() {
+
+		const params = {
+			globalId: globalId,
+			counters: storage( new StorageBufferAttribute( 1, 1 ), sampleCountersStruct ),
+			sampleCountTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadOnly(),
+		};
+
+		const fn = wgslTagFn/* wgsl */`
+			fn compute( globalId: vec3u ) -> void {
+
+				let targetDimensions = textureDimensions( ${ params.sampleCountTarget } );
+				if ( globalId.x >= targetDimensions.x || globalId.y >= targetDimensions.y ) {
+
+					return;
+
+				}
+
+				// pixels without a camera ray aren't part of the image, so skip rather than count zero
+				let combinedField = textureLoad( ${ params.sampleCountTarget }, globalId.xy ).r;
+				if ( ( combinedField & ${ SAMPLE_DISPATCHED_FLAG }u ) == 0u ) {
+
+					return;
+
+				}
+
+				let samples = combinedField & ${ SAMPLE_COUNT_MASK }u;
+
+				atomicMin( &${ params.counters }.minSamples, samples );
+				atomicMax( &${ params.counters }.maxSamples, samples );
+				atomicAdd( &${ params.counters }.pixelCount, 1u );
+
+				// the sum overflowed if it came out smaller than the value already there, so carry
+				// the wrap into "totalHi"
+				let previousTotal = atomicAdd( &${ params.counters }.totalLo, samples );
+				if ( previousTotal + samples < previousTotal ) {
+
+					atomicAdd( &${ params.counters }.totalHi, 1u );
+
+				}
+
+			}
+		`;
+
+		super( fn( params ) );
+
+		this.defineUniformAccessors( params );
+
+	}
+
+}

package/src/webgpu/compute/ZeroOutBufferKernel.js

@@ -0,0 +1,35 @@
+import { IndirectStorageBufferAttribute } from 'three/webgpu';
+import { ComputeKernel } from './ComputeKernel.js';
+import { storage, globalId } from 'three/tsl';
+import { wgslTagFn } from 'three-mesh-bvh/webgpu';
+
+export class ZeroOutBufferKernel extends ComputeKernel {
+
+	constructor( options = {} ) {
+
+		const {
+			type = 'u32',
+		} = options;
+
+		const params = {
+			globalId: globalId,
+			outputTarget: storage( new IndirectStorageBufferAttribute( 1, 1 ), type ),
+		};
+
+		const fn = wgslTagFn/* wgsl */`
+			fn compute( globalId: vec3u ) -> void {
+
+				${ params.outputTarget }[ globalId.x ] = 0;
+
+			}
+		`;
+
+		super( fn( params ) );
+
+		this.defineUniformAccessors( {
+			target: params.outputTarget,
+		} );
+
+	}
+
+}

package/src/webgpu/compute/ZeroOutKernel.js

@@ -0,0 +1,31 @@
+import { StorageTexture } from 'three/webgpu';
+import { ComputeKernel } from './ComputeKernel.js';
+import { textureStore, globalId } from 'three/tsl';
+import { wgslTagFn } from 'three-mesh-bvh/webgpu';
+
+export class ZeroOutKernel extends ComputeKernel {
+
+	constructor() {
+
+		const params = {
+			globalId: globalId,
+			outputTarget: textureStore( new StorageTexture( 1, 1 ) ).toWriteOnly(),
+		};
+
+		const fn = wgslTagFn/* wgsl */`
+			fn compute( globalId: vec3u ) -> void {
+
+				textureStore( ${ params.outputTarget }, globalId.xy, vec4( 0, 0, 0, 1 ) );
+
+			}
+		`;
+
+		super( fn( params ) );
+
+		this.defineUniformAccessors( {
+			target: params.outputTarget,
+		} );
+
+	}
+
+}

package/src/webgpu/compute/wavefront/LogicKernel.js

@@ -0,0 +1,349 @@
+import { StorageBufferAttribute, StorageTexture } from 'three/webgpu';
+import { ComputeKernel } from '../ComputeKernel.js';
+import { uniform, storage, textureStore, globalId } from 'three/tsl';
+import { proxy, proxyFn, wgslTagFn } from 'three-mesh-bvh/webgpu';
+import { misHeuristicFn, weightedAlphaBlendFn } from '../../nodes/sampling.wgsl.js';
+import { clampPathContributionFunc, isTerminatingScatterFunc } from '../../nodes/utils.wgsl.js';
+import { TRANSMISSIVE_BACKGROUND_ENVIRONMENT, TRANSMISSIVE_BACKGROUND_OVERLAY, TRANSMISSIVE_BACKGROUND_TRANSPARENT } from '../../constants.js';
+import {
+	rngInit, rand2, rand3,
+	RNG_INDEX_BACKGROUND_SAMPLE,
+	RNG_INDEX_DIRECT_LIGHT_SAMPLE,
+} from '../../nodes/random.wgsl.js';
+import { ENVIRONMENT_LIGHT_TYPE, LIGHT_FAR_DISTANCE, isMISWeightLightFn } from '../../nodes/lights.wgsl.js';
+import { lightRecordStruct, scatterRecordStruct } from '../../nodes/structs.wgsl.js';
+import { rayDataStruct, intersectionResultStruct } from './structs.js';
+import { SAMPLE_COUNT_MASK, SAMPLE_DISPATCHED_FLAG } from '../../constants.js';
+
+// Path logic over the persistent slot pool: resolves the previous frame's shadow and bounce trace
+// results, accumulates emission and NEE / forward MIS contributions, terminates finished paths into
+// the output, and samples the next NEE light for MaterialKernel. Touches no BVH data.
+export class LogicKernel extends ComputeKernel {
+
+	constructor( ) {
+
+		const params = {
+			envInfo: { value: null },
+			backgroundInfo: { value: null },
+			lightsInfo: { value: null },
+
+			// targets
+			prevOutputTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadOnly(),
+			outputTarget: textureStore( new StorageTexture( 1, 1 ) ).toWriteOnly(),
+			sampleCountTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadWrite(),
+
+			// settings
+			misEnabled: uniform( 1, 'uint' ),
+			rayCount: uniform( 0, 'uint' ),
+			maxBounces: uniform( 5, 'uint' ),
+			clampDirect: uniform( 0 ),
+			clampIndirect: uniform( 10 ),
+			transmissiveBackground: uniform( TRANSMISSIVE_BACKGROUND_OVERLAY ),
+
+			rayDataStorage: storage( new StorageBufferAttribute( 1, 1 ), rayDataStruct ),
+			rayIntersectionsStorage: storage( new StorageBufferAttribute( 1, 1 ), intersectionResultStruct ),
+			shadowRayIntersectionsStorage: storage( new StorageBufferAttribute( 1, 1 ), intersectionResultStruct ),
+
+			globalId: globalId,
+		};
+
+		// environment + background resources pulled off their providers (embedded functions)
+		const envTotalSumNode = proxy( 'envInfo.value.totalSumNode', params );
+		const envIntensityNode = proxy( 'envInfo.value.intensityNode', params );
+		const sampleEnvColor = proxy( 'envInfo.value.sampleColor', params );
+		const sampleEnvDir = proxy( 'envInfo.value.sampleDir', params );
+		const getEnvDirPdf = proxy( 'envInfo.value.getDirPdf', params );
+		const sampleBackground = proxy( 'backgroundInfo.value.sampleColor', params );
+
+		// analytic scene lights pulled off the lightsInfo provider (LightsInfoNode)
+		const lightsCountNode = proxy( 'lightsInfo.value.countNode', params );
+		const randomLightSampleFn = proxyFn( 'lightsInfo.value.randomLightSample', params );
+		const intersectLightAtIndexFn = proxyFn( 'lightsInfo.value.intersectLightAtIndex', params );
+
+		const fn = wgslTagFn/* wgsl */`
+
+			fn compute(
+				// settings
+				misEnabled: u32,
+				rayCount: u32,
+				maxBounces: u32,
+				clampDirect: f32,
+				clampIndirect: f32,
+				transmissiveBackground: u32,
+
+				globalId: vec3u
+			) -> void {
+
+				let rayDataStorage = &${ params.rayDataStorage };
+				let rayIntersectionsStorage = &${ params.rayIntersectionsStorage };
+				let shadowRayIntersectionsStorage = &${ params.shadowRayIntersectionsStorage };
+
+				// bound by "rayCount" rather than the pool length. See MaterialKernel
+				let index = globalId.x;
+				if ( index >= rayCount ) {
+
+					return;
+
+				}
+
+				// skip slots that have never spawned a ray
+				let input = ( *rayDataStorage )[ index ];
+				if ( input.rayIntersectionIndex < 0 ) {
+
+					return;
+
+				}
+
+				let indexUV = vec2u( input.pixelIndex >> 16, input.pixelIndex & 0xFFFF );
+				${ rngInit }( indexUV, input.seed, input.currentBounce + input.alphaDepth );
+
+				// one-sample NEE selection normalization (lights + env), matched with the megakernel
+				let envActive = ${ envTotalSumNode } > 0.0 && ${ envIntensityNode } > 0.0;
+				let lightsCount = ${ lightsCountNode };
+				var lightsDenom = f32( lightsCount );
+				if ( envActive ) {
+
+					lightsDenom += 1.0;
+
+				}
+
+				var resultColor = input.resultColor;
+				var throughputColor = input.throughputColor;
+
+				// resolve the previous surface's NEE shadow ray (pre-scatter throughput). The index
+				// is negative when no shadow ray was enqueued last frame
+				if ( input.shadowRayIntersectionIndex >= 0 && input.lightPdf > 0.0 ) {
+
+					let shadowHit = ( *shadowRayIntersectionsStorage )[ u32( input.shadowRayIntersectionIndex ) ];
+					let occluded = shadowHit.objectIndex >= 0;
+					if ( ! occluded ) {
+
+						// env + area lights are also bsdf-sampled, so MIS-weight them; punctual take full weight
+						let misWeight = select( 1.0, ${ misHeuristicFn }( input.lightPdf, input.lightBsdfPdf ), ${ isMISWeightLightFn }( input.lightType ) );
+						let directLight = throughputColor * input.lightEmission * input.lightBsdf * misWeight / input.lightPdf;
+						let contribution = ${ clampPathContributionFunc }( directLight, input.currentBounce, clampDirect, clampIndirect );
+						resultColor += vec4f( contribution, 0.0 );
+
+					}
+
+				}
+
+				// emission gathered at the previous surface (pre-scatter throughput). Found light clamps
+				// at the depth of the segment that found it, as in Cycles, so one less than the current
+				let emission = ${ clampPathContributionFunc }( throughputColor * input.emission, max( input.currentBounce, 1u ) - 1u, clampDirect, clampIndirect );
+				resultColor += vec4f( emission, 0.0 );
+
+				// reconstruct the scatter record staged by MaterialKernel
+				var scatterRec: ${ scatterRecordStruct };
+				scatterRec.color = input.scatterColor;
+				scatterRec.pdf = input.scatterPdf;
+
+				// the bounce limit, russian roulette, and terminating scatter checks all ran in
+				// MaterialKernel, which stages a zeroed pdf and skips the bounce trace when they fire
+				var isTerminated = all( throughputColor == vec3f( 0.0 ) ) || input.currentBounce >= maxBounces || ${ isTerminatingScatterFunc }( scatterRec );
+
+				if ( ! isTerminated ) {
+
+					// apply the scatter across the traced segment
+					throughputColor *= scatterRec.color / scatterRec.pdf;
+
+					let hitResult = ( *rayIntersectionsStorage )[ u32( input.rayIntersectionIndex ) ];
+					let didHit = hitResult.objectIndex >= 0;
+					let surfaceDist = select( ${ LIGHT_FAR_DISTANCE }, hitResult.dist, didHit );
+
+					// forward hits: a bsdf-sampled segment that lands on a area light. MIS-weighted
+					// only when NEE is also sampling the lights. The camera segment is skipped.
+					if ( input.currentBounce > 0u ) {
+
+						for ( var li = 0u; li < lightsCount; li ++ ) {
+
+							var lightRec: ${ lightRecordStruct };
+							if ( ${ intersectLightAtIndexFn }( input.origin, input.direction, li, &lightRec ) && lightRec.dist < surfaceDist ) {
+
+								var misWeight = 1.0;
+								if ( misEnabled != 0u ) {
+
+									let lightPdf = lightRec.pdf / lightsDenom;
+									misWeight = ${ misHeuristicFn }( input.scatterPdf, lightPdf );
+
+								}
+
+								let lightHit = ${ clampPathContributionFunc }( lightRec.emission * throughputColor * misWeight, input.currentBounce, clampDirect, clampIndirect );
+								resultColor += vec4f( lightHit, 0.0 );
+
+							}
+
+						}
+
+					}
+
+					if ( didHit ) {
+
+						// stage the hit for MaterialKernel
+						rayDataStorage[ index ].barycoord = hitResult.barycoord;
+						rayDataStorage[ index ].normal = hitResult.normal;
+						rayDataStorage[ index ].side = hitResult.side;
+						rayDataStorage[ index ].indices = hitResult.indices;
+						rayDataStorage[ index ].objectIndex = hitResult.objectIndex;
+						rayDataStorage[ index ].dist = hitResult.dist;
+
+						// next event estimation: pick one light or the environment with a single sample.
+						// MaterialKernel evaluates the bsdf and enqueues the shadow ray.
+						// TODO: importance-sample the selection by light intensity and solid angle
+						var lightPdf = 0.0;
+						if ( misEnabled != 0u && lightsDenom > 0.0 ) {
+
+							let ruv = ${ rand3 }( ${ RNG_INDEX_DIRECT_LIGHT_SAMPLE } );
+							let lightIndex = min( u32( ruv.x * lightsDenom ), u32( lightsDenom ) - 1u );
+							var lightRec: ${ lightRecordStruct };
+							if ( envActive && lightIndex == lightsCount ) {
+
+								// the environment, sampled from its CDF, as a light of kind ENVIRONMENT
+								let envSample = ${ sampleEnvDir }( ruv.yz );
+								lightRec.direction = envSample.direction;
+								lightRec.emission = envSample.color;
+								lightRec.pdf = envSample.pdf;
+								lightRec.dist = ${ LIGHT_FAR_DISTANCE };
+								lightRec.lightType = ${ ENVIRONMENT_LIGHT_TYPE };
+
+							} else {
+
+								lightRec = ${ randomLightSampleFn }( lightIndex, hitResult.position, ruv.yz );
+
+							}
+
+							lightPdf = lightRec.pdf / lightsDenom;
+							rayDataStorage[ index ].lightDirection = lightRec.direction;
+							rayDataStorage[ index ].lightEmission = lightRec.emission;
+							rayDataStorage[ index ].lightDist = lightRec.dist;
+							rayDataStorage[ index ].lightType = lightRec.lightType;
+
+						}
+
+						rayDataStorage[ index ].lightPdf = lightPdf;
+
+					} else {
+
+						// the segment escaped the scene: gather the environment for opaque paths, or
+						// the background for camera segments and fully transmissive paths
+						if ( input.currentBounce > 0u && input.isFullyTransmissive == 0u ) {
+
+							var misWeight = 1.0;
+							if ( misEnabled != 0u && envActive ) {
+
+								// match the env pdf scaling used by the NEE selection so the two estimators balance
+								let envPdf = ${ getEnvDirPdf }( input.direction ) / lightsDenom;
+								misWeight = ${ misHeuristicFn }( input.scatterPdf, envPdf );
+
+							}
+
+							let environment = ${ sampleEnvColor }( input.direction ).rgb * throughputColor * misWeight;
+							let contribution = ${ clampPathContributionFunc }( environment, input.currentBounce, clampDirect, clampIndirect );
+							resultColor += vec4f( contribution, 0.0 );
+
+						} else {
+
+							// hit the background
+							// support multiple transparent background blending techniques
+							let rng = ${ rand2 }( ${ RNG_INDEX_BACKGROUND_SAMPLE } );
+							let bg = ${ sampleBackground }( input.direction, rng );
+							if ( input.currentBounce == 0u ) {
+
+								// sample the background directly if this is the primary ray
+								let background = ${ clampPathContributionFunc }( bg.a * bg.rgb, input.currentBounce, clampDirect, clampIndirect );
+								resultColor = vec4f( background, bg.a );
+
+							} else {
+
+								// transmissive ray handling
+								let env = ${ sampleEnvColor }( input.direction );
+								let avg = saturate( dot( throughputColor, vec3f( 1.0 / 3.0 ) ) );
+								let transparency = ( 1.0 - bg.a ) * avg;
+
+								var misWeight = 1.0;
+								if ( misEnabled != 0u && envActive ) {
+
+									let envPdf = ${ getEnvDirPdf }( input.direction );
+									misWeight = ${ misHeuristicFn }( input.scatterPdf, envPdf );
+
+								}
+
+								if ( transmissiveBackground == ${ TRANSMISSIVE_BACKGROUND_ENVIRONMENT }u ) {
+
+									// display the env map through transmissive surfaces
+									let background = ${ clampPathContributionFunc }( env.rgb * throughputColor * misWeight, input.currentBounce, clampDirect, clampIndirect );
+									resultColor = vec4f(
+										resultColor.rgb + background,
+										1.0,
+									);
+
+								} else if ( transmissiveBackground == ${ TRANSMISSIVE_BACKGROUND_TRANSPARENT }u ) {
+
+									// fade the background by the throughput color average
+									let background = ${ clampPathContributionFunc }( bg.a * bg.rgb * throughputColor * misWeight, input.currentBounce, clampDirect, clampIndirect );
+									resultColor = vec4f(
+										resultColor.rgb + background,
+										1.0 - transparency,
+									);
+
+								} else {
+
+									// fade the background by the throughput color average, mixing in env lighting
+									var light = mix( env.rgb, bg.rgb, bg.a ) * misWeight;
+									let background = ${ clampPathContributionFunc }( light * throughputColor, input.currentBounce, clampDirect, clampIndirect );
+									resultColor = vec4f(
+										resultColor.rgb + background,
+										1.0 - transparency,
+									);
+
+								}
+
+							}
+
+						}
+
+						isTerminated = true;
+
+					}
+
+				}
+
+				// the color rows are stored top down to match a rasterized render target
+				let colorIndex = vec2u( indexUV.x, textureDimensions( ${ params.outputTarget } ).y - 1u - indexUV.y );
+				let storedSamples = textureLoad( ${ params.sampleCountTarget }, indexUV ).r & ${ SAMPLE_COUNT_MASK }u;
+
+				if ( isTerminated ) {
+
+					// blend the finished sample into the output and free the slot for a new camera ray
+					let sampleCount = storedSamples + 1;
+					let prevColor = textureLoad( ${ params.prevOutputTarget }, colorIndex );
+					let blendedColor = ${ weightedAlphaBlendFn }( prevColor, resultColor, 1.0 / f32( sampleCount ) );
+					textureStore( ${ params.sampleCountTarget }, indexUV, vec4( ${ SAMPLE_DISPATCHED_FLAG }u | sampleCount ) );
+					textureStore( ${ params.outputTarget }, colorIndex, blendedColor );
+
+					rayDataStorage[ index ].objectIndex = - 1;
+
+				} else {
+
+					// show the path so far, which the first sample overwrites
+					if ( storedSamples == 0u ) {
+
+						textureStore( ${ params.outputTarget }, colorIndex, resultColor );
+
+					}
+
+					rayDataStorage[ index ].resultColor = resultColor;
+					rayDataStorage[ index ].throughputColor = throughputColor;
+
+				}
+
+			}
+		`;
+
+		super( fn( params ) );
+
+		this.defineUniformAccessors( params );
+
+	}
+
+}