three-gpu-pathtracer 0.0.26 → 0.0.28

package/src/webgpu/compute/ComputeKernel.js

@@ -46,6 +46,14 @@ export class ComputeKernel {
 			}
 
 			const node = parameters[ key ];
+
+			// name storage buffers so equivalent kernels generate identical WGSL and hit the pipeline cache
+			if ( node.isStorageBufferNode && ! node.name ) {
+
+				node.setName( `kernel_${ key }` );
+
+			}
+
 			if ( 'value' in node ) {
 
 				Object.defineProperty( this, key, {
@@ -73,11 +81,21 @@ export class ComputeKernel {
 
 		// forward the call parameters so they remain accessible via "computeNode.parameters"
 		node.parameters = this._fn.parameters;
+
+		// release the renderer's pipeline and bindings for the kernel being replaced
+		this.kernel?.dispose();
 		this.kernel = node.computeKernel( [ x, y, z ] );
 		return this;
 
 	}
 
+	dispose() {
+
+		this.kernel?.dispose();
+		this.kernel = null;
+
+	}
+
 	getDispatchSize( tx = 1, ty = 1, tz = 1, target = [] ) {
 
 		const [ wgx, wgy, wgz ] = this.workgroupSize;

package/src/webgpu/compute/PathTracerMegaKernel.js

@@ -1,9 +1,12 @@
-import { DataTexture, Vector2, StorageTexture } from 'three/webgpu';
+import { DataTexture, Vector2 } from 'three/webgpu';
 import { ComputeKernel } from './ComputeKernel.js';
+import { createAuxTarget, createOutputTarget, createSampleCountTarget } from './targets.js';
 import { texture, sampler, uniform, globalId, textureStore } from 'three/tsl';
-import { rngInit, rngNextBounce, rand1, rand2, rand3, RNG_INDEX_RAY_JITTER, RNG_INDEX_BACKGROUND_SAMPLE, RNG_INDEX_DIRECT_LIGHT_SAMPLE, RNG_INDEX_RUSSIAN_ROULETTE, RNG_INDEX_DISPERSION_WAVELENGTH, RNG_INDEX_ALPHA_TEST } from '../nodes/random.wgsl.js';
-import { misHeuristicFn, weightedAlphaBlendFn } from '../nodes/sampling.wgsl.js';
-import { proxy, proxyFn, wgslTagFn, rayStruct } from 'three-mesh-bvh/webgpu';
+import { blendAuxFn, isAuxSurfaceFn, AUX_LAYERS } from '../nodes/aux.wgsl.js';
+import { rngInit, rngNextBounce, rand1, rand2, rand3, RNG_INDEX_RAY_JITTER, RNG_INDEX_BACKGROUND_SAMPLE, RNG_INDEX_DIRECT_LIGHT_SAMPLE, RNG_INDEX_RUSSIAN_ROULETTE, RNG_INDEX_DISPERSION_WAVELENGTH, RNG_INDEX_ALPHA_TEST, RNG_INDEX_MEDIUM_DISTANCE, RNG_INDEX_MEDIUM_DIRECTION } from '../nodes/random.wgsl.js';
+import { misHeuristicFn, weightedAlphaBlendFn, sampleHenyeyGreensteinFn } from '../nodes/sampling.wgsl.js';
+import { mediumStruct, enterMediumFn, mediumChannelPdfFn, sampleMediumDistanceFn, mediumDistancePdfFn, MEDIUM_THROUGHPUT_EPSILON } from '../nodes/medium.wgsl.js';
+import { proxy, proxyFn, contextProxy, contextProxyFn, wgslTagFn, rayStruct } from 'three-mesh-bvh/webgpu';
 import { clampPathContributionFunc, isTerminatingScatterFunc, offsetRayOriginFunc } from '../nodes/utils.wgsl.js';
 import { lightRecordStruct } from '../nodes/structs.wgsl.js';
 import { ENVIRONMENT_LIGHT_TYPE, LIGHT_FAR_DISTANCE, LIGHT_EPSILON, isMISWeightLightFn } from '../nodes/lights.wgsl.js';
@@ -15,15 +18,14 @@ export class PathTracerMegaKernel extends ComputeKernel {
 	constructor( ) {
 
 		const params = {
-			bvhData: { value: null },
 			material: { value: null },
 			envInfo: { 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(),
+			prevOutputTarget: textureStore( createOutputTarget() ).toReadOnly(),
+			outputTarget: textureStore( createOutputTarget() ).toWriteOnly(),
+			sampleCountTarget: textureStore( createSampleCountTarget() ).toReadWrite(),
 
 			// tiles
 			offset: uniform( new Vector2() ),
@@ -33,6 +35,7 @@ export class PathTracerMegaKernel extends ComputeKernel {
 			seed: uniform( 0 ),
 			maxBounces: uniform( 5 ),
 			maxTransparentBounces: uniform( 5, 'uint' ),
+			maxSubsurfaceBounces: uniform( 256, 'uint' ),
 			misEnabled: uniform( 1, 'uint' ),
 			maxSamples: uniform( 0, 'uint' ),
 			filterGlossy: uniform( 1 ),
@@ -42,6 +45,10 @@ export class PathTracerMegaKernel extends ComputeKernel {
 			backgroundInfo: { value: null },
 
 			transmissiveBackground: uniform( TRANSMISSIVE_BACKGROUND_OVERLAY ),
+			accumulateAux: uniform( 0, 'uint' ),
+
+			// the per pixel first hit albedo and normal, as running means like the color
+			auxTarget: textureStore( createAuxTarget() ).toReadWrite(),
 
 			textures: texture( new DataTexture() ),
 			textureSampler: sampler( new DataTexture() ),
@@ -50,13 +57,16 @@ export class PathTracerMegaKernel extends ComputeKernel {
 			globalId: globalId,
 		};
 
-		const raycastOutput = proxy( 'bvhData.value.fns.raycastFirstHit.outputType', params );
-		const raycastFirstHitFn = proxyFn( 'bvhData.value.fns.raycastFirstHit', params );
-		const sampleTrianglePointFn = proxyFn( 'bvhData.value.fns.sampleTrianglePoint', params );
-		const getSurfaceRecordFn = proxyFn( 'bvhData.value.fns.getSurfaceRecord', params );
-		const getCameraRayFn = proxyFn( 'bvhData.value.fns.getCameraRay', params );
+		const raycastOutput = contextProxy( 'bvhData.fns.raycastFirstHit.outputType' );
+		const raycastFirstHitFn = contextProxyFn( 'bvhData.fns.raycastFirstHit' );
+		const raycastShadowFn = contextProxyFn( 'bvhData.fns.raycastShadow' );
+		const sampleTrianglePointFn = contextProxyFn( 'bvhData.fns.sampleTrianglePoint' );
+		const getSurfaceRecordFn = contextProxyFn( 'bvhData.fns.getSurfaceRecord' );
+		const getCameraRayFn = contextProxyFn( 'bvhData.fns.getCameraRay' );
 		const bsdfSampleFn = proxyFn( 'material.value.bsdfSample', params );
 		const bsdfEvalPdfFn = proxyFn( 'material.value.bsdfEvalPdf', params );
+		const materialsStorage = contextProxy( 'bvhData.storage.materials' );
+		const transformsStorage = contextProxy( 'bvhData.storage.transforms' );
 
 		// environment resources
 		const envTotalSumNode = proxy( 'envInfo.value.totalSumNode', params );
@@ -64,12 +74,17 @@ export class PathTracerMegaKernel extends ComputeKernel {
 		const sampleEnvColor = proxy( 'envInfo.value.sampleColor', params );
 		const sampleEnvDir = proxy( 'envInfo.value.sampleDir', params );
 		const getEnvDirPdf = proxy( 'envInfo.value.getDirPdf', params );
+		const getEnvWeight = proxy( 'envInfo.value.getWeight', 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 getLightWeightFn = proxyFn( 'lightsInfo.value.getLightWeight', params );
+		const getLightsWeightFn = proxyFn( 'lightsInfo.value.getLightsWeight', params );
+		const selectLightFn = proxyFn( 'lightsInfo.value.selectLight', params );
+		const isLightVisibleToCameraFn = proxyFn( 'lightsInfo.value.isLightVisibleToCamera', params );
 
 		const shader = wgslTagFn/* wgsl */`
 
@@ -86,6 +101,7 @@ export class PathTracerMegaKernel extends ComputeKernel {
 				seed: u32,
 				maxBounces: u32,
 				maxTransparentBounces: u32,
+				maxSubsurfaceBounces: u32,
 				misEnabled: u32,
 				maxSamples: u32,
 				filterGlossy: f32,
@@ -93,12 +109,10 @@ export class PathTracerMegaKernel extends ComputeKernel {
 				clampIndirect: f32,
 
 				transmissiveBackground: u32,
+				accumulateAux: u32,
 
 			) -> void {
 
-				let transforms = &${ proxy( 'bvhData.value.storage.transforms', params ) };
-				let materials = &${ proxy( 'bvhData.value.storage.materials', params ) };
-
 				// make sure we don't bleed over the edge of our tile
 				if ( globalId.x >= tileSize.x || globalId.y >= tileSize.y ) {
 
@@ -131,6 +145,13 @@ export class PathTracerMegaKernel extends ComputeKernel {
 				var ray: ${ rayStruct };
 				if ( ! ${ getCameraRayFn }( jitteredUv, vec2f( targetDimensions ), &ray ) ) {
 
+					// a declined pixel never takes a sample, so with a cap it counts as finished
+					if ( maxSamples != 0u ) {
+
+						textureStore( ${ params.sampleCountTarget }, indexUV, vec4( maxSamples ) );
+
+					}
+
 					return;
 
 				}
@@ -140,59 +161,154 @@ export class PathTracerMegaKernel extends ComputeKernel {
 				var resultColor = vec4f( 0, 0, 0, 1 );
 				var throughputColor = vec3f( 1.0 );
 				var bsdfPdf = 0.0;
+
+				// Alpha pass throughs only move where the BVH traversal starts, "rayStart" along the ray, so the
+				// ray keeps the origin it scattered from, which its MIS weights are measured from
+				var rayStart = 0.0;
+				var traceOrigin = ray.origin;
 				var isFullyTransmissive = true;
 				var minPdf = 1.0;
 				var transparentBounce = 0u;
+
+				// the scattering volume the path is in. Media don't nest
+				var medium: ${ mediumStruct };
+				var subsurfaceBounces = 0u;
+
+				// no light is sampled at a scatter event, so whatever the next segment reaches takes full weight
+				var mediumScattered = false;
+
 				// A negative value marks a sampled wavelength whose RGB reconstruction weight has not been applied yet.
 				var dispersionWavelength = - mix( ${ DISPERSION_MIN_WAVELENGTH }.0, ${ DISPERSION_MAX_WAVELENGTH }.0, ${ rand1 }( ${ RNG_INDEX_DISPERSION_WAVELENGTH } ) );
 
-				// one-sample NEE selects between the analytic lights and the environment -
-				// lightsDenom is the number of options
+				// one-sample NEE chooses between the analytic lights and the environment
 				let envActive = ${ envTotalSumNode } > 0.0 && ${ envIntensityNode } > 0.0;
 				let lightsCount = ${ lightsCountNode };
-				var lightsDenom = f32( lightsCount );
+				var envWeight = 0.0;
 				if ( envActive ) {
 
-					lightsDenom += 1.0;
+					envWeight = ${ getEnvWeight }();
 
 				}
 
+				// added last, since a camera ray miss overwrites the pixel with the background
+				var lightHits = vec3f( 0.0 );
+				var cameraHitLight = false;
+
+				// The first surface the path scatters off fills the denoiser's auxiliary images. Until one is found
+				// the last mirror or transmission hit stands in, for paths that reach the environment. The albedo
+				// is tinted by the throughput of the surfaces looked past to reach it
+				var auxFound = false;
+				var auxAlbedo = vec3f( 0.0 );
+				var auxNormal = vec3f( 0.0 );
+
 				for ( var bounce = 0u; bounce < maxBounces; bounce ++ ) {
 
+					// Inside a medium, sample a scatter event before the next surface and only trace up to it
+					var channelPdf = vec3f( 0.0 );
+					var scatterDist = 0.0;
+					var traceRay = ray;
+					traceRay.origin = traceOrigin;
+					if ( medium.inside ) {
+
+						channelPdf = ${ mediumChannelPdfFn }( throughputColor, medium.albedo );
+						scatterDist = ${ sampleMediumDistanceFn }( medium, channelPdf, ${ rand2 }( ${ RNG_INDEX_MEDIUM_DISTANCE } ) );
+						traceRay.maxDist = select( scatterDist, min( ray.maxDist, scatterDist ), ray.maxDist > 0.0 );
+
+					}
+
 					var hitResult: ${ raycastOutput };
-					let didHit = ${ raycastFirstHitFn }( ray, &hitResult );
+					let didHit = ${ raycastFirstHitFn }( traceRay, &hitResult );
 					let surfaceDist = select( ${ LIGHT_FAR_DISTANCE }, hitResult.dist, didHit );
 
-					// forward hits: a bsdf-sampled ray that lands on a area light. MIS-weighted
-					// only when NEE is also sampling the lights
-					if ( bounce > 0u ) {
+					let segmentInMedium = medium.inside;
+					if ( medium.inside ) {
 
-						for ( var li = 0u; li < lightsCount; li ++ ) {
+						let scattered = scatterDist < surfaceDist;
+						let segmentDist = select( surfaceDist, scatterDist, scattered );
+						let transmittance = exp( - medium.extinction * segmentDist );
+						let distancePdf = dot( channelPdf, ${ mediumDistancePdfFn }( medium, segmentDist, scattered ) );
+						if ( distancePdf <= 0.0 ) {
 
-							var lightRec: ${ lightRecordStruct };
-							if ( ${ intersectLightAtIndexFn }( ray.origin, ray.direction, li, &lightRec ) && lightRec.dist < surfaceDist ) {
+							break;
 
-								var misWeight = 1.0;
-								if ( misEnabled != 0u ) {
+						}
 
-									let lightPdf = lightRec.pdf / lightsDenom;
-									misWeight = ${ misHeuristicFn }( bsdfPdf, lightPdf );
+						if ( ! scattered ) {
 
-								}
+							// reaching the surface, weighted by the probability of no scatter event before it
+							throughputColor *= transmittance / distancePdf;
+
+						} else {
 
-								let lightHit = ${ clampPathContributionFunc }( lightRec.emission * throughputColor * misWeight, bounce, clampDirect, clampIndirect );
-								resultColor += vec4f( lightHit, 0.0 );
+							throughputColor *= medium.albedo * medium.extinction * transmittance / distancePdf;
+							ray.origin = traceOrigin + ray.direction * scatterDist;
+							traceOrigin = ray.origin;
+							rayStart = 0.0;
+							ray.direction = ${ sampleHenyeyGreensteinFn }( ray.direction, medium.anisotropy, ${ rand2 }( ${ RNG_INDEX_MEDIUM_DIRECTION } ) );
+							ray.maxDist = 0.0;
+							mediumScattered = true;
+							isFullyTransmissive = false;
+
+							// subsurface bounces are capped separately from bounces
+							subsurfaceBounces ++;
+							let maxThroughput = max( max( throughputColor.r, throughputColor.g ), throughputColor.b );
+							if ( subsurfaceBounces >= maxSubsurfaceBounces || maxThroughput < ${ MEDIUM_THROUGHPUT_EPSILON } ) {
+
+								break;
 
 							}
 
+							${ rngNextBounce }();
+							bounce --;
+							continue;
+
+						}
+
+					}
+
+					// NEE's total weight at the vertex this ray left, for MIS. Camera rays have none, and
+					// neither do scatter events in a medium
+					let isMISWeighted = misEnabled != 0u && bounce > 0u && ! mediumScattered;
+					var lightTotalWeight = 0.0;
+					if ( isMISWeighted ) {
+
+						lightTotalWeight = ${ getLightsWeightFn }( ray.origin ) + envWeight;
+
+					}
+
+					// forward hits on area lights along the traced segment. Camera rays only see visibleToCameraRays
+					// lights, at full weight
+					for ( var li = 0u; li < lightsCount; li ++ ) {
+
+						if ( bounce == 0u && ! ${ isLightVisibleToCameraFn }( li ) ) {
+
+							continue;
+
+						}
+
+						var lightRec: ${ lightRecordStruct };
+						if ( ${ intersectLightAtIndexFn }( ray.origin, ray.direction, li, &lightRec ) && lightRec.dist > rayStart && lightRec.dist < rayStart + surfaceDist ) {
+
+							var misWeight = 1.0;
+							if ( isMISWeighted ) {
+
+								let selectionPdf = ${ getLightWeightFn }( li, ray.origin ) / lightTotalWeight;
+								misWeight = ${ misHeuristicFn }( bsdfPdf, lightRec.pdf * selectionPdf );
+
+							}
+
+							lightHits += ${ clampPathContributionFunc }( lightRec.emission * throughputColor * misWeight, bounce, clampDirect, clampIndirect );
+							cameraHitLight = cameraHitLight || bounce == 0u;
+
 						}
 
 					}
 
 					if ( didHit ) {
 
-						let objectInfo = transforms[ hitResult.objectIndex ];
-						var materialInfo = ( *materials )[ objectInfo.materialIndex ];
+						// whole structs are read from the buffer, not a pointer alias, which Safari fails to compile
+						let objectInfo = ${ transformsStorage }[ hitResult.objectIndex ];
+						var materialInfo = ${ materialsStorage }[ objectInfo.materialIndex ];
 
 						// a matte surface hit by the camera ray renders as a fully transparent
 						if ( materialInfo.matte != 0 && bounce == 0u ) {
@@ -220,7 +336,7 @@ export class PathTracerMegaKernel extends ComputeKernel {
 						var surface = ${ getSurfaceRecordFn }( materialInfo, vertexData, hitResult.side, hitResult.normal, view, blurRoughness );
 
 						// Stochastically pass through partially transparent surfaces by restarting
-						// the ray at the hit point, advancing the rng but not the bounce count.
+						// the traversal at the hit point, advancing the rng but not the bounce count.
 						if ( ${ rand1 }( ${ RNG_INDEX_ALPHA_TEST } ) > surface.opacity ) {
 
 							// stop once the transparent bounces run out
@@ -230,7 +346,8 @@ export class PathTracerMegaKernel extends ComputeKernel {
 
 							}
 
-							ray.origin = ${ offsetRayOriginFunc }( vertexData.position.xyz, ray.direction, hitResult.normal );
+							traceOrigin = ${ offsetRayOriginFunc }( vertexData.position.xyz, ray.direction, hitResult.normal );
+							rayStart += hitResult.dist;
 							${ rngNextBounce }();
 							transparentBounce ++;
 							bounce --;
@@ -252,8 +369,9 @@ export class PathTracerMegaKernel extends ComputeKernel {
 
 						}
 
-						// attenuate the light transmitted through the volume when exiting a backface
-						if ( hitResult.side < 0.0 && materialInfo.transmission > 0.0 ) {
+						// attenuate the light transmitted through the volume when exiting a backface. A scattering
+						// medium has already applied its transmittance above
+						if ( hitResult.side < 0.0 && materialInfo.transmission > 0.0 && ! segmentInMedium ) {
 
 							throughputColor *= ${ transmissionAttenuationFunc }( hitResult.dist, materialInfo.attenuationColor, materialInfo.attenuationDistance );
 
@@ -263,30 +381,31 @@ export class PathTracerMegaKernel extends ComputeKernel {
 						let emission = ${ clampPathContributionFunc }( throughputColor * surface.emission, bounce, clampDirect, clampIndirect );
 						resultColor += vec4f( emission, 0.0 );
 
-						// next event estimation: draw one light or the environment, each with
-						// probability 1 / lightsDenom
-						if ( misEnabled != 0u && lightsDenom > 0.0 ) {
+						// next event estimation: choose one light or the environment by its estimated
+						// contribution
+						if ( misEnabled != 0u ) {
 
-							// pick one light or the environment with a single sample
-							// TODO: importance-sample the selection by light intensity and solid angle
 							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 };
+							var selectionPdf = 0.0;
+							let lightIndex = ${ selectLightFn }( vertexData.position.xyz, envWeight, ruv.x, &selectionPdf );
+							if ( selectionPdf > 0.0 ) {
 
-							} else {
+								var lightRec: ${ lightRecordStruct };
+								if ( envActive && lightIndex == lightsCount ) {
 
-								lightRec = ${ randomLightSampleFn }( lightIndex, vertexData.position.xyz, ruv.yz );
+									// 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, vertexData.position.xyz, ruv.yz );
+
+								}
 
 								if ( lightRec.pdf > 0.0 ) {
 
@@ -296,21 +415,22 @@ export class PathTracerMegaKernel extends ComputeKernel {
 										var shadowRay: ${ rayStruct };
 										shadowRay.origin = ${ offsetRayOriginFunc }( vertexData.position.xyz, lightRec.direction, hitResult.normal );
 										shadowRay.direction = lightRec.direction;
-										shadowRay.maxDist = lightRec.dist - ${ LIGHT_EPSILON };
+										shadowRay.maxDist = lightRec.dist * ( 1.0 - ${ LIGHT_EPSILON } );
 
-										// opaque occlusion up to the light distance. A shadow-specific any hit
-										// traversal could support tinted shadows from transmissive and partially
-										// opaque objects
+										// opaque occlusion up to the light distance, attenuated by the partially
+										// transparent surfaces passed
 										var shadowHit: ${ raycastOutput };
-										let occluded = ${ raycastFirstHitFn }( shadowRay, &shadowHit );
+										let occluded = ${ raycastShadowFn }( shadowRay, &shadowHit );
 										if ( ! occluded ) {
 
-											var lightPdf = lightRec.pdf;
-											lightPdf /= lightsDenom;
+											let lightPdf = lightRec.pdf * selectionPdf;
 
-											// env + area lights are also bsdf-sampled, so MIS-weight them - punctual lights take full weight
-											let misWeight = select( 1.0, ${ misHeuristicFn }( lightPdf, evalRec.pdf ), ${ isMISWeightLightFn }( lightRec.lightType ) );
-											let directLight = throughputColor * lightRec.emission * evalRec.color * misWeight / lightPdf;
+											// env + area lights are also found by bsdf rays, so MIS-weight them. Punctual lights, and lights
+											// past the bounce limit, take full weight. Bsdf rays stopped by faces the shadow ray skipped can't
+											// find the light, so their share of the weight moves to the light sample
+											let isMISWeighted = ${ isMISWeightLightFn }( lightRec.lightType ) && bounce + 1u < maxBounces;
+											let misWeight = select( 1.0, mix( ${ misHeuristicFn }( lightPdf, evalRec.pdf ), 1.0, shadowHit.skippedOcclusion ), isMISWeighted );
+											let directLight = throughputColor * lightRec.emission * evalRec.color * misWeight * ( 1.0 - shadowHit.occlusion ) / lightPdf;
 											let contribution = ${ clampPathContributionFunc }( directLight, bounce + 1u, clampDirect, clampIndirect );
 											resultColor += vec4f( contribution, 0.0 );
 
@@ -320,9 +440,20 @@ export class PathTracerMegaKernel extends ComputeKernel {
 
 								}
 
+							}
+
 						}
 
 						let scatterRec = ${ bsdfSampleFn }( view, surface );
+
+						if ( ! auxFound ) {
+
+							auxFound = ${ isAuxSurfaceFn }( surface, scatterRec );
+							auxAlbedo = saturate( throughputColor * surface.color );
+							auxNormal = surface.normal;
+
+						}
+
 						if ( ${ isTerminatingScatterFunc }( scatterRec ) ) {
 
 							break;
@@ -353,6 +484,28 @@ export class PathTracerMegaKernel extends ComputeKernel {
 						throughputColor *= scatterRec.color;
 						throughputColor /= scatterRec.pdf;
 						bsdfPdf = scatterRec.pdf;
+						mediumScattered = false;
+
+						// As in "KHR_materials_scatter", the medium's surface bsdf is its boundary: the walk ends when
+						// light refracts out through it, and continues inside when it reflects back
+						if ( segmentInMedium && hitResult.objectIndex == medium.object && hitResult.side < 0.0 && scatterRec.isTransmissive ) {
+
+							medium.inside = false;
+
+						}
+
+						// refracting into a front face of a scattering material starts a random walk inside it
+						if ( scatterRec.isTransmissive && ! surface.thinWall && hitResult.side > 0.0 ) {
+
+							let entered = ${ enterMediumFn }( materialInfo, hitResult.objectIndex );
+
+							if ( entered.inside ) {
+
+								medium = entered;
+
+							}
+
+						}
 
 						// exit if our throughput is 0.0
 						if ( all( throughputColor == vec3f( 0.0 ) ) ) {
@@ -364,16 +517,18 @@ export class PathTracerMegaKernel extends ComputeKernel {
 						ray.origin = ${ offsetRayOriginFunc }( vertexData.position.xyz, scatterRec.direction, hitResult.normal );
 						ray.direction = scatterRec.direction;
 						ray.maxDist = 0.0;
+						traceOrigin = ray.origin;
+						rayStart = 0.0;
 
 					} else {
 
 						if ( bounce > 0u && ! isFullyTransmissive ) {
 
 							var misWeight = 1.0;
-							if ( misEnabled != 0u && envActive ) {
+							if ( isMISWeighted && envActive ) {
 
 								// match the env pdf scaling used by the NEE selection so the two estimators balance
-								let envPdf = ${ getEnvDirPdf }( ray.direction ) / lightsDenom;
+								let envPdf = ${ getEnvDirPdf }( ray.direction ) * envWeight / lightTotalWeight;
 								misWeight = ${ misHeuristicFn }( bsdfPdf, envPdf );
 
 							}
@@ -450,6 +605,14 @@ export class PathTracerMegaKernel extends ComputeKernel {
 
 				}
 
+				// visible lights are opaque against a transparent background
+				resultColor += vec4f( lightHits, 0.0 );
+				if ( cameraHitLight ) {
+
+					resultColor.a = 1.0;
+
+				}
+
 				let nextSampleCount = sampleCount + 1;
 				// store the color rows top down to match a rasterized render target
 				let colorIndex = vec2u( indexUV.x, targetDimensions.y - 1u - indexUV.y );
@@ -459,6 +622,25 @@ export class PathTracerMegaKernel extends ComputeKernel {
 				textureStore( ${ params.sampleCountTarget }, indexUV, vec4( nextSampleCount ) );
 				textureStore( ${ params.outputTarget }, colorIndex, blendedColor );
 
+				// the auxiliary images are blended the same way, in the same row order
+				if ( accumulateAux != 0u ) {
+
+					var auxPacked: vec3u;
+					for ( var layer = 0u; layer < ${ AUX_LAYERS }u; layer ++ ) {
+
+						auxPacked[ layer ] = textureLoad( ${ params.auxTarget }, colorIndex, layer ).r;
+
+					}
+
+					auxPacked = ${ blendAuxFn }( auxPacked, auxAlbedo, auxNormal, 1.0 / f32( nextSampleCount ) );
+					for ( var layer = 0u; layer < ${ AUX_LAYERS }u; layer ++ ) {
+
+						textureStore( ${ params.auxTarget }, colorIndex, layer, vec4u( auxPacked[ layer ] ) );
+
+					}
+
+				}
+
 			}`;
 
 		super( shader( params ) );

package/src/webgpu/compute/TallySampleCountsKernel.js

@@ -37,11 +37,11 @@ export class PrimeSampleCountersKernel extends ComputeKernel {
 		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 );
+				atomicStore( &${ params.counters }[ 0 ].minSamples, 0xFFFFFFFFu );
+				atomicStore( &${ params.counters }[ 0 ].maxSamples, 0u );
+				atomicStore( &${ params.counters }[ 0 ].totalLo, 0u );
+				atomicStore( &${ params.counters }[ 0 ].totalHi, 0u );
+				atomicStore( &${ params.counters }[ 0 ].pixelCount, 0u );
 
 			}
 		`;
@@ -86,16 +86,16 @@ export class TallySampleCountsKernel extends ComputeKernel {
 
 				let samples = combinedField & ${ SAMPLE_COUNT_MASK }u;
 
-				atomicMin( &${ params.counters }.minSamples, samples );
-				atomicMax( &${ params.counters }.maxSamples, samples );
-				atomicAdd( &${ params.counters }.pixelCount, 1u );
+				atomicMin( &${ params.counters }[ 0 ].minSamples, samples );
+				atomicMax( &${ params.counters }[ 0 ].maxSamples, samples );
+				atomicAdd( &${ params.counters }[ 0 ].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 );
+				let previousTotal = atomicAdd( &${ params.counters }[ 0 ].totalLo, samples );
 				if ( previousTotal + samples < previousTotal ) {
 
-					atomicAdd( &${ params.counters }.totalHi, 1u );
+					atomicAdd( &${ params.counters }[ 0 ].totalHi, 1u );
 
 				}
 

package/src/webgpu/compute/UnpackAuxKernel.js

@@ -0,0 +1,48 @@
+import { ComputeKernel } from './ComputeKernel.js';
+import { textureStore, globalId } from 'three/tsl';
+import { wgslTagFn } from 'three-mesh-bvh/webgpu';
+import { createAuxImageTarget, createAuxTarget } from './targets.js';
+import { unpackAuxAlbedoFn, unpackAuxNormalFn, AUX_LAYERS } from '../nodes/aux.wgsl.js';
+
+// Unpacks the auxiliary target into separate albedo and normal images
+export class UnpackAuxKernel extends ComputeKernel {
+
+	constructor() {
+
+		const params = {
+			auxTarget: textureStore( createAuxTarget() ).toReadOnly(),
+			albedoTarget: textureStore( createAuxImageTarget() ).toWriteOnly(),
+			normalTarget: textureStore( createAuxImageTarget() ).toWriteOnly(),
+			globalId: globalId,
+		};
+
+		const fn = wgslTagFn/* wgsl */`
+			fn compute( globalId: vec3u ) -> void {
+
+				let dims = textureDimensions( ${ params.auxTarget } );
+				if ( globalId.x >= dims.x || globalId.y >= dims.y ) {
+
+					return;
+
+				}
+
+				var packed: vec3u;
+				for ( var layer = 0u; layer < ${ AUX_LAYERS }u; layer ++ ) {
+
+					packed[ layer ] = textureLoad( ${ params.auxTarget }, globalId.xy, layer ).r;
+
+				}
+
+				textureStore( ${ params.albedoTarget }, globalId.xy, vec4f( ${ unpackAuxAlbedoFn }( packed ), 1.0 ) );
+				textureStore( ${ params.normalTarget }, globalId.xy, vec4f( ${ unpackAuxNormalFn }( packed ), 1.0 ) );
+
+			}
+		`;
+
+		super( fn( params ) );
+
+		this.defineUniformAccessors( params );
+
+	}
+
+}

package/src/webgpu/compute/ZeroOutKernel.js

@@ -12,10 +12,12 @@ export class ZeroOutKernel extends ComputeKernel {
 			outputTarget: textureStore( new StorageTexture( 1, 1 ) ).toWriteOnly(),
 		};
 
+		// resolved when the shader is built, so an array texture target is cleared a layer per z
+		const coord = () => params.outputTarget.value.isArrayTexture ? 'globalId.xy, globalId.z' : 'globalId.xy';
 		const fn = wgslTagFn/* wgsl */`
 			fn compute( globalId: vec3u ) -> void {
 
-				textureStore( ${ params.outputTarget }, globalId.xy, vec4( 0, 0, 0, 1 ) );
+				textureStore( ${ params.outputTarget }, ${ coord }, vec4( 0, 0, 0, 1 ) );
 
 			}
 		`;