Compute Shaders: Add support for TLAS, scene-level raycasting

package.json

@@ -46,7 +46,7 @@
     "rollup": "^2.70.0",
     "simple-git": "^3.10.0",
     "three": "^0.183.1",
-    "three-mesh-bvh": "^0.9.5",
+    "three-mesh-bvh": "^0.9.8",
     "typescript": "^5.9.2",
     "vite": "^6.2.2",
     "yargs": "^17.5.1"

src/webgpu/MegaKernelPathTracer.js

@@ -8,7 +8,6 @@ function* renderTask() {
 		renderer,
 		camera,
 		kernel,
-		geometry,
 		bounces,
 
 		tiles,
@@ -22,13 +21,6 @@ function* renderTask() {
 	kernel.outputTarget = outputTarget;
 	kernel.sampleCountTarget = sampleCountTarget;
 
-	kernel.geom_index = geometry.index;
-	kernel.geom_position = geometry.position;
-	kernel.geom_normals = geometry.normal;
-	kernel.geom_material_index = geometry.materialIndex;
-	kernel.bvh = geometry.bvh;
-	kernel.materials = geometry.materials;
-
 	kernel.bounces = bounces;
 	kernel.inverseProjectionMatrix.copy( camera.projectionMatrixInverse );
 	kernel.cameraToModelMatrix.copy( camera.matrixWorld );
@@ -77,17 +69,6 @@ export class MegaKernelPathTracer {
 		this.bounces = 7;
 		this.tiles = new Vector2( 2, 2 );
 
-		// geometry fields
-		this.geometry = {
-			bvh: null,
-			index: null,
-			position: null,
-			normal: null,
-
-			materialIndex: null,
-			materials: null,
-		};
-
 		// targets
 		this.outputTarget = new StorageTexture( 1, 1, );
 		this.outputTarget.format = RGBAFormat;
@@ -118,23 +99,10 @@ export class MegaKernelPathTracer {
 
 	}
 
-	setGeometryData( geometry ) {
-
-		for ( const propName in geometry ) {
-
-			const prop = this.geometry[ propName ];
-			if ( prop === undefined ) {
-
-				console.error( `Invalid property name in geometry data: ${propName}` );
-				continue;
+	setBVHData( bvhData ) {
 
-			}
-
-			// TODO: cannot dispose at the moment
-			// prop.dispose();
-			this.geometry[ propName ] = geometry[ propName ];
-
-		}
+		this.kernel.bvhData = bvhData;
+		this.reset();
 
 	}
 
@@ -223,7 +191,6 @@ export class MegaKernelPathTracer {
 		this.samples = 0;
 		this._task = null;
 
-
 		const { width, height } = sampleCountTarget;
 		const dispatchSize = sampleCountClearKernel.getDispatchSize( width, height );
 
@@ -240,7 +207,7 @@ export class MegaKernelPathTracer {
 
 	update() {
 
-		if ( ! this.camera ) {
+		if ( ! this.camera || ! this.kernel ) {
 
 			return;
 

src/webgpu/WaveFrontPathTracer.js

@@ -225,23 +225,11 @@ export class WaveFrontPathTracer {
 
 	}
 
-	setGeometryData( geometry ) {
+	setBVHData( bvhData ) {
 
-		for ( const propName in geometry ) {
-
-			const prop = this.geometry[ propName ];
-			if ( prop === undefined ) {
-
-				console.error( `Invalid property name in geometry data: ${propName}` );
-				continue;
-
-			}
-
-			// TODO: cannot dispose at the moment
-			// prop.dispose();
-			this.geometry[ propName ] = geometry[ propName ];
-
-		}
+		this.rayIntersectionKernel.bvhData = bvhData;
+		this.hitProcessKernel.bvhData = bvhData;
+		this.reset();
 
 	}
 

src/webgpu/WebGPUPathTracer.js

@@ -1,9 +1,11 @@
-import { Color, StorageBufferAttribute, PerspectiveCamera, Scene, Vector2, Clock } from 'three/webgpu';
-import { PathTracingSceneGenerator } from '../core/PathTracingSceneGenerator.js';
+import { Vector2, Scene, PerspectiveCamera } from 'three/webgpu';
+import { MeshBVH, SAH } from 'three-mesh-bvh';
 import { FullScreenQuad } from 'three/examples/jsm/postprocessing/Pass.js';
 import { RenderToScreenNodeMaterial } from './materials/RenderToScreenMaterial.js';
 import { MegaKernelPathTracer } from './MegaKernelPathTracer.js';
 import { WaveFrontPathTracer } from './WaveFrontPathTracer.js';
+import { ObjectBVH } from './lib/ObjectBVH.js';
+import { PathtracerBVHComputeData } from './lib/PathtracerBVHComputeData.js';
 
 const _resolution = new Vector2();
 export class WebGPUPathTracer {
@@ -24,19 +26,16 @@ export class WebGPUPathTracer {
 
 		this._pathTracer.dispose();
 		this._pathTracer = value ? new MegaKernelPathTracer( this._renderer ) : new WaveFrontPathTracer( this._renderer );
-		this._generator = new PathTracingSceneGenerator();
+		this._pathTracer.setBVHData( this._bvhData );
+		this.setCamera( this.camera );
 
 	}
 
 	constructor( renderer ) {
 
 		// members
 		this._renderer = renderer;
-		this._generator = new PathTracingSceneGenerator();
-		// this._pathTracer = new MegaKernelPathTracer( renderer );
-		this._pathTracer = new WaveFrontPathTracer( renderer );
-		this._queueReset = false;
-		this._clock = new Clock();
+		this._pathTracer = new MegaKernelPathTracer( renderer );
 
 		// options
 		this.renderScale = 1;
@@ -54,11 +53,25 @@ export class WebGPUPathTracer {
 		scene.updateMatrixWorld( true );
 		camera.updateMatrixWorld();
 
-		const generator = this._generator;
-		generator.setObjects( scene );
+		// Build BVH for each mesh geometry
+		scene.traverse( child => {
 
-		const result = generator.generate();
-		return this._updateFromResults( scene, camera, result );
+			if ( child.isMesh && ! child.geometry.boundsTree ) {
+
+				child.geometry.boundsTree = new MeshBVH( child.geometry, { strategy: SAH, maxLeafSize: 5 } );
+
+			}
+
+		} );
+
+		// Build TLAS and compute functions
+		const objectBVH = new ObjectBVH( scene, { strategy: SAH } );
+		const bvhData = new PathtracerBVHComputeData( objectBVH );
+		bvhData.update();
+
+		this._bvhData = bvhData;
+		this._pathTracer.setBVHData( bvhData );
+		this.setCamera( camera );
 
 	}
 
@@ -85,80 +98,6 @@ export class WebGPUPathTracer {
 
 	}
 
-	_updateFromResults( scene, camera, results ) {
-
-		const {
-			materials,
-			geometry,
-			bvh,
-			bvhChanged,
-			needsMaterialIndexUpdate,
-		} = results;
-
-		const pathTracer = this._pathTracer;
-
-		const newGeometryData = {};
-
-		if ( bvhChanged ) {
-
-			// dereference a new index attribute if we're using indirect storage
-			const dereferencedIndexAttr = geometry.index.clone();
-			const indirectBuffer = bvh._indirectBuffer;
-			if ( indirectBuffer ) {
-
-				dereferenceIndex( geometry, indirectBuffer, dereferencedIndexAttr );
-
-			}
-
-			const newIndex = new StorageBufferAttribute( dereferencedIndexAttr.array, 3 );
-			newIndex.name = 'Geometry Index';
-			newGeometryData.index = newIndex;
-
-			const newPosition = new StorageBufferAttribute( geometry.attributes.position.array, 3 );
-			newPosition.name = 'Geometry Positions';
-			newGeometryData.position = newPosition;
-
-			const newNormals = new StorageBufferAttribute( geometry.attributes.normal.array, 3 );
-			newNormals.name = 'Geometry Normals';
-			newGeometryData.normal = newNormals;
-
-			const newBvhRoots = new StorageBufferAttribute( new Float32Array( bvh._roots[ 0 ] ), 8 );
-			newBvhRoots.name = 'BVH Roots';
-			newGeometryData.bvh = newBvhRoots;
-
-		}
-
-		if ( needsMaterialIndexUpdate ) {
-
-			const newMaterialIndex = new StorageBufferAttribute( geometry.attributes.materialIndex.array, 1 );
-			newMaterialIndex.name = 'Material Index';
-			newGeometryData.materialIndex = newMaterialIndex;
-
-		}
-
-		const newMaterialsData = new Float32Array( materials.length * 3 );
-		const defaultColor = new Color();
-		for ( let i = 0; i < materials.length; i ++ ) {
-
-			const material = materials[ i ];
-			const color = material.color ?? defaultColor;
-			// Make sure those are in linear-sRGB space
-			newMaterialsData[ 3 * i + 0 ] = color.r;
-			newMaterialsData[ 3 * i + 1 ] = color.g;
-			newMaterialsData[ 3 * i + 2 ] = color.b;
-
-		}
-
-		const newMaterialsBuffer = new StorageBufferAttribute( newMaterialsData, 3 );
-		newMaterialsBuffer.name = 'Material Data';
-		newGeometryData.materials = newMaterialsBuffer;
-
-		pathTracer.setGeometryData( newGeometryData );
-
-		this.setCamera( camera );
-
-	}
-
 	renderSample() {
 
 		if ( ! this._renderer._initialized ) {
@@ -216,22 +155,3 @@ export class WebGPUPathTracer {
 	}
 
 }
-
-// TODO: Expose in three-mesh-bvh?
-function dereferenceIndex( geometry, indirectBuffer, target ) {
-
-	const unpacked = target.array;
-	const indexArray = geometry.index ? geometry.index.array : null;
-	for ( let i = 0, l = indirectBuffer.length; i < l; i ++ ) {
-
-		const i3 = 3 * i;
-		const v3 = 3 * indirectBuffer[ i ];
-		for ( let c = 0; c < 3; c ++ ) {
-
-			unpacked[ i3 + c ] = indexArray ? indexArray[ v3 + c ] : v3 + c;
-
-		}
-
-	}
-
-}