WebGPUPathTracer: Add support for opacity during BVH traversal

src/webgpu/WaveFrontPathTracer.js

@@ -29,13 +29,13 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 		// queues
 		this.rayQueue = new IndirectStorageBufferAttribute( MAX_RAY_COUNT, queuedRayStruct.getLength() );
 		this.rayQueue.name = 'Ray Queue';
-		this.rayQueueSize = new IndirectStorageBufferAttribute( 2, 1 );
-		this.rayQueueSize.name = 'Ray Queue Size';
 
 		this.hitQueue = new IndirectStorageBufferAttribute( MAX_HIT_COUNT, queuedHitStruct.getLength() );
 		this.hitQueue.name = 'Hit Queue';
-		this.hitQueueSize = new IndirectStorageBufferAttribute( 2, 1 );
-		this.hitQueueSize.name = 'Hit Queue Size';
+
+		// [0] ray head, [1] ray count, [2] hit head, [3] hit count
+		this.queueSizes = new IndirectStorageBufferAttribute( 4, 1 );
+		this.queueSizes.name = 'Queue Sizes';
 
 		// dispatches
 		this.tileIndexBuffer = new IndirectStorageBufferAttribute( 2, 1 );
@@ -160,8 +160,7 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 
 			hitProcessDispatch,
 			rayGenerationDispatch,
-			rayQueueSize,
-			hitQueueSize,
+			queueSizes,
 			tileIndexBuffer,
 		} = this;
 
@@ -182,11 +181,8 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 
 		// clear queues
 		// TODO: why do we need to se the work group size here?
-		zeroDispatchKernel.target = rayQueueSize;
-		renderer.compute( zeroDispatchKernel.kernel, [ rayQueueSize.count ] );
-
-		zeroDispatchKernel.target = hitQueueSize;
-		renderer.compute( zeroDispatchKernel.kernel, [ hitQueueSize.count ] );
+		zeroDispatchKernel.target = queueSizes;
+		renderer.compute( zeroDispatchKernel.kernel, [ queueSizes.count ] );
 
 		// clear dispatch sizes
 		zeroDispatchKernel.target = hitProcessDispatch;
@@ -211,11 +207,9 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 			sampleCountTarget,
 
 			rayQueue,
-			rayQueueSize,
-			rayGenerationDispatch,
-
 			hitQueue,
-			hitQueueSize,
+			queueSizes,
+			rayGenerationDispatch,
 			// hitProcessDispatch,
 			tileIndexBuffer,
 
@@ -224,7 +218,6 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 			rayIntersectionKernel,
 			updateRayQueueParamsKernel,
 			hitProcessKernel,
-			zeroDispatchKernel,
 
 			lowResMode
 		} = this;
@@ -259,7 +252,7 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 				primeRayGenerationDispatchKernel.tileCount.copy( tiles );
 				primeRayGenerationDispatchKernel.tileSize.copy( tileSize );
 				primeRayGenerationDispatchKernel.rayQueue = rayQueue;
-				primeRayGenerationDispatchKernel.rayQueueSize = rayQueueSize;
+				primeRayGenerationDispatchKernel.queueSizes = queueSizes;
 				primeRayGenerationDispatchKernel.outputTileIndex = tileIndexBuffer;
 				primeRayGenerationDispatchKernel.outputDispatch = rayGenerationDispatch;
 
@@ -270,7 +263,7 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 				enqueueRaysKernel.tileIndexBuffer = tileIndexBuffer;
 				enqueueRaysKernel.tileSize.copy( tileSize );
 				enqueueRaysKernel.rayQueue = rayQueue;
-				enqueueRaysKernel.rayQueueSize = rayQueueSize;
+				enqueueRaysKernel.queueSizes = queueSizes;
 				enqueueRaysKernel.sampleCountTarget = sampleCountTarget;
 
 				for ( let i = 0; i < tiles.x * tiles.y; i ++ ) {
@@ -287,15 +280,14 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 				const intersectDispatch = rayIntersectionKernel.getDispatchSize( RAYS_TO_PROCESS, 1, 1 );
 				rayIntersectionKernel.sampleCountTarget = sampleCountTarget;
 				rayIntersectionKernel.rayQueue = rayQueue;
-				rayIntersectionKernel.rayQueueSize = rayQueueSize;
+				rayIntersectionKernel.queueSizes = queueSizes;
 				rayIntersectionKernel.hitQueue = hitQueue;
-				rayIntersectionKernel.hitQueueSize = hitQueueSize;
 				renderer.compute( rayIntersectionKernel.kernel, intersectDispatch );
 
 				// mark the rays as consumed
 				const processed = intersectDispatch[ 0 ] * rayIntersectionKernel.workgroupSize[ 0 ];
 				updateRayQueueParamsKernel.processed = processed;
-				updateRayQueueParamsKernel.rayQueueSize = rayQueueSize;
+				updateRayQueueParamsKernel.queueSizes = queueSizes;
 				renderer.compute( updateRayQueueParamsKernel.kernel, [ 1, 1, 1 ] );
 
 				// TODO: we should use an indirect dispatch here to only kick off the number of threads
@@ -304,15 +296,10 @@ export class WaveFrontPathTracer extends PathTracerBackend {
 				hitProcessKernel.sampleCountTarget = sampleCountTarget;
 				hitProcessKernel.bounces = bounces;
 				hitProcessKernel.rayQueue = rayQueue;
-				hitProcessKernel.rayQueueSize = rayQueueSize;
+				hitProcessKernel.queueSizes = queueSizes;
 				hitProcessKernel.hitQueue = hitQueue;
-				hitProcessKernel.hitQueueSize = hitQueueSize;
 				renderer.compute( hitProcessKernel.kernel, hitProcessKernel.getDispatchSize( processed, 1, 1 ) );
-
-				// TODO: for some reason we need to call "setWorkgroupSize" here? Is it because work group size
-				// is cached per parameters and resets?
-				zeroDispatchKernel.target = hitQueueSize;
-				renderer.compute( zeroDispatchKernel.kernel, [ hitQueueSize.count, 1, 1 ] );
+				// Note: hit queue size ([2] and [3]) is reset at the top of the next iteration by PrimeRayGenerationDispatchKernel
 
 				// Step 4: connect to lights
 				// TODO

src/webgpu/WebGPUPathTracer.js

@@ -133,6 +133,7 @@ export class WebGPUPathTracer {
 		// Build TLAS and compute functions
 		const bvhData = new PathtracerBVHComputeData( scene );
 		bvhData.update();
+		bvhData.useTransparencyRaycastFn();
 
 		this.textureArray.setTextures( this._renderer, bvhData.textures );
 		this._pathTracer.setTextures( this.textureArray.texture );

src/webgpu/compute/wavefront/PrimeRayGenerationDispatchKernel.js

@@ -17,7 +17,7 @@ export class PrimeRayGenerationDispatchKernel extends ComputeKernel {
 			tileOffset: uniform( 1 ),
 
 			rayQueue: storage( new IndirectStorageBufferAttribute( 1, queuedRayStruct.getLength() ), queuedRayStruct ).toReadOnly(),
-			rayQueueSize: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ),
+			queueSizes: storage( new IndirectStorageBufferAttribute( 4, 1 ), 'u32' ),
 
 			outputTileIndex: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ).setName( 'outputTileIndex' ),
 			outputDispatch: storage( new IndirectStorageBufferAttribute( 3, 1 ), 'u32' ).setName( 'outputDispatch' ),
@@ -33,24 +33,28 @@ export class PrimeRayGenerationDispatchKernel extends ComputeKernel {
 			) -> void {
 
 				let rayQueue = &${ params.rayQueue };
-				let rayQueueSize = &${ params.rayQueueSize };
+				let queueSizes = &${ params.queueSizes };
 
 				let outputTileIndex = &${ params.outputTileIndex };
 				let outputDispatch = &${ params.outputDispatch };
 
+				// reset hit queue size from previous iteration
+				queueSizes[ 2 ] = 0u;
+				queueSizes[ 3 ] = 0u;
+
 				// keep the queue index small
 			    let queueCapacity = arrayLength( rayQueue );
-				if ( rayQueueSize[ 0 ] >= queueCapacity ) {
+				if ( queueSizes[ 0 ] >= queueCapacity ) {
 
 					// uint division results in a floored value
-					let offset = rayQueueSize[ 0 ] / queueCapacity;
-					rayQueueSize[ 0 ] = rayQueueSize[ 0 ] - queueCapacity * offset;
-					rayQueueSize[ 1 ] = rayQueueSize[ 1 ] - queueCapacity * offset;
+					let offset = queueSizes[ 0 ] / queueCapacity;
+					queueSizes[ 0 ] = queueSizes[ 0 ] - queueCapacity * offset;
+					queueSizes[ 1 ] = queueSizes[ 1 ] - queueCapacity * offset;
 
 				}
 
 				// calculate the amount of elements in the queue
-			    var queueSize = rayQueueSize[ 1 ] - rayQueueSize[ 0 ];
+			    var queueSize = queueSizes[ 1 ] - queueSizes[ 0 ];
 
 				// calculate the overhead of space in the queue and how much space we need to run a new tile
 				let overhead = queueCapacity - queueSize;

src/webgpu/compute/wavefront/ProcessHitsKernel.js

@@ -24,10 +24,8 @@ export class ProcessHitsKernel extends ComputeKernel {
 
 			// rays
 			rayQueue: storage( new IndirectStorageBufferAttribute( 1, queuedRayStruct.getLength() ), queuedRayStruct ),
-			rayQueueSize: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ).toAtomic(),
-
 			hitQueue: storage( new IndirectStorageBufferAttribute( 1, queuedHitStruct.getLength() ), queuedHitStruct ),
-			hitQueueSize: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ),
+			queueSizes: storage( new IndirectStorageBufferAttribute( 4, 1 ), 'u32' ).toAtomic(),
 
 			textures: texture( new DataTexture( ) ),
 			textureSampler: sampler( new DataTexture( ) ),
@@ -51,18 +49,16 @@ export class ProcessHitsKernel extends ComputeKernel {
 			) -> void {
 
 				let rayQueue = &${ params.rayQueue };
-				let rayQueueSize = &${ params.rayQueueSize };
-
 				let hitQueue = &${ params.hitQueue };
-				let hitQueueSize = &${ params.hitQueueSize };
+				let queueSizes = &${ params.queueSizes };
 
 				let materials = &${ proxy( 'bvhData.value.storage.materials', params ) };
 				let transforms = &${ proxy( 'bvhData.value.storage.transforms', params ) };
 
 				// skip any rays invocations beyond the ray count
 				let hitQueueCapacity = arrayLength( hitQueue );
-				let hitIndex = ( globalId.x + hitQueueSize[ 0 ] );
-				if ( hitIndex >= hitQueueSize[ 1 ] ) {
+				let hitIndex = ( globalId.x + atomicLoad( &queueSizes[ 2 ] ) );
+				if ( hitIndex >= atomicLoad( &queueSizes[ 3 ] ) ) {
 
 					return;
 
@@ -102,7 +98,7 @@ export class ProcessHitsKernel extends ComputeKernel {
 				} else {
 
 					let rayQueueCapacity = arrayLength( rayQueue );
-					let index = atomicAdd( &rayQueueSize[ 1 ], 1 ) % rayQueueCapacity;
+					let index = atomicAdd( &queueSizes[ 1 ], 1 ) % rayQueueCapacity;
 					let resultColor = input.resultColor + vec4f( input.throughputColor * surface.emission, 0.0 );
 					rayQueue[ index ].origin = vertexData.position.xyz;
 					rayQueue[ index ].direction = scatterRec.direction;

src/webgpu/compute/wavefront/RayGenerationKernel.js

@@ -21,7 +21,7 @@ export class RayGenerationKernel extends ComputeKernel {
 			tileSize: uniform( new Vector2() ),
 
 			rayQueue: storage( new IndirectStorageBufferAttribute( 1, queuedRayStruct.getLength() ), queuedRayStruct ),
-			rayQueueSize: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ).toAtomic(),
+			queueSizes: storage( new IndirectStorageBufferAttribute( 4, 1 ), 'u32' ).toAtomic(),
 
 			sampleCountTarget: textureStore( new StorageTexture() ).toReadWrite(),
 
@@ -40,7 +40,7 @@ export class RayGenerationKernel extends ComputeKernel {
 			) -> void {
 
 				let rayQueue = &${ params.rayQueue };
-				let rayQueueSize = &${ params.rayQueueSize };
+				let queueSizes = &${ params.queueSizes };
 				let tileIndexBuffer = &${ params.tileIndexBuffer };
 
 				// don't overstep the edge of the tile
@@ -78,7 +78,7 @@ export class RayGenerationKernel extends ComputeKernel {
 
 				// get the ray index
 				let queueCapacity = arrayLength( rayQueue );
-				let index = atomicAdd( &rayQueueSize[ 1 ], 1 ) % queueCapacity;
+				let index = atomicAdd( &queueSizes[ 1 ], 1 ) % queueCapacity;
 
 				${ pcgInit }( indexUV, seed );
 

src/webgpu/compute/wavefront/RayIntersectionKernel.js

@@ -20,10 +20,8 @@ export class RayIntersectionKernel extends ComputeKernel {
 
 			// rays
 			rayQueue: storage( new IndirectStorageBufferAttribute( 1, queuedRayStruct.getLength() ), queuedRayStruct ).toReadOnly(),
-			rayQueueSize: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ).toReadOnly(),
-
 			hitQueue: storage( new IndirectStorageBufferAttribute( 1, queuedHitStruct.getLength() ), queuedHitStruct ),
-			hitQueueSize: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ).toAtomic(),
+			queueSizes: storage( new IndirectStorageBufferAttribute( 4, 1 ), 'u32' ).toAtomic(),
 
 			// environment
 			envMap: texture( new DataTexture() ),
@@ -62,10 +60,8 @@ export class RayIntersectionKernel extends ComputeKernel {
 			) -> void {
 
 				let rayQueue = &${ params.rayQueue };
-				let rayQueueSize = &${ params.rayQueueSize };
-
 				let hitQueue = &${ params.hitQueue };
-				let hitQueueSize = &${ params.hitQueueSize };
+				let queueSizes = &${ params.queueSizes };
 
 				let envInfo = EnvironmentInfo(
 					envMapRotation,
@@ -81,8 +77,8 @@ export class RayIntersectionKernel extends ComputeKernel {
 
 				// skip any rays invocations beyond the ray count
 				let queueCapacity = arrayLength( rayQueue );
-				let rayIndex = ( globalId.x + rayQueueSize[ 0 ] );
-				if ( rayIndex >= rayQueueSize[ 1 ] ) {
+				let rayIndex = ( globalId.x + atomicLoad( &queueSizes[ 0 ] ) );
+				if ( rayIndex >= atomicLoad( &queueSizes[ 1 ] ) ) {
 
 					return;
 
@@ -102,7 +98,7 @@ export class RayIntersectionKernel extends ComputeKernel {
 				if ( ${ raycastFirstHitFn }( ray, &hitResult ) ) {
 
 					// TODO: we process all of these materials immediately to push to the ray queue
-					let index = atomicAdd( &hitQueueSize[ 1 ], 1 );
+					let index = atomicAdd( &queueSizes[ 3 ], 1 );
 					hitQueue[ index ].view = - input.direction;
 					hitQueue[ index ].indices = hitResult.indices.xyz;
 					hitQueue[ index ].barycoord = hitResult.barycoord;

src/webgpu/compute/wavefront/UpdateRayQueueParamsKernel.js

@@ -9,21 +9,21 @@ export class UpdateRayQueueParamsKernel extends ComputeKernel {
 
 		const params = {
 			processed: uniform( 0 ),
-			rayQueueSize: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ),
+			queueSizes: storage( new IndirectStorageBufferAttribute( 4, 1 ), 'u32' ),
 		};
 
 		const fn = wgslTagFn/* wgsl */`
 			fn compute( processed: u32 ) -> void {
 
-				let rayQueueSize = &${ params.rayQueueSize };
-			    var queueSize = rayQueueSize[ 1 ] - rayQueueSize[ 0 ];
+				let queueSizes = &${ params.queueSizes };
+			    var queueSize = queueSizes[ 1 ] - queueSizes[ 0 ];
 				if ( processed > queueSize ) {
 
-					rayQueueSize[ 0 ] = rayQueueSize[ 1 ];
+					queueSizes[ 0 ] = queueSizes[ 1 ];
 
 				} else {
 
-					rayQueueSize[ 0 ] = rayQueueSize[ 0 ] + processed;
+					queueSizes[ 0 ] = queueSizes[ 0 ] + processed;
 
 				}
 

src/webgpu/lib/BVHComputeData.js

@@ -116,7 +116,7 @@ const transformStruct = new StructTypeNode( {
 	_alignment1: 'uint',
 }, 'TransformStruct' );
 
-const intersectionResultStruct = new StructTypeNode( {
+export const intersectionResultStruct = new StructTypeNode( {
 	indices: 'vec4u',
 	normal: 'vec3f',
 	didHit: 'bool',
@@ -166,7 +166,7 @@ function getTotalBVHByteLength( bvh ) {
 
 }
 
-const intersectsTriangle = wgslTagFn/* wgsl */ `
+export const intersectsTriangle = wgslTagFn/* wgsl */ `
 	// fn
 	fn intersectsTriangle( ray: ${ rayStruct }, a: vec3f, b: vec3f, c: vec3f ) -> ${ intersectionResultStruct } {