From 545087871c69e60971a0e9c5fa965c8421462daf Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Tue, 17 Feb 2026 17:36:39 +0900 Subject: [PATCH 01/84] Add wrappers for full-scene compute raycasting --- src/webgpu/lib/BVHComputeFns.js | 213 ++++++++++++ src/webgpu/lib/ObjectBVH.js | 600 ++++++++++++++++++++++++++++++++ 2 files changed, 813 insertions(+) create mode 100644 src/webgpu/lib/BVHComputeFns.js create mode 100644 src/webgpu/lib/ObjectBVH.js diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js new file mode 100644 index 000000000..2fc5c6ee3 --- /dev/null +++ b/src/webgpu/lib/BVHComputeFns.js @@ -0,0 +1,213 @@ +import { Vector4 } from 'three'; +import { StorageBufferAttribute } from 'three/src/Three.WebGPU.Nodes.js'; +import { wgsl, wgslFn } from 'three/tsl'; +import { intersectTriangles, intersectsBounds } from 'three-mesh-bvh/webgpu'; + +const _vec = /* @__PURE__ */ new Vector4(); + +export const constants = wgsl( /* wgsl */` + const BVH_STACK_DEPTH = 60u; + const INFINITY = 1e20; + const TRI_INTERSECT_EPSILON = 1e-5; +` ); + +export const rayStruct = wgsl( /* wgsl */` + struct Ray { + origin: vec3f, + _alignment0: u32, + + direction: vec3f, + _alignment1: u32, + }; +` ); + +export const bvhNodeBoundsStruct = wgsl( /* wgsl */` + struct BVHBoundingBox { + min: array, + max: array, + } +` ); + +export const bvhNodeStruct = wgsl( /* wgsl */` + struct BVHNode { + bounds: BVHBoundingBox, + rightChildOrTriangleOffset: u32, + splitAxisOrTriangleCount: u32, + }; +`, [ bvhNodeBoundsStruct ] ); + +export const intersectionResultStruct = wgsl( /* wgsl */` + struct IntersectionResult { + indices: vec4u, + normal: vec3f, + didHit: bool, + barycoord: vec3f, + side: f32, + dist: f32, + }; +` ); + +export const transformStruct = wgsl( /* wgsl */` + struct TransformStruct { + matrixWorld: mat4x4f, + nodeOffset: u32, + } +` ); + +export class BVHComputeFns { + + constructor( bvh, options = {} ) { + + const { + name = 'bvh_', + attributes = [ 'position', 'uv0' ], + } = options; + + this.name = name; + this.attributes = attributes; + + this.bvh = bvh; + this.storageBufferAttributes = { + index: null, + attributes: null, + nodes: null, + transforms: null, + }; + + this.attributesStruct = null; + this.raycastFirstHitFn = null; + this.update(); + + } + + update() { + + const { attributes, name } = this; + + // construct the attribute struct + let attributesStructSize = 0; + const attributeStructContent = attributes + .map( key => { + + attributesStructSize += 16; + return `${ key }: vec4,`; + + } ).join( '\n' ); + + const attributeStruct = wgsl( /* wgsl */` + struct ${ name }GeometryStruct { + ${ attributeStructContent } + } + ` ); + + // TODO: gather the BVHs, geometries, and meshes + // TODO: handle the case where an "ObjectBVH" vs "MeshBVH" are passed + // TODO: handle batched / instanced meshes + const bvhs = []; + const geometries = []; + const meshes = []; + + // gather the size of the geometry and BVH buffers + let indexLength = 0; + let attributesLength = 0; + geometries.forEach( ( _, i ) => { + + const geometry = geometries[ i ]; + indexLength += geometry.index ? geometry.index.count : geometry.attributes.position.count; + attributesLength += geometry.attributes.position.count; + + } ); + + // initialize the geometry attributes + const attributesOffset = 0; + const indexOffset = 0; + const indexBuffer = new Uint32Array( indexLength ); + const attributesBuffer = new Float32Array( attributesLength ); + geometries.forEach( ( _, i ) => { + + // TODO: fill out based on bvh "sub range" + const geometry = geometries[ i ]; + appendGeometryData( geometry ); + + + } ); + + // TODO: fill out the BVH buffers accounting for geometry offsets + + this.storageBufferAttributes.index = new StorageBufferAttribute( indexBuffer, 1 ); + this.storageBufferAttributes.attributes = new StorageBufferAttribute( attributesBuffer, attributesStructSize ); + this.attributesStruct = attributeStruct; + this.raycastFirstHitFn = wgslFn( /* wgsl */` + fn ${ name }RaycastFirstHit( ray: Ray ) -> IntersectionResult { + + } + `, [ intersectTriangles, intersectsBounds, rayStruct, bvhNodeStruct, intersectionResultStruct, constants ] ); + + function appendGeometryData( geometry ) { + + // TODO: handle "indirect" case + if ( geometry.index ) { + + // TODO: need to handle the "sub range" case? For Batched Mesh? + for ( let i = 0, l = geometry.index.count; i < l; i ++ ) { + + indexBuffer[ i + indexOffset ] = geometry.index.getX() + attributesOffset; + + } + + indexOffset += geometry.index.count; + + } else { + + for ( let i = 0, l = geometry.attributes.position.count; i < l; i ++ ) { + + indexBuffer[ i + indexOffset ] = i + attributesOffset; + + } + + indexOffset += geometry.attributes.position.count; + + } + + attributes.forEach( ( key, ai ) => { + + const attr = geometry.attributes[ key ]; + if ( ! attr ) { + + if ( key === 'color' ) { + + _vec.set( 1, 1, 1, 1 ); + + } else { + + _vec.set( 0, 0, 0, 1 ); + + } + + } else { + + _vec.fromBufferAttribute( attr, i ); + switch ( attr.itemSize ) { + + case 3: + _vec.w = 0; + case 2: + _vec.z = 0; + case 1: + _vec.y = 0; + + } + + } + + _vec.toArray( attributesBuffer, attributesOffset * attributesStructSize + ai * 4 ); + + } ); + + attributesOffset += geometry.attributes.position.count; + + } + + } + +} diff --git a/src/webgpu/lib/ObjectBVH.js b/src/webgpu/lib/ObjectBVH.js new file mode 100644 index 000000000..f7f602b2b --- /dev/null +++ b/src/webgpu/lib/ObjectBVH.js @@ -0,0 +1,600 @@ +import { Box3, BufferGeometry, Matrix4, Mesh, Vector3, Ray, Sphere } from 'three'; +import { BVH, INTERSECTED, NOT_INTERSECTED } from 'three-mesh-bvh'; + +const _geometry = /* @__PURE__ */ new BufferGeometry(); +const _matrix = /* @__PURE__ */ new Matrix4(); +const _inverseMatrix = /* @__PURE__ */ new Matrix4(); +const _box = /* @__PURE__ */ new Box3(); +const _sphere = /* @__PURE__ */ new Sphere(); +const _vec = /* @__PURE__ */ new Vector3(); +const _ray = /* @__PURE__ */ new Ray(); +const _mesh = /* @__PURE__ */ new Mesh(); +const _geometryRange = {}; + +// TODO: account for a "custom" object? Not necessary here? Create a more abstract foundation for this case? +export function objectAcceleratedRaycast( raycaster, intersects ) { + + if ( this.boundsTree ) { + + this.boundsTree.raycast( raycaster, intersects ); + return false; + + } + +} + +export class ObjectBVH extends BVH { + + constructor( root, options = {} ) { + + options = { + precise: false, + includeInstances: true, + matrixWorld: Array.isArray( root ) ? new Matrix4() : root.matrixWorld, + maxLeafSize: 1, + ...options, + }; + + super(); + + // collect all the leaf node objects in the geometries + const objectSet = new Set(); + collectObjects( root, objectSet ); + + // calculate the number of bits required for the primary id, leaving the remainder + // for the instanceId count + const objects = Array.from( objectSet ); + const idBits = Math.ceil( Math.log2( objects.length ) ); + const idMask = constructIdMask( idBits ); + + this.objects = objects; + this.idBits = idBits; + this.idMask = idMask; + this.primitiveBuffer = null; + this.primitiveBufferStride = 1; + + // settings + this.precise = options.precise; + this.includeInstances = options.includeInstances; + this.matrixWorld = options.matrixWorld; + + this.init( options ); + + } + + init( options ) { + + const { objects, idBits } = this; + this.primitiveBuffer = new Uint32Array( this._countPrimitives( objects ) ); + this._fillPrimitiveBuffer( objects, idBits, this.primitiveBuffer ); + + super.init( options ); + + } + + writePrimitiveBounds( i, targetBuffer, writeOffset ) { + + // TODO: it would be best to cache this matrix inversion + const { primitiveBuffer } = this; + _inverseMatrix.copy( this.matrixWorld ).invert(); + + this._getPrimitiveBoundingBox( primitiveBuffer[ i ], _inverseMatrix, _box ); + const { min, max } = _box; + + targetBuffer[ writeOffset + 0 ] = min.x; + targetBuffer[ writeOffset + 1 ] = min.y; + targetBuffer[ writeOffset + 2 ] = min.z; + targetBuffer[ writeOffset + 3 ] = max.x; + targetBuffer[ writeOffset + 4 ] = max.y; + targetBuffer[ writeOffset + 5 ] = max.z; + + } + + getRootRanges() { + + return [ { offset: 0, count: this.primitiveBuffer.length } ]; + + } + + shapecast( callbacks ) { + + return super.shapecast( { + ...callbacks, + + intersectsPrimitive: callbacks.intersectsObject, + scratchPrimitive: null, + iterate: iterateOverObjects, + } ); + + } + + // TODO: this is out of sync with the MeshBVH raycast signature. + raycast( raycaster, intersects = [] ) { + + const { matrixWorld, includeInstances } = this; + const { firstHitOnly } = raycaster; + const localIntersects = []; + + // transform the ray into the local bvh frame + _inverseMatrix.copy( matrixWorld ).invert(); + _ray.copy( raycaster.ray ).applyMatrix4( _inverseMatrix ); + + let closestDistance = Infinity; + let closestHit = null; + + this.shapecast( { + boundsTraverseOrder: box => { + + return box.distanceToPoint( _ray.origin ); + + }, + intersectsBounds: box => { + + if ( firstHitOnly ) { + + if ( ! _ray.intersectBox( box, _vec ) ) { + + return NOT_INTERSECTED; + + } + + // early out if the box is further than the closest raycast + _vec.applyMatrix4( matrixWorld ); + return raycaster.ray.origin.distanceTo( _vec ) < closestDistance ? INTERSECTED : NOT_INTERSECTED; + + } else { + + return _ray.intersectsBox( box ) ? INTERSECTED : NOT_INTERSECTED; + + } + + }, + intersectsObject( object, instanceId ) { + + // skip non visible objects + if ( ! object.visible ) { + + return; + + } + + if ( object.isInstancedMesh && includeInstances ) { + + // raycast the instance + _mesh.geometry = object.geometry; + _mesh.material = object.material; + + object.getMatrixAt( instanceId, _mesh.matrixWorld ); + _mesh.matrixWorld.premultiply( object.matrixWorld ); + _mesh.raycast( raycaster, localIntersects ); + + localIntersects.forEach( hit => { + + hit.object = object; + hit.instanceId = instanceId; + + } ); + + _mesh.material = null; + + } else if ( object.isBatchedMesh && includeInstances ) { + + if ( ! object.getVisibleAt( instanceId ) ) { + + return; + + } + + // extract the geometry & material + const geometryId = object.getGeometryIdAt( instanceId ); + const geometryRange = object.getGeometryRangeAt( geometryId, _geometryRange ); + + _geometry.index = object.geometry.index; + _geometry.attributes.position = object.geometry.attributes.position; + _geometry.setDrawRange( geometryRange.start, geometryRange.count ); + + _mesh.geometry = _geometry; + _mesh.material = object.material; + + // perform a raycast against the proxy mesh + object.getMatrixAt( instanceId, _mesh.matrixWorld ); + _mesh.matrixWorld.premultiply( object.matrixWorld ); + _mesh.raycast( raycaster, localIntersects ); + + // fix up the fields + localIntersects.forEach( hit => { + + hit.object = object; + hit.batchId = instanceId; + + } ); + + _mesh.material = null; + _geometry.index = null; + _geometry.attributes.position = null; + _geometry.setDrawRange( 0, Infinity ); + + } else { + + object.raycast( raycaster, localIntersects ); + + } + + // find the closest hit to track + if ( firstHitOnly ) { + + localIntersects.forEach( hit => { + + if ( hit.distance < closestDistance ) { + + closestDistance = hit.distance; + closestHit = hit; + + } + + } ); + + } else { + + intersects.push( ...localIntersects ); + + } + + }, + } ); + + // save the closest hit only if firstHitOnly = true + if ( firstHitOnly && closestHit ) { + + intersects.push( closestHit ); + + } + + return intersects; + + } + + // get the bounding box of a primitive node accounting for the bvh options + _getPrimitiveBoundingBox( compositeId, inverseMatrixWorld, target ) { + + const { objects, idMask, idBits, precise, includeInstances } = this; + const id = getObjectId( compositeId, idMask ); + const instanceId = getInstanceId( compositeId, idBits, idMask ); + const object = objects[ id ]; + + if ( ! includeInstances && ( object.isInstancedMesh || object.isBatchedMesh ) ) { + + // if we're not using instances then just account for the overall bounds of the BatchedMesh and InstancedMesh + if ( ! object.boundingBox ) { + + object.computeBoundingBox(); + + } + + if ( ! object.boundingSphere ) { + + object.computeBoundingSphere(); + + } + + _matrix + .copy( object.matrixWorld ) + .premultiply( inverseMatrixWorld ); + + _sphere + .copy( object.boundingSphere ) + .applyMatrix4( _matrix ); + + target + .copy( object.boundingBox ) + .applyMatrix4( _matrix ); + + shrinkToSphere( target, _sphere ); + + } else if ( precise ) { + + // calculate precise bounds if necessary by calculating the bounds of all vertices + // in the bvh frame + if ( object.isInstancedMesh ) { + + object + .getMatrixAt( instanceId, _matrix ); + + _matrix + .premultiply( object.matrixWorld ) + .premultiply( inverseMatrixWorld ); + + getPreciseBounds( object.geometry, _matrix, target ); + + } else if ( object.isBatchedMesh ) { + + const geometryId = object.getGeometryIdAt( instanceId ); + const geometryRange = object.getGeometryRangeAt( geometryId, _geometryRange ); + + _geometry.index = object.geometry.index; + _geometry.attributes.position = object.geometry.attributes.position; + _geometry.setDrawRange( geometryRange.start, geometryRange.count ); + + object + .getMatrixAt( instanceId, _matrix ); + + _matrix + .premultiply( object.matrixWorld ) + .premultiply( inverseMatrixWorld ); + + getPreciseBounds( _geometry, _matrix, target ); + + } else { + + _matrix + .copy( object.matrixWorld ) + .premultiply( inverseMatrixWorld ); + + target.setFromObject( object, true ).applyMatrix4( inverseMatrixWorld ); + + } + + } else { + + // otherwise use the fast path of extracting the cached, AABB bounds and transforming them + // into the local BVH frame + if ( object.isInstancedMesh ) { + + if ( ! object.geometry.boundingBox ) { + + object.geometry.computeBoundingBox(); + + } + + if ( ! object.geometry.boundingSphere ) { + + object.geometry.computeBoundingSphere(); + + } + + object + .getMatrixAt( instanceId, _matrix ); + + _matrix + .premultiply( object.matrixWorld ) + .premultiply( inverseMatrixWorld ); + + _sphere + .copy( object.geometry.boundingSphere ) + .applyMatrix4( _matrix ); + + target + .copy( object.geometry.boundingBox ) + .applyMatrix4( _matrix ); + + shrinkToSphere( target, _sphere ); + + } else if ( object.isBatchedMesh ) { + + const geometryId = object.getGeometryIdAt( instanceId ); + + object + .getMatrixAt( instanceId, _matrix ); + + _matrix + .premultiply( object.matrixWorld ) + .premultiply( inverseMatrixWorld ); + + object + .getBoundingSphereAt( geometryId, _sphere ) + .applyMatrix4( _matrix ); + + object + .getBoundingBoxAt( geometryId, target ) + .applyMatrix4( _matrix ); + + shrinkToSphere( target, _sphere ); + + } else { + + target + .setFromObject( object, false ) + .applyMatrix4( inverseMatrixWorld ); + + } + + } + + } + + // counts the total number of primitives required by the objects in given array of objects + _countPrimitives( objects ) { + + const { includeInstances } = this; + let total = 0; + objects.forEach( object => { + + if ( object.isInstancedMesh && includeInstances ) { + + total += object.count; + + } else if ( object.isBatchedMesh && includeInstances ) { + + total += object.instanceCount; + + } else { + + total ++; + + } + + } ); + + return total; + + } + + _fillPrimitiveBuffer( objects, idBits, target ) { + + const { includeInstances } = this; + let index = 0; + objects.forEach( ( object, i ) => { + + if ( object.isInstancedMesh && includeInstances ) { + + const count = object.count; + for ( let c = 0; c < count; c ++ ) { + + target[ index ] = ( c << idBits ) | i; + index ++; + + } + + } else if ( object.isBatchedMesh && includeInstances ) { + + const { instanceCount, maxInstanceCount } = object; + let instance = 0; + let iter = 0; + // TODO: use a better check here, like "maxInstanceCount" + while ( instance < instanceCount && iter < maxInstanceCount ) { + + iter ++; + + // TODO: it would be better to have a consistent way of querying whether an + // instance were active + try { + + object.getVisibleAt( instance ); + + target[ index ] = ( instance << idBits ) | i; + instance ++; + index ++; + + } catch { + + // + + } + + } + + } else { + + target[ index ] = i; + index ++; + + } + + } ); + + } + +} + +// id functions +// construct a mask with the given number of bits set to 1 +function constructIdMask( idBits ) { + + let mask = 0; + for ( let i = 0; i < idBits; i ++ ) { + + mask = mask << 1 | 1; + + } + + return mask; + +} + +// extract the primary object id given the provided mask +function getObjectId( id, idMask ) { + + return id & idMask; + +} + +// extract the instance id given the mask and number of bits to shift +function getInstanceId( id, idBits, idMask ) { + + return ( id & ( ~ idMask ) ) >> idBits; + +} + +// traverse the full scene and collect all leaves +function collectObjects( root, objectSet = new Set() ) { + + if ( Array.isArray( root ) ) { + + root.forEach( object => collectObjects( object, objectSet ) ); + + } else { + + root.traverse( child => { + + if ( child.isMesh || child.isLine || child.isPoints ) { + + objectSet.add( child ); + + } + + } ); + + } + +} + +// calculate precise box bounds of the given geometry in the given frame +function getPreciseBounds( geometry, matrix, target ) { + + target.makeEmpty(); + + const drawRange = geometry.drawRange; + const indexAttr = geometry.index; + const posAttr = geometry.attributes.position; + const start = drawRange.start; + const vertCount = indexAttr ? indexAttr.count : posAttr.count; + const count = Math.min( vertCount - start, drawRange.count ); + for ( let i = start, l = start + count; i < l; i ++ ) { + + let vi = i; + if ( indexAttr ) { + + vi = indexAttr.getX( vi ); + + } + + _vec.fromBufferAttribute( posAttr, vi ).applyMatrix4( matrix ); + target.expandByPoint( _vec ); + + } + + return target; + +} + +// iterator helper for raycasting +function iterateOverObjects( offset, count, bvh, callback, contained, depth, /* scratch */ ) { + + const { primitiveBuffer, objects, idMask, idBits } = bvh; + for ( let i = offset, l = count + offset; i < l; i ++ ) { + + const compositeId = primitiveBuffer[ i ]; + const id = getObjectId( compositeId, idMask ); + const instanceId = getInstanceId( compositeId, idBits, idMask ); + const object = objects[ id ]; + if ( callback( object, instanceId, contained, depth ) ) { + + return true; + + } + + } + + return false; + +} + +function shrinkToSphere( box, sphere ) { + + _vec.copy( sphere.center ).addScalar( - sphere.radius ); + box.min.max( _vec ); + + _vec.copy( sphere.center ).addScalar( sphere.radius ); + box.max.min( _vec ); + +} From 80d929700b450ca3486a5b93451d29e01801217b Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 00:32:30 +0900 Subject: [PATCH 02/84] Add functions --- src/webgpu/lib/BVHComputeFns.js | 195 ++++++++++++++++++++++++++------ 1 file changed, 162 insertions(+), 33 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 2fc5c6ee3..5a55f5060 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -60,7 +60,7 @@ export class BVHComputeFns { const { name = 'bvh_', - attributes = [ 'position', 'uv0' ], + attributes = [ 'position', 'normal', 'uv0' ], } = options; this.name = name; @@ -78,11 +78,114 @@ export class BVHComputeFns { this.raycastFirstHitFn = null; this.update(); + this.getBVH = ( object, id = - 1 ) => { + + if ( object.isInstancedMesh ) { + + return object.geometry.boundsTree; + + } else if ( object.isBatchedMesh ) { + + const geometryId = object.getGeometryIdAt( id ); + return object.boundsTrees[ geometryId ]; + + } else { + + return object.geometry.boundsTree; + + } + + }; + } update() { - const { attributes, name } = this; + const { attributes, name, bvh } = this; + + // TODO: add support for materials? Optional? Custom callback? + // TODO: gather the BVHs, geometries, and meshes + // TODO: handle the case where an "ObjectBVH" vs "MeshBVH" are passed + // TODO: handle batched / instanced meshes + // TODO: how to handle skinned meshes? + const bvhs = []; + const geometries = []; + const meshes = []; + const geometryOffsets = []; + + // TODO: find the total BVH node size first, then append BVH and geometry data + let nodeLength = 0; + nodeLength += bvh._roots[ 0 ].byteLength; + bvh.primitiveBuffer.forEach( compositeId => { + + const object = bvh.getObjectFromId( compositeId ); + const instanceId = bvh.getInstanceFromId( compositeId ); + const bvh = this.getBVH( object, instanceId ); + if ( ! bvhs.includes( bvh ) ) { + + bvhs.push( bvh ); + nodeLength += bvh._roots.reduce( ( v, root ) => v + root.byteLength, 0 ); + + // write associated geometry data, save offsets for use later + if ( object.isBatchedMesh ) { + + const geometryId = object.getGeometryIdAt( instanceId ); + const geometryInfo = object.getGeometryRangeAt( geometryId ); + const offset = appendGeometryData( object.geometry, geometryInfo ); + geometryOffsets.push( offset ); + + } else { + + const offset = appendGeometryData( object.geometry ); + geometryOffsets.push( offset ); + + } + + + } + + } ); + + // TODO: build the bvh node buffers + const nodeBuffer = new ArrayBuffer( nodeLength ); + const nodeBuffer16 = new Uint16Array( nodeBuffer ); + const nodeBuffer32 = new Uint32Array( nodeBuffer ); + const nodeBufferFloat = new Float32Array( nodeBuffer ); + const bvhNodeOffsets = []; + const nodeWriteOffset = 0; + bvhs.forEach( bvh => { + + const BYTES_PER_NODE = 6 * 4 + 4 + 4; + const UINT32_PER_NODE = BYTES_PER_NODE / 4; + const IS_LEAFNODE_FLAG = 0xFFFF; + const LEAFNODE_MASK_32 = IS_LEAFNODE_FLAG << 16; + + bvh._roots.forEach( root => { + + const rootBuffer16 = new Uint16Array( root ); + const rootBuffer32 = new Uint32Array( root ); + const rootBufferFloat = new Float32Array( root ); + for ( let i = 0, l = root.byteSize / BYTES_PER_NODE; i < l; i ++ ) { + + // write bounds + nodeBufferFloat.set( new Float32Array( root, i * BYTES_PER_NODE, 6 ), nodeWriteOffset * UINT32_PER_NODE ); + + // TODO: write remaining data + // - leaf + // - primitive leaf + // - primitive offset + // - primitive count + // - child + // - split + + } + + } ); + + } ); + + // TODO: build the geometry and transforms + // construct the attribute struct let attributesStructSize = 0; @@ -100,12 +203,6 @@ export class BVHComputeFns { } ` ); - // TODO: gather the BVHs, geometries, and meshes - // TODO: handle the case where an "ObjectBVH" vs "MeshBVH" are passed - // TODO: handle batched / instanced meshes - const bvhs = []; - const geometries = []; - const meshes = []; // gather the size of the geometry and BVH buffers let indexLength = 0; @@ -143,64 +240,94 @@ export class BVHComputeFns { } `, [ intersectTriangles, intersectsBounds, rayStruct, bvhNodeStruct, intersectionResultStruct, constants ] ); - function appendGeometryData( geometry ) { + function appendGeometryData( geometry, offsets = {} ) { + + const result = indexOffset; // TODO: handle "indirect" case if ( geometry.index ) { - // TODO: need to handle the "sub range" case? For Batched Mesh? - for ( let i = 0, l = geometry.index.count; i < l; i ++ ) { + let { + indexCount = - 1, + indexStart = - 1, + } = offsets; + + if ( indexCount === - 1 ) { + + indexStart = 0; + indexCount = geometry.index.count; - indexBuffer[ i + indexOffset ] = geometry.index.getX() + attributesOffset; + } + + for ( let i = 0; i < indexCount; i ++ ) { + + indexBuffer[ i + indexOffset ] = geometry.index.getX( i - indexStart ) + attributesOffset; } - indexOffset += geometry.index.count; + indexOffset += indexCount; } else { - for ( let i = 0, l = geometry.attributes.position.count; i < l; i ++ ) { + const indexCount = geometry.attributes.position.count; + for ( let i = 0; i < indexCount; i ++ ) { indexBuffer[ i + indexOffset ] = i + attributesOffset; } - indexOffset += geometry.attributes.position.count; + indexOffset += indexCount; } - attributes.forEach( ( key, ai ) => { + let { + vertexStart = - 1, + vertexCount = - 1, + } = offsets; + + if ( vertexCount === - 1 ) { + + vertexStart = 0; + vertexCount = geometry.attributes.position.count; + + } + + attributes.forEach( ( key, interleavedOffset ) => { const attr = geometry.attributes[ key ]; - if ( ! attr ) { + for ( let i = 0; i < vertexCount; i ++ ) { - if ( key === 'color' ) { + if ( ! attr ) { - _vec.set( 1, 1, 1, 1 ); + if ( key === 'color' ) { - } else { + _vec.set( 1, 1, 1, 1 ); - _vec.set( 0, 0, 0, 1 ); + } else { - } + _vec.set( 0, 0, 0, 1 ); - } else { + } + + } else { - _vec.fromBufferAttribute( attr, i ); - switch ( attr.itemSize ) { + _vec.fromBufferAttribute( attr, i - vertexStart ); + switch ( attr.itemSize ) { - case 3: - _vec.w = 0; - case 2: - _vec.z = 0; - case 1: - _vec.y = 0; + case 3: + _vec.w = 0; + case 2: + _vec.z = 0; + case 1: + _vec.y = 0; + + } } - } + _vec.toArray( attributesBuffer, attributesOffset * attributesStructSize + interleavedOffset * 4 ); - _vec.toArray( attributesBuffer, attributesOffset * attributesStructSize + ai * 4 ); + } } ); @@ -208,6 +335,8 @@ export class BVHComputeFns { } + return indexOffset; + } } From 728d428b5db4475fc91b26af6b5d74365c28a6e0 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 00:48:25 +0900 Subject: [PATCH 03/84] Update Fns --- src/webgpu/lib/BVHComputeFns.js | 71 +++++++++++++++++++-------------- 1 file changed, 42 insertions(+), 29 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 5a55f5060..747ecbb27 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -115,12 +115,15 @@ export class BVHComputeFns { // TODO: find the total BVH node size first, then append BVH and geometry data let nodeLength = 0; + let transformCount = 0; nodeLength += bvh._roots[ 0 ].byteLength; bvh.primitiveBuffer.forEach( compositeId => { const object = bvh.getObjectFromId( compositeId ); const instanceId = bvh.getInstanceFromId( compositeId ); const bvh = this.getBVH( object, instanceId ); + transformCount += bvh._roots.length; + if ( ! bvhs.includes( bvh ) ) { bvhs.push( bvh ); @@ -153,34 +156,10 @@ export class BVHComputeFns { const nodeBufferFloat = new Float32Array( nodeBuffer ); const bvhNodeOffsets = []; const nodeWriteOffset = 0; - bvhs.forEach( bvh => { - - const BYTES_PER_NODE = 6 * 4 + 4 + 4; - const UINT32_PER_NODE = BYTES_PER_NODE / 4; - const IS_LEAFNODE_FLAG = 0xFFFF; - const LEAFNODE_MASK_32 = IS_LEAFNODE_FLAG << 16; - - bvh._roots.forEach( root => { - - const rootBuffer16 = new Uint16Array( root ); - const rootBuffer32 = new Uint32Array( root ); - const rootBufferFloat = new Float32Array( root ); - for ( let i = 0, l = root.byteSize / BYTES_PER_NODE; i < l; i ++ ) { - - // write bounds - nodeBufferFloat.set( new Float32Array( root, i * BYTES_PER_NODE, 6 ), nodeWriteOffset * UINT32_PER_NODE ); + appendBVHData( bvh, true, 0 ); + bvhs.forEach( ( bvh, i ) => { - // TODO: write remaining data - // - leaf - // - primitive leaf - // - primitive offset - // - primitive count - // - child - // - split - - } - - } ); + bvhNodeOffsets.push( appendBVHData( bvh, true, geometryOffsets[ i ] ) ); } ); @@ -240,6 +219,41 @@ export class BVHComputeFns { } `, [ intersectTriangles, intersectsBounds, rayStruct, bvhNodeStruct, intersectionResultStruct, constants ] ); + function appendBVHData( bvh, tlas = false, geometryOffset = 0 ) { + + const BYTES_PER_NODE = 6 * 4 + 4 + 4; + const UINT32_PER_NODE = BYTES_PER_NODE / 4; + const IS_LEAFNODE_FLAG = 0xFFFF; + const LEAFNODE_MASK_32 = IS_LEAFNODE_FLAG << 16; + const result = []; + + bvh._roots.forEach( root => { + + const rootBuffer16 = new Uint16Array( root ); + const rootBuffer32 = new Uint32Array( root ); + const rootBufferFloat = new Float32Array( root ); + result.push( nodeWriteOffset ); + for ( let i = 0, l = root.byteSize / BYTES_PER_NODE; i < l; i ++ ) { + + // write bounds + nodeBufferFloat.set( new Float32Array( root, i * BYTES_PER_NODE, 6 ), nodeWriteOffset * UINT32_PER_NODE ); + + // TODO: write remaining data + // - leaf + // - primitive leaf + // - primitive offset + // - primitive count + // - child + // - split + + } + + } ); + + return result; + + } + function appendGeometryData( geometry, offsets = {} ) { const result = indexOffset; @@ -332,11 +346,10 @@ export class BVHComputeFns { } ); attributesOffset += geometry.attributes.position.count; + return result; } - return indexOffset; - } } From 7455a3432d2b8d3fb1a61907f35fd4c5d4bd1dd5 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 10:10:10 +0900 Subject: [PATCH 04/84] Updates --- src/webgpu/lib/BVHComputeFns.js | 204 +++++++++++++++++++++----------- src/webgpu/lib/ObjectBVH.js | 36 ++++++ 2 files changed, 168 insertions(+), 72 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 747ecbb27..dba04ba55 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -1,4 +1,4 @@ -import { Vector4 } from 'three'; +import { Matrix4, Vector4 } from 'three'; import { StorageBufferAttribute } from 'three/src/Three.WebGPU.Nodes.js'; import { wgsl, wgslFn } from 'three/tsl'; import { intersectTriangles, intersectsBounds } from 'three-mesh-bvh/webgpu'; @@ -110,12 +110,14 @@ export class BVHComputeFns { // TODO: how to handle skinned meshes? const bvhs = []; const geometries = []; - const meshes = []; const geometryOffsets = []; + const transformBVHs = []; // TODO: find the total BVH node size first, then append BVH and geometry data let nodeLength = 0; let transformCount = 0; + let indexLength = 0; + let attributesLength = 0; nodeLength += bvh._roots[ 0 ].byteLength; bvh.primitiveBuffer.forEach( compositeId => { @@ -134,32 +136,72 @@ export class BVHComputeFns { const geometryId = object.getGeometryIdAt( instanceId ); const geometryInfo = object.getGeometryRangeAt( geometryId ); - const offset = appendGeometryData( object.geometry, geometryInfo ); - geometryOffsets.push( offset ); + geometries.push( { + geometry: object.geometry, + range: geometryInfo, + } ); + + indexLength += range.indexCount === - 1 ? range.vertexCount : range.indexCount; + attributesLength += range.vertexCount; } else { - const offset = appendGeometryData( object.geometry ); - geometryOffsets.push( offset ); + indexLength += geometry.index ? geometry.index.count : geometry.attributes.position.count; + attributesLength += geometry.attributes.position.count; + geometries.push( { + geometry: object.geometry, + range: undefined, + } ); } - } + transformBVHs.push( bvhs.indexOf( bvh ) ); + + } ); + + // initialize the geometry attributes + const attributesOffset = 0; + const indexOffset = 0; + const indexBuffer = new Uint32Array( indexLength ); + const attributesBuffer = new Float32Array( attributesLength ); + geometries.forEach( ( { geometry, range } ) => { + + const offset = appendGeometryData( geometry, range ); + geometryOffsets.push( offset ); + } ); - // TODO: build the bvh node buffers const nodeBuffer = new ArrayBuffer( nodeLength ); const nodeBuffer16 = new Uint16Array( nodeBuffer ); const nodeBuffer32 = new Uint32Array( nodeBuffer ); const nodeBufferFloat = new Float32Array( nodeBuffer ); const bvhNodeOffsets = []; const nodeWriteOffset = 0; - appendBVHData( bvh, true, 0 ); + appendBVHData( bvh, 0, true ); bvhs.forEach( ( bvh, i ) => { - bvhNodeOffsets.push( appendBVHData( bvh, true, geometryOffsets[ i ] ) ); + bvhNodeOffsets.push( appendBVHData( bvh, geometryOffsets[ i ] ), false ); + + } ); + + let transformWriteOffset = 0; + const transformBuffer = new Uint32Array( 17 * transformCount ); + bvh.primitiveBuffer.forEach( ( compositeId, i ) => { + + const matrix = new Matrix4(); + bvh.getObjectMatrix( compositeId, matrix ); + + const objectBvh = bvhs[ transformBVHs[ i ] ]; + const bvhOffset = bvhNodeOffsets[ i ]; + objectBvh._roots.forEach( ( root, ri ) => { + + matrix.toArray( transformBuffer, transformWriteOffset * 17 ); + transformBuffer[ transformWriteOffset * 17 + 16 ] = bvhOffset[ ri ]; + transformWriteOffset ++; + + } ); } ); @@ -182,34 +224,6 @@ export class BVHComputeFns { } ` ); - - // gather the size of the geometry and BVH buffers - let indexLength = 0; - let attributesLength = 0; - geometries.forEach( ( _, i ) => { - - const geometry = geometries[ i ]; - indexLength += geometry.index ? geometry.index.count : geometry.attributes.position.count; - attributesLength += geometry.attributes.position.count; - - } ); - - // initialize the geometry attributes - const attributesOffset = 0; - const indexOffset = 0; - const indexBuffer = new Uint32Array( indexLength ); - const attributesBuffer = new Float32Array( attributesLength ); - geometries.forEach( ( _, i ) => { - - // TODO: fill out based on bvh "sub range" - const geometry = geometries[ i ]; - appendGeometryData( geometry ); - - - } ); - - // TODO: fill out the BVH buffers accounting for geometry offsets - this.storageBufferAttributes.index = new StorageBufferAttribute( indexBuffer, 1 ); this.storageBufferAttributes.attributes = new StorageBufferAttribute( attributesBuffer, attributesStructSize ); this.attributesStruct = attributeStruct; @@ -219,32 +233,51 @@ export class BVHComputeFns { } `, [ intersectTriangles, intersectsBounds, rayStruct, bvhNodeStruct, intersectionResultStruct, constants ] ); - function appendBVHData( bvh, tlas = false, geometryOffset = 0 ) { + function appendBVHData( bvh, geometryOffsets, tlas = false ) { const BYTES_PER_NODE = 6 * 4 + 4 + 4; const UINT32_PER_NODE = BYTES_PER_NODE / 4; const IS_LEAFNODE_FLAG = 0xFFFF; - const LEAFNODE_MASK_32 = IS_LEAFNODE_FLAG << 16; const result = []; - bvh._roots.forEach( root => { + bvh._roots.forEach( ( root, rootIndex ) => { const rootBuffer16 = new Uint16Array( root ); const rootBuffer32 = new Uint32Array( root ); - const rootBufferFloat = new Float32Array( root ); result.push( nodeWriteOffset ); for ( let i = 0, l = root.byteSize / BYTES_PER_NODE; i < l; i ++ ) { + const r32 = i * UINT32_PER_NODE; + const r16 = r32 * 2; + const n32 = nodeWriteOffset * UINT32_PER_NODE; + const n16 = n32 * 2; + // write bounds - nodeBufferFloat.set( new Float32Array( root, i * BYTES_PER_NODE, 6 ), nodeWriteOffset * UINT32_PER_NODE ); + nodeBufferFloat.set( new Float32Array( root, i * BYTES_PER_NODE, 6 ), n32 ); + + const isLeaf = IS_LEAFNODE_FLAG === rootBuffer16[ r32 + 15 ]; + if ( isLeaf ) { + + nodeBuffer32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ]; + nodeBuffer16[ n16 + 14 ] = rootBuffer16[ r16 + 14 ] + geometryOffsets[ rootIndex ]; + if ( tlas ) { - // TODO: write remaining data - // - leaf - // - primitive leaf - // - primitive offset - // - primitive count - // - child - // - split + // 0xFFFF == mesh leaf, 0xFF00 == TLAS leaf + nodeBuffer16[ n16 + 15 ] = 0xFF00; + + } else { + + nodeBuffer16[ n16 + 15 ] = IS_LEAFNODE_FLAG; + + } + + } else { + + nodeBuffer32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ]; + nodeBuffer32[ n32 + 7 ] = rootBuffer32[ r32 + 7 ]; + + + } } @@ -256,43 +289,70 @@ export class BVHComputeFns { function appendGeometryData( geometry, offsets = {} ) { - const result = indexOffset; + const result = []; + const groups = geometry.groups.length === 0 ? [ { start: 0, count: Infinity } ] : geometry.groups; + + groups.forEach( group => { - // TODO: handle "indirect" case - if ( geometry.index ) { + result.push( indexOffset ); - let { - indexCount = - 1, - indexStart = - 1, - } = offsets; + // TODO: handle "indirect" case + // TODO: validate the write offsets here + if ( geometry.index ) { - if ( indexCount === - 1 ) { + let { + indexCount = - 1, + indexStart = - 1, + } = offsets; - indexStart = 0; - indexCount = geometry.index.count; + if ( indexCount === - 1 ) { - } + indexStart = 0; + indexCount = geometry.index.count; - for ( let i = 0; i < indexCount; i ++ ) { + } - indexBuffer[ i + indexOffset ] = geometry.index.getX( i - indexStart ) + attributesOffset; + const indexEnd = Math.min( indexStart + indexCount, group.start + group.count ); - } + indexStart = Math.max( indexStart, group.start ); + indexCount = indexEnd - indexStart; - indexOffset += indexCount; + for ( let i = 0; i < indexCount; i ++ ) { - } else { + indexBuffer[ i + indexOffset ] = geometry.index.getX( i - indexStart ) + attributesOffset; - const indexCount = geometry.attributes.position.count; - for ( let i = 0; i < indexCount; i ++ ) { + } - indexBuffer[ i + indexOffset ] = i + attributesOffset; + indexOffset += indexCount; - } + } else { - indexOffset += indexCount; + let { + vertexStart = - 1, + vertexCount = - 1, + } = offsets; - } + if ( vertexCount === - 1 ) { + + vertexStart = 0; + vertexCount = geometry.attributes.position.count; + + } + + const indexStart = Math.max( vertexStart, group.start ); + const indexEnd = Math.min( group.start + group.end, vertexStart, vertexCount ); + const indexCount = indexEnd - indexStart; + for ( let i = 0; i < indexCount; i ++ ) { + + indexBuffer[ i + indexOffset ] = i - indexStart + attributesOffset; + + } + + indexOffset += indexCount; + + } + + } ); let { vertexStart = - 1, diff --git a/src/webgpu/lib/ObjectBVH.js b/src/webgpu/lib/ObjectBVH.js index f7f602b2b..eecea45b5 100644 --- a/src/webgpu/lib/ObjectBVH.js +++ b/src/webgpu/lib/ObjectBVH.js @@ -62,6 +62,42 @@ export class ObjectBVH extends BVH { } + getObjectMatrix( compositeId, target ) { + + const { idMask, idBits, objects, matrixWorld } = this; + const id = getObjectId( compositeId, idMask ); + const instanceId = getInstanceId( compositeId, idBits, idMask ); + const object = objects[ id ]; + if ( object.isInstancedMesh || object.isBatchedMesh ) { + + object.getMatrixAt( instanceId, target ).premultiply( object.matrixWorld ); + + } else { + + target.copy( object.matrixWorld ); + + } + + _inverseMatrix.copy( matrixWorld ).invert(); + target.premultiply( _inverseMatrix ); + + } + + getObjectFromId( compositeId ) { + + const { idMask, objects } = this; + const id = getObjectId( compositeId, idMask ); + return objects[ id ]; + + } + + getInstanceFromId( compositeId ) { + + const { idMask, idBits } = this; + return getInstanceId( compositeId, idBits, idMask ); + + } + init( options ) { const { objects, idBits } = this; From 9def446a936299eb455ebd7b8b531ede15794ca0 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 10:15:40 +0900 Subject: [PATCH 05/84] Updates --- src/webgpu/lib/BVHComputeFns.js | 57 ++++++++++++++++++--------------- 1 file changed, 31 insertions(+), 26 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index dba04ba55..9c4c5336e 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -104,50 +104,52 @@ export class BVHComputeFns { const { attributes, name, bvh } = this; // TODO: add support for materials? Optional? Custom callback? - // TODO: gather the BVHs, geometries, and meshes - // TODO: handle the case where an "ObjectBVH" vs "MeshBVH" are passed - // TODO: handle batched / instanced meshes // TODO: how to handle skinned meshes? const bvhs = []; const geometries = []; const geometryOffsets = []; const transformBVHs = []; - // TODO: find the total BVH node size first, then append BVH and geometry data - let nodeLength = 0; - let transformCount = 0; - let indexLength = 0; - let attributesLength = 0; - nodeLength += bvh._roots[ 0 ].byteLength; + // accumulate the sizes of the bvh nodes buffer, number of objects, and geometry buffers + let bvhNodesBufferLength = 0; + let objectTransformsCount = 0; + let indexBufferLength = 0; + let attributesBufferLength = 0; + bvhNodesBufferLength += bvh._roots[ 0 ].byteLength; bvh.primitiveBuffer.forEach( compositeId => { const object = bvh.getObjectFromId( compositeId ); const instanceId = bvh.getInstanceFromId( compositeId ); const bvh = this.getBVH( object, instanceId ); - transformCount += bvh._roots.length; + // add a new transform for each bvh root + objectTransformsCount += bvh._roots.length; + + // if we haven't added this bvh, yet if ( ! bvhs.includes( bvh ) ) { + // increase the buffer size bvhs.push( bvh ); - nodeLength += bvh._roots.reduce( ( v, root ) => v + root.byteLength, 0 ); + bvhNodesBufferLength += bvh._roots.reduce( ( v, root ) => v + root.byteLength, 0 ); - // write associated geometry data, save offsets for use later + // save the geometry info to write later and increment the buffer sizes if ( object.isBatchedMesh ) { const geometryId = object.getGeometryIdAt( instanceId ); - const geometryInfo = object.getGeometryRangeAt( geometryId ); + const range = object.getGeometryRangeAt( geometryId ); geometries.push( { geometry: object.geometry, - range: geometryInfo, + range: range, } ); - indexLength += range.indexCount === - 1 ? range.vertexCount : range.indexCount; - attributesLength += range.vertexCount; + indexBufferLength += range.indexCount === - 1 ? range.vertexCount : range.indexCount; + attributesBufferLength += range.vertexCount; } else { - indexLength += geometry.index ? geometry.index.count : geometry.attributes.position.count; - attributesLength += geometry.attributes.position.count; + const geometry = object.geometry; + indexBufferLength += geometry.index ? geometry.index.count : geometry.attributes.position.count; + attributesBufferLength += geometry.attributes.position.count; geometries.push( { geometry: object.geometry, range: undefined, @@ -157,15 +159,16 @@ export class BVHComputeFns { } + // save the index of the bvh associated with this transform transformBVHs.push( bvhs.indexOf( bvh ) ); } ); - // initialize the geometry attributes + // write the geometry buffer attributes const attributesOffset = 0; const indexOffset = 0; - const indexBuffer = new Uint32Array( indexLength ); - const attributesBuffer = new Float32Array( attributesLength ); + const indexBuffer = new Uint32Array( indexBufferLength ); + const attributesBuffer = new Float32Array( attributesBufferLength ); geometries.forEach( ( { geometry, range } ) => { const offset = appendGeometryData( geometry, range ); @@ -173,7 +176,8 @@ export class BVHComputeFns { } ); - const nodeBuffer = new ArrayBuffer( nodeLength ); + // write the bvh data + const nodeBuffer = new ArrayBuffer( bvhNodesBufferLength ); const nodeBuffer16 = new Uint16Array( nodeBuffer ); const nodeBuffer32 = new Uint32Array( nodeBuffer ); const nodeBufferFloat = new Float32Array( nodeBuffer ); @@ -186,8 +190,9 @@ export class BVHComputeFns { } ); + // write the transforms let transformWriteOffset = 0; - const transformBuffer = new Uint32Array( 17 * transformCount ); + const transformBuffer = new Uint32Array( 17 * objectTransformsCount ); bvh.primitiveBuffer.forEach( ( compositeId, i ) => { const matrix = new Matrix4(); @@ -205,10 +210,8 @@ export class BVHComputeFns { } ); - // TODO: build the geometry and transforms - - // construct the attribute struct + // TODO: need to include materials here let attributesStructSize = 0; const attributeStructContent = attributes .map( key => { @@ -224,6 +227,8 @@ export class BVHComputeFns { } ` ); + this.storageBufferAttributes.transforms = new StorageBufferAttribute( transformBuffer, 17 ); + this.storageBufferAttributes.nodes = new StorageBufferAttribute( nodeBuffer32, 8 ); this.storageBufferAttributes.index = new StorageBufferAttribute( indexBuffer, 1 ); this.storageBufferAttributes.attributes = new StorageBufferAttribute( attributesBuffer, attributesStructSize ); this.attributesStruct = attributeStruct; From d9999f7af2d531a773b0d7d96222d12417e2b4b7 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 11:05:00 +0900 Subject: [PATCH 06/84] Small fixes --- src/webgpu/lib/BVHComputeFns.js | 79 +++++++++++++-------------------- 1 file changed, 32 insertions(+), 47 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 9c4c5336e..e80f9d1c5 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -120,17 +120,17 @@ export class BVHComputeFns { const object = bvh.getObjectFromId( compositeId ); const instanceId = bvh.getInstanceFromId( compositeId ); - const bvh = this.getBVH( object, instanceId ); + const meshBvh = this.getBVH( object, instanceId ); // add a new transform for each bvh root - objectTransformsCount += bvh._roots.length; + objectTransformsCount += meshBvh._roots.length; // if we haven't added this bvh, yet - if ( ! bvhs.includes( bvh ) ) { + if ( ! bvhs.includes( meshBvh ) ) { // increase the buffer size - bvhs.push( bvh ); - bvhNodesBufferLength += bvh._roots.reduce( ( v, root ) => v + root.byteLength, 0 ); + bvhs.push( meshBvh ); + bvhNodesBufferLength += meshBvh._roots.reduce( ( v, root ) => v + root.byteLength, 0 ); // save the geometry info to write later and increment the buffer sizes if ( object.isBatchedMesh ) { @@ -160,13 +160,13 @@ export class BVHComputeFns { } // save the index of the bvh associated with this transform - transformBVHs.push( bvhs.indexOf( bvh ) ); + transformBVHs.push( bvhs.indexOf( meshBvh ) ); } ); // write the geometry buffer attributes - const attributesOffset = 0; - const indexOffset = 0; + let attributesOffset = 0; + let indexOffset = 0; const indexBuffer = new Uint32Array( indexBufferLength ); const attributesBuffer = new Float32Array( attributesBufferLength ); geometries.forEach( ( { geometry, range } ) => { @@ -182,7 +182,7 @@ export class BVHComputeFns { const nodeBuffer32 = new Uint32Array( nodeBuffer ); const nodeBufferFloat = new Float32Array( nodeBuffer ); const bvhNodeOffsets = []; - const nodeWriteOffset = 0; + let nodeWriteOffset = 0; appendBVHData( bvh, 0, true ); bvhs.forEach( ( bvh, i ) => { @@ -284,6 +284,8 @@ export class BVHComputeFns { } + nodeWriteOffset ++; + } } ); @@ -297,6 +299,15 @@ export class BVHComputeFns { const result = []; const groups = geometry.groups.length === 0 ? [ { start: 0, count: Infinity } ] : geometry.groups; + let vertexStart = 0; + let vertexCount = geometry.attributes.position.count; + if ( offsets ) { + + vertexStart = offsets.vertexStart; + vertexCount = offsets.vertexCount; + + } + groups.forEach( group => { result.push( indexOffset ); @@ -305,26 +316,18 @@ export class BVHComputeFns { // TODO: validate the write offsets here if ( geometry.index ) { - let { - indexCount = - 1, - indexStart = - 1, - } = offsets; - - if ( indexCount === - 1 ) { + let indexStart = Math.max( 0, group.start ); + let indexCount = Math.min( geometry.index.count, group.start + group.count ) - indexStart; + if ( offsets ) { - indexStart = 0; - indexCount = geometry.index.count; + indexStart = offsets.indexStart; + indexCount = geometry.indexCount; } - const indexEnd = Math.min( indexStart + indexCount, group.start + group.count ); - - indexStart = Math.max( indexStart, group.start ); - indexCount = indexEnd - indexStart; - for ( let i = 0; i < indexCount; i ++ ) { - indexBuffer[ i + indexOffset ] = geometry.index.getX( i - indexStart ) + attributesOffset; + indexBuffer[ i + indexOffset ] = geometry.index.getX( i + indexStart ) - vertexStart + attributesOffset; } @@ -332,24 +335,18 @@ export class BVHComputeFns { } else { - let { - vertexStart = - 1, - vertexCount = - 1, - } = offsets; - - if ( vertexCount === - 1 ) { + let indexStart = Math.max( 0, group.start ); + let indexCount = Math.min( geometry.attributes.position.count, group.start + group.count ) - indexStart; + if ( offsets ) { - vertexStart = 0; - vertexCount = geometry.attributes.position.count; + indexStart = offsets.vertexStart; + indexCount = offsets.vertexCount; } - const indexStart = Math.max( vertexStart, group.start ); - const indexEnd = Math.min( group.start + group.end, vertexStart, vertexCount ); - const indexCount = indexEnd - indexStart; for ( let i = 0; i < indexCount; i ++ ) { - indexBuffer[ i + indexOffset ] = i - indexStart + attributesOffset; + indexBuffer[ i + indexOffset ] = i + indexStart + attributesOffset; } @@ -359,18 +356,6 @@ export class BVHComputeFns { } ); - let { - vertexStart = - 1, - vertexCount = - 1, - } = offsets; - - if ( vertexCount === - 1 ) { - - vertexStart = 0; - vertexCount = geometry.attributes.position.count; - - } - attributes.forEach( ( key, interleavedOffset ) => { const attr = geometry.attributes[ key ]; From 4b096b6247cd478922811a0ea3357a311e3a489e Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 11:07:04 +0900 Subject: [PATCH 07/84] Fix member name, paren --- src/webgpu/lib/BVHComputeFns.js | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index e80f9d1c5..cd66889bd 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -186,7 +186,7 @@ export class BVHComputeFns { appendBVHData( bvh, 0, true ); bvhs.forEach( ( bvh, i ) => { - bvhNodeOffsets.push( appendBVHData( bvh, geometryOffsets[ i ] ), false ); + bvhNodeOffsets.push( appendBVHData( bvh, geometryOffsets[ i ], false ) ); } ); @@ -250,7 +250,7 @@ export class BVHComputeFns { const rootBuffer16 = new Uint16Array( root ); const rootBuffer32 = new Uint32Array( root ); result.push( nodeWriteOffset ); - for ( let i = 0, l = root.byteSize / BYTES_PER_NODE; i < l; i ++ ) { + for ( let i = 0, l = root.byteLength / BYTES_PER_NODE; i < l; i ++ ) { const r32 = i * UINT32_PER_NODE; const r16 = r32 * 2; From 43838cefb6593f77e3baf50201faf0ecd9c2af83 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 11:11:16 +0900 Subject: [PATCH 08/84] Fix buffer types --- src/webgpu/lib/BVHComputeFns.js | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index cd66889bd..2fbb951b9 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -192,7 +192,9 @@ export class BVHComputeFns { // write the transforms let transformWriteOffset = 0; - const transformBuffer = new Uint32Array( 17 * objectTransformsCount ); + const transformArrayBuffer = new ArrayBuffer( 17 * 4 * objectTransformsCount ); + const transformBufferF32 = new Float32Array( transformArrayBuffer ); + const transformBufferU32 = new Uint32Array( transformArrayBuffer ); bvh.primitiveBuffer.forEach( ( compositeId, i ) => { const matrix = new Matrix4(); @@ -202,8 +204,8 @@ export class BVHComputeFns { const bvhOffset = bvhNodeOffsets[ i ]; objectBvh._roots.forEach( ( root, ri ) => { - matrix.toArray( transformBuffer, transformWriteOffset * 17 ); - transformBuffer[ transformWriteOffset * 17 + 16 ] = bvhOffset[ ri ]; + matrix.toArray( transformBufferF32, transformWriteOffset * 17 ); + transformBufferU32[ transformWriteOffset * 17 + 16 ] = bvhOffset[ ri ]; transformWriteOffset ++; } ); @@ -227,7 +229,7 @@ export class BVHComputeFns { } ` ); - this.storageBufferAttributes.transforms = new StorageBufferAttribute( transformBuffer, 17 ); + this.storageBufferAttributes.transforms = new StorageBufferAttribute( transformBufferF32, 17 ); this.storageBufferAttributes.nodes = new StorageBufferAttribute( nodeBuffer32, 8 ); this.storageBufferAttributes.index = new StorageBufferAttribute( indexBuffer, 1 ); this.storageBufferAttributes.attributes = new StorageBufferAttribute( attributesBuffer, attributesStructSize ); @@ -263,15 +265,16 @@ export class BVHComputeFns { const isLeaf = IS_LEAFNODE_FLAG === rootBuffer16[ r32 + 15 ]; if ( isLeaf ) { - nodeBuffer32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ]; - nodeBuffer16[ n16 + 14 ] = rootBuffer16[ r16 + 14 ] + geometryOffsets[ rootIndex ]; + nodeBuffer16[ n16 + 14 ] = rootBuffer16[ r16 + 14 ]; if ( tlas ) { // 0xFFFF == mesh leaf, 0xFF00 == TLAS leaf + nodeBuffer32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ]; nodeBuffer16[ n16 + 15 ] = 0xFF00; } else { + nodeBuffer32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ] + geometryOffsets[ rootIndex ]; nodeBuffer16[ n16 + 15 ] = IS_LEAFNODE_FLAG; } From 0f717209b1a633f251847ee28c0a3e8088c66883 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 11:13:22 +0900 Subject: [PATCH 09/84] Fix switch statement --- src/webgpu/lib/BVHComputeFns.js | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 2fbb951b9..ea61e66b7 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -381,12 +381,18 @@ export class BVHComputeFns { _vec.fromBufferAttribute( attr, i - vertexStart ); switch ( attr.itemSize ) { - case 3: + case 1: + _vec.y = 0; + _vec.z = 0; _vec.w = 0; + break; case 2: _vec.z = 0; - case 1: - _vec.y = 0; + _vec.w = 0; + break; + case 3: + _vec.w = 0; + break; } From 6e5cd11a34a0eb8ffff49080504117c86376387e Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 11:18:28 +0900 Subject: [PATCH 10/84] Other fixes --- src/webgpu/lib/BVHComputeFns.js | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index ea61e66b7..cc3c3f313 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -201,7 +201,7 @@ export class BVHComputeFns { bvh.getObjectMatrix( compositeId, matrix ); const objectBvh = bvhs[ transformBVHs[ i ] ]; - const bvhOffset = bvhNodeOffsets[ i ]; + const bvhOffset = bvhNodeOffsets[ transformBVHs[ i ] ]; objectBvh._roots.forEach( ( root, ri ) => { matrix.toArray( transformBufferF32, transformWriteOffset * 17 ); @@ -219,7 +219,7 @@ export class BVHComputeFns { .map( key => { attributesStructSize += 16; - return `${ key }: vec4,`; + return `${ key }: vec4f,`; } ).join( '\n' ); @@ -378,7 +378,7 @@ export class BVHComputeFns { } else { - _vec.fromBufferAttribute( attr, i - vertexStart ); + _vec.fromBufferAttribute( attr, i + vertexStart ); switch ( attr.itemSize ) { case 1: From 4ddfd785c7a77bc8a6cd9dcfd577c7448b50a03b Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 11:25:47 +0900 Subject: [PATCH 11/84] Fixes --- src/webgpu/lib/BVHComputeFns.js | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index cc3c3f313..0db61881a 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -218,7 +218,7 @@ export class BVHComputeFns { const attributeStructContent = attributes .map( key => { - attributesStructSize += 16; + attributesStructSize += 4; return `${ key }: vec4f,`; } ).join( '\n' ); @@ -262,7 +262,7 @@ export class BVHComputeFns { // write bounds nodeBufferFloat.set( new Float32Array( root, i * BYTES_PER_NODE, 6 ), n32 ); - const isLeaf = IS_LEAFNODE_FLAG === rootBuffer16[ r32 + 15 ]; + const isLeaf = IS_LEAFNODE_FLAG === rootBuffer16[ r16 + 15 ]; if ( isLeaf ) { nodeBuffer16[ n16 + 14 ] = rootBuffer16[ r16 + 14 ]; From 01118fd2e161b40c4abf10742e1fbcf240ecfa07 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 11:29:29 +0900 Subject: [PATCH 12/84] More fixes --- src/webgpu/lib/BVHComputeFns.js | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 0db61881a..ea6189093 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -168,7 +168,7 @@ export class BVHComputeFns { let attributesOffset = 0; let indexOffset = 0; const indexBuffer = new Uint32Array( indexBufferLength ); - const attributesBuffer = new Float32Array( attributesBufferLength ); + const attributesBuffer = new Float32Array( attributesBufferLength * attributes.length * 4 ); geometries.forEach( ( { geometry, range } ) => { const offset = appendGeometryData( geometry, range ); @@ -297,7 +297,7 @@ export class BVHComputeFns { } - function appendGeometryData( geometry, offsets = {} ) { + function appendGeometryData( geometry, offsets = null ) { const result = []; const groups = geometry.groups.length === 0 ? [ { start: 0, count: Infinity } ] : geometry.groups; From 2a85e75880567cf49a605d1d16b1c38474f5d38d Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 11:37:59 +0900 Subject: [PATCH 13/84] Fix variable use --- src/webgpu/lib/BVHComputeFns.js | 31 ++++++++++++++++--------------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index ea6189093..062fc6ae8 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -76,7 +76,6 @@ export class BVHComputeFns { this.attributesStruct = null; this.raycastFirstHitFn = null; - this.update(); this.getBVH = ( object, id = - 1 ) => { @@ -97,6 +96,8 @@ export class BVHComputeFns { }; + this.update(); + } update() { @@ -164,6 +165,17 @@ export class BVHComputeFns { } ); + // construct the attribute struct + // TODO: need to include materials here + let attributesStructSize = 0; + const attributeStructContent = attributes + .map( key => { + + attributesStructSize += 4; + return `${ key }: vec4f,`; + + } ).join( '\n' ); + // write the geometry buffer attributes let attributesOffset = 0; let indexOffset = 0; @@ -212,17 +224,6 @@ export class BVHComputeFns { } ); - // construct the attribute struct - // TODO: need to include materials here - let attributesStructSize = 0; - const attributeStructContent = attributes - .map( key => { - - attributesStructSize += 4; - return `${ key }: vec4f,`; - - } ).join( '\n' ); - const attributeStruct = wgsl( /* wgsl */` struct ${ name }GeometryStruct { ${ attributeStructContent } @@ -324,7 +325,7 @@ export class BVHComputeFns { if ( offsets ) { indexStart = offsets.indexStart; - indexCount = geometry.indexCount; + indexCount = offsets.indexCount; } @@ -398,13 +399,13 @@ export class BVHComputeFns { } - _vec.toArray( attributesBuffer, attributesOffset * attributesStructSize + interleavedOffset * 4 ); + _vec.toArray( attributesBuffer, ( attributesOffset + i ) * attributesStructSize + interleavedOffset * 4 ); } } ); - attributesOffset += geometry.attributes.position.count; + attributesOffset += vertexCount; return result; } From cc4e5ca43f2369fcf9731d9ecd76ed28a68678cf Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 11:42:11 +0900 Subject: [PATCH 14/84] Fix stride --- src/webgpu/lib/BVHComputeFns.js | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 062fc6ae8..76c15d496 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -116,7 +116,7 @@ export class BVHComputeFns { let objectTransformsCount = 0; let indexBufferLength = 0; let attributesBufferLength = 0; - bvhNodesBufferLength += bvh._roots[ 0 ].byteLength; + bvhNodesBufferLength += bvh._roots.reduce( ( v, root ) => v + root.byteLength, 0 ); bvh.primitiveBuffer.forEach( compositeId => { const object = bvh.getObjectFromId( compositeId ); @@ -203,8 +203,10 @@ export class BVHComputeFns { } ); // write the transforms + // stride is 20 floats (80 bytes) to match WGSL struct alignment: mat4x4f (64) + u32 (4) + 12 bytes padding to align to 16 + const TRANSFORM_STRIDE = 20; let transformWriteOffset = 0; - const transformArrayBuffer = new ArrayBuffer( 17 * 4 * objectTransformsCount ); + const transformArrayBuffer = new ArrayBuffer( TRANSFORM_STRIDE * 4 * objectTransformsCount ); const transformBufferF32 = new Float32Array( transformArrayBuffer ); const transformBufferU32 = new Uint32Array( transformArrayBuffer ); bvh.primitiveBuffer.forEach( ( compositeId, i ) => { @@ -216,8 +218,8 @@ export class BVHComputeFns { const bvhOffset = bvhNodeOffsets[ transformBVHs[ i ] ]; objectBvh._roots.forEach( ( root, ri ) => { - matrix.toArray( transformBufferF32, transformWriteOffset * 17 ); - transformBufferU32[ transformWriteOffset * 17 + 16 ] = bvhOffset[ ri ]; + matrix.toArray( transformBufferF32, transformWriteOffset * TRANSFORM_STRIDE ); + transformBufferU32[ transformWriteOffset * TRANSFORM_STRIDE + 16 ] = bvhOffset[ ri ]; transformWriteOffset ++; } ); @@ -230,7 +232,7 @@ export class BVHComputeFns { } ` ); - this.storageBufferAttributes.transforms = new StorageBufferAttribute( transformBufferF32, 17 ); + this.storageBufferAttributes.transforms = new StorageBufferAttribute( transformBufferF32, TRANSFORM_STRIDE ); this.storageBufferAttributes.nodes = new StorageBufferAttribute( nodeBuffer32, 8 ); this.storageBufferAttributes.index = new StorageBufferAttribute( indexBuffer, 1 ); this.storageBufferAttributes.attributes = new StorageBufferAttribute( attributesBuffer, attributesStructSize ); From 8193940d5e2058ca03faa43dac795d4fb2f4fde1 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 12:11:13 +0900 Subject: [PATCH 15/84] Add indirect buffer support --- src/webgpu/lib/BVHComputeFns.js | 95 ++++++++++++++++++++++----------- 1 file changed, 64 insertions(+), 31 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 76c15d496..e6276b195 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -51,9 +51,32 @@ export const transformStruct = wgsl( /* wgsl */` struct TransformStruct { matrixWorld: mat4x4f, nodeOffset: u32, + _alignment0: u32, + _alignment1: u32, + _alignment2: u32, } ` ); +function dereferenceIndex( indexAttr, indirectBuffer ) { + + const indexArray = indexAttr ? indexAttr.array : null; + const result = new Uint32Array( indirectBuffer.length * 3 ); + 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 ++ ) { + + result[ i3 + c ] = indexArray ? indexArray[ v3 + c ] : v3 + c; + + } + + } + + return result; + +} + export class BVHComputeFns { constructor( bvh, options = {} ) { @@ -134,6 +157,8 @@ export class BVHComputeFns { bvhNodesBufferLength += meshBvh._roots.reduce( ( v, root ) => v + root.byteLength, 0 ); // save the geometry info to write later and increment the buffer sizes + const indirectBuffer = meshBvh._indirectBuffer || null; + if ( object.isBatchedMesh ) { const geometryId = object.getGeometryIdAt( instanceId ); @@ -141,6 +166,7 @@ export class BVHComputeFns { geometries.push( { geometry: object.geometry, range: range, + indirectBuffer: indirectBuffer, } ); indexBufferLength += range.indexCount === - 1 ? range.vertexCount : range.indexCount; @@ -154,6 +180,7 @@ export class BVHComputeFns { geometries.push( { geometry: object.geometry, range: undefined, + indirectBuffer: indirectBuffer, } ); } @@ -181,9 +208,9 @@ export class BVHComputeFns { let indexOffset = 0; const indexBuffer = new Uint32Array( indexBufferLength ); const attributesBuffer = new Float32Array( attributesBufferLength * attributes.length * 4 ); - geometries.forEach( ( { geometry, range } ) => { + geometries.forEach( ( { geometry, range, indirectBuffer } ) => { - const offset = appendGeometryData( geometry, range ); + const offset = appendGeometryData( geometry, range, indirectBuffer ); geometryOffsets.push( offset ); } ); @@ -300,10 +327,9 @@ export class BVHComputeFns { } - function appendGeometryData( geometry, offsets = null ) { + function appendGeometryData( geometry, offsets = null, indirectBuffer = null ) { const result = []; - const groups = geometry.groups.length === 0 ? [ { start: 0, count: Infinity } ] : geometry.groups; let vertexStart = 0; let vertexCount = geometry.attributes.position.count; @@ -314,53 +340,60 @@ export class BVHComputeFns { } - groups.forEach( group => { + // write indices — when an indirect buffer is present, dereference it to + // resolve the BVH's triangle order into a flat index buffer + result.push( indexOffset ); - result.push( indexOffset ); + if ( indirectBuffer ) { - // TODO: handle "indirect" case - // TODO: validate the write offsets here - if ( geometry.index ) { + const dereferencedIndex = dereferenceIndex( geometry.index, indirectBuffer ); + for ( let i = 0; i < dereferencedIndex.length; i ++ ) { - let indexStart = Math.max( 0, group.start ); - let indexCount = Math.min( geometry.index.count, group.start + group.count ) - indexStart; - if ( offsets ) { + indexBuffer[ i + indexOffset ] = dereferencedIndex[ i ] - vertexStart + attributesOffset; - indexStart = offsets.indexStart; - indexCount = offsets.indexCount; + } - } + indexOffset += dereferencedIndex.length; - for ( let i = 0; i < indexCount; i ++ ) { + } else if ( geometry.index ) { - indexBuffer[ i + indexOffset ] = geometry.index.getX( i + indexStart ) - vertexStart + attributesOffset; + let indexStart = 0; + let indexCount = geometry.index.count; + if ( offsets ) { - } + indexStart = offsets.indexStart; + indexCount = offsets.indexCount; - indexOffset += indexCount; + } - } else { + for ( let i = 0; i < indexCount; i ++ ) { - let indexStart = Math.max( 0, group.start ); - let indexCount = Math.min( geometry.attributes.position.count, group.start + group.count ) - indexStart; - if ( offsets ) { + indexBuffer[ i + indexOffset ] = geometry.index.getX( i + indexStart ) - vertexStart + attributesOffset; - indexStart = offsets.vertexStart; - indexCount = offsets.vertexCount; + } - } + indexOffset += indexCount; + + } else { - for ( let i = 0; i < indexCount; i ++ ) { + let indexStart = 0; + let indexCount = geometry.attributes.position.count; + if ( offsets ) { - indexBuffer[ i + indexOffset ] = i + indexStart + attributesOffset; + indexStart = offsets.vertexStart; + indexCount = offsets.vertexCount; - } + } + + for ( let i = 0; i < indexCount; i ++ ) { - indexOffset += indexCount; + indexBuffer[ i + indexOffset ] = i + indexStart + attributesOffset; } - } ); + indexOffset += indexCount; + + } attributes.forEach( ( key, interleavedOffset ) => { From 65f986c132a1180c09ade51ae6f669ac0f1882b4 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 12:43:00 +0900 Subject: [PATCH 16/84] comments --- src/webgpu/lib/BVHComputeFns.js | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index e6276b195..68ac416b6 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -3,6 +3,11 @@ import { StorageBufferAttribute } from 'three/src/Three.WebGPU.Nodes.js'; import { wgsl, wgslFn } from 'three/tsl'; import { intersectTriangles, intersectsBounds } from 'three-mesh-bvh/webgpu'; +// TODO: separate update functions into utilities +// TODO: add ability to easily update a single matrix / scene rearrangement +// TODO: add material support w/ function to easily update material +// TODO: add skinned mesh bvh support +// TODO: add overrideable functions for custom implementations (custom attributes, transform fields) const _vec = /* @__PURE__ */ new Vector4(); export const constants = wgsl( /* wgsl */` @@ -127,8 +132,6 @@ export class BVHComputeFns { const { attributes, name, bvh } = this; - // TODO: add support for materials? Optional? Custom callback? - // TODO: how to handle skinned meshes? const bvhs = []; const geometries = []; const geometryOffsets = []; @@ -193,7 +196,6 @@ export class BVHComputeFns { } ); // construct the attribute struct - // TODO: need to include materials here let attributesStructSize = 0; const attributeStructContent = attributes .map( key => { From 380b5503822bbdef605447880303188118097a77 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 15:53:22 +0900 Subject: [PATCH 17/84] Add intersection function --- src/webgpu/lib/BVHComputeFns.js | 237 ++++++++++++++++++++++++-------- 1 file changed, 182 insertions(+), 55 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 68ac416b6..b644a0ea2 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -1,7 +1,15 @@ import { Matrix4, Vector4 } from 'three'; import { StorageBufferAttribute } from 'three/src/Three.WebGPU.Nodes.js'; -import { wgsl, wgslFn } from 'three/tsl'; -import { intersectTriangles, intersectsBounds } from 'three-mesh-bvh/webgpu'; +import { storage, wgsl, wgslFn } from 'three/tsl'; +import { + intersectsTriangle, + intersectsBounds, + rayStruct, + bvhNodeBoundsStruct, + bvhNodeStruct, + intersectionResultStruct, + constants, +} from 'three-mesh-bvh/webgpu'; // TODO: separate update functions into utilities // TODO: add ability to easily update a single matrix / scene rearrangement @@ -9,52 +17,15 @@ import { intersectTriangles, intersectsBounds } from 'three-mesh-bvh/webgpu'; // TODO: add skinned mesh bvh support // TODO: add overrideable functions for custom implementations (custom attributes, transform fields) const _vec = /* @__PURE__ */ new Vector4(); +const _matrix = /* @__PURE__ */ new Matrix4(); +const _inverseMatrix = /* @__PURE__ */ new Matrix4(); -export const constants = wgsl( /* wgsl */` - const BVH_STACK_DEPTH = 60u; - const INFINITY = 1e20; - const TRI_INTERSECT_EPSILON = 1e-5; -` ); - -export const rayStruct = wgsl( /* wgsl */` - struct Ray { - origin: vec3f, - _alignment0: u32, - - direction: vec3f, - _alignment1: u32, - }; -` ); - -export const bvhNodeBoundsStruct = wgsl( /* wgsl */` - struct BVHBoundingBox { - min: array, - max: array, - } -` ); - -export const bvhNodeStruct = wgsl( /* wgsl */` - struct BVHNode { - bounds: BVHBoundingBox, - rightChildOrTriangleOffset: u32, - splitAxisOrTriangleCount: u32, - }; -`, [ bvhNodeBoundsStruct ] ); - -export const intersectionResultStruct = wgsl( /* wgsl */` - struct IntersectionResult { - indices: vec4u, - normal: vec3f, - didHit: bool, - barycoord: vec3f, - side: f32, - dist: f32, - }; -` ); +export { rayStruct, bvhNodeBoundsStruct, bvhNodeStruct, intersectionResultStruct, constants }; export const transformStruct = wgsl( /* wgsl */` struct TransformStruct { matrixWorld: mat4x4f, + inverseMatrixWorld: mat4x4f, nodeOffset: u32, _alignment0: u32, _alignment1: u32, @@ -232,23 +203,25 @@ export class BVHComputeFns { } ); // write the transforms - // stride is 20 floats (80 bytes) to match WGSL struct alignment: mat4x4f (64) + u32 (4) + 12 bytes padding to align to 16 - const TRANSFORM_STRIDE = 20; + // stride is 36 floats (144 bytes) to match WGSL struct alignment: + // mat4x4f (64) + mat4x4f (64) + u32 (4) + 12 bytes padding to align to 16 + const TRANSFORM_STRIDE = 36; let transformWriteOffset = 0; const transformArrayBuffer = new ArrayBuffer( TRANSFORM_STRIDE * 4 * objectTransformsCount ); const transformBufferF32 = new Float32Array( transformArrayBuffer ); const transformBufferU32 = new Uint32Array( transformArrayBuffer ); bvh.primitiveBuffer.forEach( ( compositeId, i ) => { - const matrix = new Matrix4(); - bvh.getObjectMatrix( compositeId, matrix ); + bvh.getObjectMatrix( compositeId, _matrix ); + _inverseMatrix.copy( _matrix ).invert(); const objectBvh = bvhs[ transformBVHs[ i ] ]; const bvhOffset = bvhNodeOffsets[ transformBVHs[ i ] ]; objectBvh._roots.forEach( ( root, ri ) => { - matrix.toArray( transformBufferF32, transformWriteOffset * TRANSFORM_STRIDE ); - transformBufferU32[ transformWriteOffset * TRANSFORM_STRIDE + 16 ] = bvhOffset[ ri ]; + _matrix.toArray( transformBufferF32, transformWriteOffset * TRANSFORM_STRIDE ); + _inverseMatrix.toArray( transformBufferF32, transformWriteOffset * TRANSFORM_STRIDE + 16 ); + transformBufferU32[ transformWriteOffset * TRANSFORM_STRIDE + 32 ] = bvhOffset[ ri ]; transformWriteOffset ++; } ); @@ -261,16 +234,169 @@ export class BVHComputeFns { } ` ); - this.storageBufferAttributes.transforms = new StorageBufferAttribute( transformBufferF32, TRANSFORM_STRIDE ); - this.storageBufferAttributes.nodes = new StorageBufferAttribute( nodeBuffer32, 8 ); - this.storageBufferAttributes.index = new StorageBufferAttribute( indexBuffer, 1 ); - this.storageBufferAttributes.attributes = new StorageBufferAttribute( attributesBuffer, attributesStructSize ); + const nodesStorage = storage( new StorageBufferAttribute( nodeBuffer32, 8 ), `${ name }BVHNode` ).toReadOnly().setName( `${ name }nodes` ); + const transformsStorage = storage( new StorageBufferAttribute( transformBufferF32, TRANSFORM_STRIDE ), `${ name }TransformStruct` ).toReadOnly().setName( `${ name }transforms` ); + const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ name }index` ); + const attributesStorage = storage( new StorageBufferAttribute( attributesBuffer, attributesStructSize ), `${ name }GeometryStruct` ).toReadOnly().setName( `${ name }attributes` ); + + this.storageBufferAttributes.transforms = transformsStorage; + this.storageBufferAttributes.nodes = nodesStorage; + this.storageBufferAttributes.index = indexStorage; + this.storageBufferAttributes.attributes = attributesStorage; this.attributesStruct = attributeStruct; this.raycastFirstHitFn = wgslFn( /* wgsl */` fn ${ name }RaycastFirstHit( ray: Ray ) -> IntersectionResult { + var bestHit: IntersectionResult; + bestHit.didHit = false; + bestHit.dist = INFINITY; + + // TLAS traversal + var tlasPointer: i32 = 0; + var tlasStack: array; + tlasStack[ 0 ] = 0u; + + loop { + + if ( tlasPointer < 0 || tlasPointer >= i32( BVH_STACK_DEPTH ) ) { + + break; + + } + + let currNodeIndex = tlasStack[ tlasPointer ]; + let node = ${ name }nodes[ currNodeIndex ]; + tlasPointer = tlasPointer - 1; + + var boundsHitDist: f32 = 0.0; + if ( ! intersectsBounds( ray, node.bounds, &boundsHitDist ) || boundsHitDist > bestHit.dist ) { + + continue; + + } + + let boundsInfox = node.splitAxisOrTriangleCount; + let boundsInfoy = node.rightChildOrTriangleOffset; + let isLeaf = ( boundsInfox & 0xffff0000u ) != 0u; + + if ( isLeaf ) { + + // TLAS leaf — iterate over object transforms + let count = boundsInfox & 0x0000ffffu; + let offset = boundsInfoy; + + for ( var t = offset; t < offset + count; t = t + 1u ) { + + let transform = ${ name }transforms[ t ]; + + // Transform ray into object local space + var localRay: Ray; + localRay.origin = ( transform.inverseMatrixWorld * vec4f( ray.origin, 1.0 ) ).xyz; + localRay.direction = ( transform.inverseMatrixWorld * vec4f( ray.direction, 0.0 ) ).xyz; + + // BLAS traversal + var blasPointer: i32 = 0; + var blasStack: array; + blasStack[ 0 ] = transform.nodeOffset; + + loop { + + if ( blasPointer < 0 || blasPointer >= i32( BVH_STACK_DEPTH ) ) { + + break; + + } + + let blasNodeIndex = blasStack[ blasPointer ]; + let blasNode = ${ name }nodes[ blasNodeIndex ]; + blasPointer = blasPointer - 1; + + var blasBoundsHitDist: f32 = 0.0; + if ( ! intersectsBounds( localRay, blasNode.bounds, &blasBoundsHitDist ) || blasBoundsHitDist > bestHit.dist ) { + + continue; + + } + + let blasInfox = blasNode.splitAxisOrTriangleCount; + let blasInfoy = blasNode.rightChildOrTriangleOffset; + let isBlasLeaf = ( blasInfox & 0xffff0000u ) != 0u; + + if ( isBlasLeaf ) { + + let triCount = blasInfox & 0x0000ffffu; + let triOffset = blasInfoy; + + for ( var ti = triOffset; ti < triOffset + triCount; ti = ti + 1u ) { + + let i0 = ${ name }index[ ti * 3u ]; + let i1 = ${ name }index[ ti * 3u + 1u ]; + let i2 = ${ name }index[ ti * 3u + 2u ]; + + let a = ${ name }attributes[ i0 ].position.xyz; + let b = ${ name }attributes[ i1 ].position.xyz; + let c = ${ name }attributes[ i2 ].position.xyz; + + var triResult = intersectsTriangle( localRay, a, b, c ); + + if ( triResult.didHit && triResult.dist < bestHit.dist ) { + + bestHit = triResult; + bestHit.indices = vec4u( i0, i1, i2, ti ); + + // Transform normal to world space: normal matrix = transpose( inverse ) + bestHit.normal = normalize( ( transpose( transform.inverseMatrixWorld ) * vec4f( bestHit.normal, 0.0 ) ).xyz ); + + } + + } + + } else { + + let leftIndex = blasNodeIndex + 1u; + let splitAxis = blasInfox & 0x0000ffffu; + let rightIndex = blasNodeIndex + blasInfoy; + + let leftToRight = localRay.direction[ splitAxis ] >= 0.0; + let c1 = select( rightIndex, leftIndex, leftToRight ); + let c2 = select( leftIndex, rightIndex, leftToRight ); + + blasPointer = blasPointer + 1; + blasStack[ blasPointer ] = c2; + + blasPointer = blasPointer + 1; + blasStack[ blasPointer ] = c1; + + } + + } + + } + + } else { + + let leftIndex = currNodeIndex + 1u; + let splitAxis = boundsInfox & 0x0000ffffu; + let rightIndex = currNodeIndex + boundsInfoy; + + let leftToRight = ray.direction[ splitAxis ] >= 0.0; + let c1 = select( rightIndex, leftIndex, leftToRight ); + let c2 = select( leftIndex, rightIndex, leftToRight ); + + tlasPointer = tlasPointer + 1; + tlasStack[ tlasPointer ] = c2; + + tlasPointer = tlasPointer + 1; + tlasStack[ tlasPointer ] = c1; + + } + + } + + return bestHit; + } - `, [ intersectTriangles, intersectsBounds, rayStruct, bvhNodeStruct, intersectionResultStruct, constants ] ); + `, [ nodesStorage, transformsStorage, indexStorage, attributesStorage, intersectsTriangle, intersectsBounds, rayStruct, bvhNodeStruct, intersectionResultStruct, constants, transformStruct, attributeStruct ] ); function appendBVHData( bvh, geometryOffsets, tlas = false ) { @@ -344,7 +470,8 @@ export class BVHComputeFns { // write indices — when an indirect buffer is present, dereference it to // resolve the BVH's triangle order into a flat index buffer - result.push( indexOffset ); + // store as triangle offset (not vertex-index offset) to match BVH leaf format + result.push( indexOffset / 3 ); if ( indirectBuffer ) { From 38137040a9e10e6674cb2e327af0fa5c82c86c4f Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 16:04:29 +0900 Subject: [PATCH 18/84] Separate functions --- src/webgpu/lib/BVHComputeFns.js | 340 ++++++++++++++++++-------------- 1 file changed, 187 insertions(+), 153 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index b644a0ea2..6291f7960 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -53,6 +53,192 @@ function dereferenceIndex( indexAttr, indirectBuffer ) { } +function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct ) { + + const intersectFirstHitFn = wgslFn( /* wgsl */` + fn ${ name }IntersectFirstHit( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> IntersectionResult { + + var bestHit: IntersectionResult; + bestHit.didHit = false; + bestHit.dist = bestDist; + + var pointer: i32 = 0; + var stack: array; + stack[ 0 ] = rootNodeIndex; + + loop { + + if ( pointer < 0 || pointer >= i32( BVH_STACK_DEPTH ) ) { + + break; + + } + + let nodeIndex = stack[ pointer ]; + let node = ${ name }nodes[ nodeIndex ]; + pointer = pointer - 1; + + var boundsHitDist: f32 = 0.0; + if ( ! intersectsBounds( ray, node.bounds, &boundsHitDist ) || boundsHitDist > bestHit.dist ) { + + continue; + + } + + let infoX = node.splitAxisOrTriangleCount; + let infoY = node.rightChildOrTriangleOffset; + let isLeaf = ( infoX & 0xffff0000u ) != 0u; + + if ( isLeaf ) { + + let triCount = infoX & 0x0000ffffu; + let triOffset = infoY; + + for ( var ti = triOffset; ti < triOffset + triCount; ti = ti + 1u ) { + + let i0 = ${ name }index[ ti * 3u ]; + let i1 = ${ name }index[ ti * 3u + 1u ]; + let i2 = ${ name }index[ ti * 3u + 2u ]; + + let a = ${ name }attributes[ i0 ].position.xyz; + let b = ${ name }attributes[ i1 ].position.xyz; + let c = ${ name }attributes[ i2 ].position.xyz; + + var triResult = intersectsTriangle( ray, a, b, c ); + + if ( triResult.didHit && triResult.dist < bestHit.dist ) { + + bestHit = triResult; + bestHit.indices = vec4u( i0, i1, i2, ti ); + + } + + } + + } else { + + let leftIndex = nodeIndex + 1u; + let splitAxis = infoX & 0x0000ffffu; + let rightIndex = nodeIndex + infoY; + + let leftToRight = ray.direction[ splitAxis ] >= 0.0; + let c1 = select( rightIndex, leftIndex, leftToRight ); + let c2 = select( leftIndex, rightIndex, leftToRight ); + + pointer = pointer + 1; + stack[ pointer ] = c2; + + pointer = pointer + 1; + stack[ pointer ] = c1; + + } + + } + + return bestHit; + + } + `, [ + nodesStorage, indexStorage, attributesStorage, + intersectsTriangle, intersectsBounds, + rayStruct, bvhNodeStruct, intersectionResultStruct, constants, + attributeStruct, + ] ); + + return wgslFn( /* wgsl */` + fn ${ name }RaycastFirstHit( ray: Ray ) -> IntersectionResult { + + var bestHit: IntersectionResult; + bestHit.didHit = false; + bestHit.dist = INFINITY; + + var tlasPointer: i32 = 0; + var tlasStack: array; + tlasStack[ 0 ] = 0u; + + loop { + + if ( tlasPointer < 0 || tlasPointer >= i32( BVH_STACK_DEPTH ) ) { + + break; + + } + + let currNodeIndex = tlasStack[ tlasPointer ]; + let node = ${ name }nodes[ currNodeIndex ]; + tlasPointer = tlasPointer - 1; + + var boundsHitDist: f32 = 0.0; + if ( ! intersectsBounds( ray, node.bounds, &boundsHitDist ) || boundsHitDist > bestHit.dist ) { + + continue; + + } + + let infoX = node.splitAxisOrTriangleCount; + let infoY = node.rightChildOrTriangleOffset; + let isLeaf = ( infoX & 0xffff0000u ) != 0u; + + if ( isLeaf ) { + + let count = infoX & 0x0000ffffu; + let offset = infoY; + + for ( var t = offset; t < offset + count; t = t + 1u ) { + + let transform = ${ name }transforms[ t ]; + + // Transform ray into object local space + var localRay: Ray; + localRay.origin = ( transform.inverseMatrixWorld * vec4f( ray.origin, 1.0 ) ).xyz; + localRay.direction = ( transform.inverseMatrixWorld * vec4f( ray.direction, 0.0 ) ).xyz; + + let blasHit = ${ name }IntersectFirstHit( localRay, transform.nodeOffset, bestHit.dist ); + + if ( blasHit.didHit && blasHit.dist < bestHit.dist ) { + + bestHit = blasHit; + + // Transform normal to world space: normal matrix = transpose( inverse ) + bestHit.normal = normalize( ( transpose( transform.inverseMatrixWorld ) * vec4f( bestHit.normal, 0.0 ) ).xyz ); + + } + + } + + } else { + + let leftIndex = currNodeIndex + 1u; + let splitAxis = infoX & 0x0000ffffu; + let rightIndex = currNodeIndex + infoY; + + let leftToRight = ray.direction[ splitAxis ] >= 0.0; + let c1 = select( rightIndex, leftIndex, leftToRight ); + let c2 = select( leftIndex, rightIndex, leftToRight ); + + tlasPointer = tlasPointer + 1; + tlasStack[ tlasPointer ] = c2; + + tlasPointer = tlasPointer + 1; + tlasStack[ tlasPointer ] = c1; + + } + + } + + return bestHit; + + } + `, [ + intersectFirstHitFn, + nodesStorage, transformsStorage, + intersectsBounds, + rayStruct, bvhNodeStruct, intersectionResultStruct, constants, + transformStruct, + ] ); + +} + export class BVHComputeFns { constructor( bvh, options = {} ) { @@ -244,159 +430,7 @@ export class BVHComputeFns { this.storageBufferAttributes.index = indexStorage; this.storageBufferAttributes.attributes = attributesStorage; this.attributesStruct = attributeStruct; - this.raycastFirstHitFn = wgslFn( /* wgsl */` - fn ${ name }RaycastFirstHit( ray: Ray ) -> IntersectionResult { - - var bestHit: IntersectionResult; - bestHit.didHit = false; - bestHit.dist = INFINITY; - - // TLAS traversal - var tlasPointer: i32 = 0; - var tlasStack: array; - tlasStack[ 0 ] = 0u; - - loop { - - if ( tlasPointer < 0 || tlasPointer >= i32( BVH_STACK_DEPTH ) ) { - - break; - - } - - let currNodeIndex = tlasStack[ tlasPointer ]; - let node = ${ name }nodes[ currNodeIndex ]; - tlasPointer = tlasPointer - 1; - - var boundsHitDist: f32 = 0.0; - if ( ! intersectsBounds( ray, node.bounds, &boundsHitDist ) || boundsHitDist > bestHit.dist ) { - - continue; - - } - - let boundsInfox = node.splitAxisOrTriangleCount; - let boundsInfoy = node.rightChildOrTriangleOffset; - let isLeaf = ( boundsInfox & 0xffff0000u ) != 0u; - - if ( isLeaf ) { - - // TLAS leaf — iterate over object transforms - let count = boundsInfox & 0x0000ffffu; - let offset = boundsInfoy; - - for ( var t = offset; t < offset + count; t = t + 1u ) { - - let transform = ${ name }transforms[ t ]; - - // Transform ray into object local space - var localRay: Ray; - localRay.origin = ( transform.inverseMatrixWorld * vec4f( ray.origin, 1.0 ) ).xyz; - localRay.direction = ( transform.inverseMatrixWorld * vec4f( ray.direction, 0.0 ) ).xyz; - - // BLAS traversal - var blasPointer: i32 = 0; - var blasStack: array; - blasStack[ 0 ] = transform.nodeOffset; - - loop { - - if ( blasPointer < 0 || blasPointer >= i32( BVH_STACK_DEPTH ) ) { - - break; - - } - - let blasNodeIndex = blasStack[ blasPointer ]; - let blasNode = ${ name }nodes[ blasNodeIndex ]; - blasPointer = blasPointer - 1; - - var blasBoundsHitDist: f32 = 0.0; - if ( ! intersectsBounds( localRay, blasNode.bounds, &blasBoundsHitDist ) || blasBoundsHitDist > bestHit.dist ) { - - continue; - - } - - let blasInfox = blasNode.splitAxisOrTriangleCount; - let blasInfoy = blasNode.rightChildOrTriangleOffset; - let isBlasLeaf = ( blasInfox & 0xffff0000u ) != 0u; - - if ( isBlasLeaf ) { - - let triCount = blasInfox & 0x0000ffffu; - let triOffset = blasInfoy; - - for ( var ti = triOffset; ti < triOffset + triCount; ti = ti + 1u ) { - - let i0 = ${ name }index[ ti * 3u ]; - let i1 = ${ name }index[ ti * 3u + 1u ]; - let i2 = ${ name }index[ ti * 3u + 2u ]; - - let a = ${ name }attributes[ i0 ].position.xyz; - let b = ${ name }attributes[ i1 ].position.xyz; - let c = ${ name }attributes[ i2 ].position.xyz; - - var triResult = intersectsTriangle( localRay, a, b, c ); - - if ( triResult.didHit && triResult.dist < bestHit.dist ) { - - bestHit = triResult; - bestHit.indices = vec4u( i0, i1, i2, ti ); - - // Transform normal to world space: normal matrix = transpose( inverse ) - bestHit.normal = normalize( ( transpose( transform.inverseMatrixWorld ) * vec4f( bestHit.normal, 0.0 ) ).xyz ); - - } - - } - - } else { - - let leftIndex = blasNodeIndex + 1u; - let splitAxis = blasInfox & 0x0000ffffu; - let rightIndex = blasNodeIndex + blasInfoy; - - let leftToRight = localRay.direction[ splitAxis ] >= 0.0; - let c1 = select( rightIndex, leftIndex, leftToRight ); - let c2 = select( leftIndex, rightIndex, leftToRight ); - - blasPointer = blasPointer + 1; - blasStack[ blasPointer ] = c2; - - blasPointer = blasPointer + 1; - blasStack[ blasPointer ] = c1; - - } - - } - - } - - } else { - - let leftIndex = currNodeIndex + 1u; - let splitAxis = boundsInfox & 0x0000ffffu; - let rightIndex = currNodeIndex + boundsInfoy; - - let leftToRight = ray.direction[ splitAxis ] >= 0.0; - let c1 = select( rightIndex, leftIndex, leftToRight ); - let c2 = select( leftIndex, rightIndex, leftToRight ); - - tlasPointer = tlasPointer + 1; - tlasStack[ tlasPointer ] = c2; - - tlasPointer = tlasPointer + 1; - tlasStack[ tlasPointer ] = c1; - - } - - } - - return bestHit; - - } - `, [ nodesStorage, transformsStorage, indexStorage, attributesStorage, intersectsTriangle, intersectsBounds, rayStruct, bvhNodeStruct, intersectionResultStruct, constants, transformStruct, attributeStruct ] ); + this.raycastFirstHitFn = buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct ); function appendBVHData( bvh, geometryOffsets, tlas = false ) { From 86f60b48e34e3449141f3310764cd797d91d04f3 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 16:15:54 +0900 Subject: [PATCH 19/84] Updates --- src/webgpu/lib/BVHComputeFns.js | 39 +++++++++++++++++++-------------- 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 6291f7960..188bffefe 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -5,24 +5,32 @@ import { intersectsTriangle, intersectsBounds, rayStruct, - bvhNodeBoundsStruct, bvhNodeStruct, intersectionResultStruct, constants, } from 'three-mesh-bvh/webgpu'; +// TODO: clean up "update" function, separate it into chunks // TODO: separate update functions into utilities -// TODO: add ability to easily update a single matrix / scene rearrangement +// TODO: add ability to easily update a single matrix / scene rearrangement (partial update) // TODO: add material support w/ function to easily update material // TODO: add skinned mesh bvh support // TODO: add overrideable functions for custom implementations (custom attributes, transform fields) +// TODO: see if it's possible to replace function contents and dependencies in-place so that +// a node fn can be updated without regenerating all other materials. +// TODO: see if we can reference wgslFn names directly rather than constructing them inline over and over +// and / or use local variable definitions for the pointers to clean up the code +// TODO: see if there's a "build" step that can be leveraged fro nodes +// NEXT: Get a basic version working with megakernel + const _vec = /* @__PURE__ */ new Vector4(); const _matrix = /* @__PURE__ */ new Matrix4(); const _inverseMatrix = /* @__PURE__ */ new Matrix4(); -export { rayStruct, bvhNodeBoundsStruct, bvhNodeStruct, intersectionResultStruct, constants }; - -export const transformStruct = wgsl( /* wgsl */` +// stride is 36 floats (144 bytes) to match WGSL struct alignment: +// mat4x4f (64) + mat4x4f (64) + u32 (4) + 12 bytes padding to align to 16 +const TRANSFORM_STRUCT_SIZE = 36; +const transformStruct = wgsl( /* wgsl */` struct TransformStruct { matrixWorld: mat4x4f, inverseMatrixWorld: mat4x4f, @@ -55,8 +63,8 @@ function dereferenceIndex( indexAttr, indirectBuffer ) { function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct ) { - const intersectFirstHitFn = wgslFn( /* wgsl */` - fn ${ name }IntersectFirstHit( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> IntersectionResult { + const geometryRaycastFirstHitFn = wgslFn( /* wgsl */` + fn ${ name }RaycastGeometryFirstHit( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> IntersectionResult { var bestHit: IntersectionResult; bestHit.didHit = false; @@ -193,7 +201,7 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto localRay.origin = ( transform.inverseMatrixWorld * vec4f( ray.origin, 1.0 ) ).xyz; localRay.direction = ( transform.inverseMatrixWorld * vec4f( ray.direction, 0.0 ) ).xyz; - let blasHit = ${ name }IntersectFirstHit( localRay, transform.nodeOffset, bestHit.dist ); + let blasHit = ${ name }RaycastGeometryFirstHit( localRay, transform.nodeOffset, bestHit.dist ); if ( blasHit.didHit && blasHit.dist < bestHit.dist ) { @@ -230,7 +238,7 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto } `, [ - intersectFirstHitFn, + geometryRaycastFirstHitFn, nodesStorage, transformsStorage, intersectsBounds, rayStruct, bvhNodeStruct, intersectionResultStruct, constants, @@ -389,11 +397,8 @@ export class BVHComputeFns { } ); // write the transforms - // stride is 36 floats (144 bytes) to match WGSL struct alignment: - // mat4x4f (64) + mat4x4f (64) + u32 (4) + 12 bytes padding to align to 16 - const TRANSFORM_STRIDE = 36; let transformWriteOffset = 0; - const transformArrayBuffer = new ArrayBuffer( TRANSFORM_STRIDE * 4 * objectTransformsCount ); + const transformArrayBuffer = new ArrayBuffer( TRANSFORM_STRUCT_SIZE * 4 * objectTransformsCount ); const transformBufferF32 = new Float32Array( transformArrayBuffer ); const transformBufferU32 = new Uint32Array( transformArrayBuffer ); bvh.primitiveBuffer.forEach( ( compositeId, i ) => { @@ -405,9 +410,9 @@ export class BVHComputeFns { const bvhOffset = bvhNodeOffsets[ transformBVHs[ i ] ]; objectBvh._roots.forEach( ( root, ri ) => { - _matrix.toArray( transformBufferF32, transformWriteOffset * TRANSFORM_STRIDE ); - _inverseMatrix.toArray( transformBufferF32, transformWriteOffset * TRANSFORM_STRIDE + 16 ); - transformBufferU32[ transformWriteOffset * TRANSFORM_STRIDE + 32 ] = bvhOffset[ ri ]; + _matrix.toArray( transformBufferF32, transformWriteOffset * TRANSFORM_STRUCT_SIZE ); + _inverseMatrix.toArray( transformBufferF32, transformWriteOffset * TRANSFORM_STRUCT_SIZE + 16 ); + transformBufferU32[ transformWriteOffset * TRANSFORM_STRUCT_SIZE + 32 ] = bvhOffset[ ri ]; transformWriteOffset ++; } ); @@ -421,7 +426,7 @@ export class BVHComputeFns { ` ); const nodesStorage = storage( new StorageBufferAttribute( nodeBuffer32, 8 ), `${ name }BVHNode` ).toReadOnly().setName( `${ name }nodes` ); - const transformsStorage = storage( new StorageBufferAttribute( transformBufferF32, TRANSFORM_STRIDE ), `${ name }TransformStruct` ).toReadOnly().setName( `${ name }transforms` ); + const transformsStorage = storage( new StorageBufferAttribute( transformBufferF32, TRANSFORM_STRUCT_SIZE ), `${ name }TransformStruct` ).toReadOnly().setName( `${ name }transforms` ); const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ name }index` ); const attributesStorage = storage( new StorageBufferAttribute( attributesBuffer, attributesStructSize ), `${ name }GeometryStruct` ).toReadOnly().setName( `${ name }attributes` ); From f0bfafea08af18015f565f2beaefbdab2400372b Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 18 Feb 2026 17:26:25 +0900 Subject: [PATCH 20/84] comments --- src/webgpu/lib/BVHComputeFns.js | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 188bffefe..2d14437ac 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -14,6 +14,7 @@ import { // TODO: separate update functions into utilities // TODO: add ability to easily update a single matrix / scene rearrangement (partial update) // TODO: add material support w/ function to easily update material +// - add a callback for writing a property for a geometry to a range // TODO: add skinned mesh bvh support // TODO: add overrideable functions for custom implementations (custom attributes, transform fields) // TODO: see if it's possible to replace function contents and dependencies in-place so that @@ -21,6 +22,7 @@ import { // TODO: see if we can reference wgslFn names directly rather than constructing them inline over and over // and / or use local variable definitions for the pointers to clean up the code // TODO: see if there's a "build" step that can be leveraged fro nodes +// TODO: allow for "slotting" a new type of callback (eg distance, etc) so multiple types of queries can be made // NEXT: Get a basic version working with megakernel const _vec = /* @__PURE__ */ new Vector4(); From 7003d3cfb2d6e19c57d73a667b54a2151ebfa5b0 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 10:29:02 +0900 Subject: [PATCH 21/84] update three-mesh-bvh --- package-lock.json | 14 +++++++------- package.json | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/package-lock.json b/package-lock.json index 3fad42a7a..30f76b0a2 100644 --- a/package-lock.json +++ b/package-lock.json @@ -25,7 +25,7 @@ "rollup": "^2.70.0", "simple-git": "^3.10.0", "three": "^0.181.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" @@ -3657,9 +3657,9 @@ "license": "MIT" }, "node_modules/three-mesh-bvh": { - "version": "0.9.5", - "resolved": "https://registry.npmjs.org/three-mesh-bvh/-/three-mesh-bvh-0.9.5.tgz", - "integrity": "sha512-MYpwzUWDxPAKGhSBFin9E/7K4AAHyIm4IfMZQ/3+Z/jq/swa2dAhXx0yUNDd9mjlhLuzXkMBTGDZioL2GSlIfQ==", + "version": "0.9.8", + "resolved": "https://registry.npmjs.org/three-mesh-bvh/-/three-mesh-bvh-0.9.8.tgz", + "integrity": "sha512-YphYvdXEZSXdz6iNdWJo1RB6qvSCRyiXPEVSvNU6xVWbLDOdSrfEIsJOpgFOnefdmVEvZ6M+sY0cjh9gl7MvdA==", "dev": true, "license": "MIT", "peerDependencies": { @@ -6432,9 +6432,9 @@ "dev": true }, "three-mesh-bvh": { - "version": "0.9.5", - "resolved": "https://registry.npmjs.org/three-mesh-bvh/-/three-mesh-bvh-0.9.5.tgz", - "integrity": "sha512-MYpwzUWDxPAKGhSBFin9E/7K4AAHyIm4IfMZQ/3+Z/jq/swa2dAhXx0yUNDd9mjlhLuzXkMBTGDZioL2GSlIfQ==", + "version": "0.9.8", + "resolved": "https://registry.npmjs.org/three-mesh-bvh/-/three-mesh-bvh-0.9.8.tgz", + "integrity": "sha512-YphYvdXEZSXdz6iNdWJo1RB6qvSCRyiXPEVSvNU6xVWbLDOdSrfEIsJOpgFOnefdmVEvZ6M+sY0cjh9gl7MvdA==", "dev": true, "requires": {} }, diff --git a/package.json b/package.json index 0238598c9..6e0a622cd 100644 --- a/package.json +++ b/package.json @@ -46,7 +46,7 @@ "rollup": "^2.70.0", "simple-git": "^3.10.0", "three": "^0.181.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" From a9be061e1222d8a48e1bf3911bafd106fb1dff83 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 10:54:34 +0900 Subject: [PATCH 22/84] Adjust WebGPUPathtracer --- src/webgpu/WebGPUPathTracer.js | 128 +++++++-------------------------- 1 file changed, 24 insertions(+), 104 deletions(-) diff --git a/src/webgpu/WebGPUPathTracer.js b/src/webgpu/WebGPUPathTracer.js index a38ddeae4..2b62aa3e9 100644 --- a/src/webgpu/WebGPUPathTracer.js +++ b/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, Clock, Scene, PerspectiveCamera } from 'three/webgpu'; +import { MeshBVH } 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 { BVHComputeFns } from './lib/BVHComputeFns.js'; const _resolution = new Vector2(); export class WebGPUPathTracer { @@ -24,7 +26,8 @@ export class WebGPUPathTracer { this._pathTracer.dispose(); this._pathTracer = value ? new MegaKernelPathTracer( this._renderer ) : new WaveFrontPathTracer( this._renderer ); - this._generator = new PathTracingSceneGenerator(); + this._pathTracer.setBVHComputeFns( this._bvhComputeFns ); + this.setCamera( this.camera ); } @@ -32,10 +35,7 @@ export class WebGPUPathTracer { // members this._renderer = renderer; - this._generator = new PathTracingSceneGenerator(); - // this._pathTracer = new MegaKernelPathTracer( renderer ); - this._pathTracer = new WaveFrontPathTracer( renderer ); - this._queueReset = false; + this._pathTracer = new MegaKernelPathTracer( renderer ); this._clock = new Clock(); // options @@ -54,11 +54,24 @@ 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 ); + + } + + } ); + + // Build TLAS and compute functions + const objectBVH = new ObjectBVH( scene ); + const bvhComputeFns = new BVHComputeFns( objectBVH ); + + this._bvhComputeFns = bvhComputeFns; + this._pathTracer.setBVHComputeFns( bvhComputeFns ); + 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; - - } - - } - -} From 41fd20881476c401d50314478320133e7d1df1a6 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 10:55:54 +0900 Subject: [PATCH 23/84] Pass bvh compute fns to megakernelpathtracer --- src/webgpu/MegaKernelPathTracer.js | 45 ++++++------------------------ 1 file changed, 8 insertions(+), 37 deletions(-) diff --git a/src/webgpu/MegaKernelPathTracer.js b/src/webgpu/MegaKernelPathTracer.js index 1a390d6dc..9aedce818 100644 --- a/src/webgpu/MegaKernelPathTracer.js +++ b/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,16 +69,8 @@ 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, - }; + // bvh data + this.bvhComputeFns = null; // targets this.outputTarget = new StorageTexture( 1, 1, ); @@ -112,29 +96,17 @@ export class MegaKernelPathTracer { this.sampleCountTarget.generateMipmaps = false; // kernels - this.kernel = new PathTracerMegaKernel().setWorkgroupSize( 8, 8, 1 ); + this.kernel = null; this.sampleCountClearKernel = new ZeroOutKernel( { textureType: 'r32uint' } ).setWorkgroupSize( 8, 8, 1 ); this.outputTargetClearKernel = new ZeroOutKernel( { textureType: 'rgba32float' } ).setWorkgroupSize( 8, 8, 1 ); } - setGeometryData( geometry ) { - - for ( const propName in geometry ) { + setBVHComputeFns( bvhComputeFns ) { - 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.bvhComputeFns = bvhComputeFns; + this.kernel = new PathTracerMegaKernel( bvhComputeFns ).setWorkgroupSize( 8, 8, 1 ); + this._task = null; } @@ -223,7 +195,6 @@ export class MegaKernelPathTracer { this.samples = 0; this._task = null; - const { width, height } = sampleCountTarget; const dispatchSize = sampleCountClearKernel.getDispatchSize( width, height ); @@ -240,7 +211,7 @@ export class MegaKernelPathTracer { update() { - if ( ! this.camera ) { + if ( ! this.camera || ! this.kernel ) { return; From 696cf08c9d5b5242ae373e655a0c3d083358dc6c Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 10:56:37 +0900 Subject: [PATCH 24/84] Add temp dispose function --- src/webgpu/lib/BVHComputeFns.js | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 2d14437ac..fe4a28a38 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -617,4 +617,6 @@ export class BVHComputeFns { } + dispose() {} + } From 0472c1391ba3757a97c40ca254c85a5b1519cca2 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 10:56:42 +0900 Subject: [PATCH 25/84] add comment --- src/webgpu/lib/BVHComputeFns.js | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index fe4a28a38..2e27c7787 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -617,6 +617,10 @@ export class BVHComputeFns { } - dispose() {} + dispose() { + + // TODO: dispose buffers + + } } From dbaeae4a2c3e32e4019360b7e37c3f55a2850d30 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 11:31:53 +0900 Subject: [PATCH 26/84] Update --- src/webgpu/compute/PathTracerMegaKernel.js | 120 ++++++++++++++++++--- src/webgpu/lib/BVHComputeFns.js | 22 ++-- 2 files changed, 116 insertions(+), 26 deletions(-) diff --git a/src/webgpu/compute/PathTracerMegaKernel.js b/src/webgpu/compute/PathTracerMegaKernel.js index 90eca4eb9..7406bcfa3 100644 --- a/src/webgpu/compute/PathTracerMegaKernel.js +++ b/src/webgpu/compute/PathTracerMegaKernel.js @@ -1,11 +1,13 @@ -import { IndirectStorageBufferAttribute, Matrix4, Vector2, StorageTexture } from 'three/webgpu'; +import { Matrix4, Vector2, StorageTexture } from 'three/webgpu'; +import { ndcToCameraRay } from 'three-mesh-bvh/webgpu'; import { ComputeKernel } from './ComputeKernel.js'; -import { uniform, storage, globalId, textureStore } from 'three/tsl'; -import megakernelShader from '../nodes/megakernel.wgsl.js'; +import { uniform, globalId, textureStore, wgslFn } from 'three/tsl'; +import { pcgRand3, pcgInit } from '../nodes/random.wgsl.js'; +import { lambertBsdfFunc } from '../nodes/sampling.wgsl.js'; export class PathTracerMegaKernel extends ComputeKernel { - constructor() { + constructor( bvhComputeFns ) { const megakernelShaderParams = { prevOutputTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadOnly(), @@ -14,7 +16,6 @@ export class PathTracerMegaKernel extends ComputeKernel { offset: uniform( new Vector2() ), tileSize: uniform( new Vector2() ), - smoothNormals: uniform( 1 ), seed: uniform( 0 ), bounces: uniform( 5 ), @@ -22,20 +23,109 @@ export class PathTracerMegaKernel extends ComputeKernel { inverseProjectionMatrix: uniform( new Matrix4() ), cameraToModelMatrix: uniform( new Matrix4() ), - // bvh and geometry definition - geom_index: storage( new IndirectStorageBufferAttribute( 1, 3 ), 'vec3u' ).toReadOnly(), - geom_position: storage( new IndirectStorageBufferAttribute( 1, 3 ), 'vec3f' ).toReadOnly(), - geom_normals: storage( new IndirectStorageBufferAttribute( 1, 3 ), 'vec3f' ).toReadOnly(), - geom_material_index: storage( new IndirectStorageBufferAttribute( 1, 1 ), 'u32' ).toReadOnly(), - bvh: storage( new IndirectStorageBufferAttribute(), 'BVHNode' ).toReadOnly(), // TODO: fill this in - - materials: storage( new IndirectStorageBufferAttribute(), 'Material' ).toReadOnly(), // TODO: fill this in - // compute variables globalId: globalId, }; - super( megakernelShader( megakernelShaderParams ) ); + const { raycastFirstHitFn, name } = bvhComputeFns; + const shader = wgslFn( /* wgsl */` + + fn compute( + + // indices and target + globalId: vec3u, + prevOutputTarget: texture_storage_2d, + outputTarget: texture_storage_2d, + sampleCountTarget: texture_storage_2d, + + // tiles + offset: vec2u, + tileSize: vec2u, + + // settings + inverseProjectionMatrix: mat4x4f, + cameraToModelMatrix: mat4x4f, + seed: u32, + bounces: u32, + + ) -> void { + + // make sure we don't bleed over the edge of our tile + if ( globalId.x >= tileSize.x || globalId.y >= tileSize.y ) { + + return; + + } + + // to screen coordinates + let indexUV = offset + globalId.xy; + let targetDimensions = textureDimensions( outputTarget ); + if ( indexUV.x >= targetDimensions.x || indexUV.y >= targetDimensions.y ) { + + return; + + } + + let uv = vec2f( indexUV ) / vec2f( targetDimensions ); + let ndc = uv * 2.0 - vec2f( 1.0 ); + + pcgInitialize( indexUV, seed ); + + // scene ray + var jitter = 2.0 * ( pcgRand2() - vec2( 0.5 ) ) / vec2f( targetDimensions.xy ); + var ray = ndcToCameraRay( ndc + jitter, cameraToModelMatrix * inverseProjectionMatrix ); + + var resultColor = vec3f( 0.0 ); + var throughputColor = vec3f( 1.0 ); + + for ( var bounce = 0u; bounce < bounces; bounce ++ ) { + + let hitResult = ${ name }RaycastFirstHit( ray ); + if ( hitResult.didHit ) { + + resultColor = vec3f( 1, 0, 0 ); + + } else { + + resultColor = vec3f( 0, 0, 1 ); + + } + + break; + + // if ( hitResult.didHit ) { + + // let hitPosition = ray.origin + ray.direction * hitResult.dist; + // let scatterRec = bsdfEval( hitResult.normal, - ray.direction ); + + // // white diffuse surface + // throughputColor *= scatterRec.value / scatterRec.pdf; + + // ray.origin = hitPosition; + // ray.direction = scatterRec.direction; + + // } else { + + // let background = vec3f( 0.5 ); + // resultColor += background * throughputColor; + // break; + + // } + + } + + let sampleCount = textureLoad( sampleCountTarget, indexUV ).r + 1; + var color = textureLoad( prevOutputTarget, indexUV ).xyz; + color += ( resultColor - color.xyz ) / f32( sampleCount ); + + textureStore( sampleCountTarget, indexUV, vec4( sampleCount ) ); + textureStore( outputTarget, indexUV, vec4( color, 1.0 ) ); + + } + + `, [ raycastFirstHitFn, ndcToCameraRay, pcgRand3, pcgInit, lambertBsdfFunc ] ); + + super( shader( megakernelShaderParams ) ); this.defineUniformAccessors( megakernelShaderParams ); diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 2e27c7787..0b65aaf4f 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -85,7 +85,7 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto } let nodeIndex = stack[ pointer ]; - let node = ${ name }nodes[ nodeIndex ]; + let node = ${ name }nodes.value[ nodeIndex ]; pointer = pointer - 1; var boundsHitDist: f32 = 0.0; @@ -106,13 +106,13 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto for ( var ti = triOffset; ti < triOffset + triCount; ti = ti + 1u ) { - let i0 = ${ name }index[ ti * 3u ]; - let i1 = ${ name }index[ ti * 3u + 1u ]; - let i2 = ${ name }index[ ti * 3u + 2u ]; + let i0 = ${ name }index.value[ ti * 3u ]; + let i1 = ${ name }index.value[ ti * 3u + 1u ]; + let i2 = ${ name }index.value[ ti * 3u + 2u ]; - let a = ${ name }attributes[ i0 ].position.xyz; - let b = ${ name }attributes[ i1 ].position.xyz; - let c = ${ name }attributes[ i2 ].position.xyz; + let a = ${ name }attributes.value[ i0 ].position.xyz; + let b = ${ name }attributes.value[ i1 ].position.xyz; + let c = ${ name }attributes.value[ i2 ].position.xyz; var triResult = intersectsTriangle( ray, a, b, c ); @@ -175,7 +175,7 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto } let currNodeIndex = tlasStack[ tlasPointer ]; - let node = ${ name }nodes[ currNodeIndex ]; + let node = ${ name }nodes.value[ currNodeIndex ]; tlasPointer = tlasPointer - 1; var boundsHitDist: f32 = 0.0; @@ -196,7 +196,7 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto for ( var t = offset; t < offset + count; t = t + 1u ) { - let transform = ${ name }transforms[ t ]; + let transform = ${ name }transforms.value[ t ]; // Transform ray into object local space var localRay: Ray; @@ -427,8 +427,8 @@ export class BVHComputeFns { } ` ); - const nodesStorage = storage( new StorageBufferAttribute( nodeBuffer32, 8 ), `${ name }BVHNode` ).toReadOnly().setName( `${ name }nodes` ); - const transformsStorage = storage( new StorageBufferAttribute( transformBufferF32, TRANSFORM_STRUCT_SIZE ), `${ name }TransformStruct` ).toReadOnly().setName( `${ name }transforms` ); + const nodesStorage = storage( new StorageBufferAttribute( nodeBuffer32, 8 ), 'BVHNode' ).toReadOnly().setName( `${ name }nodes` ); + const transformsStorage = storage( new StorageBufferAttribute( transformBufferF32, TRANSFORM_STRUCT_SIZE ), 'TransformStruct' ).toReadOnly().setName( `${ name }transforms` ); const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ name }index` ); const attributesStorage = storage( new StorageBufferAttribute( attributesBuffer, attributesStructSize ), `${ name }GeometryStruct` ).toReadOnly().setName( `${ name }attributes` ); From ee41ab158682aa121870526f7597434b34fb8f67 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 12:14:22 +0900 Subject: [PATCH 27/84] Get a basic version working --- src/webgpu/compute/PathTracerMegaKernel.js | 38 ++++++++++++---------- src/webgpu/lib/BVHComputeFns.js | 1 - 2 files changed, 20 insertions(+), 19 deletions(-) diff --git a/src/webgpu/compute/PathTracerMegaKernel.js b/src/webgpu/compute/PathTracerMegaKernel.js index 7406bcfa3..79d4441cf 100644 --- a/src/webgpu/compute/PathTracerMegaKernel.js +++ b/src/webgpu/compute/PathTracerMegaKernel.js @@ -81,36 +81,38 @@ export class PathTracerMegaKernel extends ComputeKernel { for ( var bounce = 0u; bounce < bounces; bounce ++ ) { let hitResult = ${ name }RaycastFirstHit( ray ); - if ( hitResult.didHit ) { + // resultColor = hitResult.normal; + // break; + // if ( hitResult.didHit ) { - resultColor = vec3f( 1, 0, 0 ); + // resultColor = vec3f( 1, 0, 0 ); - } else { + // } else { - resultColor = vec3f( 0, 0, 1 ); + // resultColor = vec3f( 0, 0, 1 ); - } + // } - break; + // break; - // if ( hitResult.didHit ) { + if ( hitResult.didHit ) { - // let hitPosition = ray.origin + ray.direction * hitResult.dist; - // let scatterRec = bsdfEval( hitResult.normal, - ray.direction ); + let hitPosition = ray.origin + ray.direction * hitResult.dist; + let scatterRec = bsdfEval( hitResult.normal, - ray.direction ); - // // white diffuse surface - // throughputColor *= scatterRec.value / scatterRec.pdf; + // white diffuse surface + throughputColor *= hitResult.normal * scatterRec.value / scatterRec.pdf; - // ray.origin = hitPosition; - // ray.direction = scatterRec.direction; + ray.origin = hitPosition; + ray.direction = scatterRec.direction; - // } else { + } else { - // let background = vec3f( 0.5 ); - // resultColor += background * throughputColor; - // break; + let background = vec3f( 0.5 ); + resultColor += background * throughputColor; + break; - // } + } } diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 0b65aaf4f..ab3a15c16 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -204,7 +204,6 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto localRay.direction = ( transform.inverseMatrixWorld * vec4f( ray.direction, 0.0 ) ).xyz; let blasHit = ${ name }RaycastGeometryFirstHit( localRay, transform.nodeOffset, bestHit.dist ); - if ( blasHit.didHit && blasHit.dist < bestHit.dist ) { bestHit = blasHit; From 361a54abf374951c12620bfd4d70e091f13be1d1 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 12:41:56 +0900 Subject: [PATCH 28/84] Fix multiroot bvhs --- src/webgpu/lib/BVHComputeFns.js | 31 +++++++++++++++++++++++-------- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index ab3a15c16..105ac9e2e 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -390,7 +390,7 @@ export class BVHComputeFns { const nodeBufferFloat = new Float32Array( nodeBuffer ); const bvhNodeOffsets = []; let nodeWriteOffset = 0; - appendBVHData( bvh, 0, true ); + appendBVHData( bvh, 0, true, bvhs, transformBVHs ); bvhs.forEach( ( bvh, i ) => { bvhNodeOffsets.push( appendBVHData( bvh, geometryOffsets[ i ], false ) ); @@ -438,15 +438,16 @@ export class BVHComputeFns { this.attributesStruct = attributeStruct; this.raycastFirstHitFn = buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct ); - function appendBVHData( bvh, geometryOffsets, tlas = false ) { + function appendBVHData( bvh, geometryOffset, tlas = false, bvhs = null, transformBVHs = null ) { const BYTES_PER_NODE = 6 * 4 + 4 + 4; const UINT32_PER_NODE = BYTES_PER_NODE / 4; const IS_LEAFNODE_FLAG = 0xFFFF; const result = []; - bvh._roots.forEach( ( root, rootIndex ) => { + bvh._roots.forEach( root => { + let tlasOffset = 0; const rootBuffer16 = new Uint16Array( root ); const rootBuffer32 = new Uint32Array( root ); result.push( nodeWriteOffset ); @@ -463,16 +464,31 @@ export class BVHComputeFns { const isLeaf = IS_LEAFNODE_FLAG === rootBuffer16[ r16 + 15 ]; if ( isLeaf ) { - nodeBuffer16[ n16 + 14 ] = rootBuffer16[ r16 + 14 ]; if ( tlas ) { + const count = rootBuffer16[ r16 + 14 ]; + const offset = rootBuffer32[ r32 + 6 ]; + + // each root is expanded into a separate transform so we need to expand + // the embedded offsets and counts. + let rootsCount = 0; + for ( let o = offset, l = offset + count; o < l; o ++ ) { + + rootsCount += bvhs[ transformBVHs[ o ] ]._roots.length; + + } + // 0xFFFF == mesh leaf, 0xFF00 == TLAS leaf - nodeBuffer32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ]; + nodeBuffer32[ n32 + 6 ] = tlasOffset; // rootBuffer32[ r32 + 6 ]; + nodeBuffer16[ n16 + 14 ] = rootsCount; //rootBuffer16[ r16 + 14 ]; nodeBuffer16[ n16 + 15 ] = 0xFF00; + tlasOffset += rootsCount; + } else { - nodeBuffer32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ] + geometryOffsets[ rootIndex ]; + nodeBuffer32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ] + geometryOffset; + nodeBuffer16[ n16 + 14 ] = rootBuffer16[ r16 + 14 ]; nodeBuffer16[ n16 + 15 ] = IS_LEAFNODE_FLAG; } @@ -497,7 +513,6 @@ export class BVHComputeFns { function appendGeometryData( geometry, offsets = null, indirectBuffer = null ) { - const result = []; let vertexStart = 0; let vertexCount = geometry.attributes.position.count; @@ -511,7 +526,7 @@ export class BVHComputeFns { // write indices — when an indirect buffer is present, dereference it to // resolve the BVH's triangle order into a flat index buffer // store as triangle offset (not vertex-index offset) to match BVH leaf format - result.push( indexOffset / 3 ); + const result = indexOffset / 3; if ( indirectBuffer ) { From fa4463e54a586a082584e52a57b18635746d90f5 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 13:24:44 +0900 Subject: [PATCH 29/84] Cleanup --- src/webgpu/lib/BVHComputeFns.js | 155 ++++++++++++++++++-------------- 1 file changed, 87 insertions(+), 68 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 105ac9e2e..fabe0bfbf 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -25,6 +25,7 @@ import { // TODO: allow for "slotting" a new type of callback (eg distance, etc) so multiple types of queries can be made // NEXT: Get a basic version working with megakernel +const _def = /* @__PURE__ */ new Vector4(); const _vec = /* @__PURE__ */ new Vector4(); const _matrix = /* @__PURE__ */ new Matrix4(); const _inverseMatrix = /* @__PURE__ */ new Matrix4(); @@ -296,12 +297,14 @@ export class BVHComputeFns { update() { + const self = this; const { attributes, name, bvh } = this; + // collect the BVHs const bvhs = []; - const geometries = []; + const geometryInfo = []; const geometryOffsets = []; - const transformBVHs = []; + const transformBVHIndices = []; // accumulate the sizes of the bvh nodes buffer, number of objects, and geometry buffers let bvhNodesBufferLength = 0; @@ -332,7 +335,7 @@ export class BVHComputeFns { const geometryId = object.getGeometryIdAt( instanceId ); const range = object.getGeometryRangeAt( geometryId ); - geometries.push( { + geometryInfo.push( { geometry: object.geometry, range: range, indirectBuffer: indirectBuffer, @@ -346,7 +349,7 @@ export class BVHComputeFns { const geometry = object.geometry; indexBufferLength += geometry.index ? geometry.index.count : geometry.attributes.position.count; attributesBufferLength += geometry.attributes.position.count; - geometries.push( { + geometryInfo.push( { geometry: object.geometry, range: undefined, indirectBuffer: indirectBuffer, @@ -357,46 +360,48 @@ export class BVHComputeFns { } // save the index of the bvh associated with this transform - transformBVHs.push( bvhs.indexOf( meshBvh ) ); + transformBVHIndices.push( bvhs.indexOf( meshBvh ) ); } ); - // construct the attribute struct - let attributesStructSize = 0; - const attributeStructContent = attributes - .map( key => { - - attributesStructSize += 4; - return `${ key }: vec4f,`; + // - } ).join( '\n' ); + // construct the attribute struct + const attributesStructSize = 4 * attributes.length; + const attributeStruct = wgsl( /* wgsl */` + struct ${ name }GeometryStruct { + ${ attributes.map( key => `${ key }: vec4f,` ).join( '\n' ) } + } + ` ); // write the geometry buffer attributes let attributesOffset = 0; let indexOffset = 0; const indexBuffer = new Uint32Array( indexBufferLength ); const attributesBuffer = new Float32Array( attributesBufferLength * attributes.length * 4 ); - geometries.forEach( ( { geometry, range, indirectBuffer } ) => { + geometryInfo.forEach( ( { geometry, range, indirectBuffer } ) => { const offset = appendGeometryData( geometry, range, indirectBuffer ); geometryOffsets.push( offset ); } ); + // + // write the bvh data - const nodeBuffer = new ArrayBuffer( bvhNodesBufferLength ); - const nodeBuffer16 = new Uint16Array( nodeBuffer ); - const nodeBuffer32 = new Uint32Array( nodeBuffer ); - const nodeBufferFloat = new Float32Array( nodeBuffer ); + const transformBVHs = transformBVHIndices.map( i => bvhs[ i ] ); + const bvhNodesBuffer = new ArrayBuffer( bvhNodesBufferLength ); const bvhNodeOffsets = []; let nodeWriteOffset = 0; - appendBVHData( bvh, 0, true, bvhs, transformBVHs ); + appendBVHData( bvh, 0, transformBVHs, bvhNodesBuffer, true ); bvhs.forEach( ( bvh, i ) => { - bvhNodeOffsets.push( appendBVHData( bvh, geometryOffsets[ i ], false ) ); + bvhNodeOffsets.push( appendBVHData( bvh, geometryOffsets[ i ], transformBVHs, bvhNodesBuffer, false ) ); } ); + // + // write the transforms let transformWriteOffset = 0; const transformArrayBuffer = new ArrayBuffer( TRANSFORM_STRUCT_SIZE * 4 * objectTransformsCount ); @@ -407,8 +412,8 @@ export class BVHComputeFns { bvh.getObjectMatrix( compositeId, _matrix ); _inverseMatrix.copy( _matrix ).invert(); - const objectBvh = bvhs[ transformBVHs[ i ] ]; - const bvhOffset = bvhNodeOffsets[ transformBVHs[ i ] ]; + const objectBvh = bvhs[ transformBVHIndices[ i ] ]; + const bvhOffset = bvhNodeOffsets[ transformBVHIndices[ i ] ]; objectBvh._roots.forEach( ( root, ri ) => { _matrix.toArray( transformBufferF32, transformWriteOffset * TRANSFORM_STRUCT_SIZE ); @@ -420,13 +425,10 @@ export class BVHComputeFns { } ); - const attributeStruct = wgsl( /* wgsl */` - struct ${ name }GeometryStruct { - ${ attributeStructContent } - } - ` ); + // - const nodesStorage = storage( new StorageBufferAttribute( nodeBuffer32, 8 ), 'BVHNode' ).toReadOnly().setName( `${ name }nodes` ); + // set up the storage buffers + const nodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), 'BVHNode' ).toReadOnly().setName( `${ name }nodes` ); const transformsStorage = storage( new StorageBufferAttribute( transformBufferF32, TRANSFORM_STRUCT_SIZE ), 'TransformStruct' ).toReadOnly().setName( `${ name }transforms` ); const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ name }index` ); const attributesStorage = storage( new StorageBufferAttribute( attributesBuffer, attributesStructSize ), `${ name }GeometryStruct` ).toReadOnly().setName( `${ name }attributes` ); @@ -438,7 +440,11 @@ export class BVHComputeFns { this.attributesStruct = attributeStruct; this.raycastFirstHitFn = buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct ); - function appendBVHData( bvh, geometryOffset, tlas = false, bvhs = null, transformBVHs = null ) { + function appendBVHData( bvh, geometryOffset, transformBVHs, target, tlas = false ) { + + const targetU16 = new Uint16Array( target ); + const targetU32 = new Uint32Array( target ); + const targetF32 = new Float32Array( target ); const BYTES_PER_NODE = 6 * 4 + 4 + 4; const UINT32_PER_NODE = BYTES_PER_NODE / 4; @@ -459,7 +465,7 @@ export class BVHComputeFns { const n16 = n32 * 2; // write bounds - nodeBufferFloat.set( new Float32Array( root, i * BYTES_PER_NODE, 6 ), n32 ); + targetF32.set( new Float32Array( root, i * BYTES_PER_NODE, 6 ), n32 ); const isLeaf = IS_LEAFNODE_FLAG === rootBuffer16[ r16 + 15 ]; if ( isLeaf ) { @@ -474,29 +480,29 @@ export class BVHComputeFns { let rootsCount = 0; for ( let o = offset, l = offset + count; o < l; o ++ ) { - rootsCount += bvhs[ transformBVHs[ o ] ]._roots.length; + rootsCount += transformBVHs[ o ]._roots.length; } // 0xFFFF == mesh leaf, 0xFF00 == TLAS leaf - nodeBuffer32[ n32 + 6 ] = tlasOffset; // rootBuffer32[ r32 + 6 ]; - nodeBuffer16[ n16 + 14 ] = rootsCount; //rootBuffer16[ r16 + 14 ]; - nodeBuffer16[ n16 + 15 ] = 0xFF00; + targetU32[ n32 + 6 ] = tlasOffset; // rootBuffer32[ r32 + 6 ]; + targetU16[ n16 + 14 ] = rootsCount; //rootBuffer16[ r16 + 14 ]; + targetU16[ n16 + 15 ] = 0xFF00; tlasOffset += rootsCount; } else { - nodeBuffer32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ] + geometryOffset; - nodeBuffer16[ n16 + 14 ] = rootBuffer16[ r16 + 14 ]; - nodeBuffer16[ n16 + 15 ] = IS_LEAFNODE_FLAG; + targetU32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ] + geometryOffset; + targetU16[ n16 + 14 ] = rootBuffer16[ r16 + 14 ]; + targetU16[ n16 + 15 ] = IS_LEAFNODE_FLAG; } } else { - nodeBuffer32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ]; - nodeBuffer32[ n32 + 7 ] = rootBuffer32[ r32 + 7 ]; + targetU32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ]; + targetU32[ n32 + 7 ] = rootBuffer32[ r32 + 7 ]; } @@ -513,7 +519,6 @@ export class BVHComputeFns { function appendGeometryData( geometry, offsets = null, indirectBuffer = null ) { - let vertexStart = 0; let vertexCount = geometry.attributes.position.count; if ( offsets ) { @@ -579,48 +584,48 @@ export class BVHComputeFns { } - attributes.forEach( ( key, interleavedOffset ) => { + const groups = geometry.groups.length === 0 ? [ { start: vertexStart, count: vertexCount } ] : geometry.groups; + groups.forEach( ( { start, count }, groupIndex ) => { - const attr = geometry.attributes[ key ]; - for ( let i = 0; i < vertexCount; i ++ ) { + attributes.forEach( ( key, interleavedOffset ) => { - if ( ! attr ) { + const attr = geometry.attributes[ key ]; + self.getDefaultAttributeValue( key, groupIndex, _def ); - if ( key === 'color' ) { + for ( let i = 0; i < count; i ++ ) { - _vec.set( 1, 1, 1, 1 ); + if ( attr ) { - } else { + _vec.fromBufferAttribute( attr, i + start ); - _vec.set( 0, 0, 0, 1 ); + switch ( attr.itemSize ) { - } + case 1: + _vec.y = _def.y; + _vec.z = _def.z; + _vec.w = _def.w; + break; + case 2: + _vec.z = _def.z; + _vec.w = _def.w; + break; + case 3: + _vec.w = _def.w; + break; - } else { + } + + } else { - _vec.fromBufferAttribute( attr, i + vertexStart ); - switch ( attr.itemSize ) { - - case 1: - _vec.y = 0; - _vec.z = 0; - _vec.w = 0; - break; - case 2: - _vec.z = 0; - _vec.w = 0; - break; - case 3: - _vec.w = 0; - break; + _vec.copy( _def ); } - } + _vec.toArray( attributesBuffer, ( attributesOffset + i ) * attributesStructSize + interleavedOffset * 4 ); - _vec.toArray( attributesBuffer, ( attributesOffset + i ) * attributesStructSize + interleavedOffset * 4 ); + } - } + } ); } ); @@ -631,6 +636,20 @@ export class BVHComputeFns { } + getDefaultAttributeValue( key, groupIndex, target ) { + + if ( key === 'color' ) { + + target.set( 1, 1, 1, 1 ); + + } else { + + target.set( 0, 0, 0, 1 ); + + } + + } + dispose() { // TODO: dispose buffers From f0eff98080d15ef2ac40d5ec81a9d01ebd79d8cc Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 13:31:47 +0900 Subject: [PATCH 30/84] Updates --- src/webgpu/lib/BVHComputeFns.js | 43 ++++++++++++++++++--------------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index fabe0bfbf..6e700f608 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -23,7 +23,6 @@ import { // and / or use local variable definitions for the pointers to clean up the code // TODO: see if there's a "build" step that can be leveraged fro nodes // TODO: allow for "slotting" a new type of callback (eg distance, etc) so multiple types of queries can be made -// NEXT: Get a basic version working with megakernel const _def = /* @__PURE__ */ new Vector4(); const _vec = /* @__PURE__ */ new Vector4(); @@ -272,25 +271,6 @@ export class BVHComputeFns { this.attributesStruct = null; this.raycastFirstHitFn = null; - this.getBVH = ( object, id = - 1 ) => { - - if ( object.isInstancedMesh ) { - - return object.geometry.boundsTree; - - } else if ( object.isBatchedMesh ) { - - const geometryId = object.getGeometryIdAt( id ); - return object.boundsTrees[ geometryId ]; - - } else { - - return object.geometry.boundsTree; - - } - - }; - this.update(); } @@ -636,12 +616,35 @@ export class BVHComputeFns { } + getBVH( object, id ) { + + if ( object.isInstancedMesh ) { + + return object.geometry.boundsTree; + + } else if ( object.isBatchedMesh ) { + + const geometryId = object.getGeometryIdAt( id ); + return object.boundsTrees[ geometryId ]; + + } else { + + return object.geometry.boundsTree; + + } + + } + getDefaultAttributeValue( key, groupIndex, target ) { if ( key === 'color' ) { target.set( 1, 1, 1, 1 ); + } else if ( key === 'groupIndex' ) { + + target.set( groupIndex, 0, 0, 1 ); + } else { target.set( 0, 0, 0, 1 ); From 6a49319b9abd2bad02202a2a8a05601548213d57 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 14:02:16 +0900 Subject: [PATCH 31/84] Add output for object index --- src/webgpu/WebGPUPathTracer.js | 4 +- src/webgpu/lib/BVHComputeFns.js | 74 +++++++++++++++++++++++++++++---- 2 files changed, 69 insertions(+), 9 deletions(-) diff --git a/src/webgpu/WebGPUPathTracer.js b/src/webgpu/WebGPUPathTracer.js index 2b62aa3e9..e6bb01a18 100644 --- a/src/webgpu/WebGPUPathTracer.js +++ b/src/webgpu/WebGPUPathTracer.js @@ -1,5 +1,5 @@ import { Vector2, Clock, Scene, PerspectiveCamera } from 'three/webgpu'; -import { MeshBVH } from 'three-mesh-bvh'; +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'; @@ -59,7 +59,7 @@ export class WebGPUPathTracer { if ( child.isMesh && ! child.geometry.boundsTree ) { - child.geometry.boundsTree = new MeshBVH( child.geometry ); + child.geometry.boundsTree = new MeshBVH( child.geometry, { strategy: SAH, maxLeafSize: 1 } ); } diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 6e700f608..ccf8c0270 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -2,11 +2,9 @@ import { Matrix4, Vector4 } from 'three'; import { StorageBufferAttribute } from 'three/src/Three.WebGPU.Nodes.js'; import { storage, wgsl, wgslFn } from 'three/tsl'; import { - intersectsTriangle, intersectsBounds, rayStruct, bvhNodeStruct, - intersectionResultStruct, constants, } from 'three-mesh-bvh/webgpu'; @@ -63,7 +61,67 @@ function dereferenceIndex( indexAttr, indirectBuffer ) { } -function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct ) { +const intersectionResultStruct = wgsl( /* wgsl */` + struct IntersectionResult { + didHit: bool, + indices: vec4u, + normal: vec3f, + barycoord: vec3f, + side: f32, + dist: f32, + objectIndex: u32, + }; +` ); + +const intersectsTriangle = wgslFn( /* wgsl */ ` + + fn intersectsTriangle( ray: Ray, a: vec3f, b: vec3f, c: vec3f ) -> IntersectionResult { + + var result: IntersectionResult; + result.didHit = false; + + let edge1 = b - a; + let edge2 = c - a; + let n = cross( edge1, edge2 ); + + let det = - dot( ray.direction, n ); + + if ( abs( det ) < TRI_INTERSECT_EPSILON ) { + + return result; + + } + + let invdet = 1.0 / det; + + let AO = ray.origin - a; + let DAO = cross( AO, ray.direction ); + + let u = dot( edge2, DAO ) * invdet; + let v = -dot( edge1, DAO ) * invdet; + let t = dot( AO, n ) * invdet; + + let w = 1.0 - u - v; + + if ( u < - TRI_INTERSECT_EPSILON || v < - TRI_INTERSECT_EPSILON || w < - TRI_INTERSECT_EPSILON || t < TRI_INTERSECT_EPSILON ) { + + return result; + + } + + result.didHit = true; + result.barycoord = vec3f( w, u, v ); + result.dist = t; + result.side = sign( det ); + result.normal = result.side * normalize( n ); + + return result; + + } + +`, [ rayStruct, intersectionResultStruct, constants ] ); + +function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ) { const geometryRaycastFirstHitFn = wgslFn( /* wgsl */` fn ${ name }RaycastGeometryFirstHit( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> IntersectionResult { @@ -207,6 +265,7 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto if ( blasHit.didHit && blasHit.dist < bestHit.dist ) { bestHit = blasHit; + bestHit.objectIndex = t; // Transform normal to world space: normal matrix = transpose( inverse ) bestHit.normal = normalize( ( transpose( transform.inverseMatrixWorld ) * vec4f( bestHit.normal, 0.0 ) ).xyz ); @@ -358,7 +417,7 @@ export class BVHComputeFns { let attributesOffset = 0; let indexOffset = 0; const indexBuffer = new Uint32Array( indexBufferLength ); - const attributesBuffer = new Float32Array( attributesBufferLength * attributes.length * 4 ); + const attributesBuffer = new ArrayBuffer( attributesBufferLength * attributes.length * 4 * 4 ); geometryInfo.forEach( ( { geometry, range, indirectBuffer } ) => { const offset = appendGeometryData( geometry, range, indirectBuffer ); @@ -411,14 +470,14 @@ export class BVHComputeFns { const nodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), 'BVHNode' ).toReadOnly().setName( `${ name }nodes` ); const transformsStorage = storage( new StorageBufferAttribute( transformBufferF32, TRANSFORM_STRUCT_SIZE ), 'TransformStruct' ).toReadOnly().setName( `${ name }transforms` ); const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ name }index` ); - const attributesStorage = storage( new StorageBufferAttribute( attributesBuffer, attributesStructSize ), `${ name }GeometryStruct` ).toReadOnly().setName( `${ name }attributes` ); + const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributesStructSize ), `${ name }GeometryStruct` ).toReadOnly().setName( `${ name }attributes` ); this.storageBufferAttributes.transforms = transformsStorage; this.storageBufferAttributes.nodes = nodesStorage; this.storageBufferAttributes.index = indexStorage; this.storageBufferAttributes.attributes = attributesStorage; this.attributesStruct = attributeStruct; - this.raycastFirstHitFn = buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct ); + this.raycastFirstHitFn = buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ); function appendBVHData( bvh, geometryOffset, transformBVHs, target, tlas = false ) { @@ -564,6 +623,7 @@ export class BVHComputeFns { } + const attributesBufferF32 = new Float32Array( attributesBuffer ); const groups = geometry.groups.length === 0 ? [ { start: vertexStart, count: vertexCount } ] : geometry.groups; groups.forEach( ( { start, count }, groupIndex ) => { @@ -601,7 +661,7 @@ export class BVHComputeFns { } - _vec.toArray( attributesBuffer, ( attributesOffset + i ) * attributesStructSize + interleavedOffset * 4 ); + _vec.toArray( attributesBufferF32, ( attributesOffset + i ) * attributesStructSize + interleavedOffset * 4 ); } From 4363359a6185e0866bc803932899c3695091bcf8 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 14:18:20 +0900 Subject: [PATCH 32/84] Updates --- src/webgpu/lib/BVHComputeFns.js | 61 ++++++++++++++------------------- 1 file changed, 26 insertions(+), 35 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index ccf8c0270..78f240cb7 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -624,48 +624,43 @@ export class BVHComputeFns { } const attributesBufferF32 = new Float32Array( attributesBuffer ); - const groups = geometry.groups.length === 0 ? [ { start: vertexStart, count: vertexCount } ] : geometry.groups; - groups.forEach( ( { start, count }, groupIndex ) => { + attributes.forEach( ( key, interleavedOffset ) => { - attributes.forEach( ( key, interleavedOffset ) => { + const attr = geometry.attributes[ key ]; + self.getDefaultAttributeValue( key, _def ); - const attr = geometry.attributes[ key ]; - self.getDefaultAttributeValue( key, groupIndex, _def ); + for ( let i = 0; i < vertexCount; i ++ ) { - for ( let i = 0; i < count; i ++ ) { + if ( attr ) { - if ( attr ) { + _vec.fromBufferAttribute( attr, i + vertexStart ); - _vec.fromBufferAttribute( attr, i + start ); + switch ( attr.itemSize ) { - switch ( attr.itemSize ) { - - case 1: - _vec.y = _def.y; - _vec.z = _def.z; - _vec.w = _def.w; - break; - case 2: - _vec.z = _def.z; - _vec.w = _def.w; - break; - case 3: - _vec.w = _def.w; - break; - - } - - } else { - - _vec.copy( _def ); + case 1: + _vec.y = _def.y; + _vec.z = _def.z; + _vec.w = _def.w; + break; + case 2: + _vec.z = _def.z; + _vec.w = _def.w; + break; + case 3: + _vec.w = _def.w; + break; } - _vec.toArray( attributesBufferF32, ( attributesOffset + i ) * attributesStructSize + interleavedOffset * 4 ); + } else { + + _vec.copy( _def ); } - } ); + _vec.toArray( attributesBufferF32, ( attributesOffset + i ) * attributesStructSize + interleavedOffset * 4 ); + + } } ); @@ -695,16 +690,12 @@ export class BVHComputeFns { } - getDefaultAttributeValue( key, groupIndex, target ) { + getDefaultAttributeValue( key, target ) { if ( key === 'color' ) { target.set( 1, 1, 1, 1 ); - } else if ( key === 'groupIndex' ) { - - target.set( groupIndex, 0, 0, 1 ); - } else { target.set( 0, 0, 0, 1 ); From 86e5ea07f78972a25068ccb1328f20787d24150f Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 14:48:08 +0900 Subject: [PATCH 33/84] clean up --- src/webgpu/lib/BVHComputeFns.js | 73 ++++++++++++++++----------------- 1 file changed, 36 insertions(+), 37 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 78f240cb7..86b3c6a93 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -340,10 +340,8 @@ export class BVHComputeFns { const { attributes, name, bvh } = this; // collect the BVHs - const bvhs = []; const geometryInfo = []; - const geometryOffsets = []; - const transformBVHIndices = []; + const transformInfo = []; // accumulate the sizes of the bvh nodes buffer, number of objects, and geometry buffers let bvhNodesBufferLength = 0; @@ -361,24 +359,33 @@ export class BVHComputeFns { objectTransformsCount += meshBvh._roots.length; // if we haven't added this bvh, yet - if ( ! bvhs.includes( meshBvh ) ) { + if ( ! geometryInfo.find( info => info.bvh === meshBvh ) ) { // increase the buffer size - bvhs.push( meshBvh ); bvhNodesBufferLength += meshBvh._roots.reduce( ( v, root ) => v + root.byteLength, 0 ); // save the geometry info to write later and increment the buffer sizes const indirectBuffer = meshBvh._indirectBuffer || null; + const info = { + index: geometryInfo.length, + bvh: meshBvh, + geometry: object.geometry, + range: undefined, + indirectBuffer: indirectBuffer, + + bvhBufferOffsets: null, + indexBufferOffset: null, + + + }; + + geometryInfo.push( info ); if ( object.isBatchedMesh ) { const geometryId = object.getGeometryIdAt( instanceId ); const range = object.getGeometryRangeAt( geometryId ); - geometryInfo.push( { - geometry: object.geometry, - range: range, - indirectBuffer: indirectBuffer, - } ); + info.range = range; indexBufferLength += range.indexCount === - 1 ? range.vertexCount : range.indexCount; attributesBufferLength += range.vertexCount; @@ -388,18 +395,17 @@ export class BVHComputeFns { const geometry = object.geometry; indexBufferLength += geometry.index ? geometry.index.count : geometry.attributes.position.count; attributesBufferLength += geometry.attributes.position.count; - geometryInfo.push( { - geometry: object.geometry, - range: undefined, - indirectBuffer: indirectBuffer, - } ); } } // save the index of the bvh associated with this transform - transformBVHIndices.push( bvhs.indexOf( meshBvh ) ); + transformInfo.push( { + data: geometryInfo.find( info => object.geometry === info.geometry ), + object, + instanceId, + } ); } ); @@ -413,29 +419,22 @@ export class BVHComputeFns { } ` ); - // write the geometry buffer attributes + // write the geometry buffer attributes & bvh data let attributesOffset = 0; let indexOffset = 0; + let nodeWriteOffset = 0; const indexBuffer = new Uint32Array( indexBufferLength ); const attributesBuffer = new ArrayBuffer( attributesBufferLength * attributes.length * 4 * 4 ); - geometryInfo.forEach( ( { geometry, range, indirectBuffer } ) => { - - const offset = appendGeometryData( geometry, range, indirectBuffer ); - geometryOffsets.push( offset ); - - } ); - - // - - // write the bvh data - const transformBVHs = transformBVHIndices.map( i => bvhs[ i ] ); const bvhNodesBuffer = new ArrayBuffer( bvhNodesBufferLength ); - const bvhNodeOffsets = []; - let nodeWriteOffset = 0; - appendBVHData( bvh, 0, transformBVHs, bvhNodesBuffer, true ); - bvhs.forEach( ( bvh, i ) => { + appendBVHData( bvh, 0, transformInfo, bvhNodesBuffer, true ); + geometryInfo.forEach( info => { + + const { geometry, range, indirectBuffer } = info; + const indexOffset = appendGeometryData( geometry, range, indirectBuffer ); + const bvhNodeOffsets = appendBVHData( info.bvh, indexOffset, transformInfo, bvhNodesBuffer, false ); - bvhNodeOffsets.push( appendBVHData( bvh, geometryOffsets[ i ], transformBVHs, bvhNodesBuffer, false ) ); + info.indexBufferOffset = indexOffset; + info.bvhNodeOffsets = bvhNodeOffsets; } ); @@ -451,8 +450,8 @@ export class BVHComputeFns { bvh.getObjectMatrix( compositeId, _matrix ); _inverseMatrix.copy( _matrix ).invert(); - const objectBvh = bvhs[ transformBVHIndices[ i ] ]; - const bvhOffset = bvhNodeOffsets[ transformBVHIndices[ i ] ]; + const objectBvh = transformInfo[ i ].data.bvh; + const bvhOffset = transformInfo[ i ].data.bvhNodeOffsets; objectBvh._roots.forEach( ( root, ri ) => { _matrix.toArray( transformBufferF32, transformWriteOffset * TRANSFORM_STRUCT_SIZE ); @@ -479,7 +478,7 @@ export class BVHComputeFns { this.attributesStruct = attributeStruct; this.raycastFirstHitFn = buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ); - function appendBVHData( bvh, geometryOffset, transformBVHs, target, tlas = false ) { + function appendBVHData( bvh, geometryOffset, transformInfo, target, tlas = false ) { const targetU16 = new Uint16Array( target ); const targetU32 = new Uint32Array( target ); @@ -519,7 +518,7 @@ export class BVHComputeFns { let rootsCount = 0; for ( let o = offset, l = offset + count; o < l; o ++ ) { - rootsCount += transformBVHs[ o ]._roots.length; + rootsCount += transformInfo[ o ].data.bvh._roots.length; } From 42f30e28447b52f358e50a08f8340454c9f595fd Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 15:30:17 +0900 Subject: [PATCH 34/84] More cleanup --- src/webgpu/lib/BVHComputeFns.js | 164 +++++++++++++++----------------- 1 file changed, 77 insertions(+), 87 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 86b3c6a93..5864d01e1 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -8,6 +8,10 @@ import { constants, } from 'three-mesh-bvh/webgpu'; +const BYTES_PER_NODE = 6 * 4 + 4 + 4; +const UINT32_PER_NODE = BYTES_PER_NODE / 4; +const IS_LEAFNODE_FLAG = 0xFFFF; + // TODO: clean up "update" function, separate it into chunks // TODO: separate update functions into utilities // TODO: add ability to easily update a single matrix / scene rearrangement (partial update) @@ -61,6 +65,12 @@ function dereferenceIndex( indexAttr, indirectBuffer ) { } +function getTotalBVHByteLength( bvh ) { + + return bvh._roots.reduce( ( v, root ) => v + root.byteLength, 0 ); + +} + const intersectionResultStruct = wgsl( /* wgsl */` struct IntersectionResult { didHit: bool, @@ -344,67 +354,70 @@ export class BVHComputeFns { const transformInfo = []; // accumulate the sizes of the bvh nodes buffer, number of objects, and geometry buffers - let bvhNodesBufferLength = 0; - let objectTransformsCount = 0; + let bvhNodesBufferLength = getTotalBVHByteLength( bvh ); let indexBufferLength = 0; let attributesBufferLength = 0; - bvhNodesBufferLength += bvh._roots.reduce( ( v, root ) => v + root.byteLength, 0 ); bvh.primitiveBuffer.forEach( compositeId => { const object = bvh.getObjectFromId( compositeId ); const instanceId = bvh.getInstanceFromId( compositeId ); const meshBvh = this.getBVH( object, instanceId ); - // add a new transform for each bvh root - objectTransformsCount += meshBvh._roots.length; - // if we haven't added this bvh, yet if ( ! geometryInfo.find( info => info.bvh === meshBvh ) ) { - // increase the buffer size - bvhNodesBufferLength += meshBvh._roots.reduce( ( v, root ) => v + root.byteLength, 0 ); + // TODO: account for indirect buffer here? // save the geometry info to write later and increment the buffer sizes - const indirectBuffer = meshBvh._indirectBuffer || null; const info = { index: geometryInfo.length, bvh: meshBvh, geometry: object.geometry, - range: undefined, - indirectBuffer: indirectBuffer, + range: { + start: 0, + count: 0, + vertexStart: 0, + vertexCount: 0, + }, bvhBufferOffsets: null, indexBufferOffset: null, - }; - geometryInfo.push( info ); - if ( object.isBatchedMesh ) { const geometryId = object.getGeometryIdAt( instanceId ); const range = object.getGeometryRangeAt( geometryId ); - info.range = range; - - indexBufferLength += range.indexCount === - 1 ? range.vertexCount : range.indexCount; - attributesBufferLength += range.vertexCount; + Object.assign( info.range, range ); } else { const geometry = object.geometry; - indexBufferLength += geometry.index ? geometry.index.count : geometry.attributes.position.count; - attributesBufferLength += geometry.attributes.position.count; + info.range.count = geometry.index ? geometry.index.count : geometry.attributes.position.count, + info.range.vertexCount = geometry.attributes.position.count; } + // increase the buffer sizes for bvh and geometry + bvhNodesBufferLength += getTotalBVHByteLength( meshBvh ); + indexBufferLength += info.range.count; + attributesBufferLength += info.range.vertexCount; + geometryInfo.push( info ); + } // save the index of the bvh associated with this transform - transformInfo.push( { - data: geometryInfo.find( info => object.geometry === info.geometry ), - object, - instanceId, + meshBvh._roots.forEach( ( root, i ) => { + + transformInfo.push( { + data: geometryInfo.find( info => object.geometry === info.geometry ), + root: i, + object, + instanceId, + compositeId, + } ); + } ); } ); @@ -426,67 +439,67 @@ export class BVHComputeFns { const indexBuffer = new Uint32Array( indexBufferLength ); const attributesBuffer = new ArrayBuffer( attributesBufferLength * attributes.length * 4 * 4 ); const bvhNodesBuffer = new ArrayBuffer( bvhNodesBufferLength ); - appendBVHData( bvh, 0, transformInfo, bvhNodesBuffer, true ); + + // append TLAS data + appendBVHData( bvh, 0, transformInfo, 0, bvhNodesBuffer, true ); + nodeWriteOffset += getTotalBVHByteLength( bvh ) / BYTES_PER_NODE; geometryInfo.forEach( info => { - const { geometry, range, indirectBuffer } = info; - const indexOffset = appendGeometryData( geometry, range, indirectBuffer ); - const bvhNodeOffsets = appendBVHData( info.bvh, indexOffset, transformInfo, bvhNodesBuffer, false ); + // append bvh data + const bvhNodeOffsets = appendBVHData( info.bvh, indexOffset / 3, transformInfo, nodeWriteOffset, bvhNodesBuffer, false ); + info.bvhNodeOffsets = bvhNodeOffsets; + // append geometry data + appendGeometryData( info.geometry, info.range, indexOffset, attributesOffset, info.bvh._indirectBuffer ); info.indexBufferOffset = indexOffset; - info.bvhNodeOffsets = bvhNodeOffsets; + + // step the write offsets forward + indexOffset += info.range.count; + attributesOffset += info.range.vertexCount; + nodeWriteOffset += getTotalBVHByteLength( info.bvh ) / BYTES_PER_NODE; } ); // // write the transforms - let transformWriteOffset = 0; - const transformArrayBuffer = new ArrayBuffer( TRANSFORM_STRUCT_SIZE * 4 * objectTransformsCount ); + const transformArrayBuffer = new ArrayBuffer( TRANSFORM_STRUCT_SIZE * 4 * transformInfo.length ); const transformBufferF32 = new Float32Array( transformArrayBuffer ); const transformBufferU32 = new Uint32Array( transformArrayBuffer ); - bvh.primitiveBuffer.forEach( ( compositeId, i ) => { + transformInfo.forEach( ( info, i ) => { + const { compositeId, data, root } = info; bvh.getObjectMatrix( compositeId, _matrix ); _inverseMatrix.copy( _matrix ).invert(); - const objectBvh = transformInfo[ i ].data.bvh; - const bvhOffset = transformInfo[ i ].data.bvhNodeOffsets; - objectBvh._roots.forEach( ( root, ri ) => { - - _matrix.toArray( transformBufferF32, transformWriteOffset * TRANSFORM_STRUCT_SIZE ); - _inverseMatrix.toArray( transformBufferF32, transformWriteOffset * TRANSFORM_STRUCT_SIZE + 16 ); - transformBufferU32[ transformWriteOffset * TRANSFORM_STRUCT_SIZE + 32 ] = bvhOffset[ ri ]; - transformWriteOffset ++; - - } ); + const { bvhNodeOffsets } = data; + _matrix.toArray( transformBufferF32, i * TRANSFORM_STRUCT_SIZE ); + _inverseMatrix.toArray( transformBufferF32, i * TRANSFORM_STRUCT_SIZE + 16 ); + transformBufferU32[ i * TRANSFORM_STRUCT_SIZE + 32 ] = bvhNodeOffsets[ root ]; } ); // // set up the storage buffers - const nodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), 'BVHNode' ).toReadOnly().setName( `${ name }nodes` ); + const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), 'BVHNode' ).toReadOnly().setName( `${ name }nodes` ); const transformsStorage = storage( new StorageBufferAttribute( transformBufferF32, TRANSFORM_STRUCT_SIZE ), 'TransformStruct' ).toReadOnly().setName( `${ name }transforms` ); const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ name }index` ); const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributesStructSize ), `${ name }GeometryStruct` ).toReadOnly().setName( `${ name }attributes` ); this.storageBufferAttributes.transforms = transformsStorage; - this.storageBufferAttributes.nodes = nodesStorage; + this.storageBufferAttributes.nodes = bvhNodesStorage; this.storageBufferAttributes.index = indexStorage; this.storageBufferAttributes.attributes = attributesStorage; this.attributesStruct = attributeStruct; - this.raycastFirstHitFn = buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ); + this.raycastFirstHitFn = buildRaycastFirstHitFn( name, bvhNodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ); - function appendBVHData( bvh, geometryOffset, transformInfo, target, tlas = false ) { + function appendBVHData( bvh, geometryOffset, transformInfo, nodeWriteOffset, target, tlas = false ) { const targetU16 = new Uint16Array( target ); const targetU32 = new Uint32Array( target ); const targetF32 = new Float32Array( target ); - const BYTES_PER_NODE = 6 * 4 + 4 + 4; - const UINT32_PER_NODE = BYTES_PER_NODE / 4; - const IS_LEAFNODE_FLAG = 0xFFFF; const result = []; bvh._roots.forEach( root => { @@ -555,22 +568,9 @@ export class BVHComputeFns { } - function appendGeometryData( geometry, offsets = null, indirectBuffer = null ) { - - let vertexStart = 0; - let vertexCount = geometry.attributes.position.count; - if ( offsets ) { - - vertexStart = offsets.vertexStart; - vertexCount = offsets.vertexCount; - - } - - // write indices — when an indirect buffer is present, dereference it to - // resolve the BVH's triangle order into a flat index buffer - // store as triangle offset (not vertex-index offset) to match BVH leaf format - const result = indexOffset / 3; + function appendIndexData( geometry, range, indexOffset, attributesOffset, indirectBuffer = null ) { + const { start, count, vertexStart } = range; if ( indirectBuffer ) { const dereferencedIndex = dereferenceIndex( geometry.index, indirectBuffer ); @@ -584,43 +584,34 @@ export class BVHComputeFns { } else if ( geometry.index ) { - let indexStart = 0; - let indexCount = geometry.index.count; - if ( offsets ) { - - indexStart = offsets.indexStart; - indexCount = offsets.indexCount; - - } - - for ( let i = 0; i < indexCount; i ++ ) { + for ( let i = 0; i < count; i ++ ) { - indexBuffer[ i + indexOffset ] = geometry.index.getX( i + indexStart ) - vertexStart + attributesOffset; + indexBuffer[ i + indexOffset ] = geometry.index.getX( i + start ) - vertexStart + attributesOffset; } - indexOffset += indexCount; + indexOffset += count; } else { - let indexStart = 0; - let indexCount = geometry.attributes.position.count; - if ( offsets ) { + for ( let i = 0; i < count; i ++ ) { - indexStart = offsets.vertexStart; - indexCount = offsets.vertexCount; + indexBuffer[ i + indexOffset ] = i + start + attributesOffset; } - for ( let i = 0; i < indexCount; i ++ ) { + indexOffset += count; - indexBuffer[ i + indexOffset ] = i + indexStart + attributesOffset; + } - } - indexOffset += indexCount; + } - } + function appendGeometryData( geometry, range, indexOffset, attributesOffset, indirectBuffer = null ) { + + appendIndexData( geometry, range, indexOffset, attributesOffset, indirectBuffer ); + + const { vertexStart, vertexCount } = range; const attributesBufferF32 = new Float32Array( attributesBuffer ); attributes.forEach( ( key, interleavedOffset ) => { @@ -664,7 +655,6 @@ export class BVHComputeFns { } ); attributesOffset += vertexCount; - return result; } From c3816e8ffb048be295231aea0d0cda408e3e6453 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 15:44:44 +0900 Subject: [PATCH 35/84] Cleanup --- src/webgpu/lib/BVHComputeFns.js | 73 ++++++++++++++++++++------------- 1 file changed, 45 insertions(+), 28 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 5864d01e1..d82757cef 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -1,5 +1,5 @@ import { Matrix4, Vector4 } from 'three'; -import { StorageBufferAttribute } from 'three/src/Three.WebGPU.Nodes.js'; +import { StorageBufferAttribute } from 'three/webgpu'; import { storage, wgsl, wgslFn } from 'three/tsl'; import { intersectsBounds, @@ -450,7 +450,8 @@ export class BVHComputeFns { info.bvhNodeOffsets = bvhNodeOffsets; // append geometry data - appendGeometryData( info.geometry, info.range, indexOffset, attributesOffset, info.bvh._indirectBuffer ); + appendIndexData( info.geometry, info.bvh._indirectBuffer, info.range, attributesOffset, indexOffset, indexBuffer ); + appendGeometryData( info.geometry, info.range, attributesOffset, attributesBuffer ); info.indexBufferOffset = indexOffset; // step the write offsets forward @@ -464,18 +465,10 @@ export class BVHComputeFns { // write the transforms const transformArrayBuffer = new ArrayBuffer( TRANSFORM_STRUCT_SIZE * 4 * transformInfo.length ); - const transformBufferF32 = new Float32Array( transformArrayBuffer ); - const transformBufferU32 = new Uint32Array( transformArrayBuffer ); transformInfo.forEach( ( info, i ) => { - const { compositeId, data, root } = info; - bvh.getObjectMatrix( compositeId, _matrix ); - _inverseMatrix.copy( _matrix ).invert(); - - const { bvhNodeOffsets } = data; - _matrix.toArray( transformBufferF32, i * TRANSFORM_STRUCT_SIZE ); - _inverseMatrix.toArray( transformBufferF32, i * TRANSFORM_STRUCT_SIZE + 16 ); - transformBufferU32[ i * TRANSFORM_STRUCT_SIZE + 32 ] = bvhNodeOffsets[ root ]; + _inverseMatrix.copy( bvh.matrixWorld ).invert(); + appendTransformData( info, _inverseMatrix, i, transformArrayBuffer ); } ); @@ -483,7 +476,7 @@ export class BVHComputeFns { // set up the storage buffers const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), 'BVHNode' ).toReadOnly().setName( `${ name }nodes` ); - const transformsStorage = storage( new StorageBufferAttribute( transformBufferF32, TRANSFORM_STRUCT_SIZE ), 'TransformStruct' ).toReadOnly().setName( `${ name }transforms` ); + const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), TRANSFORM_STRUCT_SIZE ), 'TransformStruct' ).toReadOnly().setName( `${ name }transforms` ); const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ name }index` ); const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributesStructSize ), `${ name }GeometryStruct` ).toReadOnly().setName( `${ name }attributes` ); @@ -494,6 +487,34 @@ export class BVHComputeFns { this.attributesStruct = attributeStruct; this.raycastFirstHitFn = buildRaycastFirstHitFn( name, bvhNodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ); + function appendTransformData( info, premultiplyMatrix, writeOffset, target ) { + + const transformBufferF32 = new Float32Array( target ); + const transformBufferU32 = new Uint32Array( target ); + + const { object, instanceId, root, data } = info; + const { bvhNodeOffsets } = data; + + if ( object.isInstancedMesh || object.isBatchedMesh ) { + + object.getMatrixAt( instanceId, _matrix ).premultiply( object.matrixWorld ); + + } else { + + _matrix.copy( object.matrixWorld ); + + } + + _matrix.premultiply( premultiplyMatrix ); + _matrix.toArray( transformBufferF32, writeOffset * TRANSFORM_STRUCT_SIZE ); + + _matrix.invert(); + _matrix.toArray( transformBufferF32, writeOffset * TRANSFORM_STRUCT_SIZE + 16 ); + + transformBufferU32[ writeOffset * TRANSFORM_STRUCT_SIZE + 32 ] = bvhNodeOffsets[ root ]; + + } + function appendBVHData( bvh, geometryOffset, transformInfo, nodeWriteOffset, target, tlas = false ) { const targetU16 = new Uint16Array( target ); @@ -568,7 +589,7 @@ export class BVHComputeFns { } - function appendIndexData( geometry, range, indexOffset, attributesOffset, indirectBuffer = null ) { + function appendIndexData( geometry, indirectBuffer, range, valueOffset, writeOffset, target ) { const { start, count, vertexStart } = range; if ( indirectBuffer ) { @@ -576,44 +597,40 @@ export class BVHComputeFns { const dereferencedIndex = dereferenceIndex( geometry.index, indirectBuffer ); for ( let i = 0; i < dereferencedIndex.length; i ++ ) { - indexBuffer[ i + indexOffset ] = dereferencedIndex[ i ] - vertexStart + attributesOffset; + target[ i + writeOffset ] = dereferencedIndex[ i ] - vertexStart + valueOffset; } - indexOffset += dereferencedIndex.length; + writeOffset += dereferencedIndex.length; } else if ( geometry.index ) { for ( let i = 0; i < count; i ++ ) { - indexBuffer[ i + indexOffset ] = geometry.index.getX( i + start ) - vertexStart + attributesOffset; + target[ i + writeOffset ] = geometry.index.getX( i + start ) - vertexStart + valueOffset; } - indexOffset += count; + writeOffset += count; } else { for ( let i = 0; i < count; i ++ ) { - indexBuffer[ i + indexOffset ] = i + start + attributesOffset; + target[ i + writeOffset ] = i + start + valueOffset; } - indexOffset += count; + writeOffset += count; } - } - function appendGeometryData( geometry, range, indexOffset, attributesOffset, indirectBuffer = null ) { - - appendIndexData( geometry, range, indexOffset, attributesOffset, indirectBuffer ); + function appendGeometryData( geometry, range, writeOffset, target ) { const { vertexStart, vertexCount } = range; - - const attributesBufferF32 = new Float32Array( attributesBuffer ); + const attributesBufferF32 = new Float32Array( target ); attributes.forEach( ( key, interleavedOffset ) => { const attr = geometry.attributes[ key ]; @@ -648,13 +665,13 @@ export class BVHComputeFns { } - _vec.toArray( attributesBufferF32, ( attributesOffset + i ) * attributesStructSize + interleavedOffset * 4 ); + _vec.toArray( attributesBufferF32, ( writeOffset + i ) * attributesStructSize + interleavedOffset * 4 ); } } ); - attributesOffset += vertexCount; + writeOffset += vertexCount; } From e73117adfe44e9c159c7bd0158b4f886e9865590 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 17:47:41 +0900 Subject: [PATCH 36/84] Cleanup --- src/webgpu/compute/PathTracerMegaKernel.js | 18 +- src/webgpu/lib/BVHComputeFns.js | 220 +++++++++++---------- 2 files changed, 122 insertions(+), 116 deletions(-) diff --git a/src/webgpu/compute/PathTracerMegaKernel.js b/src/webgpu/compute/PathTracerMegaKernel.js index 79d4441cf..a6e126819 100644 --- a/src/webgpu/compute/PathTracerMegaKernel.js +++ b/src/webgpu/compute/PathTracerMegaKernel.js @@ -27,7 +27,7 @@ export class PathTracerMegaKernel extends ComputeKernel { globalId: globalId, }; - const { raycastFirstHitFn, name } = bvhComputeFns; + const { fns, name } = bvhComputeFns; const shader = wgslFn( /* wgsl */` fn compute( @@ -81,20 +81,6 @@ export class PathTracerMegaKernel extends ComputeKernel { for ( var bounce = 0u; bounce < bounces; bounce ++ ) { let hitResult = ${ name }RaycastFirstHit( ray ); - // resultColor = hitResult.normal; - // break; - // if ( hitResult.didHit ) { - - // resultColor = vec3f( 1, 0, 0 ); - - // } else { - - // resultColor = vec3f( 0, 0, 1 ); - - // } - - // break; - if ( hitResult.didHit ) { let hitPosition = ray.origin + ray.direction * hitResult.dist; @@ -125,7 +111,7 @@ export class PathTracerMegaKernel extends ComputeKernel { } - `, [ raycastFirstHitFn, ndcToCameraRay, pcgRand3, pcgInit, lambertBsdfFunc ] ); + `, [ fns.raycastFirstHit, ndcToCameraRay, pcgRand3, pcgInit, lambertBsdfFunc ] ); super( shader( megakernelShaderParams ) ); diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index d82757cef..249e2f525 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -1,5 +1,5 @@ import { Matrix4, Vector4 } from 'three'; -import { StorageBufferAttribute } from 'three/webgpu'; +import { CodeNode, StorageBufferAttribute } from 'three/webgpu'; import { storage, wgsl, wgslFn } from 'three/tsl'; import { intersectsBounds, @@ -23,27 +23,61 @@ const IS_LEAFNODE_FLAG = 0xFFFF; // a node fn can be updated without regenerating all other materials. // TODO: see if we can reference wgslFn names directly rather than constructing them inline over and over // and / or use local variable definitions for the pointers to clean up the code -// TODO: see if there's a "build" step that can be leveraged fro nodes +// TODO: see if there's a "build" step that can be leveraged for nodes to make integration more simple // TODO: allow for "slotting" a new type of callback (eg distance, etc) so multiple types of queries can be made +// - add a "shapecast" style function with functions and return types that can be slotted in +// NEXT: add support for custom transform fields const _def = /* @__PURE__ */ new Vector4(); const _vec = /* @__PURE__ */ new Vector4(); const _matrix = /* @__PURE__ */ new Matrix4(); const _inverseMatrix = /* @__PURE__ */ new Matrix4(); +// a more structured "struct" node that bookkeeps the struct name, byte size +class WGSLStructNode extends CodeNode { + + get uintSize() { + + return this.byteSize / 4; + + } + + constructor( name, byteSize, fields, includes = [] ) { + + const content = Object + .entries( fields ) + .map( ( [ name, type ] ) => { + + return `${ name }: ${ type },`; + + } ).join( '\n' ); + + const code = /* wgsl */` + struct ${ name } { + ${ content } + } + `; + + super( code, includes, 'wgsl' ); + this.name = name; + this.byteSize = byteSize; + + } + +} + +const wgslStruct = ( ...args ) => new WGSLStructNode( ...args ); + // stride is 36 floats (144 bytes) to match WGSL struct alignment: // mat4x4f (64) + mat4x4f (64) + u32 (4) + 12 bytes padding to align to 16 -const TRANSFORM_STRUCT_SIZE = 36; -const transformStruct = wgsl( /* wgsl */` - struct TransformStruct { - matrixWorld: mat4x4f, - inverseMatrixWorld: mat4x4f, - nodeOffset: u32, - _alignment0: u32, - _alignment1: u32, - _alignment2: u32, - } -` ); +const transformStruct = wgslStruct( 'TransformStruct', 36 * 4, { + matrixWorld: 'mat4x4f', + inverseMatrixWorld: 'mat4x4f', + nodeOffset: 'u32', + _alignment0: 'u32', + _alignment1: 'u32', + _alignment2: 'u32', +} ); function dereferenceIndex( indexAttr, indirectBuffer ) { @@ -71,17 +105,15 @@ function getTotalBVHByteLength( bvh ) { } -const intersectionResultStruct = wgsl( /* wgsl */` - struct IntersectionResult { - didHit: bool, - indices: vec4u, - normal: vec3f, - barycoord: vec3f, - side: f32, - dist: f32, - objectIndex: u32, - }; -` ); +const intersectionResultStruct = wgslStruct( 'IntersectionResult', 16 * 4, { + indices: 'vec4u', + normal: 'vec3f', + didHit: 'bool', + barycoord: 'vec3f', + objectIndex: 'u32', + side: 'f32', + dist: 'f32', +} ); const intersectsTriangle = wgslFn( /* wgsl */ ` @@ -330,15 +362,22 @@ export class BVHComputeFns { this.attributes = attributes; this.bvh = bvh; - this.storageBufferAttributes = { + this.storage = { index: null, attributes: null, nodes: null, transforms: null, }; - this.attributesStruct = null; - this.raycastFirstHitFn = null; + this.structs = { + attributes: null, + transform: null, + intersection: null, + }; + + this.fns = { + raycastFirstHit: null, + }; this.update(); @@ -361,46 +400,26 @@ export class BVHComputeFns { const object = bvh.getObjectFromId( compositeId ); const instanceId = bvh.getInstanceFromId( compositeId ); - const meshBvh = this.getBVH( object, instanceId ); + const range = { start: 0, count: 0, vertexStart: 0, vertexCount: 0 }; + const primBvh = this.getBVH( object, instanceId, range ); // if we haven't added this bvh, yet - if ( ! geometryInfo.find( info => info.bvh === meshBvh ) ) { - - // TODO: account for indirect buffer here? + if ( ! geometryInfo.find( info => info.bvh === primBvh ) ) { // save the geometry info to write later and increment the buffer sizes const info = { index: geometryInfo.length, - bvh: meshBvh, - geometry: object.geometry, - range: { - start: 0, - count: 0, - vertexStart: 0, - vertexCount: 0, - }, + bvh: primBvh, + geometry: primBvh.geometry, + range: range, bvhBufferOffsets: null, indexBufferOffset: null, }; - if ( object.isBatchedMesh ) { - - const geometryId = object.getGeometryIdAt( instanceId ); - const range = object.getGeometryRangeAt( geometryId ); - Object.assign( info.range, range ); - - } else { - - const geometry = object.geometry; - info.range.count = geometry.index ? geometry.index.count : geometry.attributes.position.count, - info.range.vertexCount = geometry.attributes.position.count; - - } - // increase the buffer sizes for bvh and geometry - bvhNodesBufferLength += getTotalBVHByteLength( meshBvh ); + bvhNodesBufferLength += getTotalBVHByteLength( primBvh ); indexBufferLength += info.range.count; attributesBufferLength += info.range.vertexCount; geometryInfo.push( info ); @@ -408,7 +427,7 @@ export class BVHComputeFns { } // save the index of the bvh associated with this transform - meshBvh._roots.forEach( ( root, i ) => { + primBvh._roots.forEach( ( root, i ) => { transformInfo.push( { data: geometryInfo.find( info => object.geometry === info.geometry ), @@ -464,11 +483,11 @@ export class BVHComputeFns { // // write the transforms - const transformArrayBuffer = new ArrayBuffer( TRANSFORM_STRUCT_SIZE * 4 * transformInfo.length ); + const transformArrayBuffer = new ArrayBuffer( transformStruct.byteSize * transformInfo.length ); transformInfo.forEach( ( info, i ) => { _inverseMatrix.copy( bvh.matrixWorld ).invert(); - appendTransformData( info, _inverseMatrix, i, transformArrayBuffer ); + this.writeTransformData( info, _inverseMatrix, i, transformArrayBuffer ); } ); @@ -476,44 +495,18 @@ export class BVHComputeFns { // set up the storage buffers const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), 'BVHNode' ).toReadOnly().setName( `${ name }nodes` ); - const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), TRANSFORM_STRUCT_SIZE ), 'TransformStruct' ).toReadOnly().setName( `${ name }transforms` ); + const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), transformStruct.uintSize ), 'TransformStruct' ).toReadOnly().setName( `${ name }transforms` ); const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ name }index` ); const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributesStructSize ), `${ name }GeometryStruct` ).toReadOnly().setName( `${ name }attributes` ); - this.storageBufferAttributes.transforms = transformsStorage; - this.storageBufferAttributes.nodes = bvhNodesStorage; - this.storageBufferAttributes.index = indexStorage; - this.storageBufferAttributes.attributes = attributesStorage; - this.attributesStruct = attributeStruct; - this.raycastFirstHitFn = buildRaycastFirstHitFn( name, bvhNodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ); - - function appendTransformData( info, premultiplyMatrix, writeOffset, target ) { - - const transformBufferF32 = new Float32Array( target ); - const transformBufferU32 = new Uint32Array( target ); - - const { object, instanceId, root, data } = info; - const { bvhNodeOffsets } = data; - - if ( object.isInstancedMesh || object.isBatchedMesh ) { - - object.getMatrixAt( instanceId, _matrix ).premultiply( object.matrixWorld ); - - } else { - - _matrix.copy( object.matrixWorld ); - - } - - _matrix.premultiply( premultiplyMatrix ); - _matrix.toArray( transformBufferF32, writeOffset * TRANSFORM_STRUCT_SIZE ); - - _matrix.invert(); - _matrix.toArray( transformBufferF32, writeOffset * TRANSFORM_STRUCT_SIZE + 16 ); - - transformBufferU32[ writeOffset * TRANSFORM_STRUCT_SIZE + 32 ] = bvhNodeOffsets[ root ]; - - } + this.storage.transforms = transformsStorage; + this.storage.nodes = bvhNodesStorage; + this.storage.index = indexStorage; + this.storage.attributes = attributesStorage; + this.structs.attributes = attributeStruct; + this.structs.transform = transformStruct; + this.structs.intersectionResult = intersectionResultStruct; + this.fns.raycastFirstHit = buildRaycastFirstHitFn( name, bvhNodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ); function appendBVHData( bvh, geometryOffset, transformInfo, nodeWriteOffset, target, tlas = false ) { @@ -522,10 +515,9 @@ export class BVHComputeFns { const targetF32 = new Float32Array( target ); const result = []; - + let tlasOffset = 0; bvh._roots.forEach( root => { - let tlasOffset = 0; const rootBuffer16 = new Uint16Array( root ); const rootBuffer32 = new Uint32Array( root ); result.push( nodeWriteOffset ); @@ -576,7 +568,6 @@ export class BVHComputeFns { targetU32[ n32 + 6 ] = rootBuffer32[ r32 + 6 ]; targetU32[ n32 + 7 ] = rootBuffer32[ r32 + 7 ]; - } nodeWriteOffset ++; @@ -677,19 +668,48 @@ export class BVHComputeFns { } - getBVH( object, id ) { + writeTransformData( info, premultiplyMatrix, writeOffset, targetBuffer ) { - if ( object.isInstancedMesh ) { + const transformBufferF32 = new Float32Array( targetBuffer ); + const transformBufferU32 = new Uint32Array( targetBuffer ); - return object.geometry.boundsTree; + const { object, instanceId, root, data } = info; + const { bvhNodeOffsets } = data; + + if ( object.isInstancedMesh || object.isBatchedMesh ) { + + object.getMatrixAt( instanceId, _matrix ).premultiply( object.matrixWorld ); + + } else { + + _matrix.copy( object.matrixWorld ); + + } + + _matrix.premultiply( premultiplyMatrix ); + _matrix.toArray( transformBufferF32, writeOffset * transformStruct.uintSize ); + + _matrix.invert(); + _matrix.toArray( transformBufferF32, writeOffset * transformStruct.uintSize + 16 ); + + transformBufferU32[ writeOffset * transformStruct.uintSize + 32 ] = bvhNodeOffsets[ root ]; + + } + + getBVH( object, instanceId, rangeTarget ) { - } else if ( object.isBatchedMesh ) { + if ( object.isBatchedMesh ) { - const geometryId = object.getGeometryIdAt( id ); + const geometryId = object.getGeometryIdAt( instanceId ); + const range = object.getGeometryRangeAt( geometryId ); + Object.assign( rangeTarget, range ); return object.boundsTrees[ geometryId ]; } else { + const geometry = object.geometry; + rangeTarget.count = geometry.index ? geometry.index.count : geometry.attributes.position.count; + rangeTarget.vertexCount = geometry.attributes.position.count; return object.geometry.boundsTree; } From 9e035951ee8be99d0dae0f8afff30fe815384157 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Thu, 19 Feb 2026 18:24:02 +0900 Subject: [PATCH 37/84] Move attributes to a wgslstruct definition --- src/webgpu/lib/BVHComputeFns.js | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 249e2f525..25abf85ac 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -444,12 +444,11 @@ export class BVHComputeFns { // // construct the attribute struct - const attributesStructSize = 4 * attributes.length; - const attributeStruct = wgsl( /* wgsl */` - struct ${ name }GeometryStruct { - ${ attributes.map( key => `${ key }: vec4f,` ).join( '\n' ) } - } - ` ); + const attributeStruct = wgslStruct( + `${ name }GeometryStruct`, + 4 * 4 * attributes.length, + attributes.reduce( ( o, key ) => ( { ...o, [ key ]: 'vec4f' } ), {} ), + ); // write the geometry buffer attributes & bvh data let attributesOffset = 0; @@ -495,14 +494,13 @@ export class BVHComputeFns { // set up the storage buffers const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), 'BVHNode' ).toReadOnly().setName( `${ name }nodes` ); - const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), transformStruct.uintSize ), 'TransformStruct' ).toReadOnly().setName( `${ name }transforms` ); + const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), transformStruct.uintSize ), transformStruct.name ).toReadOnly().setName( `${ name }transforms` ); const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ name }index` ); - const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributesStructSize ), `${ name }GeometryStruct` ).toReadOnly().setName( `${ name }attributes` ); + const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributeStruct.uintSize ), attributeStruct.name ).toReadOnly().setName( `${ name }attributes` ); this.storage.transforms = transformsStorage; this.storage.nodes = bvhNodesStorage; this.storage.index = indexStorage; - this.storage.attributes = attributesStorage; this.structs.attributes = attributeStruct; this.structs.transform = transformStruct; this.structs.intersectionResult = intersectionResultStruct; @@ -656,7 +654,7 @@ export class BVHComputeFns { } - _vec.toArray( attributesBufferF32, ( writeOffset + i ) * attributesStructSize + interleavedOffset * 4 ); + _vec.toArray( attributesBufferF32, ( writeOffset + i ) * attributeStruct.uintSize + interleavedOffset * 4 ); } From 809ef0adfa2ce9d94b7df062486ec0ab1b813051 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Fri, 20 Feb 2026 12:46:14 +0900 Subject: [PATCH 38/84] Cleanup --- src/webgpu/lib/BVHComputeFns.js | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 25abf85ac..1ca12cea7 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -1,6 +1,6 @@ import { Matrix4, Vector4 } from 'three'; import { CodeNode, StorageBufferAttribute } from 'three/webgpu'; -import { storage, wgsl, wgslFn } from 'three/tsl'; +import { storage, wgslFn } from 'three/tsl'; import { intersectsBounds, rayStruct, @@ -166,7 +166,7 @@ const intersectsTriangle = wgslFn( /* wgsl */ ` function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ) { const geometryRaycastFirstHitFn = wgslFn( /* wgsl */` - fn ${ name }RaycastGeometryFirstHit( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> IntersectionResult { + fn ${ name }RaycastFirstHit_blas( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> IntersectionResult { var bestHit: IntersectionResult; bestHit.didHit = false; @@ -303,7 +303,7 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto localRay.origin = ( transform.inverseMatrixWorld * vec4f( ray.origin, 1.0 ) ).xyz; localRay.direction = ( transform.inverseMatrixWorld * vec4f( ray.direction, 0.0 ) ).xyz; - let blasHit = ${ name }RaycastGeometryFirstHit( localRay, transform.nodeOffset, bestHit.dist ); + let blasHit = ${ name }RaycastFirstHit_blas( localRay, transform.nodeOffset, bestHit.dist ); if ( blasHit.didHit && blasHit.dist < bestHit.dist ) { bestHit = blasHit; @@ -360,8 +360,8 @@ export class BVHComputeFns { this.name = name; this.attributes = attributes; - this.bvh = bvh; + this.storage = { index: null, attributes: null, @@ -370,9 +370,8 @@ export class BVHComputeFns { }; this.structs = { + transform: transformStruct, attributes: null, - transform: null, - intersection: null, }; this.fns = { @@ -501,9 +500,8 @@ export class BVHComputeFns { this.storage.transforms = transformsStorage; this.storage.nodes = bvhNodesStorage; this.storage.index = indexStorage; + this.storage.attributes = attributesStorage; this.structs.attributes = attributeStruct; - this.structs.transform = transformStruct; - this.structs.intersectionResult = intersectionResultStruct; this.fns.raycastFirstHit = buildRaycastFirstHitFn( name, bvhNodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ); function appendBVHData( bvh, geometryOffset, transformInfo, nodeWriteOffset, target, tlas = false ) { @@ -716,16 +714,20 @@ export class BVHComputeFns { getDefaultAttributeValue( key, target ) { - if ( key === 'color' ) { + switch ( key ) { + case 'position': + case 'color': target.set( 1, 1, 1, 1 ); + break; - } else { - - target.set( 0, 0, 0, 1 ); + default: + target.set( 0, 0, 0, 0 ); } + return target; + } dispose() { From 80e8959c922fe4c58e0d62a5cc8ddbeea433517d Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Fri, 20 Feb 2026 13:25:07 +0900 Subject: [PATCH 39/84] Add material support --- src/webgpu/WebGPUPathTracer.js | 6 +- src/webgpu/compute/PathTracerMegaKernel.js | 12 +++- src/webgpu/lib/BVHComputeFns.js | 51 ++-------------- src/webgpu/lib/PathtracerBVHComputeFns.js | 70 ++++++++++++++++++++++ src/webgpu/lib/WGSLStructNode.js | 36 +++++++++++ 5 files changed, 124 insertions(+), 51 deletions(-) create mode 100644 src/webgpu/lib/PathtracerBVHComputeFns.js create mode 100644 src/webgpu/lib/WGSLStructNode.js diff --git a/src/webgpu/WebGPUPathTracer.js b/src/webgpu/WebGPUPathTracer.js index e6bb01a18..dc20f96f7 100644 --- a/src/webgpu/WebGPUPathTracer.js +++ b/src/webgpu/WebGPUPathTracer.js @@ -5,7 +5,7 @@ import { RenderToScreenNodeMaterial } from './materials/RenderToScreenMaterial.j import { MegaKernelPathTracer } from './MegaKernelPathTracer.js'; import { WaveFrontPathTracer } from './WaveFrontPathTracer.js'; import { ObjectBVH } from './lib/ObjectBVH.js'; -import { BVHComputeFns } from './lib/BVHComputeFns.js'; +import { PathtracerBVHComputeFns } from './lib/PathtracerBVHComputeFns.js'; const _resolution = new Vector2(); export class WebGPUPathTracer { @@ -59,7 +59,7 @@ export class WebGPUPathTracer { if ( child.isMesh && ! child.geometry.boundsTree ) { - child.geometry.boundsTree = new MeshBVH( child.geometry, { strategy: SAH, maxLeafSize: 1 } ); + child.geometry.boundsTree = new MeshBVH( child.geometry, { strategy: SAH, maxLeafSize: 5 } ); } @@ -67,7 +67,7 @@ export class WebGPUPathTracer { // Build TLAS and compute functions const objectBVH = new ObjectBVH( scene ); - const bvhComputeFns = new BVHComputeFns( objectBVH ); + const bvhComputeFns = new PathtracerBVHComputeFns( objectBVH ); this._bvhComputeFns = bvhComputeFns; this._pathTracer.setBVHComputeFns( bvhComputeFns ); diff --git a/src/webgpu/compute/PathTracerMegaKernel.js b/src/webgpu/compute/PathTracerMegaKernel.js index a6e126819..d8c409846 100644 --- a/src/webgpu/compute/PathTracerMegaKernel.js +++ b/src/webgpu/compute/PathTracerMegaKernel.js @@ -27,7 +27,7 @@ export class PathTracerMegaKernel extends ComputeKernel { globalId: globalId, }; - const { fns, name } = bvhComputeFns; + const { fns, storage, structs, name } = bvhComputeFns; const shader = wgslFn( /* wgsl */` fn compute( @@ -86,8 +86,11 @@ export class PathTracerMegaKernel extends ComputeKernel { let hitPosition = ray.origin + ray.direction * hitResult.dist; let scatterRec = bsdfEval( hitResult.normal, - ray.direction ); + let transform = ${ name }transforms.value[ hitResult.objectIndex ]; + let material = ${ name }materials.value[ transform.materialIndex ]; + // white diffuse surface - throughputColor *= hitResult.normal * scatterRec.value / scatterRec.pdf; + throughputColor *= material.albedo * scatterRec.value / scatterRec.pdf; ray.origin = hitPosition; ray.direction = scatterRec.direction; @@ -111,7 +114,10 @@ export class PathTracerMegaKernel extends ComputeKernel { } - `, [ fns.raycastFirstHit, ndcToCameraRay, pcgRand3, pcgInit, lambertBsdfFunc ] ); + `, [ + fns.raycastFirstHit, storage.materials, structs.material, + ndcToCameraRay, pcgRand3, pcgInit, lambertBsdfFunc, + ] ); super( shader( megakernelShaderParams ) ); diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 1ca12cea7..75223152d 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -1,5 +1,5 @@ import { Matrix4, Vector4 } from 'three'; -import { CodeNode, StorageBufferAttribute } from 'three/webgpu'; +import { StorageBufferAttribute } from 'three/webgpu'; import { storage, wgslFn } from 'three/tsl'; import { intersectsBounds, @@ -7,13 +7,12 @@ import { bvhNodeStruct, constants, } from 'three-mesh-bvh/webgpu'; +import { wgslStruct } from './WGSLStructNode.js'; const BYTES_PER_NODE = 6 * 4 + 4 + 4; const UINT32_PER_NODE = BYTES_PER_NODE / 4; const IS_LEAFNODE_FLAG = 0xFFFF; -// TODO: clean up "update" function, separate it into chunks -// TODO: separate update functions into utilities // TODO: add ability to easily update a single matrix / scene rearrangement (partial update) // TODO: add material support w/ function to easily update material // - add a callback for writing a property for a geometry to a range @@ -26,48 +25,12 @@ const IS_LEAFNODE_FLAG = 0xFFFF; // TODO: see if there's a "build" step that can be leveraged for nodes to make integration more simple // TODO: allow for "slotting" a new type of callback (eg distance, etc) so multiple types of queries can be made // - add a "shapecast" style function with functions and return types that can be slotted in -// NEXT: add support for custom transform fields const _def = /* @__PURE__ */ new Vector4(); const _vec = /* @__PURE__ */ new Vector4(); const _matrix = /* @__PURE__ */ new Matrix4(); const _inverseMatrix = /* @__PURE__ */ new Matrix4(); -// a more structured "struct" node that bookkeeps the struct name, byte size -class WGSLStructNode extends CodeNode { - - get uintSize() { - - return this.byteSize / 4; - - } - - constructor( name, byteSize, fields, includes = [] ) { - - const content = Object - .entries( fields ) - .map( ( [ name, type ] ) => { - - return `${ name }: ${ type },`; - - } ).join( '\n' ); - - const code = /* wgsl */` - struct ${ name } { - ${ content } - } - `; - - super( code, includes, 'wgsl' ); - this.name = name; - this.byteSize = byteSize; - - } - -} - -const wgslStruct = ( ...args ) => new WGSLStructNode( ...args ); - // stride is 36 floats (144 bytes) to match WGSL struct alignment: // mat4x4f (64) + mat4x4f (64) + u32 (4) + 12 bytes padding to align to 16 const transformStruct = wgslStruct( 'TransformStruct', 36 * 4, { @@ -378,14 +341,12 @@ export class BVHComputeFns { raycastFirstHit: null, }; - this.update(); - } update() { const self = this; - const { attributes, name, bvh } = this; + const { attributes, structs, name, bvh } = this; // collect the BVHs const geometryInfo = []; @@ -481,7 +442,7 @@ export class BVHComputeFns { // // write the transforms - const transformArrayBuffer = new ArrayBuffer( transformStruct.byteSize * transformInfo.length ); + const transformArrayBuffer = new ArrayBuffer( structs.transform.byteSize * transformInfo.length ); transformInfo.forEach( ( info, i ) => { _inverseMatrix.copy( bvh.matrixWorld ).invert(); @@ -493,7 +454,7 @@ export class BVHComputeFns { // set up the storage buffers const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), 'BVHNode' ).toReadOnly().setName( `${ name }nodes` ); - const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), transformStruct.uintSize ), transformStruct.name ).toReadOnly().setName( `${ name }transforms` ); + const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), structs.transform.uintSize ), structs.transform.name ).toReadOnly().setName( `${ name }transforms` ); const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ name }index` ); const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributeStruct.uintSize ), attributeStruct.name ).toReadOnly().setName( `${ name }attributes` ); @@ -502,7 +463,7 @@ export class BVHComputeFns { this.storage.index = indexStorage; this.storage.attributes = attributesStorage; this.structs.attributes = attributeStruct; - this.fns.raycastFirstHit = buildRaycastFirstHitFn( name, bvhNodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ); + this.fns.raycastFirstHit = buildRaycastFirstHitFn( name, bvhNodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, structs.transform ); function appendBVHData( bvh, geometryOffset, transformInfo, nodeWriteOffset, target, tlas = false ) { diff --git a/src/webgpu/lib/PathtracerBVHComputeFns.js b/src/webgpu/lib/PathtracerBVHComputeFns.js new file mode 100644 index 000000000..6abf5843c --- /dev/null +++ b/src/webgpu/lib/PathtracerBVHComputeFns.js @@ -0,0 +1,70 @@ +import { StorageBufferAttribute } from 'three/webgpu'; +import { BVHComputeFns } from './BVHComputeFns.js'; +import { wgslStruct } from './WGSLStructNode.js'; +import { storage } from 'three/tsl'; + +const transformStruct = wgslStruct( 'TransformStruct', 36 * 4, { + matrixWorld: 'mat4x4f', + inverseMatrixWorld: 'mat4x4f', + nodeOffset: 'u32', + materialIndex: 'u32', + _alignment0: 'u32', + _alignment1: 'u32', +} ); + +const materialStruct = wgslStruct( 'MaterialStruct', 4 * 4, { + albedo: 'vec3f', +} ); + +export class PathtracerBVHComputeFns extends BVHComputeFns { + + constructor( ...args ) { + + super( ...args ); + + this.structs.transform = transformStruct; + this.structs.material = materialStruct; + this.materials = null; + + this.update(); + + } + + update() { + + this.materials = []; + super.update(); + + const { materials, structs, name } = this; + const materialBuffer = new ArrayBuffer( structs.material.byteSize * materials.length ); + const materialBufferF32 = new Float32Array( materialBuffer ); + materials.forEach( ( mat, i ) => { + + mat.color.toArray( materialBufferF32, i * structs.material.uintSize ); + + } ); + + const materialStorage = storage( new StorageBufferAttribute( new Uint32Array( materialBuffer ), 8 ), structs.material.name ).toReadOnly().setName( `${ name }materials` ); + this.storage.materials = materialStorage; + + } + + writeTransformData( info, premultiplyMatrix, writeOffset, targetBuffer ) { + + super.writeTransformData( info, premultiplyMatrix, writeOffset, targetBuffer ); + + const { materials } = this; + const material = info.object.material; + if ( ! materials.includes( material ) ) { + + materials.push( material ); + + } + + const index = materials.indexOf( material ); + const transformBufferU32 = new Uint32Array( targetBuffer ); + transformBufferU32[ writeOffset * transformStruct.uintSize + 33 ] = index; + + } + +} diff --git a/src/webgpu/lib/WGSLStructNode.js b/src/webgpu/lib/WGSLStructNode.js new file mode 100644 index 000000000..9d6cf6340 --- /dev/null +++ b/src/webgpu/lib/WGSLStructNode.js @@ -0,0 +1,36 @@ +import { CodeNode } from 'three/webgpu'; + +// a more structured "struct" node that bookkeeps the struct name, byte size +export class WGSLStructNode extends CodeNode { + + get uintSize() { + + return this.byteSize / 4; + + } + + constructor( name, byteSize, fields, includes = [] ) { + + const content = Object + .entries( fields ) + .map( ( [ name, type ] ) => { + + return `${ name }: ${ type },`; + + } ).join( '\n' ); + + const code = /* wgsl */` + struct ${ name } { + ${ content } + } + `; + + super( code, includes, 'wgsl' ); + this.name = name; + this.byteSize = byteSize; + + } + +} + +export const wgslStruct = ( ...args ) => new WGSLStructNode( ...args ); From c51a6c9cc5377a39abbd9568a2afee3be2f88e3b Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Fri, 20 Feb 2026 13:46:41 +0900 Subject: [PATCH 40/84] Cleanup --- src/webgpu/lib/BVHComputeFns.js | 25 +++++++++++-------------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 75223152d..6e872d6e9 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -17,9 +17,6 @@ const IS_LEAFNODE_FLAG = 0xFFFF; // TODO: add material support w/ function to easily update material // - add a callback for writing a property for a geometry to a range // TODO: add skinned mesh bvh support -// TODO: add overrideable functions for custom implementations (custom attributes, transform fields) -// TODO: see if it's possible to replace function contents and dependencies in-place so that -// a node fn can be updated without regenerating all other materials. // TODO: see if we can reference wgslFn names directly rather than constructing them inline over and over // and / or use local variable definitions for the pointers to clean up the code // TODO: see if there's a "build" step that can be leveraged for nodes to make integration more simple @@ -31,17 +28,6 @@ const _vec = /* @__PURE__ */ new Vector4(); const _matrix = /* @__PURE__ */ new Matrix4(); const _inverseMatrix = /* @__PURE__ */ new Matrix4(); -// stride is 36 floats (144 bytes) to match WGSL struct alignment: -// mat4x4f (64) + mat4x4f (64) + u32 (4) + 12 bytes padding to align to 16 -const transformStruct = wgslStruct( 'TransformStruct', 36 * 4, { - matrixWorld: 'mat4x4f', - inverseMatrixWorld: 'mat4x4f', - nodeOffset: 'u32', - _alignment0: 'u32', - _alignment1: 'u32', - _alignment2: 'u32', -} ); - function dereferenceIndex( indexAttr, indirectBuffer ) { const indexArray = indexAttr ? indexAttr.array : null; @@ -68,6 +54,17 @@ function getTotalBVHByteLength( bvh ) { } +// stride is 36 floats (144 bytes) to match WGSL struct alignment: +// mat4x4f (64) + mat4x4f (64) + u32 (4) + 12 bytes padding to align to 16 +const transformStruct = wgslStruct( 'TransformStruct', 36 * 4, { + matrixWorld: 'mat4x4f', + inverseMatrixWorld: 'mat4x4f', + nodeOffset: 'u32', + _alignment0: 'u32', + _alignment1: 'u32', + _alignment2: 'u32', +} ); + const intersectionResultStruct = wgslStruct( 'IntersectionResult', 16 * 4, { indices: 'vec4u', normal: 'vec3f', From fa67477dc57dd867cafd145687fff68d38277af1 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Fri, 20 Feb 2026 19:03:40 +0900 Subject: [PATCH 41/84] Add a node proxy mechanism --- src/webgpu/MegaKernelPathTracer.js | 10 ++--- src/webgpu/WebGPUPathTracer.js | 6 +-- src/webgpu/compute/PathTracerMegaKernel.js | 10 +++-- src/webgpu/lib/NodeProxy.js | 50 ++++++++++++++++++++++ 4 files changed, 63 insertions(+), 13 deletions(-) create mode 100644 src/webgpu/lib/NodeProxy.js diff --git a/src/webgpu/MegaKernelPathTracer.js b/src/webgpu/MegaKernelPathTracer.js index 9aedce818..f84f71e40 100644 --- a/src/webgpu/MegaKernelPathTracer.js +++ b/src/webgpu/MegaKernelPathTracer.js @@ -69,9 +69,6 @@ export class MegaKernelPathTracer { this.bounces = 7; this.tiles = new Vector2( 2, 2 ); - // bvh data - this.bvhComputeFns = null; - // targets this.outputTarget = new StorageTexture( 1, 1, ); this.outputTarget.format = RGBAFormat; @@ -96,16 +93,15 @@ export class MegaKernelPathTracer { this.sampleCountTarget.generateMipmaps = false; // kernels - this.kernel = null; + this.kernel = new PathTracerMegaKernel().setWorkgroupSize( 8, 8, 1 ); this.sampleCountClearKernel = new ZeroOutKernel( { textureType: 'r32uint' } ).setWorkgroupSize( 8, 8, 1 ); this.outputTargetClearKernel = new ZeroOutKernel( { textureType: 'rgba32float' } ).setWorkgroupSize( 8, 8, 1 ); } - setBVHComputeFns( bvhComputeFns ) { + setBVHData( bvhData ) { - this.bvhComputeFns = bvhComputeFns; - this.kernel = new PathTracerMegaKernel( bvhComputeFns ).setWorkgroupSize( 8, 8, 1 ); + this.kernel.bvhData = bvhData; this._task = null; } diff --git a/src/webgpu/WebGPUPathTracer.js b/src/webgpu/WebGPUPathTracer.js index dc20f96f7..ff6abcd10 100644 --- a/src/webgpu/WebGPUPathTracer.js +++ b/src/webgpu/WebGPUPathTracer.js @@ -26,7 +26,7 @@ export class WebGPUPathTracer { this._pathTracer.dispose(); this._pathTracer = value ? new MegaKernelPathTracer( this._renderer ) : new WaveFrontPathTracer( this._renderer ); - this._pathTracer.setBVHComputeFns( this._bvhComputeFns ); + this._pathTracer.setBVHData( this._bvhData ); this.setCamera( this.camera ); } @@ -69,8 +69,8 @@ export class WebGPUPathTracer { const objectBVH = new ObjectBVH( scene ); const bvhComputeFns = new PathtracerBVHComputeFns( objectBVH ); - this._bvhComputeFns = bvhComputeFns; - this._pathTracer.setBVHComputeFns( bvhComputeFns ); + this._bvhData = bvhComputeFns; + this._pathTracer.setBVHData( bvhComputeFns ); this.setCamera( camera ); } diff --git a/src/webgpu/compute/PathTracerMegaKernel.js b/src/webgpu/compute/PathTracerMegaKernel.js index d8c409846..afe6e5aaa 100644 --- a/src/webgpu/compute/PathTracerMegaKernel.js +++ b/src/webgpu/compute/PathTracerMegaKernel.js @@ -4,12 +4,15 @@ import { ComputeKernel } from './ComputeKernel.js'; import { uniform, globalId, textureStore, wgslFn } from 'three/tsl'; import { pcgRand3, pcgInit } from '../nodes/random.wgsl.js'; import { lambertBsdfFunc } from '../nodes/sampling.wgsl.js'; +import { proxy } from '../lib/NodeProxy.js'; export class PathTracerMegaKernel extends ComputeKernel { - constructor( bvhComputeFns ) { + constructor( name = 'bvh_' ) { const megakernelShaderParams = { + bvhData: { value: null }, + prevOutputTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadOnly(), outputTarget: textureStore( new StorageTexture( 1, 1 ) ).toWriteOnly(), sampleCountTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadWrite(), @@ -27,7 +30,6 @@ export class PathTracerMegaKernel extends ComputeKernel { globalId: globalId, }; - const { fns, storage, structs, name } = bvhComputeFns; const shader = wgslFn( /* wgsl */` fn compute( @@ -115,7 +117,9 @@ export class PathTracerMegaKernel extends ComputeKernel { } `, [ - fns.raycastFirstHit, storage.materials, structs.material, + proxy( 'bvhData.value.storage.materials', megakernelShaderParams ), + proxy( 'bvhData.value.structs.material', megakernelShaderParams ), + proxy( 'bvhData.value.fns.raycastFirstHit', megakernelShaderParams ), ndcToCameraRay, pcgRand3, pcgInit, lambertBsdfFunc, ] ); diff --git a/src/webgpu/lib/NodeProxy.js b/src/webgpu/lib/NodeProxy.js new file mode 100644 index 000000000..3a95c4da5 --- /dev/null +++ b/src/webgpu/lib/NodeProxy.js @@ -0,0 +1,50 @@ +import { Node } from 'three/webgpu'; + +export class NodeProxy extends Node { + + get node() { + + const { properties, object } = this; + let value = object; + for ( let i = 0, l = properties.length; i < l; i ++ ) { + + value = value[ properties[ i ] ]; + + } + + if ( 'functionNode' in value ) { + + return value.functionNode; + + } else { + + return value; + + } + + } + + constructor( property, object = null ) { + + super(); + this.object = object; + this.property = property; + this.properties = property.split( '.' ); + + } + + getNodeType( builder ) { + + return this.node.getNodeType( builder ); + + } + + setup( builder ) { + + return this.node; + + } + +} + +export const proxy = ( ...args ) => new NodeProxy( ...args ); From b7717825d901b8c33a0a3ec918608f7b8de63fed Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Fri, 20 Feb 2026 19:06:49 +0900 Subject: [PATCH 42/84] Get wavefront pathtracer working --- src/webgpu/MegaKernelPathTracer.js | 2 +- src/webgpu/WaveFrontPathTracer.js | 20 ++------- src/webgpu/compute/PathTracerMegaKernel.js | 16 ++++---- .../compute/wavefront/ProcessHitsKernel.js | 41 ++++++++++--------- .../wavefront/RayIntersectionKernel.js | 29 +++++-------- src/webgpu/compute/wavefront/structs.js | 2 +- src/webgpu/lib/BVHComputeFns.js | 19 +++++++++ src/webgpu/lib/WGSLStructNode.js | 1 + 8 files changed, 68 insertions(+), 62 deletions(-) diff --git a/src/webgpu/MegaKernelPathTracer.js b/src/webgpu/MegaKernelPathTracer.js index f84f71e40..58053bbb2 100644 --- a/src/webgpu/MegaKernelPathTracer.js +++ b/src/webgpu/MegaKernelPathTracer.js @@ -102,7 +102,7 @@ export class MegaKernelPathTracer { setBVHData( bvhData ) { this.kernel.bvhData = bvhData; - this._task = null; + this.reset(); } diff --git a/src/webgpu/WaveFrontPathTracer.js b/src/webgpu/WaveFrontPathTracer.js index b87ef5a34..11960ff10 100644 --- a/src/webgpu/WaveFrontPathTracer.js +++ b/src/webgpu/WaveFrontPathTracer.js @@ -231,23 +231,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(); } diff --git a/src/webgpu/compute/PathTracerMegaKernel.js b/src/webgpu/compute/PathTracerMegaKernel.js index afe6e5aaa..51eaad9d9 100644 --- a/src/webgpu/compute/PathTracerMegaKernel.js +++ b/src/webgpu/compute/PathTracerMegaKernel.js @@ -10,7 +10,7 @@ export class PathTracerMegaKernel extends ComputeKernel { constructor( name = 'bvh_' ) { - const megakernelShaderParams = { + const parameters = { bvhData: { value: null }, prevOutputTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadOnly(), @@ -85,8 +85,9 @@ export class PathTracerMegaKernel extends ComputeKernel { let hitResult = ${ name }RaycastFirstHit( ray ); if ( hitResult.didHit ) { + let vertexData = ${ name }sampleTrianglePoint( hitResult.barycoord, hitResult.indices.xyz ); let hitPosition = ray.origin + ray.direction * hitResult.dist; - let scatterRec = bsdfEval( hitResult.normal, - ray.direction ); + let scatterRec = bsdfEval( normalize( vertexData.normal.xyz ), - ray.direction ); let transform = ${ name }transforms.value[ hitResult.objectIndex ]; let material = ${ name }materials.value[ transform.materialIndex ]; @@ -117,15 +118,16 @@ export class PathTracerMegaKernel extends ComputeKernel { } `, [ - proxy( 'bvhData.value.storage.materials', megakernelShaderParams ), - proxy( 'bvhData.value.structs.material', megakernelShaderParams ), - proxy( 'bvhData.value.fns.raycastFirstHit', megakernelShaderParams ), + proxy( 'bvhData.value.storage.materials', parameters ), + proxy( 'bvhData.value.structs.material', parameters ), + proxy( 'bvhData.value.fns.raycastFirstHit', parameters ), + proxy( 'bvhData.value.fns.sampleTrianglePoint', parameters ), ndcToCameraRay, pcgRand3, pcgInit, lambertBsdfFunc, ] ); - super( shader( megakernelShaderParams ) ); + super( shader( parameters ) ); - this.defineUniformAccessors( megakernelShaderParams ); + this.defineUniformAccessors( parameters ); } diff --git a/src/webgpu/compute/wavefront/ProcessHitsKernel.js b/src/webgpu/compute/wavefront/ProcessHitsKernel.js index ebe21b62c..54b990304 100644 --- a/src/webgpu/compute/wavefront/ProcessHitsKernel.js +++ b/src/webgpu/compute/wavefront/ProcessHitsKernel.js @@ -1,17 +1,19 @@ import { IndirectStorageBufferAttribute, StorageTexture } from 'three/webgpu'; import { ComputeKernel } from '../ComputeKernel.js'; import { uniform, storage, wgslFn, textureStore, globalId } from 'three/tsl'; -import { constants, getVertexAttribute } from 'three-mesh-bvh/webgpu'; +import { constants } from 'three-mesh-bvh/webgpu'; import { pcgRand3, pcgInit } from '../../nodes/random.wgsl.js'; -import { materialStruct } from '../../nodes/structs.wgsl.js'; import { lambertBsdfFunc } from '../../nodes/sampling.wgsl.js'; import { queuedRayStruct, queuedHitStruct, QUEUED_RAY_SIZE, QUEUED_HIT_SIZE } from './structs.js'; +import { proxy } from '../../lib/NodeProxy.js'; export class ProcessHitsKernel extends ComputeKernel { - constructor() { + constructor( name = 'bvh_' ) { const parameters = { + bvhData: { value: null }, + prevOutputTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadOnly(), outputTarget: textureStore( new StorageTexture( 1, 1 ) ).toWriteOnly(), sampleCountTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadWrite(), @@ -27,11 +29,6 @@ export class ProcessHitsKernel extends ComputeKernel { hitQueue: storage( new IndirectStorageBufferAttribute( 1, QUEUED_HIT_SIZE ), 'QueuedHit' ), hitQueueSize: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ), - // bvh and geometry definition - geom_position: storage( new IndirectStorageBufferAttribute( 1, 3 ), 'vec3f' ).toReadOnly(), - geom_normals: storage( new IndirectStorageBufferAttribute( 1, 3 ), 'vec3f' ).toReadOnly(), - materials: storage( new IndirectStorageBufferAttribute(), 'Material' ).toReadOnly(), // TODO: fill in initial values - globalId: globalId, }; @@ -55,11 +52,6 @@ export class ProcessHitsKernel extends ComputeKernel { hitQueue: ptr, read_write>, hitQueueSize: ptr, read_write>, - // scene - geom_position: ptr, read>, - geom_normals: ptr, read>, - materials: ptr, read>, - globalId: vec3u ) -> void { @@ -80,10 +72,13 @@ export class ProcessHitsKernel extends ComputeKernel { pcgInitialize( indexUV, seed ); - let material = materials[ input.materialIndex ]; - let hitPosition = getVertexAttribute( input.barycoord, input.indices.xyz, geom_position ); - let hitNormal = getVertexAttribute( input.barycoord, input.indices.xyz, geom_normals ); - let scatterRec = bsdfEval( hitNormal, input.view ); + let object = ${ name }transforms.value[ input.objectIndex ]; + let material = ${ name }materials.value[ object.materialIndex ]; + var vertexData = ${ name }sampleTrianglePoint( input.barycoord, input.indices.xyz ); + vertexData.normal = normalize( transpose( object.inverseMatrixWorld ) * vertexData.normal ); + vertexData.position = object.matrixWorld * vertexData.position; + + let scatterRec = bsdfEval( vertexData.normal.xyz, input.view ); if ( input.currentBounce >= bounces ) { @@ -99,7 +94,7 @@ export class ProcessHitsKernel extends ComputeKernel { let rayQueueCapacity = arrayLength( rayQueue ); let index = atomicAdd( &rayQueueSize[ 1 ], 1 ) % rayQueueCapacity; - rayQueue[ index ].ray.origin = hitPosition; + rayQueue[ index ].ray.origin = vertexData.position.xyz; rayQueue[ index ].ray.direction = scatterRec.direction; rayQueue[ index ].pixel = indexUV; rayQueue[ index ].throughputColor = input.throughputColor * material.albedo * scatterRec.value / scatterRec.pdf; @@ -108,7 +103,15 @@ export class ProcessHitsKernel extends ComputeKernel { } } - `, [ queuedRayStruct, lambertBsdfFunc, constants, getVertexAttribute, pcgRand3, pcgInit, queuedHitStruct, materialStruct ] ); + `, [ + proxy( 'bvhData.value.structs.material', parameters ), + proxy( 'bvhData.value.structs.transform', parameters ), + proxy( 'bvhData.value.storage.materials', parameters ), + proxy( 'bvhData.value.storage.transforms', parameters ), + proxy( 'bvhData.value.fns.sampleTrianglePoint', parameters ), + queuedRayStruct, lambertBsdfFunc, constants, + pcgRand3, pcgInit, queuedHitStruct, + ] ); super( fn( parameters ) ); diff --git a/src/webgpu/compute/wavefront/RayIntersectionKernel.js b/src/webgpu/compute/wavefront/RayIntersectionKernel.js index 086caa575..4752bcb60 100644 --- a/src/webgpu/compute/wavefront/RayIntersectionKernel.js +++ b/src/webgpu/compute/wavefront/RayIntersectionKernel.js @@ -1,15 +1,18 @@ import { IndirectStorageBufferAttribute, StorageTexture } from 'three/webgpu'; import { ComputeKernel } from '../ComputeKernel.js'; import { storage, wgslFn, textureStore, globalId } from 'three/tsl'; -import { bvhIntersectFirstHit, constants } from 'three-mesh-bvh/webgpu'; +import { constants } from 'three-mesh-bvh/webgpu'; import { pcgRand3, pcgInit } from '../../nodes/random.wgsl.js'; import { queuedRayStruct, queuedHitStruct, QUEUED_RAY_SIZE, QUEUED_HIT_SIZE } from './structs.js'; +import { proxy } from '../../lib/NodeProxy.js'; export class RayIntersectionKernel extends ComputeKernel { - constructor() { + constructor( name = 'bvh_' ) { const parameters = { + bvhData: { value: null }, + prevOutputTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadOnly(), outputTarget: textureStore( new StorageTexture( 1, 1 ) ).toWriteOnly(), sampleCountTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadWrite(), @@ -21,12 +24,6 @@ export class RayIntersectionKernel extends ComputeKernel { hitQueue: storage( new IndirectStorageBufferAttribute( 1, QUEUED_HIT_SIZE ), 'QueuedHit' ), hitQueueSize: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ).toAtomic(), - // bvh and geometry definition - geom_index: storage( new IndirectStorageBufferAttribute( 1, 3 ), 'vec3u' ).toReadOnly(), - geom_position: storage( new IndirectStorageBufferAttribute( 1, 3 ), 'vec3f' ).toReadOnly(), - geom_material_index: storage( new IndirectStorageBufferAttribute( 1, 1 ), 'u32' ).toReadOnly(), - bvh: storage( new IndirectStorageBufferAttribute(), 'BVHNode' ).toReadOnly(), // TODO: fill in sizes - globalId: globalId, }; @@ -46,12 +43,6 @@ export class RayIntersectionKernel extends ComputeKernel { hitQueue: ptr, read_write>, hitQueueSize: ptr>, read_write>, - // scene - geom_position: ptr, read>, - geom_index: ptr, read>, - geom_material_index: ptr, read>, - bvh: ptr, read>, - globalId: vec3u ) -> void { @@ -73,18 +64,17 @@ export class RayIntersectionKernel extends ComputeKernel { pcgInitialize( indexUV, seed ); // run intersection - let hitResult = bvhIntersectFirstHit( geom_index, geom_position, bvh, input.ray ); + let hitResult = ${ name }RaycastFirstHit( input.ray ); if ( hitResult.didHit ) { // TODO: we process all of these materials immediately to push to the ray queue - let materialIndex = geom_material_index[ hitResult.indices.x ]; let index = atomicAdd( &hitQueueSize[ 1 ], 1 ); hitQueue[ index ].view = - input.ray.direction; hitQueue[ index ].indices = hitResult.indices.xyz; hitQueue[ index ].barycoord = hitResult.barycoord; hitQueue[ index ].pixel_x = input.pixel.x; hitQueue[ index ].pixel_y = input.pixel.y; - hitQueue[ index ].materialIndex = materialIndex; + hitQueue[ index ].objectIndex = hitResult.objectIndex; hitQueue[ index ].throughputColor = input.throughputColor; hitQueue[ index ].currentBounce = input.currentBounce;; @@ -103,7 +93,10 @@ export class RayIntersectionKernel extends ComputeKernel { } } - `, [ queuedRayStruct, bvhIntersectFirstHit, constants, pcgRand3, pcgInit, queuedHitStruct ] ); + `, [ + proxy( 'bvhData.value.fns.raycastFirstHit', parameters ), + proxy( 'bvhData.value.structs.material', parameters ), + queuedRayStruct, constants, pcgRand3, pcgInit, queuedHitStruct ] ); super( fn( parameters ) ); diff --git a/src/webgpu/compute/wavefront/structs.js b/src/webgpu/compute/wavefront/structs.js index f84308cb2..f71dfdd98 100644 --- a/src/webgpu/compute/wavefront/structs.js +++ b/src/webgpu/compute/wavefront/structs.js @@ -23,6 +23,6 @@ export const queuedHitStruct = wgsl( /* wgsl */` view: vec3f, currentBounce: u32, throughputColor: vec3f, - materialIndex: u32, + objectIndex: u32, }; ` ); diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 6e872d6e9..21072bbc4 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -462,6 +462,25 @@ export class BVHComputeFns { this.structs.attributes = attributeStruct; this.fns.raycastFirstHit = buildRaycastFirstHitFn( name, bvhNodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, structs.transform ); + const interpolateBody = Object.keys( attributeStruct.fields ) + .map( key => { + + return `result.${ key } = a0.${ key } * barycoord.x + a1.${ key } * barycoord.y + a2.${ key } * barycoord.z;`; + + } ).join( '\n' ); + this.fns.sampleTrianglePoint = wgslFn( /* wgsl */` + fn ${ name }sampleTrianglePoint( barycoord: vec3f, indices: vec3u ) -> ${ attributeStruct.name } { + + var result: ${ attributeStruct.name }; + var a0 = ${ name }attributes.value[ indices.x ]; + var a1 = ${ name }attributes.value[ indices.y ]; + var a2 = ${ name }attributes.value[ indices.z ]; + ${ interpolateBody } + return result; + + } + `, [ attributesStorage, attributeStruct ] ); + function appendBVHData( bvh, geometryOffset, transformInfo, nodeWriteOffset, target, tlas = false ) { const targetU16 = new Uint16Array( target ); diff --git a/src/webgpu/lib/WGSLStructNode.js b/src/webgpu/lib/WGSLStructNode.js index 9d6cf6340..8ead00f36 100644 --- a/src/webgpu/lib/WGSLStructNode.js +++ b/src/webgpu/lib/WGSLStructNode.js @@ -28,6 +28,7 @@ export class WGSLStructNode extends CodeNode { super( code, includes, 'wgsl' ); this.name = name; this.byteSize = byteSize; + this.fields = fields; } From d459bd1b2921f18712556895c22589502b6a041b Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Fri, 20 Feb 2026 19:37:33 +0900 Subject: [PATCH 43/84] Initial indirect bvh handling --- src/webgpu/WebGPUPathTracer.js | 7 +-- src/webgpu/lib/BVHComputeFns.js | 20 ++++++-- src/webgpu/lib/PathtracerBVHComputeFns.js | 59 +++++++++++++++++++++-- 3 files changed, 75 insertions(+), 11 deletions(-) diff --git a/src/webgpu/WebGPUPathTracer.js b/src/webgpu/WebGPUPathTracer.js index ff6abcd10..88538ad69 100644 --- a/src/webgpu/WebGPUPathTracer.js +++ b/src/webgpu/WebGPUPathTracer.js @@ -67,10 +67,11 @@ export class WebGPUPathTracer { // Build TLAS and compute functions const objectBVH = new ObjectBVH( scene ); - const bvhComputeFns = new PathtracerBVHComputeFns( objectBVH ); + const bvhData = new PathtracerBVHComputeFns( objectBVH ); + bvhData.update(); - this._bvhData = bvhComputeFns; - this._pathTracer.setBVHData( bvhComputeFns ); + this._bvhData = bvhData; + this._pathTracer.setBVHData( bvhData ); this.setCamera( camera ); } diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 21072bbc4..fc9634378 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -671,22 +671,36 @@ export class BVHComputeFns { getBVH( object, instanceId, rangeTarget ) { - if ( object.isBatchedMesh ) { + let bvh = null; + if ( object.boundsTree ) { + + // TODO + // this is a case where a mesh has morph targets and skinned meshes + + } else if ( object.isBatchedMesh ) { const geometryId = object.getGeometryIdAt( instanceId ); const range = object.getGeometryRangeAt( geometryId ); Object.assign( rangeTarget, range ); - return object.boundsTrees[ geometryId ]; + bvh = object.boundsTrees[ geometryId ]; } else { const geometry = object.geometry; rangeTarget.count = geometry.index ? geometry.index.count : geometry.attributes.position.count; rangeTarget.vertexCount = geometry.attributes.position.count; - return object.geometry.boundsTree; + bvh = object.geometry.boundsTree; } + if ( ! bvh ) { + + throw new Error( 'BVHComputeFns: BVH not found.' ); + + } + + return bvh; + } getDefaultAttributeValue( key, target ) { diff --git a/src/webgpu/lib/PathtracerBVHComputeFns.js b/src/webgpu/lib/PathtracerBVHComputeFns.js index 6abf5843c..7b384106f 100644 --- a/src/webgpu/lib/PathtracerBVHComputeFns.js +++ b/src/webgpu/lib/PathtracerBVHComputeFns.js @@ -1,7 +1,8 @@ -import { StorageBufferAttribute } from 'three/webgpu'; +import { BufferAttribute, BufferGeometry, StorageBufferAttribute } from 'three/webgpu'; import { BVHComputeFns } from './BVHComputeFns.js'; import { wgslStruct } from './WGSLStructNode.js'; import { storage } from 'three/tsl'; +import { MeshBVH, SAH } from 'three-mesh-bvh'; const transformStruct = wgslStruct( 'TransformStruct', 36 * 4, { matrixWorld: 'mat4x4f', @@ -24,15 +25,13 @@ export class PathtracerBVHComputeFns extends BVHComputeFns { this.structs.transform = transformStruct; this.structs.material = materialStruct; - this.materials = null; - - this.update(); + this.materials = []; + this.bvhMap = new Map(); } update() { - this.materials = []; super.update(); const { materials, structs, name } = this; @@ -47,6 +46,9 @@ export class PathtracerBVHComputeFns extends BVHComputeFns { const materialStorage = storage( new StorageBufferAttribute( new Uint32Array( materialBuffer ), 8 ), structs.material.name ).toReadOnly().setName( `${ name }materials` ); this.storage.materials = materialStorage; + this.bvhMap.clear(); + this.materials.length = 0; + } writeTransformData( info, premultiplyMatrix, writeOffset, targetBuffer ) { @@ -67,4 +69,51 @@ export class PathtracerBVHComputeFns extends BVHComputeFns { } + getBVH( object, instanceId, rangeTarget ) { + + const { bvhMap } = this; + const bvh = super.getBVH( object, instanceId, rangeTarget ); + if ( bvhMap.has( bvh ) ) { + + const data = bvhMap.get( bvh ); + Object.assign( rangeTarget, data.range ); + return data.bvh; + + } else if ( bvh.indirect ) { + + const proxyGeometry = new BufferGeometry(); + proxyGeometry.attributes = bvh.geometry.attributes; + + let array; + if ( bvh.geometry.index ) { + + array = bvh.geometry.index.array.slice( rangeTarget.start, rangeTarget.count + rangeTarget.start ); + + } else { + + const { start, count } = rangeTarget; + array = new Uint32Array( count ); + for ( let i = 0, l = rangeTarget.count; i < l; i ++ ) { + + array[ i ] = start + i; + + } + + } + + proxyGeometry.index = new BufferAttribute( array, 1 ); + rangeTarget.start = 0; + + const newBVH = new MeshBVH( proxyGeometry, { strategy: SAH, maxLeafSize: 5 } ); + bvhMap.set( bvh, { bvh: newBVH, range: { ...rangeTarget } } ); + return newBVH; + + } else { + + return bvh; + + } + + } + } From a1c9d14233463b83f929515fbcd852b2fea8f89e Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Fri, 20 Feb 2026 23:20:58 +0900 Subject: [PATCH 44/84] Fix SAH case --- src/webgpu/WebGPUPathTracer.js | 2 +- src/webgpu/lib/BVHComputeFns.js | 22 ++++++++++++---------- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/src/webgpu/WebGPUPathTracer.js b/src/webgpu/WebGPUPathTracer.js index 88538ad69..6d6c28bc5 100644 --- a/src/webgpu/WebGPUPathTracer.js +++ b/src/webgpu/WebGPUPathTracer.js @@ -66,7 +66,7 @@ export class WebGPUPathTracer { } ); // Build TLAS and compute functions - const objectBVH = new ObjectBVH( scene ); + const objectBVH = new ObjectBVH( scene, { strategy: SAH } ); const bvhData = new PathtracerBVHComputeFns( objectBVH ); bvhData.update(); diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index fc9634378..fc93127f2 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -384,10 +384,11 @@ export class BVHComputeFns { } // save the index of the bvh associated with this transform + const data = geometryInfo.find( info => object.geometry === info.geometry ); primBvh._roots.forEach( ( root, i ) => { transformInfo.push( { - data: geometryInfo.find( info => object.geometry === info.geometry ), + data, root: i, object, instanceId, @@ -509,24 +510,25 @@ export class BVHComputeFns { if ( tlas ) { + // 0xFFFF == mesh leaf, 0xFF00 == TLAS leaf + targetU32[ n32 + 6 ] = tlasOffset; + targetU16[ n16 + 15 ] = 0xFF00; + const count = rootBuffer16[ r16 + 14 ]; - const offset = rootBuffer32[ r32 + 6 ]; + // const offset = rootBuffer32[ r32 + 6 ]; // each root is expanded into a separate transform so we need to expand // the embedded offsets and counts. let rootsCount = 0; - for ( let o = offset, l = offset + count; o < l; o ++ ) { + for ( let o = 0; o < count; o ++ ) { - rootsCount += transformInfo[ o ].data.bvh._roots.length; + const roots = transformInfo[ tlasOffset ].data.bvh._roots.length; + tlasOffset += roots; + rootsCount += roots; } - // 0xFFFF == mesh leaf, 0xFF00 == TLAS leaf - targetU32[ n32 + 6 ] = tlasOffset; // rootBuffer32[ r32 + 6 ]; - targetU16[ n16 + 14 ] = rootsCount; //rootBuffer16[ r16 + 14 ]; - targetU16[ n16 + 15 ] = 0xFF00; - - tlasOffset += rootsCount; + targetU16[ n16 + 14 ] = rootsCount; } else { From 34acdf22b2aae7e16eea28c3e3eccd79987fcc60 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Fri, 20 Feb 2026 23:35:59 +0900 Subject: [PATCH 45/84] Cleanup --- src/webgpu/lib/BVHComputeFns.js | 31 ++++++++++++++++++------------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index fc93127f2..1bec554e8 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -346,7 +346,7 @@ export class BVHComputeFns { const { attributes, structs, name, bvh } = this; // collect the BVHs - const geometryInfo = []; + const bvhInfo = []; const transformInfo = []; // accumulate the sizes of the bvh nodes buffer, number of objects, and geometry buffers @@ -361,13 +361,12 @@ export class BVHComputeFns { const primBvh = this.getBVH( object, instanceId, range ); // if we haven't added this bvh, yet - if ( ! geometryInfo.find( info => info.bvh === primBvh ) ) { + if ( ! bvhInfo.find( info => info.bvh === primBvh ) ) { // save the geometry info to write later and increment the buffer sizes const info = { - index: geometryInfo.length, + index: bvhInfo.length, bvh: primBvh, - geometry: primBvh.geometry, range: range, bvhBufferOffsets: null, @@ -379,12 +378,12 @@ export class BVHComputeFns { bvhNodesBufferLength += getTotalBVHByteLength( primBvh ); indexBufferLength += info.range.count; attributesBufferLength += info.range.vertexCount; - geometryInfo.push( info ); + bvhInfo.push( info ); } // save the index of the bvh associated with this transform - const data = geometryInfo.find( info => object.geometry === info.geometry ); + const data = bvhInfo.find( info => primBvh === info.bvh ); primBvh._roots.forEach( ( root, i ) => { transformInfo.push( { @@ -419,15 +418,15 @@ export class BVHComputeFns { // append TLAS data appendBVHData( bvh, 0, transformInfo, 0, bvhNodesBuffer, true ); nodeWriteOffset += getTotalBVHByteLength( bvh ) / BYTES_PER_NODE; - geometryInfo.forEach( info => { + bvhInfo.forEach( info => { // append bvh data const bvhNodeOffsets = appendBVHData( info.bvh, indexOffset / 3, transformInfo, nodeWriteOffset, bvhNodesBuffer, false ); info.bvhNodeOffsets = bvhNodeOffsets; // append geometry data - appendIndexData( info.geometry, info.bvh._indirectBuffer, info.range, attributesOffset, indexOffset, indexBuffer ); - appendGeometryData( info.geometry, info.range, attributesOffset, attributesBuffer ); + appendIndexData( info.bvh, info.range, attributesOffset, indexOffset, indexBuffer ); + appendGeometryData( info.bvh, info.range, attributesOffset, attributesBuffer ); info.indexBufferOffset = indexOffset; // step the write offsets forward @@ -555,12 +554,15 @@ export class BVHComputeFns { } - function appendIndexData( geometry, indirectBuffer, range, valueOffset, writeOffset, target ) { + function appendIndexData( bvh, range, valueOffset, writeOffset, target ) { + // TODO: check if this is a skinned mesh bvh and use the mesh geometry + + const { geometry } = bvh; const { start, count, vertexStart } = range; - if ( indirectBuffer ) { + if ( bvh.indirect ) { - const dereferencedIndex = dereferenceIndex( geometry.index, indirectBuffer ); + const dereferencedIndex = dereferenceIndex( geometry.index, bvh._indirectBuffer ); for ( let i = 0; i < dereferencedIndex.length; i ++ ) { target[ i + writeOffset ] = dereferencedIndex[ i ] - vertexStart + valueOffset; @@ -593,8 +595,11 @@ export class BVHComputeFns { } - function appendGeometryData( geometry, range, writeOffset, target ) { + function appendGeometryData( bvh, range, writeOffset, target ) { + + // TODO: check if this is a skinned mesh bvh and use the mesh geometry + const { geometry } = bvh; const { vertexStart, vertexCount } = range; const attributesBufferF32 = new Float32Array( target ); attributes.forEach( ( key, interleavedOffset ) => { From 09b5fe81710b0ef87f14de22d6a211b440d6c291 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sat, 21 Feb 2026 00:02:12 +0900 Subject: [PATCH 46/84] Updates --- src/webgpu/lib/BVHComputeFns.js | 18 ++++++++++++------ src/webgpu/lib/PathtracerBVHComputeFns.js | 1 + 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeFns.js index 1bec554e8..25ae035f9 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeFns.js @@ -556,8 +556,6 @@ export class BVHComputeFns { function appendIndexData( bvh, range, valueOffset, writeOffset, target ) { - // TODO: check if this is a skinned mesh bvh and use the mesh geometry - const { geometry } = bvh; const { start, count, vertexStart } = range; if ( bvh.indirect ) { @@ -597,9 +595,8 @@ export class BVHComputeFns { function appendGeometryData( bvh, range, writeOffset, target ) { - // TODO: check if this is a skinned mesh bvh and use the mesh geometry - - const { geometry } = bvh; + // if "mesh" is present then it is assumed to be a SkinnedMeshBVH + const { geometry, mesh = null } = bvh; const { vertexStart, vertexCount } = range; const attributesBufferF32 = new Float32Array( target ); attributes.forEach( ( key, interleavedOffset ) => { @@ -611,7 +608,16 @@ export class BVHComputeFns { if ( attr ) { - _vec.fromBufferAttribute( attr, i + vertexStart ); + if ( key === 'position' && mesh ) { + + // TODO: normals and tangents need to be transformed here, as well + mesh.getVertexPosition( i + vertexStart, _vec ); + + } else { + + _vec.fromBufferAttribute( attr, i + vertexStart ); + + } switch ( attr.itemSize ) { diff --git a/src/webgpu/lib/PathtracerBVHComputeFns.js b/src/webgpu/lib/PathtracerBVHComputeFns.js index 7b384106f..05fa79ab5 100644 --- a/src/webgpu/lib/PathtracerBVHComputeFns.js +++ b/src/webgpu/lib/PathtracerBVHComputeFns.js @@ -104,6 +104,7 @@ export class PathtracerBVHComputeFns extends BVHComputeFns { proxyGeometry.index = new BufferAttribute( array, 1 ); rangeTarget.start = 0; + // TODO: need to handle SkinnedMeshBVH here const newBVH = new MeshBVH( proxyGeometry, { strategy: SAH, maxLeafSize: 5 } ); bvhMap.set( bvh, { bvh: newBVH, range: { ...rangeTarget } } ); return newBVH; From 00013d60ef421d19da159c49e284eb5c2f93eb44 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sat, 21 Feb 2026 15:37:04 +0900 Subject: [PATCH 47/84] Move NodeProxy --- src/webgpu/compute/PathTracerMegaKernel.js | 2 +- src/webgpu/compute/wavefront/ProcessHitsKernel.js | 2 +- src/webgpu/compute/wavefront/RayIntersectionKernel.js | 2 +- src/webgpu/lib/{ => nodes}/NodeProxy.js | 0 4 files changed, 3 insertions(+), 3 deletions(-) rename src/webgpu/lib/{ => nodes}/NodeProxy.js (100%) diff --git a/src/webgpu/compute/PathTracerMegaKernel.js b/src/webgpu/compute/PathTracerMegaKernel.js index 51eaad9d9..4a61933fe 100644 --- a/src/webgpu/compute/PathTracerMegaKernel.js +++ b/src/webgpu/compute/PathTracerMegaKernel.js @@ -4,7 +4,7 @@ import { ComputeKernel } from './ComputeKernel.js'; import { uniform, globalId, textureStore, wgslFn } from 'three/tsl'; import { pcgRand3, pcgInit } from '../nodes/random.wgsl.js'; import { lambertBsdfFunc } from '../nodes/sampling.wgsl.js'; -import { proxy } from '../lib/NodeProxy.js'; +import { proxy } from '../lib/nodes/NodeProxy.js'; export class PathTracerMegaKernel extends ComputeKernel { diff --git a/src/webgpu/compute/wavefront/ProcessHitsKernel.js b/src/webgpu/compute/wavefront/ProcessHitsKernel.js index 54b990304..9481f9a39 100644 --- a/src/webgpu/compute/wavefront/ProcessHitsKernel.js +++ b/src/webgpu/compute/wavefront/ProcessHitsKernel.js @@ -5,7 +5,7 @@ import { constants } from 'three-mesh-bvh/webgpu'; import { pcgRand3, pcgInit } from '../../nodes/random.wgsl.js'; import { lambertBsdfFunc } from '../../nodes/sampling.wgsl.js'; import { queuedRayStruct, queuedHitStruct, QUEUED_RAY_SIZE, QUEUED_HIT_SIZE } from './structs.js'; -import { proxy } from '../../lib/NodeProxy.js'; +import { proxy } from '../../lib/nodes/NodeProxy.js'; export class ProcessHitsKernel extends ComputeKernel { diff --git a/src/webgpu/compute/wavefront/RayIntersectionKernel.js b/src/webgpu/compute/wavefront/RayIntersectionKernel.js index 4752bcb60..8b7d468d5 100644 --- a/src/webgpu/compute/wavefront/RayIntersectionKernel.js +++ b/src/webgpu/compute/wavefront/RayIntersectionKernel.js @@ -4,7 +4,7 @@ import { storage, wgslFn, textureStore, globalId } from 'three/tsl'; import { constants } from 'three-mesh-bvh/webgpu'; import { pcgRand3, pcgInit } from '../../nodes/random.wgsl.js'; import { queuedRayStruct, queuedHitStruct, QUEUED_RAY_SIZE, QUEUED_HIT_SIZE } from './structs.js'; -import { proxy } from '../../lib/NodeProxy.js'; +import { proxy } from '../../lib/nodes/NodeProxy.js'; export class RayIntersectionKernel extends ComputeKernel { diff --git a/src/webgpu/lib/NodeProxy.js b/src/webgpu/lib/nodes/NodeProxy.js similarity index 100% rename from src/webgpu/lib/NodeProxy.js rename to src/webgpu/lib/nodes/NodeProxy.js From d20dcde288668e4179f26a41dd642b231297e76d Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sat, 21 Feb 2026 15:40:39 +0900 Subject: [PATCH 48/84] Rename computeFns --- src/webgpu/WebGPUPathTracer.js | 4 ++-- src/webgpu/lib/{BVHComputeFns.js => BVHComputeData.js} | 4 ++-- ...PathtracerBVHComputeFns.js => PathtracerBVHComputeData.js} | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) rename src/webgpu/lib/{BVHComputeFns.js => BVHComputeData.js} (99%) rename src/webgpu/lib/{PathtracerBVHComputeFns.js => PathtracerBVHComputeData.js} (96%) diff --git a/src/webgpu/WebGPUPathTracer.js b/src/webgpu/WebGPUPathTracer.js index 6d6c28bc5..a522d5a84 100644 --- a/src/webgpu/WebGPUPathTracer.js +++ b/src/webgpu/WebGPUPathTracer.js @@ -5,7 +5,7 @@ import { RenderToScreenNodeMaterial } from './materials/RenderToScreenMaterial.j import { MegaKernelPathTracer } from './MegaKernelPathTracer.js'; import { WaveFrontPathTracer } from './WaveFrontPathTracer.js'; import { ObjectBVH } from './lib/ObjectBVH.js'; -import { PathtracerBVHComputeFns } from './lib/PathtracerBVHComputeFns.js'; +import { PathtracerBVHComputeData } from './lib/PathtracerBVHComputeData.js'; const _resolution = new Vector2(); export class WebGPUPathTracer { @@ -67,7 +67,7 @@ export class WebGPUPathTracer { // Build TLAS and compute functions const objectBVH = new ObjectBVH( scene, { strategy: SAH } ); - const bvhData = new PathtracerBVHComputeFns( objectBVH ); + const bvhData = new PathtracerBVHComputeData( objectBVH ); bvhData.update(); this._bvhData = bvhData; diff --git a/src/webgpu/lib/BVHComputeFns.js b/src/webgpu/lib/BVHComputeData.js similarity index 99% rename from src/webgpu/lib/BVHComputeFns.js rename to src/webgpu/lib/BVHComputeData.js index 25ae035f9..381861aa3 100644 --- a/src/webgpu/lib/BVHComputeFns.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -309,7 +309,7 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto } -export class BVHComputeFns { +export class BVHComputeData { constructor( bvh, options = {} ) { @@ -708,7 +708,7 @@ export class BVHComputeFns { if ( ! bvh ) { - throw new Error( 'BVHComputeFns: BVH not found.' ); + throw new Error( 'BVHComputeData: BVH not found.' ); } diff --git a/src/webgpu/lib/PathtracerBVHComputeFns.js b/src/webgpu/lib/PathtracerBVHComputeData.js similarity index 96% rename from src/webgpu/lib/PathtracerBVHComputeFns.js rename to src/webgpu/lib/PathtracerBVHComputeData.js index 05fa79ab5..164ad7e72 100644 --- a/src/webgpu/lib/PathtracerBVHComputeFns.js +++ b/src/webgpu/lib/PathtracerBVHComputeData.js @@ -1,5 +1,5 @@ import { BufferAttribute, BufferGeometry, StorageBufferAttribute } from 'three/webgpu'; -import { BVHComputeFns } from './BVHComputeFns.js'; +import { BVHComputeData } from './BVHComputeData.js'; import { wgslStruct } from './WGSLStructNode.js'; import { storage } from 'three/tsl'; import { MeshBVH, SAH } from 'three-mesh-bvh'; @@ -17,7 +17,7 @@ const materialStruct = wgslStruct( 'MaterialStruct', 4 * 4, { albedo: 'vec3f', } ); -export class PathtracerBVHComputeFns extends BVHComputeFns { +export class PathtracerBVHComputeData extends BVHComputeData { constructor( ...args ) { From 9d133b979727046b9600a21f4f45d1a8bb9521a5 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sat, 21 Feb 2026 15:41:08 +0900 Subject: [PATCH 49/84] Move node to folder --- src/webgpu/lib/BVHComputeData.js | 2 +- src/webgpu/lib/PathtracerBVHComputeData.js | 2 +- src/webgpu/lib/{ => nodes}/WGSLStructNode.js | 0 3 files changed, 2 insertions(+), 2 deletions(-) rename src/webgpu/lib/{ => nodes}/WGSLStructNode.js (100%) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index 381861aa3..ec8abbc80 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -7,7 +7,7 @@ import { bvhNodeStruct, constants, } from 'three-mesh-bvh/webgpu'; -import { wgslStruct } from './WGSLStructNode.js'; +import { wgslStruct } from './nodes/WGSLStructNode.js'; const BYTES_PER_NODE = 6 * 4 + 4 + 4; const UINT32_PER_NODE = BYTES_PER_NODE / 4; diff --git a/src/webgpu/lib/PathtracerBVHComputeData.js b/src/webgpu/lib/PathtracerBVHComputeData.js index 164ad7e72..f5d16cd9d 100644 --- a/src/webgpu/lib/PathtracerBVHComputeData.js +++ b/src/webgpu/lib/PathtracerBVHComputeData.js @@ -1,6 +1,6 @@ import { BufferAttribute, BufferGeometry, StorageBufferAttribute } from 'three/webgpu'; import { BVHComputeData } from './BVHComputeData.js'; -import { wgslStruct } from './WGSLStructNode.js'; +import { wgslStruct } from './nodes/WGSLStructNode.js'; import { storage } from 'three/tsl'; import { MeshBVH, SAH } from 'three-mesh-bvh'; diff --git a/src/webgpu/lib/WGSLStructNode.js b/src/webgpu/lib/nodes/WGSLStructNode.js similarity index 100% rename from src/webgpu/lib/WGSLStructNode.js rename to src/webgpu/lib/nodes/WGSLStructNode.js From 7c97449b239dd9f61118adc58128d167fd320f8d Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sat, 21 Feb 2026 16:16:41 +0900 Subject: [PATCH 50/84] Proxy improvements --- src/webgpu/lib/nodes/NodeProxy.js | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/src/webgpu/lib/nodes/NodeProxy.js b/src/webgpu/lib/nodes/NodeProxy.js index 3a95c4da5..613fcc522 100644 --- a/src/webgpu/lib/nodes/NodeProxy.js +++ b/src/webgpu/lib/nodes/NodeProxy.js @@ -2,6 +2,12 @@ import { Node } from 'three/webgpu'; export class NodeProxy extends Node { + static get type() { + + return 'NodeProxy'; + + } + get node() { const { properties, object } = this; @@ -33,12 +39,21 @@ export class NodeProxy extends Node { } + // delegate type resolution to the target node getNodeType( builder ) { return this.node.getNodeType( builder ); } + // include the target node's cache key so the proxy invalidates when the target changes + customCacheKey() { + + return this.node.getCacheKey(); + + } + + // return the target node as the output so the builder uses it for analyze/generate setup( builder ) { return this.node; From 08a5fab23c870f8b5ac7d7766b1f79bde0636e8a Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sat, 21 Feb 2026 17:24:27 +0900 Subject: [PATCH 51/84] Add a wgslFnTagNode --- src/webgpu/lib/nodes/WGSLFnTagNode.js | 177 ++++++++++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 src/webgpu/lib/nodes/WGSLFnTagNode.js diff --git a/src/webgpu/lib/nodes/WGSLFnTagNode.js b/src/webgpu/lib/nodes/WGSLFnTagNode.js new file mode 100644 index 000000000..44ad609b1 --- /dev/null +++ b/src/webgpu/lib/nodes/WGSLFnTagNode.js @@ -0,0 +1,177 @@ +import { FunctionNode, Node } from 'three/webgpu'; + +// minimal node that outputs a raw WGSL expression verbatim when built, +// bypassing TSL's temp variable wrapping and type formatting +class RawExpression extends Node { + + constructor( code ) { + + super(); + this.code = code; + + } + + build() { + + return this.code; + + } + +} + +export class WGSLFnTagNode extends FunctionNode { + + static get type() { + + return 'WGSLFnTagNode'; + + } + + constructor( tokens, args ) { + + // extract FunctionNode dependencies for includes — only function definitions + // need to be pre-registered so their code appears before ours in the output. + // callable wrappers and FunctionCallNodes are unwrapped to the underlying FunctionNode; + // plain nodes (uniforms, storage, etc) are built inline in generate() and don't need includes. + const includes = []; + + for ( const arg of args ) { + + if ( typeof arg === 'function' && arg.functionNode ) { + + includes.push( arg.functionNode ); + + } else if ( arg && arg.isNode && arg.functionNode ) { + + includes.push( arg.functionNode ); + + } + + } + + super( '', includes, 'wgsl' ); + + this.tokens = tokens; + this.args = args; + + } + + // parse the function signature from the static template parts (tokens[0] always + // contains the full signature since interpolations only appear in the body) + getNodeFunction( builder ) { + + const { tokens } = this; + const nodeData = builder.getDataFromNode( this ); + let nodeFunction = nodeData.nodeFunction; + if ( nodeFunction === undefined ) { + + const braceIndex = tokens[ 0 ].indexOf( '{' ); + const sig = braceIndex !== - 1 + ? tokens[ 0 ].substring( 0, braceIndex ) + : tokens[ 0 ]; + + nodeFunction = builder.parser.parseFunction( sig + ' {}' ); + nodeData.nodeFunction = nodeFunction; + + } + + return nodeFunction; + + } + + generate( builder, output ) { + + const { tokens, args } = this; + + // let FunctionNode.generate handle includes, code registration, property naming, + // and type normalization (e.g. stripping "-> void" which is not valid WGSL) + const result = super.generate( builder, output ); + + // assemble the body by interleaving static tokens with resolved node names + const parts = []; + + for ( let i = 0, l = tokens.length; i < l; i ++ ) { + + parts.push( tokens[ i ] ); + + if ( i < args.length ) { + + const arg = args[ i ]; + if ( typeof arg === 'function' && arg.functionNode ) { + + // callable wrapper (from wgslFn/wgslFnTag) — resolve to function name + parts.push( arg.functionNode.build( builder, 'property' ) ); + + } else if ( arg.isNode && arg.functionNode ) { + + // FunctionCallNode — use generate() to get the inline call expression + // (build() would wrap it in a temp variable that lives outside our WGSL scope) + parts.push( arg.generate( builder ) ); + + } else { + + parts.push( arg.build( builder ) ); + + } + + } + + } + + const { type } = this.getNodeFunction( builder ); + const nodeCode = builder.getCodeFromNode( this, type ); + + // use the declaration from super (handles type normalization and name assignment), + // replace its empty body with the assembled body from the template + const declaration = nodeCode.code; + const declPrefix = declaration.substring( 0, declaration.indexOf( '{' ) + 1 ); + + const assembledCode = parts.join( '' ); + const bodyStart = assembledCode.indexOf( '{' ) + 1; + const bodyEnd = assembledCode.lastIndexOf( '}' ); + const body = assembledCode.substring( bodyStart, bodyEnd ); + + nodeCode.code = declPrefix + body + '}\n'; + + return result; + + } + +} + +const wgslFnTag = ( tokens, ...args ) => { + + const functionNode = new WGSLFnTagNode( tokens, args ); + + const fn = ( ...params ) => { + + // wrap string parameter values as raw WGSL expressions + // so they output verbatim as identifiers (e.g. local variable names) + if ( params.length === 1 && params[ 0 ] && typeof params[ 0 ] === 'object' && ! params[ 0 ].isNode ) { + + const obj = params[ 0 ]; + for ( const key in obj ) { + + if ( typeof obj[ key ] === 'string' ) { + + obj[ key ] = new RawExpression( obj[ key ] ); + + } + + } + + } + + return functionNode.call( ...params ); + + }; + + fn.functionNode = functionNode; + + return fn; + +}; + +export { + wgslFnTag, +}; From e2ab8776eef8cd1865ed610f95853fcdb3428529 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sat, 21 Feb 2026 17:29:10 +0900 Subject: [PATCH 52/84] Remove deprecated "clock" --- src/webgpu/WebGPUPathTracer.js | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/webgpu/WebGPUPathTracer.js b/src/webgpu/WebGPUPathTracer.js index a522d5a84..c6ecc24c4 100644 --- a/src/webgpu/WebGPUPathTracer.js +++ b/src/webgpu/WebGPUPathTracer.js @@ -1,4 +1,4 @@ -import { Vector2, Clock, Scene, PerspectiveCamera } from 'three/webgpu'; +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'; @@ -36,7 +36,6 @@ export class WebGPUPathTracer { // members this._renderer = renderer; this._pathTracer = new MegaKernelPathTracer( renderer ); - this._clock = new Clock(); // options this.renderScale = 1; From 13c7f7ae3c5870d2d17e58f0864decae9311c926 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sat, 21 Feb 2026 17:54:56 +0900 Subject: [PATCH 53/84] Use built-in struct definition calculation --- src/webgpu/lib/BVHComputeData.js | 33 +++++++++++----------- src/webgpu/lib/PathtracerBVHComputeData.js | 18 ++++++------ src/webgpu/lib/nodes/WGSLStructNode.js | 29 +++---------------- 3 files changed, 29 insertions(+), 51 deletions(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index ec8abbc80..d738f048a 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -56,23 +56,23 @@ function getTotalBVHByteLength( bvh ) { // stride is 36 floats (144 bytes) to match WGSL struct alignment: // mat4x4f (64) + mat4x4f (64) + u32 (4) + 12 bytes padding to align to 16 -const transformStruct = wgslStruct( 'TransformStruct', 36 * 4, { +const transformStruct = wgslStruct( 'TransformStruct', { matrixWorld: 'mat4x4f', inverseMatrixWorld: 'mat4x4f', - nodeOffset: 'u32', - _alignment0: 'u32', - _alignment1: 'u32', - _alignment2: 'u32', + nodeOffset: 'uint', + _alignment0: 'uint', + _alignment1: 'uint', + _alignment2: 'uint', } ); -const intersectionResultStruct = wgslStruct( 'IntersectionResult', 16 * 4, { +const intersectionResultStruct = wgslStruct( 'IntersectionResult', { indices: 'vec4u', normal: 'vec3f', didHit: 'bool', barycoord: 'vec3f', - objectIndex: 'u32', - side: 'f32', - dist: 'f32', + objectIndex: 'uint', + side: 'float', + dist: 'float', } ); const intersectsTriangle = wgslFn( /* wgsl */ ` @@ -403,7 +403,6 @@ export class BVHComputeData { // construct the attribute struct const attributeStruct = wgslStruct( `${ name }GeometryStruct`, - 4 * 4 * attributes.length, attributes.reduce( ( o, key ) => ( { ...o, [ key ]: 'vec4f' } ), {} ), ); @@ -439,7 +438,7 @@ export class BVHComputeData { // // write the transforms - const transformArrayBuffer = new ArrayBuffer( structs.transform.byteSize * transformInfo.length ); + const transformArrayBuffer = new ArrayBuffer( structs.transform.getLength() * transformInfo.length * 4 ); transformInfo.forEach( ( info, i ) => { _inverseMatrix.copy( bvh.matrixWorld ).invert(); @@ -451,9 +450,9 @@ export class BVHComputeData { // set up the storage buffers const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), 'BVHNode' ).toReadOnly().setName( `${ name }nodes` ); - const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), structs.transform.uintSize ), structs.transform.name ).toReadOnly().setName( `${ name }transforms` ); + const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), structs.transform.getLength() ), structs.transform.name ).toReadOnly().setName( `${ name }transforms` ); const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ name }index` ); - const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributeStruct.uintSize ), attributeStruct.name ).toReadOnly().setName( `${ name }attributes` ); + const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributeStruct.getLength() ), attributeStruct.name ).toReadOnly().setName( `${ name }attributes` ); this.storage.transforms = transformsStorage; this.storage.nodes = bvhNodesStorage; @@ -642,7 +641,7 @@ export class BVHComputeData { } - _vec.toArray( attributesBufferF32, ( writeOffset + i ) * attributeStruct.uintSize + interleavedOffset * 4 ); + _vec.toArray( attributesBufferF32, ( writeOffset + i ) * attributeStruct.getLength() + interleavedOffset * 4 ); } @@ -673,12 +672,12 @@ export class BVHComputeData { } _matrix.premultiply( premultiplyMatrix ); - _matrix.toArray( transformBufferF32, writeOffset * transformStruct.uintSize ); + _matrix.toArray( transformBufferF32, writeOffset * transformStruct.getLength() ); _matrix.invert(); - _matrix.toArray( transformBufferF32, writeOffset * transformStruct.uintSize + 16 ); + _matrix.toArray( transformBufferF32, writeOffset * transformStruct.getLength() + 16 ); - transformBufferU32[ writeOffset * transformStruct.uintSize + 32 ] = bvhNodeOffsets[ root ]; + transformBufferU32[ writeOffset * transformStruct.getLength() + 32 ] = bvhNodeOffsets[ root ]; } diff --git a/src/webgpu/lib/PathtracerBVHComputeData.js b/src/webgpu/lib/PathtracerBVHComputeData.js index f5d16cd9d..919a67623 100644 --- a/src/webgpu/lib/PathtracerBVHComputeData.js +++ b/src/webgpu/lib/PathtracerBVHComputeData.js @@ -4,16 +4,16 @@ import { wgslStruct } from './nodes/WGSLStructNode.js'; import { storage } from 'three/tsl'; import { MeshBVH, SAH } from 'three-mesh-bvh'; -const transformStruct = wgslStruct( 'TransformStruct', 36 * 4, { +const transformStruct = wgslStruct( 'TransformStruct', { matrixWorld: 'mat4x4f', inverseMatrixWorld: 'mat4x4f', - nodeOffset: 'u32', - materialIndex: 'u32', - _alignment0: 'u32', - _alignment1: 'u32', + nodeOffset: 'uint', + materialIndex: 'uint', + _alignment0: 'uint', + _alignment1: 'uint', } ); -const materialStruct = wgslStruct( 'MaterialStruct', 4 * 4, { +const materialStruct = wgslStruct( 'MaterialStruct', { albedo: 'vec3f', } ); @@ -35,11 +35,11 @@ export class PathtracerBVHComputeData extends BVHComputeData { super.update(); const { materials, structs, name } = this; - const materialBuffer = new ArrayBuffer( structs.material.byteSize * materials.length ); + const materialBuffer = new ArrayBuffer( structs.material.getLength() * materials.length * 4 ); const materialBufferF32 = new Float32Array( materialBuffer ); materials.forEach( ( mat, i ) => { - mat.color.toArray( materialBufferF32, i * structs.material.uintSize ); + mat.color.toArray( materialBufferF32, i * structs.material.getLength() ); } ); @@ -65,7 +65,7 @@ export class PathtracerBVHComputeData extends BVHComputeData { const index = materials.indexOf( material ); const transformBufferU32 = new Uint32Array( targetBuffer ); - transformBufferU32[ writeOffset * transformStruct.uintSize + 33 ] = index; + transformBufferU32[ writeOffset * transformStruct.getLength() + 33 ] = index; } diff --git a/src/webgpu/lib/nodes/WGSLStructNode.js b/src/webgpu/lib/nodes/WGSLStructNode.js index 8ead00f36..ffd1d10e0 100644 --- a/src/webgpu/lib/nodes/WGSLStructNode.js +++ b/src/webgpu/lib/nodes/WGSLStructNode.js @@ -1,33 +1,12 @@ -import { CodeNode } from 'three/webgpu'; +import { StructTypeNode } from 'three/webgpu'; // a more structured "struct" node that bookkeeps the struct name, byte size -export class WGSLStructNode extends CodeNode { +export class WGSLStructNode extends StructTypeNode { - get uintSize() { + constructor( name, fields ) { - return this.byteSize / 4; - - } - - constructor( name, byteSize, fields, includes = [] ) { - - const content = Object - .entries( fields ) - .map( ( [ name, type ] ) => { - - return `${ name }: ${ type },`; - - } ).join( '\n' ); - - const code = /* wgsl */` - struct ${ name } { - ${ content } - } - `; - - super( code, includes, 'wgsl' ); + super( fields ); this.name = name; - this.byteSize = byteSize; this.fields = fields; } From 10b9ef990b2527060ed416a6b2c7733a2f47dfe4 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sat, 21 Feb 2026 18:13:06 +0900 Subject: [PATCH 54/84] Remove unnecessary wgslstruct --- src/webgpu/lib/BVHComputeData.js | 22 +++++++++++----------- src/webgpu/lib/PathtracerBVHComputeData.js | 11 +++++------ src/webgpu/lib/nodes/WGSLFnTagNode.js | 4 ++++ src/webgpu/lib/nodes/WGSLStructNode.js | 16 ---------------- 4 files changed, 20 insertions(+), 33 deletions(-) delete mode 100644 src/webgpu/lib/nodes/WGSLStructNode.js diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index d738f048a..4d7f54345 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -1,5 +1,5 @@ import { Matrix4, Vector4 } from 'three'; -import { StorageBufferAttribute } from 'three/webgpu'; +import { StorageBufferAttribute, StructTypeNode } from 'three/webgpu'; import { storage, wgslFn } from 'three/tsl'; import { intersectsBounds, @@ -7,7 +7,6 @@ import { bvhNodeStruct, constants, } from 'three-mesh-bvh/webgpu'; -import { wgslStruct } from './nodes/WGSLStructNode.js'; const BYTES_PER_NODE = 6 * 4 + 4 + 4; const UINT32_PER_NODE = BYTES_PER_NODE / 4; @@ -56,16 +55,16 @@ function getTotalBVHByteLength( bvh ) { // stride is 36 floats (144 bytes) to match WGSL struct alignment: // mat4x4f (64) + mat4x4f (64) + u32 (4) + 12 bytes padding to align to 16 -const transformStruct = wgslStruct( 'TransformStruct', { +const transformStruct = new StructTypeNode( { matrixWorld: 'mat4x4f', inverseMatrixWorld: 'mat4x4f', nodeOffset: 'uint', _alignment0: 'uint', _alignment1: 'uint', _alignment2: 'uint', -} ); +}, 'TransformStruct' ); -const intersectionResultStruct = wgslStruct( 'IntersectionResult', { +const intersectionResultStruct = new StructTypeNode( { indices: 'vec4u', normal: 'vec3f', didHit: 'bool', @@ -73,7 +72,7 @@ const intersectionResultStruct = wgslStruct( 'IntersectionResult', { objectIndex: 'uint', side: 'float', dist: 'float', -} ); +}, 'IntersectionResult' ); const intersectsTriangle = wgslFn( /* wgsl */ ` @@ -401,9 +400,9 @@ export class BVHComputeData { // // construct the attribute struct - const attributeStruct = wgslStruct( - `${ name }GeometryStruct`, + const attributeStruct = new StructTypeNode( attributes.reduce( ( o, key ) => ( { ...o, [ key ]: 'vec4f' } ), {} ), + `${ name }GeometryStruct`, ); // write the geometry buffer attributes & bvh data @@ -461,10 +460,11 @@ export class BVHComputeData { this.structs.attributes = attributeStruct; this.fns.raycastFirstHit = buildRaycastFirstHitFn( name, bvhNodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, structs.transform ); - const interpolateBody = Object.keys( attributeStruct.fields ) - .map( key => { + const interpolateBody = attributeStruct + .membersLayout + .map( ( { name } ) => { - return `result.${ key } = a0.${ key } * barycoord.x + a1.${ key } * barycoord.y + a2.${ key } * barycoord.z;`; + return `result.${ name } = a0.${ name } * barycoord.x + a1.${ name } * barycoord.y + a2.${ name } * barycoord.z;`; } ).join( '\n' ); this.fns.sampleTrianglePoint = wgslFn( /* wgsl */` diff --git a/src/webgpu/lib/PathtracerBVHComputeData.js b/src/webgpu/lib/PathtracerBVHComputeData.js index 919a67623..75126510b 100644 --- a/src/webgpu/lib/PathtracerBVHComputeData.js +++ b/src/webgpu/lib/PathtracerBVHComputeData.js @@ -1,21 +1,20 @@ -import { BufferAttribute, BufferGeometry, StorageBufferAttribute } from 'three/webgpu'; +import { BufferAttribute, BufferGeometry, StorageBufferAttribute, StructTypeNode } from 'three/webgpu'; import { BVHComputeData } from './BVHComputeData.js'; -import { wgslStruct } from './nodes/WGSLStructNode.js'; import { storage } from 'three/tsl'; import { MeshBVH, SAH } from 'three-mesh-bvh'; -const transformStruct = wgslStruct( 'TransformStruct', { +const transformStruct = new StructTypeNode( { matrixWorld: 'mat4x4f', inverseMatrixWorld: 'mat4x4f', nodeOffset: 'uint', materialIndex: 'uint', _alignment0: 'uint', _alignment1: 'uint', -} ); +}, 'TransformStruct' ); -const materialStruct = wgslStruct( 'MaterialStruct', { +const materialStruct = new StructTypeNode( { albedo: 'vec3f', -} ); +}, 'MaterialStruct' ); export class PathtracerBVHComputeData extends BVHComputeData { diff --git a/src/webgpu/lib/nodes/WGSLFnTagNode.js b/src/webgpu/lib/nodes/WGSLFnTagNode.js index 44ad609b1..81c07c220 100644 --- a/src/webgpu/lib/nodes/WGSLFnTagNode.js +++ b/src/webgpu/lib/nodes/WGSLFnTagNode.js @@ -1,5 +1,9 @@ import { FunctionNode, Node } from 'three/webgpu'; +// TODO: allow for glsl and wgsl version +// TODO: allow for structs +// TODO: allow for arbitrary includes (support array?) + // minimal node that outputs a raw WGSL expression verbatim when built, // bypassing TSL's temp variable wrapping and type formatting class RawExpression extends Node { diff --git a/src/webgpu/lib/nodes/WGSLStructNode.js b/src/webgpu/lib/nodes/WGSLStructNode.js deleted file mode 100644 index ffd1d10e0..000000000 --- a/src/webgpu/lib/nodes/WGSLStructNode.js +++ /dev/null @@ -1,16 +0,0 @@ -import { StructTypeNode } from 'three/webgpu'; - -// a more structured "struct" node that bookkeeps the struct name, byte size -export class WGSLStructNode extends StructTypeNode { - - constructor( name, fields ) { - - super( fields ); - this.name = name; - this.fields = fields; - - } - -} - -export const wgslStruct = ( ...args ) => new WGSLStructNode( ...args ); From 2b8396d1c8e0eacfc7aef8d212ca3a19506e1ec0 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sat, 21 Feb 2026 18:19:25 +0900 Subject: [PATCH 55/84] Add lang arg --- src/webgpu/lib/nodes/WGSLFnTagNode.js | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/webgpu/lib/nodes/WGSLFnTagNode.js b/src/webgpu/lib/nodes/WGSLFnTagNode.js index 81c07c220..bf310b2f6 100644 --- a/src/webgpu/lib/nodes/WGSLFnTagNode.js +++ b/src/webgpu/lib/nodes/WGSLFnTagNode.js @@ -1,6 +1,5 @@ import { FunctionNode, Node } from 'three/webgpu'; -// TODO: allow for glsl and wgsl version // TODO: allow for structs // TODO: allow for arbitrary includes (support array?) @@ -31,7 +30,7 @@ export class WGSLFnTagNode extends FunctionNode { } - constructor( tokens, args ) { + constructor( tokens, args, lang = 'wgsl' ) { // extract FunctionNode dependencies for includes — only function definitions // need to be pre-registered so their code appears before ours in the output. @@ -53,7 +52,7 @@ export class WGSLFnTagNode extends FunctionNode { } - super( '', includes, 'wgsl' ); + super( '', includes, lang ); this.tokens = tokens; this.args = args; From 18547f41b017cccb621987f1dce67faa30c2f926 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 00:49:22 +0900 Subject: [PATCH 56/84] Add support for includes, constants --- src/webgpu/lib/nodes/WGSLFnTagNode.js | 110 ++++++++++++++++++++++---- 1 file changed, 94 insertions(+), 16 deletions(-) diff --git a/src/webgpu/lib/nodes/WGSLFnTagNode.js b/src/webgpu/lib/nodes/WGSLFnTagNode.js index bf310b2f6..51058d173 100644 --- a/src/webgpu/lib/nodes/WGSLFnTagNode.js +++ b/src/webgpu/lib/nodes/WGSLFnTagNode.js @@ -1,8 +1,5 @@ import { FunctionNode, Node } from 'three/webgpu'; -// TODO: allow for structs -// TODO: allow for arbitrary includes (support array?) - // minimal node that outputs a raw WGSL expression verbatim when built, // bypassing TSL's temp variable wrapping and type formatting class RawExpression extends Node { @@ -22,6 +19,27 @@ class RawExpression extends Node { } +// returns the StructTypeNode from either a direct StructTypeNode or a struct() callable wrapper +function getStructLayout( arg ) { + + if ( arg && arg.isNode && arg.isStructLayoutNode ) return arg; + if ( typeof arg === 'function' && arg.isStruct ) return arg.layout; + return null; + +} + +// returns the node that should be registered as an include for the given arg, +// or null if the arg doesn't represent a dependency (e.g. a string, number, or plain node) +function getIncludeNode( arg ) { + + if ( typeof arg === 'function' && arg.functionNode ) return arg.functionNode; + if ( arg && arg.isNode && arg.functionNode ) return arg.functionNode; + if ( getStructLayout( arg ) ) return getStructLayout( arg ); + if ( arg && arg.isNode && arg.isCodeNode ) return arg; + return null; + +} + export class WGSLFnTagNode extends FunctionNode { static get type() { @@ -32,21 +50,28 @@ export class WGSLFnTagNode extends FunctionNode { constructor( tokens, args, lang = 'wgsl' ) { - // extract FunctionNode dependencies for includes — only function definitions - // need to be pre-registered so their code appears before ours in the output. + // extract dependencies for includes — function definitions, struct types, + // and code nodes need to be pre-registered so their code appears before ours. // callable wrappers and FunctionCallNodes are unwrapped to the underlying FunctionNode; // plain nodes (uniforms, storage, etc) are built inline in generate() and don't need includes. + // arrays are treated as explicit include lists — each element is registered as a dependency. const includes = []; for ( const arg of args ) { - if ( typeof arg === 'function' && arg.functionNode ) { + if ( Array.isArray( arg ) ) { - includes.push( arg.functionNode ); + for ( const element of arg ) { - } else if ( arg && arg.isNode && arg.functionNode ) { + const node = getIncludeNode( element ); + if ( node ) includes.push( node ); - includes.push( arg.functionNode ); + } + + } else { + + const node = getIncludeNode( arg ); + if ( node ) includes.push( node ); } @@ -59,19 +84,57 @@ export class WGSLFnTagNode extends FunctionNode { } - // parse the function signature from the static template parts (tokens[0] always - // contains the full signature since interpolations only appear in the body) + // assemble the signature from tokens and arg names (struct types may appear + // in the signature as return types or parameter types), then parse it getNodeFunction( builder ) { - const { tokens } = this; + const { tokens, args } = this; const nodeData = builder.getDataFromNode( this ); let nodeFunction = nodeData.nodeFunction; if ( nodeFunction === undefined ) { - const braceIndex = tokens[ 0 ].indexOf( '{' ); + // reconstruct the full code with known names for struct args + // and dummy identifiers for everything else + let fullCode = ''; + for ( let i = 0, l = tokens.length; i < l; i ++ ) { + + fullCode += tokens[ i ]; + + if ( i < args.length ) { + + const arg = args[ i ]; + if ( Array.isArray( arg ) ) { + + // include array — no text output + + } else if ( typeof arg === 'string' || typeof arg === 'number' ) { + + fullCode += String( arg ); + + } else { + + const structLayout = getStructLayout( arg ); + if ( structLayout ) { + + // use getNodeType to get the correct name (may be auto-generated) + fullCode += structLayout.getNodeType( builder ); + + } else { + + fullCode += '_arg' + i; + + } + + } + + } + + } + + const braceIndex = fullCode.indexOf( '{' ); const sig = braceIndex !== - 1 - ? tokens[ 0 ].substring( 0, braceIndex ) - : tokens[ 0 ]; + ? fullCode.substring( 0, braceIndex ) + : fullCode; nodeFunction = builder.parser.parseFunction( sig + ' {}' ); nodeData.nodeFunction = nodeFunction; @@ -100,7 +163,16 @@ export class WGSLFnTagNode extends FunctionNode { if ( i < args.length ) { const arg = args[ i ]; - if ( typeof arg === 'function' && arg.functionNode ) { + if ( Array.isArray( arg ) ) { + + // include array — no text output + + } else if ( typeof arg === 'string' || typeof arg === 'number' ) { + + // raw literal — output verbatim + parts.push( String( arg ) ); + + } else if ( typeof arg === 'function' && arg.functionNode ) { // callable wrapper (from wgslFn/wgslFnTag) — resolve to function name parts.push( arg.functionNode.build( builder, 'property' ) ); @@ -111,6 +183,12 @@ export class WGSLFnTagNode extends FunctionNode { // (build() would wrap it in a temp variable that lives outside our WGSL scope) parts.push( arg.generate( builder ) ); + } else if ( getStructLayout( arg ) ) { + + // struct (StructTypeNode or struct() callable) — build to register + // the struct definition, output just the type name + parts.push( getStructLayout( arg ).build( builder ) ); + } else { parts.push( arg.build( builder ) ); From 9d02fe3a24e670154bee84cc3c016aaaced3cf4f Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 09:59:41 +0900 Subject: [PATCH 57/84] Use the tag function --- src/webgpu/lib/BVHComputeData.js | 94 +++++++++++++-------------- src/webgpu/lib/nodes/WGSLFnTagNode.js | 34 ++++++++-- 2 files changed, 77 insertions(+), 51 deletions(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index 4d7f54345..fb9ff541f 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -1,12 +1,13 @@ import { Matrix4, Vector4 } from 'three'; import { StorageBufferAttribute, StructTypeNode } from 'three/webgpu'; -import { storage, wgslFn } from 'three/tsl'; +import { storage } from 'three/tsl'; import { intersectsBounds, rayStruct, bvhNodeStruct, constants, } from 'three-mesh-bvh/webgpu'; +import { wgslFnTag } from './nodes/WGSLFnTagNode'; const BYTES_PER_NODE = 6 * 4 + 4 + 4; const UINT32_PER_NODE = BYTES_PER_NODE / 4; @@ -74,11 +75,14 @@ const intersectionResultStruct = new StructTypeNode( { dist: 'float', }, 'IntersectionResult' ); -const intersectsTriangle = wgslFn( /* wgsl */ ` +const intersectsTriangle = wgslFnTag/* wgsl */ ` + // includes + ${ [ rayStruct, constants ] } - fn intersectsTriangle( ray: Ray, a: vec3f, b: vec3f, c: vec3f ) -> IntersectionResult { + // fn + fn intersectsTriangle( ray: Ray, a: vec3f, b: vec3f, c: vec3f ) -> ${ intersectionResultStruct } { - var result: IntersectionResult; + var result: ${ intersectionResultStruct }; result.didHit = false; let edge1 = b - a; @@ -119,15 +123,18 @@ const intersectsTriangle = wgslFn( /* wgsl */ ` return result; } - -`, [ rayStruct, intersectionResultStruct, constants ] ); +`; function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ) { - const geometryRaycastFirstHitFn = wgslFn( /* wgsl */` - fn ${ name }RaycastFirstHit_blas( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> IntersectionResult { + const geometryRaycastFirstHitFn = wgslFnTag/* wgsl */` + // includes + ${ [ rayStruct, bvhNodeStruct, constants, attributeStruct ] } + + // fn + fn ${ name }RaycastFirstHit_blas( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> ${ intersectionResultStruct } { - var bestHit: IntersectionResult; + var bestHit: ${ intersectionResultStruct }; bestHit.didHit = false; bestHit.dist = bestDist; @@ -144,11 +151,11 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto } let nodeIndex = stack[ pointer ]; - let node = ${ name }nodes.value[ nodeIndex ]; + let node = ${ nodesStorage }[ nodeIndex ]; pointer = pointer - 1; var boundsHitDist: f32 = 0.0; - if ( ! intersectsBounds( ray, node.bounds, &boundsHitDist ) || boundsHitDist > bestHit.dist ) { + if ( ! ${ intersectsBounds }( ray, node.bounds, &boundsHitDist ) || boundsHitDist > bestHit.dist ) { continue; @@ -165,15 +172,15 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto for ( var ti = triOffset; ti < triOffset + triCount; ti = ti + 1u ) { - let i0 = ${ name }index.value[ ti * 3u ]; - let i1 = ${ name }index.value[ ti * 3u + 1u ]; - let i2 = ${ name }index.value[ ti * 3u + 2u ]; + let i0 = ${ indexStorage }[ ti * 3u ]; + let i1 = ${ indexStorage }[ ti * 3u + 1u ]; + let i2 = ${ indexStorage }[ ti * 3u + 2u ]; - let a = ${ name }attributes.value[ i0 ].position.xyz; - let b = ${ name }attributes.value[ i1 ].position.xyz; - let c = ${ name }attributes.value[ i2 ].position.xyz; + let a = ${ attributesStorage }[ i0 ].position.xyz; + let b = ${ attributesStorage }[ i1 ].position.xyz; + let c = ${ attributesStorage }[ i2 ].position.xyz; - var triResult = intersectsTriangle( ray, a, b, c ); + var triResult = ${ intersectsTriangle }( ray, a, b, c ); if ( triResult.didHit && triResult.dist < bestHit.dist ) { @@ -207,17 +214,16 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto return bestHit; } - `, [ - nodesStorage, indexStorage, attributesStorage, - intersectsTriangle, intersectsBounds, - rayStruct, bvhNodeStruct, intersectionResultStruct, constants, - attributeStruct, - ] ); + `; + + return wgslFnTag /* wgsl */` + // includes + ${ [ rayStruct, bvhNodeStruct, constants, transformStruct ] } - return wgslFn( /* wgsl */` - fn ${ name }RaycastFirstHit( ray: Ray ) -> IntersectionResult { + // fn + fn ${ name }RaycastFirstHit( ray: Ray ) -> ${ intersectionResultStruct } { - var bestHit: IntersectionResult; + var bestHit: ${ intersectionResultStruct }; bestHit.didHit = false; bestHit.dist = INFINITY; @@ -234,11 +240,11 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto } let currNodeIndex = tlasStack[ tlasPointer ]; - let node = ${ name }nodes.value[ currNodeIndex ]; + let node = ${ nodesStorage }[ currNodeIndex ]; tlasPointer = tlasPointer - 1; var boundsHitDist: f32 = 0.0; - if ( ! intersectsBounds( ray, node.bounds, &boundsHitDist ) || boundsHitDist > bestHit.dist ) { + if ( ! ${ intersectsBounds }( ray, node.bounds, &boundsHitDist ) || boundsHitDist > bestHit.dist ) { continue; @@ -255,14 +261,14 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto for ( var t = offset; t < offset + count; t = t + 1u ) { - let transform = ${ name }transforms.value[ t ]; + let transform = ${ transformsStorage }[ t ]; // Transform ray into object local space var localRay: Ray; localRay.origin = ( transform.inverseMatrixWorld * vec4f( ray.origin, 1.0 ) ).xyz; localRay.direction = ( transform.inverseMatrixWorld * vec4f( ray.direction, 0.0 ) ).xyz; - let blasHit = ${ name }RaycastFirstHit_blas( localRay, transform.nodeOffset, bestHit.dist ); + let blasHit = ${ geometryRaycastFirstHitFn( { ray: 'localRay', rootNodeIndex: 'transform.nodeOffset', bestDist: 'bestHit.dist' } ) }; if ( blasHit.didHit && blasHit.dist < bestHit.dist ) { bestHit = blasHit; @@ -297,14 +303,7 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto return bestHit; - } - `, [ - geometryRaycastFirstHitFn, - nodesStorage, transformsStorage, - intersectsBounds, - rayStruct, bvhNodeStruct, intersectionResultStruct, constants, - transformStruct, - ] ); + }`; } @@ -467,18 +466,19 @@ export class BVHComputeData { return `result.${ name } = a0.${ name } * barycoord.x + a1.${ name } * barycoord.y + a2.${ name } * barycoord.z;`; } ).join( '\n' ); - this.fns.sampleTrianglePoint = wgslFn( /* wgsl */` - fn ${ name }sampleTrianglePoint( barycoord: vec3f, indices: vec3u ) -> ${ attributeStruct.name } { - - var result: ${ attributeStruct.name }; - var a0 = ${ name }attributes.value[ indices.x ]; - var a1 = ${ name }attributes.value[ indices.y ]; - var a2 = ${ name }attributes.value[ indices.z ]; + this.fns.sampleTrianglePoint = wgslFnTag/* wgsl */` + // fn + fn ${ name }sampleTrianglePoint( barycoord: vec3f, indices: vec3u ) -> ${ attributeStruct } { + + var result: ${ attributeStruct }; + var a0 = ${ attributesStorage }[ indices.x ]; + var a1 = ${ attributesStorage }[ indices.y ]; + var a2 = ${ attributesStorage }[ indices.z ]; ${ interpolateBody } return result; } - `, [ attributesStorage, attributeStruct ] ); + `; function appendBVHData( bvh, geometryOffset, transformInfo, nodeWriteOffset, target, tlas = false ) { diff --git a/src/webgpu/lib/nodes/WGSLFnTagNode.js b/src/webgpu/lib/nodes/WGSLFnTagNode.js index 51058d173..cd51d4a01 100644 --- a/src/webgpu/lib/nodes/WGSLFnTagNode.js +++ b/src/webgpu/lib/nodes/WGSLFnTagNode.js @@ -28,6 +28,27 @@ function getStructLayout( arg ) { } +// replaces any string parameters on a FunctionCallNode with RawExpression nodes +// so they output as raw WGSL identifiers (e.g. local variable names) +function convertStringParams( callNode ) { + + const params = callNode.parameters; + if ( params && typeof params === 'object' && ! Array.isArray( params ) && ! params.isNode ) { + + const converted = {}; + for ( const key in params ) { + + const v = params[ key ]; + converted[ key ] = ( typeof v === 'string' ) ? new RawExpression( v ) : v; + + } + + callNode.setParameters( converted ); + + } + +} + // returns the node that should be registered as an include for the given arg, // or null if the arg doesn't represent a dependency (e.g. a string, number, or plain node) function getIncludeNode( arg ) { @@ -63,8 +84,11 @@ export class WGSLFnTagNode extends FunctionNode { for ( const element of arg ) { + // unwrap callable wrappers; accept any remaining node directly + // (storage, uniforms, etc. need to be built to register bindings) const node = getIncludeNode( element ); if ( node ) includes.push( node ); + else if ( element && element.isNode ) includes.push( element ); } @@ -132,9 +156,8 @@ export class WGSLFnTagNode extends FunctionNode { } const braceIndex = fullCode.indexOf( '{' ); - const sig = braceIndex !== - 1 - ? fullCode.substring( 0, braceIndex ) - : fullCode; + let sig = braceIndex !== - 1 ? fullCode.substring( 0, braceIndex ) : fullCode; + sig = sig.replace( /\/\/.+[\n\r]/g, '' ); nodeFunction = builder.parser.parseFunction( sig + ' {}' ); nodeData.nodeFunction = nodeFunction; @@ -180,7 +203,10 @@ export class WGSLFnTagNode extends FunctionNode { } else if ( arg.isNode && arg.functionNode ) { // FunctionCallNode — use generate() to get the inline call expression - // (build() would wrap it in a temp variable that lives outside our WGSL scope) + // (build() would wrap it in a temp variable that lives outside our WGSL scope). + // convert any string params to RawExpression — the function may have been + // created by wgslFn (which doesn't handle string-to-node conversion) + convertStringParams( arg ); parts.push( arg.generate( builder ) ); } else if ( getStructLayout( arg ) ) { From 15d21afef2d79b968dda9f8852123e6514d3d5b9 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 10:13:44 +0900 Subject: [PATCH 58/84] Pass attribute sizes --- src/webgpu/lib/BVHComputeData.js | 43 +++++++++++----------- src/webgpu/lib/PathtracerBVHComputeData.js | 15 ++++++-- 2 files changed, 32 insertions(+), 26 deletions(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index fb9ff541f..a2c39a1f9 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -125,14 +125,14 @@ const intersectsTriangle = wgslFnTag/* wgsl */ ` } `; -function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ) { +function buildRaycastFirstHitFn( prefix, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ) { const geometryRaycastFirstHitFn = wgslFnTag/* wgsl */` // includes ${ [ rayStruct, bvhNodeStruct, constants, attributeStruct ] } // fn - fn ${ name }RaycastFirstHit_blas( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> ${ intersectionResultStruct } { + fn ${ prefix }RaycastFirstHit_blas( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> ${ intersectionResultStruct } { var bestHit: ${ intersectionResultStruct }; bestHit.didHit = false; @@ -221,7 +221,7 @@ function buildRaycastFirstHitFn( name, nodesStorage, transformsStorage, indexSto ${ [ rayStruct, bvhNodeStruct, constants, transformStruct ] } // fn - fn ${ name }RaycastFirstHit( ray: Ray ) -> ${ intersectionResultStruct } { + fn ${ prefix }RaycastFirstHit( ray: Ray ) -> ${ intersectionResultStruct } { var bestHit: ${ intersectionResultStruct }; bestHit.didHit = false; @@ -312,11 +312,11 @@ export class BVHComputeData { constructor( bvh, options = {} ) { const { - name = 'bvh_', - attributes = [ 'position', 'normal', 'uv0' ], + prefix = 'bvh_', + attributes = { position: 'vec4f' }, } = options; - this.name = name; + this.prefix = prefix; this.attributes = attributes; this.bvh = bvh; @@ -341,7 +341,7 @@ export class BVHComputeData { update() { const self = this; - const { attributes, structs, name, bvh } = this; + const { attributes, structs, prefix, bvh } = this; // collect the BVHs const bvhInfo = []; @@ -399,17 +399,14 @@ export class BVHComputeData { // // construct the attribute struct - const attributeStruct = new StructTypeNode( - attributes.reduce( ( o, key ) => ( { ...o, [ key ]: 'vec4f' } ), {} ), - `${ name }GeometryStruct`, - ); + const attributeStruct = new StructTypeNode( attributes, `${ prefix }GeometryStruct` ); // write the geometry buffer attributes & bvh data let attributesOffset = 0; let indexOffset = 0; let nodeWriteOffset = 0; const indexBuffer = new Uint32Array( indexBufferLength ); - const attributesBuffer = new ArrayBuffer( attributesBufferLength * attributes.length * 4 * 4 ); + const attributesBuffer = new ArrayBuffer( attributesBufferLength * attributeStruct.getLength() * 4 ); const bvhNodesBuffer = new ArrayBuffer( bvhNodesBufferLength ); // append TLAS data @@ -447,17 +444,17 @@ export class BVHComputeData { // // set up the storage buffers - const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), 'BVHNode' ).toReadOnly().setName( `${ name }nodes` ); - const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), structs.transform.getLength() ), structs.transform.name ).toReadOnly().setName( `${ name }transforms` ); - const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ name }index` ); - const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributeStruct.getLength() ), attributeStruct.name ).toReadOnly().setName( `${ name }attributes` ); + const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), 'BVHNode' ).toReadOnly().setName( `${ prefix }nodes` ); + const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), structs.transform.getLength() ), structs.transform.name ).toReadOnly().setName( `${ prefix }transforms` ); + const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ prefix }index` ); + const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributeStruct.getLength() ), attributeStruct.name ).toReadOnly().setName( `${ prefix }attributes` ); this.storage.transforms = transformsStorage; this.storage.nodes = bvhNodesStorage; this.storage.index = indexStorage; this.storage.attributes = attributesStorage; this.structs.attributes = attributeStruct; - this.fns.raycastFirstHit = buildRaycastFirstHitFn( name, bvhNodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, structs.transform ); + this.fns.raycastFirstHit = buildRaycastFirstHitFn( prefix, bvhNodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, structs.transform ); const interpolateBody = attributeStruct .membersLayout @@ -468,7 +465,7 @@ export class BVHComputeData { } ).join( '\n' ); this.fns.sampleTrianglePoint = wgslFnTag/* wgsl */` // fn - fn ${ name }sampleTrianglePoint( barycoord: vec3f, indices: vec3u ) -> ${ attributeStruct } { + fn ${ prefix }sampleTrianglePoint( barycoord: vec3f, indices: vec3u ) -> ${ attributeStruct } { var result: ${ attributeStruct }; var a0 = ${ attributesStorage }[ indices.x ]; @@ -598,16 +595,18 @@ export class BVHComputeData { const { geometry, mesh = null } = bvh; const { vertexStart, vertexCount } = range; const attributesBufferF32 = new Float32Array( target ); - attributes.forEach( ( key, interleavedOffset ) => { + attributeStruct.membersLayout.forEach( ( { name }, interleavedOffset ) => { - const attr = geometry.attributes[ key ]; - self.getDefaultAttributeValue( key, _def ); + // TODO: we should be able to have access to memory layout offsets here via the struct + // API but it's not currently available. + const attr = geometry.attributes[ name ]; + self.getDefaultAttributeValue( name, _def ); for ( let i = 0; i < vertexCount; i ++ ) { if ( attr ) { - if ( key === 'position' && mesh ) { + if ( name === 'position' && mesh ) { // TODO: normals and tangents need to be transformed here, as well mesh.getVertexPosition( i + vertexStart, _vec ); diff --git a/src/webgpu/lib/PathtracerBVHComputeData.js b/src/webgpu/lib/PathtracerBVHComputeData.js index 75126510b..3121c17c8 100644 --- a/src/webgpu/lib/PathtracerBVHComputeData.js +++ b/src/webgpu/lib/PathtracerBVHComputeData.js @@ -18,9 +18,16 @@ const materialStruct = new StructTypeNode( { export class PathtracerBVHComputeData extends BVHComputeData { - constructor( ...args ) { - - super( ...args ); + constructor( bvh, options = {} ) { + + super( bvh, { + attributes: { + position: 'vec4f', + normal: 'vec4f', + uv0: 'vec4f', + }, + ...options, + } ); this.structs.transform = transformStruct; this.structs.material = materialStruct; @@ -33,7 +40,7 @@ export class PathtracerBVHComputeData extends BVHComputeData { super.update(); - const { materials, structs, name } = this; + const { materials, structs, prefix: name } = this; const materialBuffer = new ArrayBuffer( structs.material.getLength() * materials.length * 4 ); const materialBufferF32 = new Float32Array( materialBuffer ); materials.forEach( ( mat, i ) => { From d2995cb2847af4be7afaa4e029c33aa7777292fc Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 10:48:48 +0900 Subject: [PATCH 59/84] WGSLFnTagNode simplification --- src/webgpu/lib/nodes/WGSLFnTagNode.js | 30 +++------------------------ 1 file changed, 3 insertions(+), 27 deletions(-) diff --git a/src/webgpu/lib/nodes/WGSLFnTagNode.js b/src/webgpu/lib/nodes/WGSLFnTagNode.js index cd51d4a01..31b5969c8 100644 --- a/src/webgpu/lib/nodes/WGSLFnTagNode.js +++ b/src/webgpu/lib/nodes/WGSLFnTagNode.js @@ -22,33 +22,12 @@ class RawExpression extends Node { // returns the StructTypeNode from either a direct StructTypeNode or a struct() callable wrapper function getStructLayout( arg ) { - if ( arg && arg.isNode && arg.isStructLayoutNode ) return arg; + if ( arg && arg.isStructLayoutNode ) return arg; if ( typeof arg === 'function' && arg.isStruct ) return arg.layout; return null; } -// replaces any string parameters on a FunctionCallNode with RawExpression nodes -// so they output as raw WGSL identifiers (e.g. local variable names) -function convertStringParams( callNode ) { - - const params = callNode.parameters; - if ( params && typeof params === 'object' && ! Array.isArray( params ) && ! params.isNode ) { - - const converted = {}; - for ( const key in params ) { - - const v = params[ key ]; - converted[ key ] = ( typeof v === 'string' ) ? new RawExpression( v ) : v; - - } - - callNode.setParameters( converted ); - - } - -} - // returns the node that should be registered as an include for the given arg, // or null if the arg doesn't represent a dependency (e.g. a string, number, or plain node) function getIncludeNode( arg ) { @@ -56,7 +35,7 @@ function getIncludeNode( arg ) { if ( typeof arg === 'function' && arg.functionNode ) return arg.functionNode; if ( arg && arg.isNode && arg.functionNode ) return arg.functionNode; if ( getStructLayout( arg ) ) return getStructLayout( arg ); - if ( arg && arg.isNode && arg.isCodeNode ) return arg; + if ( arg && arg.isCodeNode ) return arg; return null; } @@ -203,10 +182,7 @@ export class WGSLFnTagNode extends FunctionNode { } else if ( arg.isNode && arg.functionNode ) { // FunctionCallNode — use generate() to get the inline call expression - // (build() would wrap it in a temp variable that lives outside our WGSL scope). - // convert any string params to RawExpression — the function may have been - // created by wgslFn (which doesn't handle string-to-node conversion) - convertStringParams( arg ); + // (build() would wrap it in a temp variable that lives outside our WGSL scope) parts.push( arg.generate( builder ) ); } else if ( getStructLayout( arg ) ) { From d1aed18c5d4f61c3322a3958968fb9c5be9159eb Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 11:37:28 +0900 Subject: [PATCH 60/84] Simplification --- src/webgpu/lib/nodes/WGSLFnTagNode.js | 112 ++++++++++++++++---------- 1 file changed, 70 insertions(+), 42 deletions(-) diff --git a/src/webgpu/lib/nodes/WGSLFnTagNode.js b/src/webgpu/lib/nodes/WGSLFnTagNode.js index 31b5969c8..d4f98af88 100644 --- a/src/webgpu/lib/nodes/WGSLFnTagNode.js +++ b/src/webgpu/lib/nodes/WGSLFnTagNode.js @@ -19,12 +19,40 @@ class RawExpression extends Node { } -// returns the StructTypeNode from either a direct StructTypeNode or a struct() callable wrapper -function getStructLayout( arg ) { +// wraps a FunctionNode so that build() returns just the function name +class PropertyRefNode extends Node { - if ( arg && arg.isStructLayoutNode ) return arg; - if ( typeof arg === 'function' && arg.isStruct ) return arg.layout; - return null; + constructor( node ) { + + super(); + this.node = node; + + } + + build( builder ) { + + return this.node.build( builder, 'property' ); + + } + +} + +// wraps a FunctionCallNode so that build() returns the inline call expression, +// bypassing TempNode's variable wrapping +class InlineCallNode extends Node { + + constructor( node ) { + + super(); + this.node = node; + + } + + build( builder ) { + + return this.node.generate( builder ); + + } } @@ -32,15 +60,26 @@ function getStructLayout( arg ) { // or null if the arg doesn't represent a dependency (e.g. a string, number, or plain node) function getIncludeNode( arg ) { - if ( typeof arg === 'function' && arg.functionNode ) return arg.functionNode; - if ( arg && arg.isNode && arg.functionNode ) return arg.functionNode; - if ( getStructLayout( arg ) ) return getStructLayout( arg ); - if ( arg && arg.isCodeNode ) return arg; + if ( typeof arg === 'function' ) { + + if ( arg.functionNode ) return arg.functionNode; + if ( arg.isStruct ) return arg.layout; + return null; + + } + + if ( arg && arg.isNode ) { + + if ( arg.functionNode ) return arg.functionNode; + if ( arg.isStructLayoutNode || arg.isCodeNode ) return arg; + + } + return null; } -export class WGSLFnTagNode extends FunctionNode { +export class WGSLTagFnNode extends FunctionNode { static get type() { @@ -50,8 +89,8 @@ export class WGSLFnTagNode extends FunctionNode { constructor( tokens, args, lang = 'wgsl' ) { - // extract dependencies for includes — function definitions, struct types, - // and code nodes need to be pre-registered so their code appears before ours. + // extract dependencies for includes from the original args — function definitions, + // struct types, and code nodes need to be pre-registered so their code appears before ours. // callable wrappers and FunctionCallNodes are unwrapped to the underlying FunctionNode; // plain nodes (uniforms, storage, etc) are built inline in generate() and don't need includes. // arrays are treated as explicit include lists — each element is registered as a dependency. @@ -80,10 +119,23 @@ export class WGSLFnTagNode extends FunctionNode { } + // normalize args so generate() can resolve them uniformly with build(): + // - callable wrappers → PropertyRefNode (emits just the function name) + // - struct callables → StructTypeNode (emits the type name via build) + // - FunctionCallNodes → InlineCallNode (emits inline call, bypassing TempNode) + const normalizedArgs = args.map( arg => { + + if ( typeof arg === 'function' && arg.functionNode ) return new PropertyRefNode( arg.functionNode ); + if ( typeof arg === 'function' && arg.isStruct ) return arg.layout; + if ( arg && arg.isNode && arg.functionNode ) return new InlineCallNode( arg ); + return arg; + + } ); + super( '', includes, lang ); this.tokens = tokens; - this.args = args; + this.args = normalizedArgs; } @@ -114,19 +166,13 @@ export class WGSLFnTagNode extends FunctionNode { fullCode += String( arg ); - } else { - - const structLayout = getStructLayout( arg ); - if ( structLayout ) { - - // use getNodeType to get the correct name (may be auto-generated) - fullCode += structLayout.getNodeType( builder ); + } else if ( arg.isStructLayoutNode ) { - } else { + fullCode += arg.getNodeType( builder ); - fullCode += '_arg' + i; + } else { - } + fullCode += '_arg' + i; } @@ -171,26 +217,8 @@ export class WGSLFnTagNode extends FunctionNode { } else if ( typeof arg === 'string' || typeof arg === 'number' ) { - // raw literal — output verbatim parts.push( String( arg ) ); - } else if ( typeof arg === 'function' && arg.functionNode ) { - - // callable wrapper (from wgslFn/wgslFnTag) — resolve to function name - parts.push( arg.functionNode.build( builder, 'property' ) ); - - } else if ( arg.isNode && arg.functionNode ) { - - // FunctionCallNode — use generate() to get the inline call expression - // (build() would wrap it in a temp variable that lives outside our WGSL scope) - parts.push( arg.generate( builder ) ); - - } else if ( getStructLayout( arg ) ) { - - // struct (StructTypeNode or struct() callable) — build to register - // the struct definition, output just the type name - parts.push( getStructLayout( arg ).build( builder ) ); - } else { parts.push( arg.build( builder ) ); @@ -224,7 +252,7 @@ export class WGSLFnTagNode extends FunctionNode { const wgslFnTag = ( tokens, ...args ) => { - const functionNode = new WGSLFnTagNode( tokens, args ); + const functionNode = new WGSLTagFnNode( tokens, args ); const fn = ( ...params ) => { From 0ce323f7c98331c6242d3f21dc32bbef74ff7439 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 11:54:44 +0900 Subject: [PATCH 61/84] Simplification --- src/webgpu/lib/BVHComputeData.js | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index a2c39a1f9..210421f5f 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -1,6 +1,6 @@ import { Matrix4, Vector4 } from 'three'; import { StorageBufferAttribute, StructTypeNode } from 'three/webgpu'; -import { storage } from 'three/tsl'; +import { storage, struct } from 'three/tsl'; import { intersectsBounds, rayStruct, @@ -125,11 +125,13 @@ const intersectsTriangle = wgslFnTag/* wgsl */ ` } `; -function buildRaycastFirstHitFn( prefix, nodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, transformStruct ) { +function buildRaycastFirstHitFn( prefix, storage, structs ) { + + const geometryRaycastFirstHitFn = wgslFnTag/* wgsl */` // includes - ${ [ rayStruct, bvhNodeStruct, constants, attributeStruct ] } + ${ [ rayStruct, bvhNodeStruct, constants, structs.attributes ] } // fn fn ${ prefix }RaycastFirstHit_blas( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> ${ intersectionResultStruct } { @@ -151,7 +153,7 @@ function buildRaycastFirstHitFn( prefix, nodesStorage, transformsStorage, indexS } let nodeIndex = stack[ pointer ]; - let node = ${ nodesStorage }[ nodeIndex ]; + let node = ${ storage.nodes }[ nodeIndex ]; pointer = pointer - 1; var boundsHitDist: f32 = 0.0; @@ -172,13 +174,13 @@ function buildRaycastFirstHitFn( prefix, nodesStorage, transformsStorage, indexS for ( var ti = triOffset; ti < triOffset + triCount; ti = ti + 1u ) { - let i0 = ${ indexStorage }[ ti * 3u ]; - let i1 = ${ indexStorage }[ ti * 3u + 1u ]; - let i2 = ${ indexStorage }[ ti * 3u + 2u ]; + let i0 = ${ storage.index }[ ti * 3u ]; + let i1 = ${ storage.index }[ ti * 3u + 1u ]; + let i2 = ${ storage.index }[ ti * 3u + 2u ]; - let a = ${ attributesStorage }[ i0 ].position.xyz; - let b = ${ attributesStorage }[ i1 ].position.xyz; - let c = ${ attributesStorage }[ i2 ].position.xyz; + let a = ${ storage.attributes }[ i0 ].position.xyz; + let b = ${ storage.attributes }[ i1 ].position.xyz; + let c = ${ storage.attributes }[ i2 ].position.xyz; var triResult = ${ intersectsTriangle }( ray, a, b, c ); @@ -240,7 +242,7 @@ function buildRaycastFirstHitFn( prefix, nodesStorage, transformsStorage, indexS } let currNodeIndex = tlasStack[ tlasPointer ]; - let node = ${ nodesStorage }[ currNodeIndex ]; + let node = ${ storage.nodes }[ currNodeIndex ]; tlasPointer = tlasPointer - 1; var boundsHitDist: f32 = 0.0; @@ -261,7 +263,7 @@ function buildRaycastFirstHitFn( prefix, nodesStorage, transformsStorage, indexS for ( var t = offset; t < offset + count; t = t + 1u ) { - let transform = ${ transformsStorage }[ t ]; + let transform = ${ storage.transforms }[ t ]; // Transform ray into object local space var localRay: Ray; @@ -454,7 +456,7 @@ export class BVHComputeData { this.storage.index = indexStorage; this.storage.attributes = attributesStorage; this.structs.attributes = attributeStruct; - this.fns.raycastFirstHit = buildRaycastFirstHitFn( prefix, bvhNodesStorage, transformsStorage, indexStorage, attributesStorage, attributeStruct, structs.transform ); + this.fns.raycastFirstHit = buildRaycastFirstHitFn( prefix, this.storage, this.structs ); const interpolateBody = attributeStruct .membersLayout From 68c8f034b45fe3e9dd4e7e2171cc1d89e05a9628 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 12:27:11 +0900 Subject: [PATCH 62/84] Fix up transforms --- src/webgpu/lib/BVHComputeData.js | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index 210421f5f..84482cc52 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -127,8 +127,6 @@ const intersectsTriangle = wgslFnTag/* wgsl */ ` function buildRaycastFirstHitFn( prefix, storage, structs ) { - - const geometryRaycastFirstHitFn = wgslFnTag/* wgsl */` // includes ${ [ rayStruct, bvhNodeStruct, constants, structs.attributes ] } @@ -218,7 +216,7 @@ function buildRaycastFirstHitFn( prefix, storage, structs ) { } `; - return wgslFnTag /* wgsl */` + return wgslFnTag/* wgsl */` // includes ${ [ rayStruct, bvhNodeStruct, constants, transformStruct ] } @@ -656,6 +654,7 @@ export class BVHComputeData { writeTransformData( info, premultiplyMatrix, writeOffset, targetBuffer ) { + const { structs } = this; const transformBufferF32 = new Float32Array( targetBuffer ); const transformBufferU32 = new Uint32Array( targetBuffer ); @@ -673,12 +672,12 @@ export class BVHComputeData { } _matrix.premultiply( premultiplyMatrix ); - _matrix.toArray( transformBufferF32, writeOffset * transformStruct.getLength() ); + _matrix.toArray( transformBufferF32, writeOffset * structs.transform.getLength() ); _matrix.invert(); - _matrix.toArray( transformBufferF32, writeOffset * transformStruct.getLength() + 16 ); + _matrix.toArray( transformBufferF32, writeOffset * structs.transform.getLength() + 16 ); - transformBufferU32[ writeOffset * transformStruct.getLength() + 32 ] = bvhNodeOffsets[ root ]; + transformBufferU32[ writeOffset * structs.transform.getLength() + 32 ] = bvhNodeOffsets[ root ]; } From a48309e7c3130e3482eafbd8f84cdb3a7006b108 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 12:31:36 +0900 Subject: [PATCH 63/84] Rename classes --- src/webgpu/lib/BVHComputeData.js | 12 ++++++------ .../lib/nodes/{WGSLFnTagNode.js => WGSLTagFnNode.js} | 6 +----- 2 files changed, 7 insertions(+), 11 deletions(-) rename src/webgpu/lib/nodes/{WGSLFnTagNode.js => WGSLTagFnNode.js} (98%) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index 84482cc52..cf51063e6 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -1,13 +1,13 @@ import { Matrix4, Vector4 } from 'three'; import { StorageBufferAttribute, StructTypeNode } from 'three/webgpu'; -import { storage, struct } from 'three/tsl'; +import { storage } from 'three/tsl'; import { intersectsBounds, rayStruct, bvhNodeStruct, constants, } from 'three-mesh-bvh/webgpu'; -import { wgslFnTag } from './nodes/WGSLFnTagNode'; +import { wgslTagFn } from './nodes/WGSLTagFnNode.js'; const BYTES_PER_NODE = 6 * 4 + 4 + 4; const UINT32_PER_NODE = BYTES_PER_NODE / 4; @@ -75,7 +75,7 @@ const intersectionResultStruct = new StructTypeNode( { dist: 'float', }, 'IntersectionResult' ); -const intersectsTriangle = wgslFnTag/* wgsl */ ` +const intersectsTriangle = wgslTagFn/* wgsl */ ` // includes ${ [ rayStruct, constants ] } @@ -127,7 +127,7 @@ const intersectsTriangle = wgslFnTag/* wgsl */ ` function buildRaycastFirstHitFn( prefix, storage, structs ) { - const geometryRaycastFirstHitFn = wgslFnTag/* wgsl */` + const geometryRaycastFirstHitFn = wgslTagFn/* wgsl */` // includes ${ [ rayStruct, bvhNodeStruct, constants, structs.attributes ] } @@ -216,7 +216,7 @@ function buildRaycastFirstHitFn( prefix, storage, structs ) { } `; - return wgslFnTag/* wgsl */` + return wgslTagFn/* wgsl */` // includes ${ [ rayStruct, bvhNodeStruct, constants, transformStruct ] } @@ -463,7 +463,7 @@ export class BVHComputeData { return `result.${ name } = a0.${ name } * barycoord.x + a1.${ name } * barycoord.y + a2.${ name } * barycoord.z;`; } ).join( '\n' ); - this.fns.sampleTrianglePoint = wgslFnTag/* wgsl */` + this.fns.sampleTrianglePoint = wgslTagFn/* wgsl */` // fn fn ${ prefix }sampleTrianglePoint( barycoord: vec3f, indices: vec3u ) -> ${ attributeStruct } { diff --git a/src/webgpu/lib/nodes/WGSLFnTagNode.js b/src/webgpu/lib/nodes/WGSLTagFnNode.js similarity index 98% rename from src/webgpu/lib/nodes/WGSLFnTagNode.js rename to src/webgpu/lib/nodes/WGSLTagFnNode.js index d4f98af88..e3b457c1c 100644 --- a/src/webgpu/lib/nodes/WGSLFnTagNode.js +++ b/src/webgpu/lib/nodes/WGSLTagFnNode.js @@ -250,7 +250,7 @@ export class WGSLTagFnNode extends FunctionNode { } -const wgslFnTag = ( tokens, ...args ) => { +export const wgslTagFn = ( tokens, ...args ) => { const functionNode = new WGSLTagFnNode( tokens, args ); @@ -282,7 +282,3 @@ const wgslFnTag = ( tokens, ...args ) => { return fn; }; - -export { - wgslFnTag, -}; From 2da88d979c863be5f85472d7b7fed822f00ba889 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 13:06:22 +0900 Subject: [PATCH 64/84] Simplification --- src/webgpu/lib/nodes/WGSLTagFnNode.js | 85 +++++++++++---------------- 1 file changed, 34 insertions(+), 51 deletions(-) diff --git a/src/webgpu/lib/nodes/WGSLTagFnNode.js b/src/webgpu/lib/nodes/WGSLTagFnNode.js index e3b457c1c..30c02176e 100644 --- a/src/webgpu/lib/nodes/WGSLTagFnNode.js +++ b/src/webgpu/lib/nodes/WGSLTagFnNode.js @@ -1,19 +1,18 @@ import { FunctionNode, Node } from 'three/webgpu'; -// minimal node that outputs a raw WGSL expression verbatim when built, -// bypassing TSL's temp variable wrapping and type formatting -class RawExpression extends Node { +// minimal node that outputs a raw WGSL expression verbatim when built +class LiteralExpression extends Node { - constructor( code ) { + constructor( literal ) { super(); - this.code = code; + this.literal = literal; } build() { - return this.code; + return this.literal; } @@ -56,8 +55,7 @@ class InlineCallNode extends Node { } -// returns the node that should be registered as an include for the given arg, -// or null if the arg doesn't represent a dependency (e.g. a string, number, or plain node) +// returns the node that should be registered as an include for the given arg function getIncludeNode( arg ) { if ( typeof arg === 'function' ) { @@ -89,24 +87,16 @@ export class WGSLTagFnNode extends FunctionNode { constructor( tokens, args, lang = 'wgsl' ) { - // extract dependencies for includes from the original args — function definitions, - // struct types, and code nodes need to be pre-registered so their code appears before ours. - // callable wrappers and FunctionCallNodes are unwrapped to the underlying FunctionNode; - // plain nodes (uniforms, storage, etc) are built inline in generate() and don't need includes. - // arrays are treated as explicit include lists — each element is registered as a dependency. + // assemble all the nodes needed for includes const includes = []; - for ( const arg of args ) { if ( Array.isArray( arg ) ) { for ( const element of arg ) { - // unwrap callable wrappers; accept any remaining node directly - // (storage, uniforms, etc. need to be built to register bindings) const node = getIncludeNode( element ); if ( node ) includes.push( node ); - else if ( element && element.isNode ) includes.push( element ); } @@ -119,10 +109,10 @@ export class WGSLTagFnNode extends FunctionNode { } - // normalize args so generate() can resolve them uniformly with build(): - // - callable wrappers → PropertyRefNode (emits just the function name) - // - struct callables → StructTypeNode (emits the type name via build) - // - FunctionCallNodes → InlineCallNode (emits inline call, bypassing TempNode) + // normalize args so generate function can resolve them with build() later: + // - callable wrappers > PropertyRefNode (emits just the function name) + // - struct callables > StructTypeNode (emits the type name via build) + // - FunctionCallNodes > InlineCallNode (emits inline call) const normalizedArgs = args.map( arg => { if ( typeof arg === 'function' && arg.functionNode ) return new PropertyRefNode( arg.functionNode ); @@ -139,8 +129,7 @@ export class WGSLTagFnNode extends FunctionNode { } - // assemble the signature from tokens and arg names (struct types may appear - // in the signature as return types or parameter types), then parse it + // assemble the signature from tokens and arg names then parse getNodeFunction( builder ) { const { tokens, args } = this; @@ -164,12 +153,19 @@ export class WGSLTagFnNode extends FunctionNode { } else if ( typeof arg === 'string' || typeof arg === 'number' ) { + // literals fullCode += String( arg ); } else if ( arg.isStructLayoutNode ) { + // struct type node fullCode += arg.getNodeType( builder ); + } else if ( arg.isStruct ) { + + // struct + fullCode += arg.layout.getNodeType( builder ); + } else { fullCode += '_arg' + i; @@ -180,11 +176,11 @@ export class WGSLTagFnNode extends FunctionNode { } - const braceIndex = fullCode.indexOf( '{' ); - let sig = braceIndex !== - 1 ? fullCode.substring( 0, braceIndex ) : fullCode; - sig = sig.replace( /\/\/.+[\n\r]/g, '' ); + // remove comments + fullCode = fullCode.replace( /\/\/.+[\n\r]/g, '' ); - nodeFunction = builder.parser.parseFunction( sig + ' {}' ); + // parse it so we have the signature defined - we will define the body content after + nodeFunction = builder.parser.parseFunction( fullCode ); nodeData.nodeFunction = nodeFunction; } @@ -193,21 +189,17 @@ export class WGSLTagFnNode extends FunctionNode { } + // get the code for the function generate( builder, output ) { const { tokens, args } = this; - // let FunctionNode.generate handle includes, code registration, property naming, - // and type normalization (e.g. stripping "-> void" which is not valid WGSL) + // rebuild the function body again because we can call "build", now const result = super.generate( builder, output ); - - // assemble the body by interleaving static tokens with resolved node names - const parts = []; - + let fullCode = ''; for ( let i = 0, l = tokens.length; i < l; i ++ ) { - parts.push( tokens[ i ] ); - + fullCode += tokens[ i ]; if ( i < args.length ) { const arg = args[ i ]; @@ -217,11 +209,11 @@ export class WGSLTagFnNode extends FunctionNode { } else if ( typeof arg === 'string' || typeof arg === 'number' ) { - parts.push( String( arg ) ); + fullCode += String( arg ); } else { - parts.push( arg.build( builder ) ); + fullCode += arg.build( builder ); } @@ -232,17 +224,8 @@ export class WGSLTagFnNode extends FunctionNode { const { type } = this.getNodeFunction( builder ); const nodeCode = builder.getCodeFromNode( this, type ); - // use the declaration from super (handles type normalization and name assignment), - // replace its empty body with the assembled body from the template - const declaration = nodeCode.code; - const declPrefix = declaration.substring( 0, declaration.indexOf( '{' ) + 1 ); - - const assembledCode = parts.join( '' ); - const bodyStart = assembledCode.indexOf( '{' ) + 1; - const bodyEnd = assembledCode.lastIndexOf( '}' ); - const body = assembledCode.substring( bodyStart, bodyEnd ); - - nodeCode.code = declPrefix + body + '}\n'; + fullCode = fullCode.replace( /\/\/.+[\n\r]/g, '' ).replace( /->\s*void/, '' ).replace( /\s+/g, ' ' ).trim(); + nodeCode.code = fullCode; return result; @@ -256,8 +239,8 @@ export const wgslTagFn = ( tokens, ...args ) => { const fn = ( ...params ) => { - // wrap string parameter values as raw WGSL expressions - // so they output verbatim as identifiers (e.g. local variable names) + // wrap string parameter values as raw WGSL expressions so they + // output verbatim as identifiers like local variable names if ( params.length === 1 && params[ 0 ] && typeof params[ 0 ] === 'object' && ! params[ 0 ].isNode ) { const obj = params[ 0 ]; @@ -265,7 +248,7 @@ export const wgslTagFn = ( tokens, ...args ) => { if ( typeof obj[ key ] === 'string' ) { - obj[ key ] = new RawExpression( obj[ key ] ); + obj[ key ] = new LiteralExpression( obj[ key ] ); } From 7fde608fa630a2ddbaa7593417dc64fc551d4a1c Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 13:58:10 +0900 Subject: [PATCH 65/84] Add comments --- src/webgpu/lib/BVHComputeData.js | 3 +++ src/webgpu/lib/nodes/WGSLTagFnNode.js | 2 ++ 2 files changed, 5 insertions(+) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index cf51063e6..0ea6a1840 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -127,6 +127,9 @@ const intersectsTriangle = wgslTagFn/* wgsl */ ` function buildRaycastFirstHitFn( prefix, storage, structs ) { + // TODO: reduce the redundancy between these functions - possibly using code snippets or + // macro-expansion-style mechanisms? + const geometryRaycastFirstHitFn = wgslTagFn/* wgsl */` // includes ${ [ rayStruct, bvhNodeStruct, constants, structs.attributes ] } diff --git a/src/webgpu/lib/nodes/WGSLTagFnNode.js b/src/webgpu/lib/nodes/WGSLTagFnNode.js index 30c02176e..a647b438a 100644 --- a/src/webgpu/lib/nodes/WGSLTagFnNode.js +++ b/src/webgpu/lib/nodes/WGSLTagFnNode.js @@ -233,6 +233,8 @@ export class WGSLTagFnNode extends FunctionNode { } +// template tag literal function version of "wgslFn" so easy interpolation of TSL nodes +// TODO: add a raw "wgsl" version for code snippets export const wgslTagFn = ( tokens, ...args ) => { const functionNode = new WGSLTagFnNode( tokens, args ); From dae28f5869c3a0ee4c2c6873b5ea275c398ab42e Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 14:12:40 +0900 Subject: [PATCH 66/84] Move structs internally --- src/webgpu/compute/PathTracerMegaKernel.js | 2 +- .../compute/wavefront/ProcessHitsKernel.js | 2 +- .../compute/wavefront/RayGenerationKernel.js | 2 +- .../wavefront/RayIntersectionKernel.js | 2 +- src/webgpu/compute/wavefront/structs.js | 2 +- src/webgpu/lib/BVHComputeData.js | 6 +- src/webgpu/lib/wgsl/common.wgsl.js | 70 +++++++++++++++++++ src/webgpu/lib/wgsl/structs.wgsl.js | 28 ++++++++ 8 files changed, 107 insertions(+), 7 deletions(-) create mode 100644 src/webgpu/lib/wgsl/common.wgsl.js create mode 100644 src/webgpu/lib/wgsl/structs.wgsl.js diff --git a/src/webgpu/compute/PathTracerMegaKernel.js b/src/webgpu/compute/PathTracerMegaKernel.js index 4a61933fe..932493452 100644 --- a/src/webgpu/compute/PathTracerMegaKernel.js +++ b/src/webgpu/compute/PathTracerMegaKernel.js @@ -1,5 +1,5 @@ import { Matrix4, Vector2, StorageTexture } from 'three/webgpu'; -import { ndcToCameraRay } from 'three-mesh-bvh/webgpu'; +import { ndcToCameraRay } from '../lib/wgsl/common.wgsl.js'; import { ComputeKernel } from './ComputeKernel.js'; import { uniform, globalId, textureStore, wgslFn } from 'three/tsl'; import { pcgRand3, pcgInit } from '../nodes/random.wgsl.js'; diff --git a/src/webgpu/compute/wavefront/ProcessHitsKernel.js b/src/webgpu/compute/wavefront/ProcessHitsKernel.js index 9481f9a39..6b798bde0 100644 --- a/src/webgpu/compute/wavefront/ProcessHitsKernel.js +++ b/src/webgpu/compute/wavefront/ProcessHitsKernel.js @@ -1,7 +1,7 @@ import { IndirectStorageBufferAttribute, StorageTexture } from 'three/webgpu'; import { ComputeKernel } from '../ComputeKernel.js'; import { uniform, storage, wgslFn, textureStore, globalId } from 'three/tsl'; -import { constants } from 'three-mesh-bvh/webgpu'; +import { constants } from '../../lib/wgsl/common.wgsl.js'; import { pcgRand3, pcgInit } from '../../nodes/random.wgsl.js'; import { lambertBsdfFunc } from '../../nodes/sampling.wgsl.js'; import { queuedRayStruct, queuedHitStruct, QUEUED_RAY_SIZE, QUEUED_HIT_SIZE } from './structs.js'; diff --git a/src/webgpu/compute/wavefront/RayGenerationKernel.js b/src/webgpu/compute/wavefront/RayGenerationKernel.js index 00ee6a873..696103f0d 100644 --- a/src/webgpu/compute/wavefront/RayGenerationKernel.js +++ b/src/webgpu/compute/wavefront/RayGenerationKernel.js @@ -2,7 +2,7 @@ import { Vector2, Matrix4 } from 'three'; import { IndirectStorageBufferAttribute, StorageTexture } from 'three/webgpu'; import { wgslFn, uniform, storage, globalId, textureStore } from 'three/tsl'; import { ComputeKernel } from '../ComputeKernel.js'; -import { ndcToCameraRay } from 'three-mesh-bvh/webgpu'; +import { ndcToCameraRay } from '../../lib/wgsl/common.wgsl.js'; import { pcgInit, pcgRand2 } from '../../nodes/random.wgsl.js'; import { QUEUED_RAY_SIZE, queuedRayStruct } from './structs.js'; diff --git a/src/webgpu/compute/wavefront/RayIntersectionKernel.js b/src/webgpu/compute/wavefront/RayIntersectionKernel.js index 8b7d468d5..83f6cfe30 100644 --- a/src/webgpu/compute/wavefront/RayIntersectionKernel.js +++ b/src/webgpu/compute/wavefront/RayIntersectionKernel.js @@ -1,7 +1,7 @@ import { IndirectStorageBufferAttribute, StorageTexture } from 'three/webgpu'; import { ComputeKernel } from '../ComputeKernel.js'; import { storage, wgslFn, textureStore, globalId } from 'three/tsl'; -import { constants } from 'three-mesh-bvh/webgpu'; +import { constants } from '../../lib/wgsl/common.wgsl.js'; import { pcgRand3, pcgInit } from '../../nodes/random.wgsl.js'; import { queuedRayStruct, queuedHitStruct, QUEUED_RAY_SIZE, QUEUED_HIT_SIZE } from './structs.js'; import { proxy } from '../../lib/nodes/NodeProxy.js'; diff --git a/src/webgpu/compute/wavefront/structs.js b/src/webgpu/compute/wavefront/structs.js index f71dfdd98..06e69ef2f 100644 --- a/src/webgpu/compute/wavefront/structs.js +++ b/src/webgpu/compute/wavefront/structs.js @@ -1,5 +1,5 @@ import { wgsl } from 'three/tsl'; -import { rayStruct } from 'three-mesh-bvh/webgpu'; +import { rayStruct } from '../../lib/wgsl/structs.wgsl.js'; export const QUEUED_RAY_SIZE = 16; diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index 0ea6a1840..8848f9413 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -3,10 +3,12 @@ import { StorageBufferAttribute, StructTypeNode } from 'three/webgpu'; import { storage } from 'three/tsl'; import { intersectsBounds, + constants, +} from './wgsl/common.wgsl.js'; +import { rayStruct, bvhNodeStruct, - constants, -} from 'three-mesh-bvh/webgpu'; +} from './wgsl/structs.wgsl.js'; import { wgslTagFn } from './nodes/WGSLTagFnNode.js'; const BYTES_PER_NODE = 6 * 4 + 4 + 4; diff --git a/src/webgpu/lib/wgsl/common.wgsl.js b/src/webgpu/lib/wgsl/common.wgsl.js new file mode 100644 index 000000000..aed172934 --- /dev/null +++ b/src/webgpu/lib/wgsl/common.wgsl.js @@ -0,0 +1,70 @@ +import { wgslFn, wgsl } from 'three/tsl'; +import { bvhNodeBoundsStruct, rayStruct } from './structs.wgsl.js'; + +export const constants = wgsl( /* wgsl */` + + const BVH_STACK_DEPTH = 60u; + const INFINITY = 1e20; + const TRI_INTERSECT_EPSILON = 1e-5; + +` ); + +export const ndcToCameraRay = wgslFn( /* wgsl*/` + + fn ndcToCameraRay( ndc: vec2f, inverseModelViewProjection: mat4x4f ) -> Ray { + + // Calculate the ray by picking the points at the near and far plane and deriving the ray + // direction from the two points. This approach works for both orthographic and perspective + // camera projection matrices. + // The returned ray direction is not normalized and extends to the camera far plane. + var homogeneous = vec4f(); + var ray = Ray(); + + homogeneous = inverseModelViewProjection * vec4f( ndc, 0.0, 1.0 ); + ray.origin = homogeneous.xyz / homogeneous.w; + + homogeneous = inverseModelViewProjection * vec4f( ndc, 1.0, 1.0 ); + ray.direction = ( homogeneous.xyz / homogeneous.w ) - ray.origin; + + return ray; + + } +`, [ rayStruct ] ); + +export const intersectsBounds = wgslFn( /* wgsl */` + + fn intersectsBounds( + ray: Ray, + bounds: BVHBoundingBox, + dist: ptr + ) -> bool { + + let boundsMin = vec3( bounds.min[0], bounds.min[1], bounds.min[2] ); + let boundsMax = vec3( bounds.max[0], bounds.max[1], bounds.max[2] ); + + let invDir = 1.0 / ray.direction; + let tMinPlane = ( boundsMin - ray.origin ) * invDir; + let tMaxPlane = ( boundsMax - ray.origin ) * invDir; + + let tMinHit = vec3f( + min( tMinPlane.x, tMaxPlane.x ), + min( tMinPlane.y, tMaxPlane.y ), + min( tMinPlane.z, tMaxPlane.z ) + ); + + let tMaxHit = vec3f( + max( tMinPlane.x, tMaxPlane.x ), + max( tMinPlane.y, tMaxPlane.y ), + max( tMinPlane.z, tMaxPlane.z ) + ); + + let t0 = max( max( tMinHit.x, tMinHit.y ), tMinHit.z ); + let t1 = min( min( tMaxHit.x, tMaxHit.y ), tMaxHit.z ); + + ( *dist ) = max( t0, 0.0 ); + + return t1 >= ( *dist ); + + } + +`, [ rayStruct, bvhNodeBoundsStruct ] ); diff --git a/src/webgpu/lib/wgsl/structs.wgsl.js b/src/webgpu/lib/wgsl/structs.wgsl.js new file mode 100644 index 000000000..040f36be3 --- /dev/null +++ b/src/webgpu/lib/wgsl/structs.wgsl.js @@ -0,0 +1,28 @@ +import { StructTypeNode } from 'three/webgpu'; + +export const rayStruct = new StructTypeNode( { + origin: 'vec3f', + direction: 'vec3f', +}, 'Ray' ); + +export const bvhNodeBoundsStruct = new StructTypeNode( { + min: 'array', + max: 'array', +}, 'BVHBoundingBox' ); +bvhNodeBoundsStruct.getLength = () => 6; + +export const bvhNodeStruct = new StructTypeNode( { + bounds: 'BVHBoundingBox', + rightChildOrTriangleOffset: 'uint', + splitAxisOrTriangleCount: 'uint', +}, 'BVHNode' ); +bvhNodeStruct.getLength = () => bvhNodeBoundsStruct.getLength() + 2; + +export const intersectionResultStruct = new StructTypeNode( { + didHit: 'bool', + indices: 'vec4u', + normal: 'vec3f', + barycoord: 'vec3f', + side: 'float', + dist: 'float', +}, 'IntersectionResult' ); From fb71782935bcb37ce02e0634ae4c8b22c6144742 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 14:21:32 +0900 Subject: [PATCH 67/84] Simplify constants --- .../compute/wavefront/ProcessHitsKernel.js | 2 +- .../compute/wavefront/RayIntersectionKernel.js | 3 +-- src/webgpu/lib/BVHComputeData.js | 18 ++++++++++-------- src/webgpu/lib/wgsl/common.wgsl.js | 14 ++++++-------- src/webgpu/nodes/structs.wgsl.js | 2 +- 5 files changed, 19 insertions(+), 20 deletions(-) diff --git a/src/webgpu/compute/wavefront/ProcessHitsKernel.js b/src/webgpu/compute/wavefront/ProcessHitsKernel.js index 6b798bde0..05da4668c 100644 --- a/src/webgpu/compute/wavefront/ProcessHitsKernel.js +++ b/src/webgpu/compute/wavefront/ProcessHitsKernel.js @@ -109,7 +109,7 @@ export class ProcessHitsKernel extends ComputeKernel { proxy( 'bvhData.value.storage.materials', parameters ), proxy( 'bvhData.value.storage.transforms', parameters ), proxy( 'bvhData.value.fns.sampleTrianglePoint', parameters ), - queuedRayStruct, lambertBsdfFunc, constants, + queuedRayStruct, lambertBsdfFunc, pcgRand3, pcgInit, queuedHitStruct, ] ); diff --git a/src/webgpu/compute/wavefront/RayIntersectionKernel.js b/src/webgpu/compute/wavefront/RayIntersectionKernel.js index 83f6cfe30..fc0cbd6b9 100644 --- a/src/webgpu/compute/wavefront/RayIntersectionKernel.js +++ b/src/webgpu/compute/wavefront/RayIntersectionKernel.js @@ -1,7 +1,6 @@ import { IndirectStorageBufferAttribute, StorageTexture } from 'three/webgpu'; import { ComputeKernel } from '../ComputeKernel.js'; import { storage, wgslFn, textureStore, globalId } from 'three/tsl'; -import { constants } from '../../lib/wgsl/common.wgsl.js'; import { pcgRand3, pcgInit } from '../../nodes/random.wgsl.js'; import { queuedRayStruct, queuedHitStruct, QUEUED_RAY_SIZE, QUEUED_HIT_SIZE } from './structs.js'; import { proxy } from '../../lib/nodes/NodeProxy.js'; @@ -96,7 +95,7 @@ export class RayIntersectionKernel extends ComputeKernel { `, [ proxy( 'bvhData.value.fns.raycastFirstHit', parameters ), proxy( 'bvhData.value.structs.material', parameters ), - queuedRayStruct, constants, pcgRand3, pcgInit, queuedHitStruct ] ); + queuedRayStruct, pcgRand3, pcgInit, queuedHitStruct ] ); super( fn( parameters ) ); diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index 8848f9413..f6874d5d3 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -79,11 +79,12 @@ const intersectionResultStruct = new StructTypeNode( { const intersectsTriangle = wgslTagFn/* wgsl */ ` // includes - ${ [ rayStruct, constants ] } + ${ [ rayStruct ] } // fn fn intersectsTriangle( ray: Ray, a: vec3f, b: vec3f, c: vec3f ) -> ${ intersectionResultStruct } { + var TRI_INTERSECT_EPSILON = ${ constants.TRI_INTERSECT_EPSILON }; var result: ${ intersectionResultStruct }; result.didHit = false; @@ -132,9 +133,10 @@ function buildRaycastFirstHitFn( prefix, storage, structs ) { // TODO: reduce the redundancy between these functions - possibly using code snippets or // macro-expansion-style mechanisms? + const { BVH_STACK_DEPTH, INFINITY } = constants; const geometryRaycastFirstHitFn = wgslTagFn/* wgsl */` // includes - ${ [ rayStruct, bvhNodeStruct, constants, structs.attributes ] } + ${ [ rayStruct, bvhNodeStruct, structs.attributes ] } // fn fn ${ prefix }RaycastFirstHit_blas( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> ${ intersectionResultStruct } { @@ -144,12 +146,12 @@ function buildRaycastFirstHitFn( prefix, storage, structs ) { bestHit.dist = bestDist; var pointer: i32 = 0; - var stack: array; + var stack: array; stack[ 0 ] = rootNodeIndex; loop { - if ( pointer < 0 || pointer >= i32( BVH_STACK_DEPTH ) ) { + if ( pointer < 0 || pointer >= i32( ${ BVH_STACK_DEPTH } ) ) { break; @@ -230,15 +232,15 @@ function buildRaycastFirstHitFn( prefix, storage, structs ) { var bestHit: ${ intersectionResultStruct }; bestHit.didHit = false; - bestHit.dist = INFINITY; + bestHit.dist = ${ INFINITY }; var tlasPointer: i32 = 0; - var tlasStack: array; + var tlasStack: array; tlasStack[ 0 ] = 0u; loop { - if ( tlasPointer < 0 || tlasPointer >= i32( BVH_STACK_DEPTH ) ) { + if ( tlasPointer < 0 || tlasPointer >= i32( ${ BVH_STACK_DEPTH } ) ) { break; @@ -449,7 +451,7 @@ export class BVHComputeData { // // set up the storage buffers - const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), 'BVHNode' ).toReadOnly().setName( `${ prefix }nodes` ); + const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), bvhNodeStruct.name ).toReadOnly().setName( `${ prefix }nodes` ); const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), structs.transform.getLength() ), structs.transform.name ).toReadOnly().setName( `${ prefix }transforms` ); const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ prefix }index` ); const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributeStruct.getLength() ), attributeStruct.name ).toReadOnly().setName( `${ prefix }attributes` ); diff --git a/src/webgpu/lib/wgsl/common.wgsl.js b/src/webgpu/lib/wgsl/common.wgsl.js index aed172934..e5003f16b 100644 --- a/src/webgpu/lib/wgsl/common.wgsl.js +++ b/src/webgpu/lib/wgsl/common.wgsl.js @@ -1,13 +1,11 @@ -import { wgslFn, wgsl } from 'three/tsl'; +import { wgslFn, uint, float } from 'three/tsl'; import { bvhNodeBoundsStruct, rayStruct } from './structs.wgsl.js'; -export const constants = wgsl( /* wgsl */` - - const BVH_STACK_DEPTH = 60u; - const INFINITY = 1e20; - const TRI_INTERSECT_EPSILON = 1e-5; - -` ); +export const constants = { + BVH_STACK_DEPTH: uint( 60 ), + INFINITY: float( 1e20 ), + TRI_INTERSECT_EPSILON: float( 1e-5 ), +}; export const ndcToCameraRay = wgslFn( /* wgsl*/` diff --git a/src/webgpu/nodes/structs.wgsl.js b/src/webgpu/nodes/structs.wgsl.js index 97cb1ba17..8d872b772 100644 --- a/src/webgpu/nodes/structs.wgsl.js +++ b/src/webgpu/nodes/structs.wgsl.js @@ -2,7 +2,7 @@ import { wgsl } from 'three/tsl'; import { rayStruct } from 'three-mesh-bvh/webgpu'; export const constants = wgsl( /* wgsl */ ` - const PI: f32 = 3.141592653589793; + const PI: f32 = 3.141592653589793; ` ); export const scatterRecordStruct = wgsl( /* wgsl */ ` From 25480cd104fd576a9e8907e36923ff22f9954bef Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 14:22:42 +0900 Subject: [PATCH 68/84] Simplify ray struct usage --- src/webgpu/lib/BVHComputeData.js | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index f6874d5d3..bb47b5dd1 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -78,11 +78,8 @@ const intersectionResultStruct = new StructTypeNode( { }, 'IntersectionResult' ); const intersectsTriangle = wgslTagFn/* wgsl */ ` - // includes - ${ [ rayStruct ] } - // fn - fn intersectsTriangle( ray: Ray, a: vec3f, b: vec3f, c: vec3f ) -> ${ intersectionResultStruct } { + fn intersectsTriangle( ray: ${ rayStruct }, a: vec3f, b: vec3f, c: vec3f ) -> ${ intersectionResultStruct } { var TRI_INTERSECT_EPSILON = ${ constants.TRI_INTERSECT_EPSILON }; var result: ${ intersectionResultStruct }; @@ -136,10 +133,10 @@ function buildRaycastFirstHitFn( prefix, storage, structs ) { const { BVH_STACK_DEPTH, INFINITY } = constants; const geometryRaycastFirstHitFn = wgslTagFn/* wgsl */` // includes - ${ [ rayStruct, bvhNodeStruct, structs.attributes ] } + ${ [ bvhNodeStruct, structs.attributes ] } // fn - fn ${ prefix }RaycastFirstHit_blas( ray: Ray, rootNodeIndex: u32, bestDist: f32 ) -> ${ intersectionResultStruct } { + fn ${ prefix }RaycastFirstHit_blas( ray: ${ rayStruct }, rootNodeIndex: u32, bestDist: f32 ) -> ${ intersectionResultStruct } { var bestHit: ${ intersectionResultStruct }; bestHit.didHit = false; From 228fbf37fb169c5d7d29ab956bbb2cb13780bc73 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Sun, 22 Feb 2026 14:27:30 +0900 Subject: [PATCH 69/84] Fix bvh compute data transform struct reference --- src/webgpu/lib/BVHComputeData.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index bb47b5dd1..c194b6fa2 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -222,7 +222,7 @@ function buildRaycastFirstHitFn( prefix, storage, structs ) { return wgslTagFn/* wgsl */` // includes - ${ [ rayStruct, bvhNodeStruct, constants, transformStruct ] } + ${ [ rayStruct, bvhNodeStruct, constants, structs.transform ] } // fn fn ${ prefix }RaycastFirstHit( ray: Ray ) -> ${ intersectionResultStruct } { From c7bdc22584737dcc0c4330a83d2b8cd61115bc71 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Mon, 23 Feb 2026 12:15:33 +0900 Subject: [PATCH 70/84] Add wgslTagCode --- src/webgpu/lib/nodes/WGSLTagFnNode.js | 150 +++++++++++++++++--------- 1 file changed, 100 insertions(+), 50 deletions(-) diff --git a/src/webgpu/lib/nodes/WGSLTagFnNode.js b/src/webgpu/lib/nodes/WGSLTagFnNode.js index a647b438a..519873c3b 100644 --- a/src/webgpu/lib/nodes/WGSLTagFnNode.js +++ b/src/webgpu/lib/nodes/WGSLTagFnNode.js @@ -1,4 +1,4 @@ -import { FunctionNode, Node } from 'three/webgpu'; +import { CodeNode, FunctionNode, Node } from 'three/webgpu'; // minimal node that outputs a raw WGSL expression verbatim when built class LiteralExpression extends Node { @@ -77,55 +77,97 @@ function getIncludeNode( arg ) { } -export class WGSLTagFnNode extends FunctionNode { +// extract dependency nodes from template args for include registration +function extractIncludes( args ) { - static get type() { + const includes = []; + for ( const arg of args ) { - return 'WGSLFnTagNode'; + if ( Array.isArray( arg ) ) { + + for ( const element of arg ) { + + const node = getIncludeNode( element ); + if ( node ) includes.push( node ); + + } + + } else { + + const node = getIncludeNode( arg ); + if ( node ) includes.push( node ); + + } } - constructor( tokens, args, lang = 'wgsl' ) { + return includes; + +} + +// normalize args so generate can resolve them uniformly with build(): +// - callable wrappers > PropertyRefNode (emits just the function name) +// - struct callables > StructTypeNode (emits the type name via build) +// - FunctionCallNodes > InlineCallNode (emits inline call) +function normalizeArgs( args ) { - // assemble all the nodes needed for includes - const includes = []; - for ( const arg of args ) { + return args.map( arg => { + if ( typeof arg === 'function' && arg.functionNode ) return new PropertyRefNode( arg.functionNode ); + if ( typeof arg === 'function' && arg.isStruct ) return arg.layout; + if ( arg && arg.isNode && arg.functionNode ) return new InlineCallNode( arg ); + return arg; + + } ); + +} + +// interleave static tokens with resolved arg values +function assembleTemplate( tokens, args, builder ) { + + let code = ''; + for ( let i = 0, l = tokens.length; i < l; i ++ ) { + + code += tokens[ i ]; + if ( i < args.length ) { + + const arg = args[ i ]; if ( Array.isArray( arg ) ) { - for ( const element of arg ) { + // include array — no text output - const node = getIncludeNode( element ); - if ( node ) includes.push( node ); + } else if ( typeof arg === 'string' || typeof arg === 'number' ) { - } + code += String( arg ); } else { - const node = getIncludeNode( arg ); - if ( node ) includes.push( node ); + code += arg.build( builder ); } } - // normalize args so generate function can resolve them with build() later: - // - callable wrappers > PropertyRefNode (emits just the function name) - // - struct callables > StructTypeNode (emits the type name via build) - // - FunctionCallNodes > InlineCallNode (emits inline call) - const normalizedArgs = args.map( arg => { + } + + return code; + +} + +export class WGSLTagFnNode extends FunctionNode { + + static get type() { + + return 'WGSLFnTagNode'; - if ( typeof arg === 'function' && arg.functionNode ) return new PropertyRefNode( arg.functionNode ); - if ( typeof arg === 'function' && arg.isStruct ) return arg.layout; - if ( arg && arg.isNode && arg.functionNode ) return new InlineCallNode( arg ); - return arg; + } - } ); + constructor( tokens, args, lang = 'wgsl' ) { - super( '', includes, lang ); + super( '', extractIncludes( args ), lang ); this.tokens = tokens; - this.args = normalizedArgs; + this.args = normalizeArgs( args ); } @@ -192,49 +234,54 @@ export class WGSLTagFnNode extends FunctionNode { // get the code for the function generate( builder, output ) { - const { tokens, args } = this; - - // rebuild the function body again because we can call "build", now const result = super.generate( builder, output ); - let fullCode = ''; - for ( let i = 0, l = tokens.length; i < l; i ++ ) { + const fullCode = assembleTemplate( this.tokens, this.args, builder ); - fullCode += tokens[ i ]; - if ( i < args.length ) { + const { type } = this.getNodeFunction( builder ); + const nodeCode = builder.getCodeFromNode( this, type ); - const arg = args[ i ]; - if ( Array.isArray( arg ) ) { + nodeCode.code = fullCode.replace( /\/\/.+[\n\r]/g, '' ).replace( /->\s*void/, '' ).replace( /\s+/g, ' ' ).trim(); - // include array — no text output + return result; - } else if ( typeof arg === 'string' || typeof arg === 'number' ) { + } - fullCode += String( arg ); +} - } else { +export class WGSLTagCodeNode extends CodeNode { - fullCode += arg.build( builder ); + static get type() { - } + return 'WGSLTagCodeNode'; - } + } - } + constructor( tokens, args, lang = 'wgsl' ) { - const { type } = this.getNodeFunction( builder ); - const nodeCode = builder.getCodeFromNode( this, type ); + super( '', extractIncludes( args ), lang ); - fullCode = fullCode.replace( /\/\/.+[\n\r]/g, '' ).replace( /->\s*void/, '' ).replace( /\s+/g, ' ' ).trim(); - nodeCode.code = fullCode; + this.tokens = tokens; + this.args = normalizeArgs( args ); - return result; + } + + generate( builder ) { + + // build includes so dependencies are registered before the parent code block + const includes = this.getIncludes( builder ); + for ( const include of includes ) { + + include.build( builder ); + + } + + return assembleTemplate( this.tokens, this.args, builder ); } } -// template tag literal function version of "wgslFn" so easy interpolation of TSL nodes -// TODO: add a raw "wgsl" version for code snippets +// template tag literal function version of "wgslFn" for easy interpolation of TSL nodes export const wgslTagFn = ( tokens, ...args ) => { const functionNode = new WGSLTagFnNode( tokens, args ); @@ -267,3 +314,6 @@ export const wgslTagFn = ( tokens, ...args ) => { return fn; }; + +// template tag literal for reusable WGSL code snippets with dependency resolution +export const wgslTagCode = ( tokens, ...args ) => new WGSLTagCodeNode( tokens, args ); From 513597684b56b111a300acd2ae34d2d00cf36928 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Mon, 23 Feb 2026 12:28:15 +0900 Subject: [PATCH 71/84] Deduplicate code --- src/webgpu/lib/BVHComputeData.js | 145 ++++++++++++------------------- 1 file changed, 55 insertions(+), 90 deletions(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index c194b6fa2..8a6f7bbb4 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -9,7 +9,7 @@ import { rayStruct, bvhNodeStruct, } from './wgsl/structs.wgsl.js'; -import { wgslTagFn } from './nodes/WGSLTagFnNode.js'; +import { wgslTagCode, wgslTagFn } from './nodes/WGSLTagFnNode.js'; const BYTES_PER_NODE = 6 * 4 + 4 + 4; const UINT32_PER_NODE = BYTES_PER_NODE / 4; @@ -127,17 +127,11 @@ const intersectsTriangle = wgslTagFn/* wgsl */ ` function buildRaycastFirstHitFn( prefix, storage, structs ) { - // TODO: reduce the redundancy between these functions - possibly using code snippets or - // macro-expansion-style mechanisms? - const { BVH_STACK_DEPTH, INFINITY } = constants; - const geometryRaycastFirstHitFn = wgslTagFn/* wgsl */` - // includes - ${ [ bvhNodeStruct, structs.attributes ] } - - // fn - fn ${ prefix }RaycastFirstHit_blas( ray: ${ rayStruct }, rootNodeIndex: u32, bestDist: f32 ) -> ${ intersectionResultStruct } { + const getFnBody = leafSnippet => { + // returns a function with a snippet inserted for the leaf intersection test + return wgslTagCode/* wgsl */` var bestHit: ${ intersectionResultStruct }; bestHit.didHit = false; bestHit.dist = bestDist; @@ -171,29 +165,7 @@ function buildRaycastFirstHitFn( prefix, storage, structs ) { if ( isLeaf ) { - let triCount = infoX & 0x0000ffffu; - let triOffset = infoY; - - for ( var ti = triOffset; ti < triOffset + triCount; ti = ti + 1u ) { - - let i0 = ${ storage.index }[ ti * 3u ]; - let i1 = ${ storage.index }[ ti * 3u + 1u ]; - let i2 = ${ storage.index }[ ti * 3u + 2u ]; - - let a = ${ storage.attributes }[ i0 ].position.xyz; - let b = ${ storage.attributes }[ i1 ].position.xyz; - let c = ${ storage.attributes }[ i2 ].position.xyz; - - var triResult = ${ intersectsTriangle }( ray, a, b, c ); - - if ( triResult.didHit && triResult.dist < bestHit.dist ) { - - bestHit = triResult; - bestHit.indices = vec4u( i0, i1, i2, ti ); - - } - - } + ${ leafSnippet } } else { @@ -216,98 +188,91 @@ function buildRaycastFirstHitFn( prefix, storage, structs ) { } return bestHit; + `; - } - `; + }; - return wgslTagFn/* wgsl */` + const blasFn = wgslTagFn/* wgsl */` // includes - ${ [ rayStruct, bvhNodeStruct, constants, structs.transform ] } + ${ [ bvhNodeStruct, structs.attributes ] } // fn - fn ${ prefix }RaycastFirstHit( ray: Ray ) -> ${ intersectionResultStruct } { + fn ${ prefix }RaycastFirstHit_blas( ray: ${ rayStruct }, rootNodeIndex: u32, bestDist: f32 ) -> ${ intersectionResultStruct } { - var bestHit: ${ intersectionResultStruct }; - bestHit.didHit = false; - bestHit.dist = ${ INFINITY }; + ${ getFnBody( wgslTagCode/* wgsl */` - var tlasPointer: i32 = 0; - var tlasStack: array; - tlasStack[ 0 ] = 0u; + let triCount = infoX & 0x0000ffffu; + let triOffset = infoY; - loop { + for ( var ti = triOffset; ti < triOffset + triCount; ti = ti + 1u ) { - if ( tlasPointer < 0 || tlasPointer >= i32( ${ BVH_STACK_DEPTH } ) ) { + let i0 = ${ storage.index }[ ti * 3u ]; + let i1 = ${ storage.index }[ ti * 3u + 1u ]; + let i2 = ${ storage.index }[ ti * 3u + 2u ]; - break; + let a = ${ storage.attributes }[ i0 ].position.xyz; + let b = ${ storage.attributes }[ i1 ].position.xyz; + let c = ${ storage.attributes }[ i2 ].position.xyz; - } + var triResult = ${ intersectsTriangle }( ray, a, b, c ); - let currNodeIndex = tlasStack[ tlasPointer ]; - let node = ${ storage.nodes }[ currNodeIndex ]; - tlasPointer = tlasPointer - 1; + if ( triResult.didHit && triResult.dist < bestHit.dist ) { - var boundsHitDist: f32 = 0.0; - if ( ! ${ intersectsBounds }( ray, node.bounds, &boundsHitDist ) || boundsHitDist > bestHit.dist ) { + bestHit = triResult; + bestHit.indices = vec4u( i0, i1, i2, ti ); - continue; + } } - let infoX = node.splitAxisOrTriangleCount; - let infoY = node.rightChildOrTriangleOffset; - let isLeaf = ( infoX & 0xffff0000u ) != 0u; - - if ( isLeaf ) { - - let count = infoX & 0x0000ffffu; - let offset = infoY; + ` ) } - for ( var t = offset; t < offset + count; t = t + 1u ) { + } + `; - let transform = ${ storage.transforms }[ t ]; + const tlasFn = wgslTagFn/* wgsl */` + // includes + ${ [ rayStruct, bvhNodeStruct, constants, structs.transform ] } - // Transform ray into object local space - var localRay: Ray; - localRay.origin = ( transform.inverseMatrixWorld * vec4f( ray.origin, 1.0 ) ).xyz; - localRay.direction = ( transform.inverseMatrixWorld * vec4f( ray.direction, 0.0 ) ).xyz; + // fn + fn ${ prefix }RaycastFirstHit( ray: Ray ) -> ${ intersectionResultStruct } { - let blasHit = ${ geometryRaycastFirstHitFn( { ray: 'localRay', rootNodeIndex: 'transform.nodeOffset', bestDist: 'bestHit.dist' } ) }; - if ( blasHit.didHit && blasHit.dist < bestHit.dist ) { + let bestDist = ${ INFINITY }; + let rootNodeIndex = 0u; - bestHit = blasHit; - bestHit.objectIndex = t; + ${ getFnBody( wgslTagCode/* wgsl */` - // Transform normal to world space: normal matrix = transpose( inverse ) - bestHit.normal = normalize( ( transpose( transform.inverseMatrixWorld ) * vec4f( bestHit.normal, 0.0 ) ).xyz ); + let count = infoX & 0x0000ffffu; + let offset = infoY; - } + for ( var t = offset; t < offset + count; t = t + 1u ) { - } + let transform = ${ storage.transforms }[ t ]; - } else { + // Transform ray into object local space + var localRay: Ray; + localRay.origin = ( transform.inverseMatrixWorld * vec4f( ray.origin, 1.0 ) ).xyz; + localRay.direction = ( transform.inverseMatrixWorld * vec4f( ray.direction, 0.0 ) ).xyz; - let leftIndex = currNodeIndex + 1u; - let splitAxis = infoX & 0x0000ffffu; - let rightIndex = currNodeIndex + infoY; + let blasHit = ${ blasFn( { ray: 'localRay', rootNodeIndex: 'transform.nodeOffset', bestDist: 'bestHit.dist' } ) }; + if ( blasHit.didHit && blasHit.dist < bestHit.dist ) { - let leftToRight = ray.direction[ splitAxis ] >= 0.0; - let c1 = select( rightIndex, leftIndex, leftToRight ); - let c2 = select( leftIndex, rightIndex, leftToRight ); + bestHit = blasHit; + bestHit.objectIndex = t; - tlasPointer = tlasPointer + 1; - tlasStack[ tlasPointer ] = c2; + // Transform normal to world space: normal matrix = transpose( inverse ) + bestHit.normal = normalize( ( transpose( transform.inverseMatrixWorld ) * vec4f( bestHit.normal, 0.0 ) ).xyz ); - tlasPointer = tlasPointer + 1; - tlasStack[ tlasPointer ] = c1; + } } - } + ` ) } - return bestHit; + } + `; - }`; + return tlasFn; } From a9db555eebb66ac50d427acf765f4d50bac9667f Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Mon, 23 Feb 2026 12:37:39 +0900 Subject: [PATCH 72/84] Add shim for allowing structs to be passed to storage nodes --- src/webgpu/lib/BVHComputeData.js | 82 ++++++++++++++++++-------------- 1 file changed, 47 insertions(+), 35 deletions(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index 8a6f7bbb4..05a4df0b8 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -11,10 +11,6 @@ import { } from './wgsl/structs.wgsl.js'; import { wgslTagCode, wgslTagFn } from './nodes/WGSLTagFnNode.js'; -const BYTES_PER_NODE = 6 * 4 + 4 + 4; -const UINT32_PER_NODE = BYTES_PER_NODE / 4; -const IS_LEAFNODE_FLAG = 0xFFFF; - // TODO: add ability to easily update a single matrix / scene rearrangement (partial update) // TODO: add material support w/ function to easily update material // - add a callback for writing a property for a geometry to a range @@ -25,11 +21,54 @@ const IS_LEAFNODE_FLAG = 0xFFFF; // TODO: allow for "slotting" a new type of callback (eg distance, etc) so multiple types of queries can be made // - add a "shapecast" style function with functions and return types that can be slotted in +// temporary shim so StructTypeNodes can be passed to storage functions +Object.defineProperty( StructTypeNode.prototype, 'layout', { + + get() { + + return this; + + } + +} ); +StructTypeNode.prototype.isStruct = true; + +// + +// structs +const transformStruct = new StructTypeNode( { + matrixWorld: 'mat4x4f', + inverseMatrixWorld: 'mat4x4f', + nodeOffset: 'uint', + _alignment0: 'uint', + _alignment1: 'uint', + _alignment2: 'uint', +}, 'TransformStruct' ); + +const intersectionResultStruct = new StructTypeNode( { + indices: 'vec4u', + normal: 'vec3f', + didHit: 'bool', + barycoord: 'vec3f', + objectIndex: 'uint', + side: 'float', + dist: 'float', +}, 'IntersectionResult' ); + +// + +// node constants +const BYTES_PER_NODE = 6 * 4 + 4 + 4; +const UINT32_PER_NODE = BYTES_PER_NODE / 4; +const IS_LEAFNODE_FLAG = 0xFFFF; + +// scratch const _def = /* @__PURE__ */ new Vector4(); const _vec = /* @__PURE__ */ new Vector4(); const _matrix = /* @__PURE__ */ new Matrix4(); const _inverseMatrix = /* @__PURE__ */ new Matrix4(); +// functions function dereferenceIndex( indexAttr, indirectBuffer ) { const indexArray = indexAttr ? indexAttr.array : null; @@ -56,27 +95,6 @@ function getTotalBVHByteLength( bvh ) { } -// stride is 36 floats (144 bytes) to match WGSL struct alignment: -// mat4x4f (64) + mat4x4f (64) + u32 (4) + 12 bytes padding to align to 16 -const transformStruct = new StructTypeNode( { - matrixWorld: 'mat4x4f', - inverseMatrixWorld: 'mat4x4f', - nodeOffset: 'uint', - _alignment0: 'uint', - _alignment1: 'uint', - _alignment2: 'uint', -}, 'TransformStruct' ); - -const intersectionResultStruct = new StructTypeNode( { - indices: 'vec4u', - normal: 'vec3f', - didHit: 'bool', - barycoord: 'vec3f', - objectIndex: 'uint', - side: 'float', - dist: 'float', -}, 'IntersectionResult' ); - const intersectsTriangle = wgslTagFn/* wgsl */ ` // fn fn intersectsTriangle( ray: ${ rayStruct }, a: vec3f, b: vec3f, c: vec3f ) -> ${ intersectionResultStruct } { @@ -125,7 +143,7 @@ const intersectsTriangle = wgslTagFn/* wgsl */ ` } `; -function buildRaycastFirstHitFn( prefix, storage, structs ) { +function buildRaycastFirstHitFn( prefix, storage ) { const { BVH_STACK_DEPTH, INFINITY } = constants; const getFnBody = leafSnippet => { @@ -193,9 +211,6 @@ function buildRaycastFirstHitFn( prefix, storage, structs ) { }; const blasFn = wgslTagFn/* wgsl */` - // includes - ${ [ bvhNodeStruct, structs.attributes ] } - // fn fn ${ prefix }RaycastFirstHit_blas( ray: ${ rayStruct }, rootNodeIndex: u32, bestDist: f32 ) -> ${ intersectionResultStruct } { @@ -231,9 +246,6 @@ function buildRaycastFirstHitFn( prefix, storage, structs ) { `; const tlasFn = wgslTagFn/* wgsl */` - // includes - ${ [ rayStruct, bvhNodeStruct, constants, structs.transform ] } - // fn fn ${ prefix }RaycastFirstHit( ray: Ray ) -> ${ intersectionResultStruct } { @@ -413,10 +425,10 @@ export class BVHComputeData { // // set up the storage buffers - const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), bvhNodeStruct.name ).toReadOnly().setName( `${ prefix }nodes` ); - const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), structs.transform.getLength() ), structs.transform.name ).toReadOnly().setName( `${ prefix }transforms` ); + const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 8 ), bvhNodeStruct ).toReadOnly().setName( `${ prefix }nodes` ); + const transformsStorage = storage( new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), structs.transform.getLength() ), structs.transform ).toReadOnly().setName( `${ prefix }transforms` ); const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( `${ prefix }index` ); - const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributeStruct.getLength() ), attributeStruct.name ).toReadOnly().setName( `${ prefix }attributes` ); + const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributeStruct.getLength() ), attributeStruct ).toReadOnly().setName( `${ prefix }attributes` ); this.storage.transforms = transformsStorage; this.storage.nodes = bvhNodesStorage; From 6e0d3ae24e261d348ab96fdc3feefb3c854e65a5 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Mon, 23 Feb 2026 13:24:49 +0900 Subject: [PATCH 73/84] Add "getShapecastFn" impl --- src/webgpu/lib/BVHComputeData.js | 140 +++++++++++++++++++++++++++++-- 1 file changed, 135 insertions(+), 5 deletions(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index 05a4df0b8..a86ee9000 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -15,11 +15,6 @@ import { wgslTagCode, wgslTagFn } from './nodes/WGSLTagFnNode.js'; // TODO: add material support w/ function to easily update material // - add a callback for writing a property for a geometry to a range // TODO: add skinned mesh bvh support -// TODO: see if we can reference wgslFn names directly rather than constructing them inline over and over -// and / or use local variable definitions for the pointers to clean up the code -// TODO: see if there's a "build" step that can be leveraged for nodes to make integration more simple -// TODO: allow for "slotting" a new type of callback (eg distance, etc) so multiple types of queries can be made -// - add a "shapecast" style function with functions and return types that can be slotted in // temporary shim so StructTypeNodes can be passed to storage functions Object.defineProperty( StructTypeNode.prototype, 'layout', { @@ -319,6 +314,141 @@ export class BVHComputeData { } + getShapecastFn( options ) { + + const { + name, + shapeStruct, + resultStruct, + + boundsOrderFn, + intersectsBoundsFn, + intersectRangeFn, + transformShapeFn, + transformResultFn, + } = options; + + const { BVH_STACK_DEPTH, INFINITY } = constants; + const getFnBody = leafSnippet => { + + // returns a function with a snippet inserted for the leaf intersection test + return wgslTagCode/* wgsl */` + var bestHit: ${ resultStruct }; + bestHit.didHit = false; + bestHit.dist = bestDist; + + var pointer: i32 = 0; + var stack: array; + stack[ 0 ] = rootNodeIndex; + + loop { + + if ( pointer < 0 || pointer >= i32( ${ BVH_STACK_DEPTH } ) ) { + + break; + + } + + let nodeIndex = stack[ pointer ]; + let node = ${ storage.nodes }[ nodeIndex ]; + pointer = pointer - 1; + + var boundsHitDist: f32 = 0.0; + if ( ! ${ intersectsBoundsFn }( shape, node.bounds, &boundsHitDist ) || boundsHitDist > bestHit.dist ) { + + continue; + + } + + let infoX = node.splitAxisOrTriangleCount; + let infoY = node.rightChildOrTriangleOffset; + let isLeaf = ( infoX & 0xffff0000u ) != 0u; + + if ( isLeaf ) { + + let count = infoX & 0x0000ffffu; + let offset = infoY; + ${ leafSnippet } + + } else { + + let leftIndex = nodeIndex + 1u; + let splitAxis = infoX & 0x0000ffffu; + let rightIndex = nodeIndex + infoY; + + let leftToRight = ${ boundsOrderFn }( splitAxis, node ); + let c1 = select( rightIndex, leftIndex, leftToRight ); + let c2 = select( leftIndex, rightIndex, leftToRight ); + + pointer = pointer + 1; + stack[ pointer ] = c2; + + pointer = pointer + 1; + stack[ pointer ] = c1; + + } + + } + + return bestHit; + `; + + }; + + const blasFn = wgslTagFn/* wgsl */` + // fn + fn ${ name }_blas( shape: ${ shapeStruct }, rootNodeIndex: u32, bestDist: f32 ) -> ${ resultStruct } { + + ${ getFnBody( wgslTagCode/* wgsl */` + + let result = ${ intersectRangeFn( { offset: 'offset', count: 'count' } ) } + if ( result.didHit && result.dist < bestHit.dist ) { + + bestHit = result; + + } + + ` ) } + + } + `; + + const tlasFn = wgslTagFn/* wgsl */` + // fn + fn ${ name }( shape: ${ shapeStruct } ) -> ${ resultStruct } { + + let bestDist = ${ INFINITY }; + let rootNodeIndex = 0u; + + ${ getFnBody( wgslTagCode/* wgsl */` + + for ( var t = offset; t < offset + count; t = t + 1u ) { + + let transform = ${ storage.transforms }[ t ]; + + // Transform shape into object local space + let localShape = ${ transformShapeFn }( shape, transform ); + let blasHit = ${ blasFn( { shape: 'localShape', rootNodeIndex: 'transform.nodeOffset', bestDist: 'bestHit.dist' } ) }; + if ( blasHit.didHit && blasHit.dist < bestHit.dist ) { + + bestHit = blasHit; + bestHit.objectIndex = t; + + ${ transformResultFn }( &bestHit, transform.inverseMatrixWorld ); + + } + + } + + ` ) } + + } + `; + + return tlasFn; + + } + update() { const self = this; From fef5d6352916a93a8a86add6e5dd9a894ca41964 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Mon, 23 Feb 2026 14:12:12 +0900 Subject: [PATCH 74/84] Switch to use the shapecast creation function for raycast --- src/webgpu/lib/BVHComputeData.js | 283 +++++++++++------------------ src/webgpu/lib/wgsl/common.wgsl.js | 4 +- 2 files changed, 105 insertions(+), 182 deletions(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index a86ee9000..fd002063f 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -1,14 +1,8 @@ import { Matrix4, Vector4 } from 'three'; import { StorageBufferAttribute, StructTypeNode } from 'three/webgpu'; -import { storage } from 'three/tsl'; -import { - intersectsBounds, - constants, -} from './wgsl/common.wgsl.js'; -import { - rayStruct, - bvhNodeStruct, -} from './wgsl/structs.wgsl.js'; +import { storage, storageBarrier } from 'three/tsl'; +import { rayIntersectsBounds, constants } from './wgsl/common.wgsl.js'; +import { rayStruct, bvhNodeStruct } from './wgsl/structs.wgsl.js'; import { wgslTagCode, wgslTagFn } from './nodes/WGSLTagFnNode.js'; // TODO: add ability to easily update a single matrix / scene rearrangement (partial update) @@ -16,7 +10,8 @@ import { wgslTagCode, wgslTagFn } from './nodes/WGSLTagFnNode.js'; // - add a callback for writing a property for a geometry to a range // TODO: add skinned mesh bvh support -// temporary shim so StructTypeNodes can be passed to storage functions +// temporary shim so StructTypeNodes can be passed to storage functions until +// this is fixed in three.js Object.defineProperty( StructTypeNode.prototype, 'layout', { get() { @@ -138,151 +133,6 @@ const intersectsTriangle = wgslTagFn/* wgsl */ ` } `; -function buildRaycastFirstHitFn( prefix, storage ) { - - const { BVH_STACK_DEPTH, INFINITY } = constants; - const getFnBody = leafSnippet => { - - // returns a function with a snippet inserted for the leaf intersection test - return wgslTagCode/* wgsl */` - var bestHit: ${ intersectionResultStruct }; - bestHit.didHit = false; - bestHit.dist = bestDist; - - var pointer: i32 = 0; - var stack: array; - stack[ 0 ] = rootNodeIndex; - - loop { - - if ( pointer < 0 || pointer >= i32( ${ BVH_STACK_DEPTH } ) ) { - - break; - - } - - let nodeIndex = stack[ pointer ]; - let node = ${ storage.nodes }[ nodeIndex ]; - pointer = pointer - 1; - - var boundsHitDist: f32 = 0.0; - if ( ! ${ intersectsBounds }( ray, node.bounds, &boundsHitDist ) || boundsHitDist > bestHit.dist ) { - - continue; - - } - - let infoX = node.splitAxisOrTriangleCount; - let infoY = node.rightChildOrTriangleOffset; - let isLeaf = ( infoX & 0xffff0000u ) != 0u; - - if ( isLeaf ) { - - ${ leafSnippet } - - } else { - - let leftIndex = nodeIndex + 1u; - let splitAxis = infoX & 0x0000ffffu; - let rightIndex = nodeIndex + infoY; - - let leftToRight = ray.direction[ splitAxis ] >= 0.0; - let c1 = select( rightIndex, leftIndex, leftToRight ); - let c2 = select( leftIndex, rightIndex, leftToRight ); - - pointer = pointer + 1; - stack[ pointer ] = c2; - - pointer = pointer + 1; - stack[ pointer ] = c1; - - } - - } - - return bestHit; - `; - - }; - - const blasFn = wgslTagFn/* wgsl */` - // fn - fn ${ prefix }RaycastFirstHit_blas( ray: ${ rayStruct }, rootNodeIndex: u32, bestDist: f32 ) -> ${ intersectionResultStruct } { - - ${ getFnBody( wgslTagCode/* wgsl */` - - let triCount = infoX & 0x0000ffffu; - let triOffset = infoY; - - for ( var ti = triOffset; ti < triOffset + triCount; ti = ti + 1u ) { - - let i0 = ${ storage.index }[ ti * 3u ]; - let i1 = ${ storage.index }[ ti * 3u + 1u ]; - let i2 = ${ storage.index }[ ti * 3u + 2u ]; - - let a = ${ storage.attributes }[ i0 ].position.xyz; - let b = ${ storage.attributes }[ i1 ].position.xyz; - let c = ${ storage.attributes }[ i2 ].position.xyz; - - var triResult = ${ intersectsTriangle }( ray, a, b, c ); - - if ( triResult.didHit && triResult.dist < bestHit.dist ) { - - bestHit = triResult; - bestHit.indices = vec4u( i0, i1, i2, ti ); - - } - - } - - ` ) } - - } - `; - - const tlasFn = wgslTagFn/* wgsl */` - // fn - fn ${ prefix }RaycastFirstHit( ray: Ray ) -> ${ intersectionResultStruct } { - - let bestDist = ${ INFINITY }; - let rootNodeIndex = 0u; - - ${ getFnBody( wgslTagCode/* wgsl */` - - let count = infoX & 0x0000ffffu; - let offset = infoY; - - for ( var t = offset; t < offset + count; t = t + 1u ) { - - let transform = ${ storage.transforms }[ t ]; - - // Transform ray into object local space - var localRay: Ray; - localRay.origin = ( transform.inverseMatrixWorld * vec4f( ray.origin, 1.0 ) ).xyz; - localRay.direction = ( transform.inverseMatrixWorld * vec4f( ray.direction, 0.0 ) ).xyz; - - let blasHit = ${ blasFn( { ray: 'localRay', rootNodeIndex: 'transform.nodeOffset', bestDist: 'bestHit.dist' } ) }; - if ( blasHit.didHit && blasHit.dist < bestHit.dist ) { - - bestHit = blasHit; - bestHit.objectIndex = t; - - // Transform normal to world space: normal matrix = transpose( inverse ) - bestHit.normal = normalize( ( transpose( transform.inverseMatrixWorld ) * vec4f( bestHit.normal, 0.0 ) ).xyz ); - - } - - } - - ` ) } - - } - `; - - return tlasFn; - -} - export class BVHComputeData { constructor( bvh, options = {} ) { @@ -328,6 +178,7 @@ export class BVHComputeData { transformResultFn, } = options; + const { storage } = this; const { BVH_STACK_DEPTH, INFINITY } = constants; const getFnBody = leafSnippet => { @@ -376,7 +227,7 @@ export class BVHComputeData { let splitAxis = infoX & 0x0000ffffu; let rightIndex = nodeIndex + infoY; - let leftToRight = ${ boundsOrderFn }( splitAxis, node ); + let leftToRight = ${ boundsOrderFn }( shape, splitAxis, node ); let c1 = select( rightIndex, leftIndex, leftToRight ); let c2 = select( leftIndex, rightIndex, leftToRight ); @@ -401,7 +252,7 @@ export class BVHComputeData { ${ getFnBody( wgslTagCode/* wgsl */` - let result = ${ intersectRangeFn( { offset: 'offset', count: 'count' } ) } + let result = ${ intersectRangeFn }( shape, offset, count, bestDist ); if ( result.didHit && result.dist < bestHit.dist ) { bestHit = result; @@ -427,14 +278,14 @@ export class BVHComputeData { let transform = ${ storage.transforms }[ t ]; // Transform shape into object local space - let localShape = ${ transformShapeFn }( shape, transform ); + let localShape = ${ transformShapeFn }( shape, transform.inverseMatrixWorld ); let blasHit = ${ blasFn( { shape: 'localShape', rootNodeIndex: 'transform.nodeOffset', bestDist: 'bestHit.dist' } ) }; if ( blasHit.didHit && blasHit.dist < bestHit.dist ) { bestHit = blasHit; bestHit.objectIndex = t; - ${ transformResultFn }( &bestHit, transform.inverseMatrixWorld ); + ${ transformResultFn }( &bestHit, transform.matrixWorld, transform.inverseMatrixWorld ); } @@ -565,28 +416,8 @@ export class BVHComputeData { this.storage.index = indexStorage; this.storage.attributes = attributesStorage; this.structs.attributes = attributeStruct; - this.fns.raycastFirstHit = buildRaycastFirstHitFn( prefix, this.storage, this.structs ); - - const interpolateBody = attributeStruct - .membersLayout - .map( ( { name } ) => { - - return `result.${ name } = a0.${ name } * barycoord.x + a1.${ name } * barycoord.y + a2.${ name } * barycoord.z;`; - - } ).join( '\n' ); - this.fns.sampleTrianglePoint = wgslTagFn/* wgsl */` - // fn - fn ${ prefix }sampleTrianglePoint( barycoord: vec3f, indices: vec3u ) -> ${ attributeStruct } { - var result: ${ attributeStruct }; - var a0 = ${ attributesStorage }[ indices.x ]; - var a1 = ${ attributesStorage }[ indices.y ]; - var a2 = ${ attributesStorage }[ indices.z ]; - ${ interpolateBody } - return result; - - } - `; + this._initFns(); function appendBVHData( bvh, geometryOffset, transformInfo, nodeWriteOffset, target, tlas = false ) { @@ -763,6 +594,98 @@ export class BVHComputeData { } + _initFns() { + + const { storage, structs, fns, prefix } = this; + + // raycast first hit + fns.raycastFirstHit = this.getShapecastFn( { + name: prefix + 'RaycastFirstHit', + shapeStruct: rayStruct, + resultStruct: intersectionResultStruct, + + boundsOrderFn: wgslTagFn/* wgsl */` + fn getBoundsOrder( ray: ${ rayStruct }, splitAxis: u32, node: ${ bvhNodeStruct } ) -> bool { + + return ray.direction[ splitAxis ] >= 0.0; + + } + `, + intersectsBoundsFn: rayIntersectsBounds, + intersectRangeFn: wgslTagFn/* wgsl */` + fn intersectRange( ray: ${ rayStruct }, offset: u32, count: u32, bestDist: f32 ) -> ${ intersectionResultStruct } { + + var bestHit: ${ intersectionResultStruct }; + bestHit.didHit = false; + bestHit.dist = bestDist; + + for ( var ti = offset; ti < offset + count; ti = ti + 1u ) { + + let i0 = ${ storage.index }[ ti * 3u ]; + let i1 = ${ storage.index }[ ti * 3u + 1u ]; + let i2 = ${ storage.index }[ ti * 3u + 2u ]; + + let a = ${ storage.attributes }[ i0 ].position.xyz; + let b = ${ storage.attributes }[ i1 ].position.xyz; + let c = ${ storage.attributes }[ i2 ].position.xyz; + + var triResult = ${ intersectsTriangle }( ray, a, b, c ); + if ( triResult.didHit && triResult.dist < bestHit.dist ) { + + bestHit = triResult; + bestHit.indices = vec4u( i0, i1, i2, ti ); + + } + + } + + return bestHit; + + } + `, + transformShapeFn: wgslTagFn/* wgsl */` + fn transformRay( ray: ${ rayStruct }, toLocal: mat4x4f ) -> ${ rayStruct } { + + var localRay: Ray; + localRay.origin = ( toLocal * vec4f( ray.origin, 1.0 ) ).xyz; + localRay.direction = ( toLocal * vec4f( ray.direction, 0.0 ) ).xyz; + return localRay; + + } + `, + transformResultFn: wgslTagFn/* wgsl */` + fn transformResult( hit: ptr, toWorld: mat4x4f, toLocal: mat4x4f ) -> void { + + hit.normal = normalize( ( transpose( toLocal ) * vec4f( hit.normal, 0.0 ) ).xyz ); + + } + `, + } ); + + const interpolateBody = structs + .attributes + .membersLayout + .map( ( { name } ) => { + + return `result.${ name } = a0.${ name } * barycoord.x + a1.${ name } * barycoord.y + a2.${ name } * barycoord.z;`; + + } ).join( '\n' ); + fns.sampleTrianglePoint = wgslTagFn/* wgsl */` + // fn + fn ${ prefix }sampleTrianglePoint( barycoord: vec3f, indices: vec3u ) -> ${ structs.attributes } { + + var result: ${ structs.attributes }; + var a0 = ${ storage.attributes }[ indices.x ]; + var a1 = ${ storage.attributes }[ indices.y ]; + var a2 = ${ storage.attributes }[ indices.z ]; + ${ interpolateBody } + return result; + + } + `; + + } + writeTransformData( info, premultiplyMatrix, writeOffset, targetBuffer ) { const { structs } = this; diff --git a/src/webgpu/lib/wgsl/common.wgsl.js b/src/webgpu/lib/wgsl/common.wgsl.js index e5003f16b..e34eb4f20 100644 --- a/src/webgpu/lib/wgsl/common.wgsl.js +++ b/src/webgpu/lib/wgsl/common.wgsl.js @@ -29,9 +29,9 @@ export const ndcToCameraRay = wgslFn( /* wgsl*/` } `, [ rayStruct ] ); -export const intersectsBounds = wgslFn( /* wgsl */` +export const rayIntersectsBounds = wgslFn( /* wgsl */` - fn intersectsBounds( + fn rayIntersectsBounds( ray: Ray, bounds: BVHBoundingBox, dist: ptr From cb31eda159662493b9a065b035dd765039c971c1 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Mon, 23 Feb 2026 14:18:20 +0900 Subject: [PATCH 75/84] Fixups --- src/webgpu/compute/PathTracerMegaKernel.js | 11 ++++++----- src/webgpu/compute/wavefront/ProcessHitsKernel.js | 9 ++++----- src/webgpu/compute/wavefront/RayIntersectionKernel.js | 4 ++-- src/webgpu/lib/BVHComputeData.js | 8 +------- 4 files changed, 13 insertions(+), 19 deletions(-) diff --git a/src/webgpu/compute/PathTracerMegaKernel.js b/src/webgpu/compute/PathTracerMegaKernel.js index 932493452..cea76dd40 100644 --- a/src/webgpu/compute/PathTracerMegaKernel.js +++ b/src/webgpu/compute/PathTracerMegaKernel.js @@ -8,7 +8,7 @@ import { proxy } from '../lib/nodes/NodeProxy.js'; export class PathTracerMegaKernel extends ComputeKernel { - constructor( name = 'bvh_' ) { + constructor() { const parameters = { bvhData: { value: null }, @@ -82,15 +82,15 @@ export class PathTracerMegaKernel extends ComputeKernel { for ( var bounce = 0u; bounce < bounces; bounce ++ ) { - let hitResult = ${ name }RaycastFirstHit( ray ); + let hitResult = bvh_RaycastFirstHit( ray ); if ( hitResult.didHit ) { - let vertexData = ${ name }sampleTrianglePoint( hitResult.barycoord, hitResult.indices.xyz ); + let vertexData = bvh_sampleTrianglePoint( hitResult.barycoord, hitResult.indices.xyz ); let hitPosition = ray.origin + ray.direction * hitResult.dist; let scatterRec = bsdfEval( normalize( vertexData.normal.xyz ), - ray.direction ); - let transform = ${ name }transforms.value[ hitResult.objectIndex ]; - let material = ${ name }materials.value[ transform.materialIndex ]; + let transform = bvh_transforms.value[ hitResult.objectIndex ]; + let material = bvh_materials.value[ transform.materialIndex ]; // white diffuse surface throughputColor *= material.albedo * scatterRec.value / scatterRec.pdf; @@ -120,6 +120,7 @@ export class PathTracerMegaKernel extends ComputeKernel { `, [ proxy( 'bvhData.value.storage.materials', parameters ), proxy( 'bvhData.value.structs.material', parameters ), + proxy( 'bvhData.value.structs.transform', parameters ), proxy( 'bvhData.value.fns.raycastFirstHit', parameters ), proxy( 'bvhData.value.fns.sampleTrianglePoint', parameters ), ndcToCameraRay, pcgRand3, pcgInit, lambertBsdfFunc, diff --git a/src/webgpu/compute/wavefront/ProcessHitsKernel.js b/src/webgpu/compute/wavefront/ProcessHitsKernel.js index 05da4668c..688c46fa8 100644 --- a/src/webgpu/compute/wavefront/ProcessHitsKernel.js +++ b/src/webgpu/compute/wavefront/ProcessHitsKernel.js @@ -1,7 +1,6 @@ import { IndirectStorageBufferAttribute, StorageTexture } from 'three/webgpu'; import { ComputeKernel } from '../ComputeKernel.js'; import { uniform, storage, wgslFn, textureStore, globalId } from 'three/tsl'; -import { constants } from '../../lib/wgsl/common.wgsl.js'; import { pcgRand3, pcgInit } from '../../nodes/random.wgsl.js'; import { lambertBsdfFunc } from '../../nodes/sampling.wgsl.js'; import { queuedRayStruct, queuedHitStruct, QUEUED_RAY_SIZE, QUEUED_HIT_SIZE } from './structs.js'; @@ -9,7 +8,7 @@ import { proxy } from '../../lib/nodes/NodeProxy.js'; export class ProcessHitsKernel extends ComputeKernel { - constructor( name = 'bvh_' ) { + constructor() { const parameters = { bvhData: { value: null }, @@ -72,9 +71,9 @@ export class ProcessHitsKernel extends ComputeKernel { pcgInitialize( indexUV, seed ); - let object = ${ name }transforms.value[ input.objectIndex ]; - let material = ${ name }materials.value[ object.materialIndex ]; - var vertexData = ${ name }sampleTrianglePoint( input.barycoord, input.indices.xyz ); + let object = bvh_transforms.value[ input.objectIndex ]; + let material = bvh_materials.value[ object.materialIndex ]; + var vertexData = bvh_sampleTrianglePoint( input.barycoord, input.indices.xyz ); vertexData.normal = normalize( transpose( object.inverseMatrixWorld ) * vertexData.normal ); vertexData.position = object.matrixWorld * vertexData.position; diff --git a/src/webgpu/compute/wavefront/RayIntersectionKernel.js b/src/webgpu/compute/wavefront/RayIntersectionKernel.js index fc0cbd6b9..1c0e28981 100644 --- a/src/webgpu/compute/wavefront/RayIntersectionKernel.js +++ b/src/webgpu/compute/wavefront/RayIntersectionKernel.js @@ -7,7 +7,7 @@ import { proxy } from '../../lib/nodes/NodeProxy.js'; export class RayIntersectionKernel extends ComputeKernel { - constructor( name = 'bvh_' ) { + constructor() { const parameters = { bvhData: { value: null }, @@ -63,7 +63,7 @@ export class RayIntersectionKernel extends ComputeKernel { pcgInitialize( indexUV, seed ); // run intersection - let hitResult = ${ name }RaycastFirstHit( input.ray ); + let hitResult = bvh_RaycastFirstHit( input.ray ); if ( hitResult.didHit ) { // TODO: we process all of these materials immediately to push to the ray queue diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index fd002063f..151d3d070 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -1,6 +1,6 @@ import { Matrix4, Vector4 } from 'three'; import { StorageBufferAttribute, StructTypeNode } from 'three/webgpu'; -import { storage, storageBarrier } from 'three/tsl'; +import { storage } from 'three/tsl'; import { rayIntersectsBounds, constants } from './wgsl/common.wgsl.js'; import { rayStruct, bvhNodeStruct } from './wgsl/structs.wgsl.js'; import { wgslTagCode, wgslTagFn } from './nodes/WGSLTagFnNode.js'; @@ -505,8 +505,6 @@ export class BVHComputeData { } - writeOffset += dereferencedIndex.length; - } else if ( geometry.index ) { for ( let i = 0; i < count; i ++ ) { @@ -515,8 +513,6 @@ export class BVHComputeData { } - writeOffset += count; - } else { for ( let i = 0; i < count; i ++ ) { @@ -525,8 +521,6 @@ export class BVHComputeData { } - writeOffset += count; - } } From 51ce75190d650dfd9fd84cf9a8b78cacd4617013 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Mon, 23 Feb 2026 15:18:59 +0900 Subject: [PATCH 76/84] Move pathtracer compute data --- src/webgpu/WebGPUPathTracer.js | 2 +- src/webgpu/{lib => nodes}/PathtracerBVHComputeData.js | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename src/webgpu/{lib => nodes}/PathtracerBVHComputeData.js (98%) diff --git a/src/webgpu/WebGPUPathTracer.js b/src/webgpu/WebGPUPathTracer.js index c6ecc24c4..f592a513b 100644 --- a/src/webgpu/WebGPUPathTracer.js +++ b/src/webgpu/WebGPUPathTracer.js @@ -5,7 +5,7 @@ import { RenderToScreenNodeMaterial } from './materials/RenderToScreenMaterial.j import { MegaKernelPathTracer } from './MegaKernelPathTracer.js'; import { WaveFrontPathTracer } from './WaveFrontPathTracer.js'; import { ObjectBVH } from './lib/ObjectBVH.js'; -import { PathtracerBVHComputeData } from './lib/PathtracerBVHComputeData.js'; +import { PathtracerBVHComputeData } from './nodes/PathtracerBVHComputeData.js'; const _resolution = new Vector2(); export class WebGPUPathTracer { diff --git a/src/webgpu/lib/PathtracerBVHComputeData.js b/src/webgpu/nodes/PathtracerBVHComputeData.js similarity index 98% rename from src/webgpu/lib/PathtracerBVHComputeData.js rename to src/webgpu/nodes/PathtracerBVHComputeData.js index 3121c17c8..ab28e40e9 100644 --- a/src/webgpu/lib/PathtracerBVHComputeData.js +++ b/src/webgpu/nodes/PathtracerBVHComputeData.js @@ -1,5 +1,5 @@ import { BufferAttribute, BufferGeometry, StorageBufferAttribute, StructTypeNode } from 'three/webgpu'; -import { BVHComputeData } from './BVHComputeData.js'; +import { BVHComputeData } from '../lib/BVHComputeData.js'; import { storage } from 'three/tsl'; import { MeshBVH, SAH } from 'three-mesh-bvh'; From 8cf5a2a18c4234e90e597a1ab30316a47e703512 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Mon, 23 Feb 2026 15:29:24 +0900 Subject: [PATCH 77/84] comments --- src/webgpu/nodes/PathtracerBVHComputeData.js | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/webgpu/nodes/PathtracerBVHComputeData.js b/src/webgpu/nodes/PathtracerBVHComputeData.js index ab28e40e9..e5ae9b306 100644 --- a/src/webgpu/nodes/PathtracerBVHComputeData.js +++ b/src/webgpu/nodes/PathtracerBVHComputeData.js @@ -16,10 +16,12 @@ const materialStruct = new StructTypeNode( { albedo: 'vec3f', }, 'MaterialStruct' ); +// Pathtracer-specific version of the BVHComputeData tht includes material mapping, property structs export class PathtracerBVHComputeData extends BVHComputeData { constructor( bvh, options = {} ) { + // TODO: once supported we should use the appropriately-sized member sizes super( bvh, { attributes: { position: 'vec4f', @@ -40,6 +42,7 @@ export class PathtracerBVHComputeData extends BVHComputeData { super.update(); + // build material storage const { materials, structs, prefix: name } = this; const materialBuffer = new ArrayBuffer( structs.material.getLength() * materials.length * 4 ); const materialBufferF32 = new Float32Array( materialBuffer ); @@ -61,6 +64,7 @@ export class PathtracerBVHComputeData extends BVHComputeData { super.writeTransformData( info, premultiplyMatrix, writeOffset, targetBuffer ); + // write material data to the transforms const { materials } = this; const material = info.object.material; if ( ! materials.includes( material ) ) { @@ -87,6 +91,7 @@ export class PathtracerBVHComputeData extends BVHComputeData { } else if ( bvh.indirect ) { + // "indirect" bvhs are not supported since they cannot be unpacked in a way tht will allow for coherent material indices const proxyGeometry = new BufferGeometry(); proxyGeometry.attributes = bvh.geometry.attributes; From e919513228797f75ccb7365df158c23e8b5ba995 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Mon, 23 Feb 2026 22:35:34 +0900 Subject: [PATCH 78/84] Fix instances --- src/webgpu/lib/BVHComputeData.js | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index 151d3d070..ad88be98e 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -691,7 +691,8 @@ export class BVHComputeData { if ( object.isInstancedMesh || object.isBatchedMesh ) { - object.getMatrixAt( instanceId, _matrix ).premultiply( object.matrixWorld ); + object.getMatrixAt( instanceId, _matrix ); + _matrix.premultiply( object.matrixWorld ); } else { From 33b2b916ac8ff6769b43edfb8a66f1d229b1f3e5 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Tue, 24 Feb 2026 10:46:11 +0900 Subject: [PATCH 79/84] Make proxy call-able --- src/webgpu/lib/nodes/NodeProxy.js | 33 ++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/src/webgpu/lib/nodes/NodeProxy.js b/src/webgpu/lib/nodes/NodeProxy.js index 613fcc522..cf922c5bb 100644 --- a/src/webgpu/lib/nodes/NodeProxy.js +++ b/src/webgpu/lib/nodes/NodeProxy.js @@ -1,5 +1,29 @@ import { Node } from 'three/webgpu'; +class ProxyCallNode extends Node { + + static get type() { + + return 'ProxyCallNode'; + + } + + constructor( proxyNode, params ) { + + super(); + this.proxyNode = proxyNode; + this.params = params; + + } + + setup() { + + return this.proxyNode.node.call( ...this.params ); + + } + +} + export class NodeProxy extends Node { static get type() { @@ -62,4 +86,11 @@ export class NodeProxy extends Node { } -export const proxy = ( ...args ) => new NodeProxy( ...args ); +export const proxy = ( ...args ) => { + + const nodeProxy = new NodeProxy( ...args ); + const fn = ( ...params ) => new ProxyCallNode( nodeProxy, params ); + fn.functionNode = nodeProxy; + return fn; + +}; From d9fc765cd7ef0af6d92682cc4f4a912d240c18c7 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Tue, 24 Feb 2026 12:45:53 +0900 Subject: [PATCH 80/84] Add glsl variants of the code tag functions --- src/webgpu/lib/nodes/WGSLTagFnNode.js | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/webgpu/lib/nodes/WGSLTagFnNode.js b/src/webgpu/lib/nodes/WGSLTagFnNode.js index 519873c3b..85a1468a0 100644 --- a/src/webgpu/lib/nodes/WGSLTagFnNode.js +++ b/src/webgpu/lib/nodes/WGSLTagFnNode.js @@ -281,10 +281,7 @@ export class WGSLTagCodeNode extends CodeNode { } -// template tag literal function version of "wgslFn" for easy interpolation of TSL nodes -export const wgslTagFn = ( tokens, ...args ) => { - - const functionNode = new WGSLTagFnNode( tokens, args ); +const getFn = functionNode => { const fn = ( ...params ) => { @@ -310,10 +307,15 @@ export const wgslTagFn = ( tokens, ...args ) => { }; fn.functionNode = functionNode; - return fn; }; -// template tag literal for reusable WGSL code snippets with dependency resolution +// template tag literal function version of "wgslFn" & "wgsl" to generate +// functions & code snippets respectively +export const wgslTagFn = ( tokens, ...args ) => getFn( new WGSLTagFnNode( tokens, args ) ); export const wgslTagCode = ( tokens, ...args ) => new WGSLTagCodeNode( tokens, args ); + +// glsl versions +export const glslTagFn = ( tokens, ...args ) => getFn( new WGSLTagFnNode( tokens, args, 'glsl' ) ); +export const glslTagCode = ( tokens, ...args ) => new WGSLTagCodeNode( tokens, args, 'glsl' ); From a6fd8ea76bdb2de246f85e37ddbcef4c09edc465 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 25 Feb 2026 12:06:22 +0900 Subject: [PATCH 81/84] Adjust proxy node --- src/webgpu/lib/nodes/NodeProxy.js | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/webgpu/lib/nodes/NodeProxy.js b/src/webgpu/lib/nodes/NodeProxy.js index cf922c5bb..ee3a9962e 100644 --- a/src/webgpu/lib/nodes/NodeProxy.js +++ b/src/webgpu/lib/nodes/NodeProxy.js @@ -88,9 +88,15 @@ export class NodeProxy extends Node { export const proxy = ( ...args ) => { + return new NodeProxy( ...args ); + +}; + +export const proxyFn = ( ...args ) => { + const nodeProxy = new NodeProxy( ...args ); const fn = ( ...params ) => new ProxyCallNode( nodeProxy, params ); fn.functionNode = nodeProxy; - return fn; + return nodeProxy; }; From 855925ad7bf7817b74dfabc0a696b63141d89b83 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 25 Feb 2026 12:08:13 +0900 Subject: [PATCH 82/84] Remove unnecessary "writeOffset" --- src/webgpu/lib/BVHComputeData.js | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/webgpu/lib/BVHComputeData.js b/src/webgpu/lib/BVHComputeData.js index ad88be98e..71434f2c4 100644 --- a/src/webgpu/lib/BVHComputeData.js +++ b/src/webgpu/lib/BVHComputeData.js @@ -582,8 +582,6 @@ export class BVHComputeData { } ); - writeOffset += vertexCount; - } } From 53ecd950d7a095f9e30a0d6d02d0facfeeffbc45 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 25 Feb 2026 12:22:52 +0900 Subject: [PATCH 83/84] Update ObjectBVH, fix BatchedMesh instance bug --- src/webgpu/lib/ObjectBVH.js | 33 ++++++--------------------------- 1 file changed, 6 insertions(+), 27 deletions(-) diff --git a/src/webgpu/lib/ObjectBVH.js b/src/webgpu/lib/ObjectBVH.js index eecea45b5..c0540ed2b 100644 --- a/src/webgpu/lib/ObjectBVH.js +++ b/src/webgpu/lib/ObjectBVH.js @@ -62,27 +62,6 @@ export class ObjectBVH extends BVH { } - getObjectMatrix( compositeId, target ) { - - const { idMask, idBits, objects, matrixWorld } = this; - const id = getObjectId( compositeId, idMask ); - const instanceId = getInstanceId( compositeId, idBits, idMask ); - const object = objects[ id ]; - if ( object.isInstancedMesh || object.isBatchedMesh ) { - - object.getMatrixAt( instanceId, target ).premultiply( object.matrixWorld ); - - } else { - - target.copy( object.matrixWorld ); - - } - - _inverseMatrix.copy( matrixWorld ).invert(); - target.premultiply( _inverseMatrix ); - - } - getObjectFromId( compositeId ) { const { idMask, objects } = this; @@ -484,10 +463,10 @@ export class ObjectBVH extends BVH { } else if ( object.isBatchedMesh && includeInstances ) { const { instanceCount, maxInstanceCount } = object; - let instance = 0; + let foundInstances = 0; let iter = 0; - // TODO: use a better check here, like "maxInstanceCount" - while ( instance < instanceCount && iter < maxInstanceCount ) { + + while ( foundInstances < instanceCount && iter < maxInstanceCount ) { iter ++; @@ -495,10 +474,10 @@ export class ObjectBVH extends BVH { // instance were active try { - object.getVisibleAt( instance ); + object.getVisibleAt( iter ); - target[ index ] = ( instance << idBits ) | i; - instance ++; + target[ index ] = ( iter << idBits ) | i; + foundInstances ++; index ++; } catch { From f5371e0ee828300543d8a60c98e88ff1475619b9 Mon Sep 17 00:00:00 2001 From: Garrett Johnson Date: Wed, 25 Feb 2026 13:18:35 +0900 Subject: [PATCH 84/84] Improve bvh update --- src/webgpu/MegaKernelPathTracer.js | 1 + src/webgpu/WaveFrontPathTracer.js | 4 ++++ src/webgpu/compute/ComputeKernel.js | 8 +++++++- 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/src/webgpu/MegaKernelPathTracer.js b/src/webgpu/MegaKernelPathTracer.js index 58053bbb2..82b1c6e03 100644 --- a/src/webgpu/MegaKernelPathTracer.js +++ b/src/webgpu/MegaKernelPathTracer.js @@ -102,6 +102,7 @@ export class MegaKernelPathTracer { setBVHData( bvhData ) { this.kernel.bvhData = bvhData; + this.kernel.needsUpdate = true; this.reset(); } diff --git a/src/webgpu/WaveFrontPathTracer.js b/src/webgpu/WaveFrontPathTracer.js index b3d1cc0c5..5aad3a6a6 100644 --- a/src/webgpu/WaveFrontPathTracer.js +++ b/src/webgpu/WaveFrontPathTracer.js @@ -228,7 +228,11 @@ export class WaveFrontPathTracer { setBVHData( bvhData ) { this.rayIntersectionKernel.bvhData = bvhData; + this.rayIntersectionKernel.needsUpdate = true; + this.hitProcessKernel.bvhData = bvhData; + this.hitProcessKernel.needsUpdate = true; + this.reset(); } diff --git a/src/webgpu/compute/ComputeKernel.js b/src/webgpu/compute/ComputeKernel.js index 520f475a2..81ba08eec 100644 --- a/src/webgpu/compute/ComputeKernel.js +++ b/src/webgpu/compute/ComputeKernel.js @@ -12,6 +12,13 @@ export class ComputeKernel { } + set needsUpdate( v ) { + + // TODO: hack to force the kernel to rebuild since "needsUpdate" is not respected + this.setWorkgroupSize( ...this.workgroupSize ); + + } + constructor( fn, options = {} ) { const { @@ -60,7 +67,6 @@ export class ComputeKernel { setWorkgroupSize( x = 64, y = 1, z = 1 ) { - // this.workgroupSize = [ x, y, z ]; this.kernel = this._fn.computeKernel( [ x, y, z ] ); return this;