Hugging Face's logo

Spaces:
SimWorld-AI
/
industrial-harbor-viewer
Running

App Files Community
industrial-harbor-viewer / node_modules /three-mesh-bvh /src /webgpu /BVHComputeData.js
KoeYe's picture
KoeYe
Add viewer app, scene data, and three.js runtime
b797337 verified
Raw History Blame Contribute Delete
23.9 kB
/** @import { Object3D, BufferGeometry, Vector4 } from 'three' */
import { Matrix4, Mesh, Group } from 'three';
import { StorageBufferAttribute, StructTypeNode } from 'three/webgpu';
import { storage } from 'three/tsl';
import { MeshBVH } from '../core/MeshBVH.js';
import { SkinnedMeshBVH } from '../core/SkinnedMeshBVH.js';
import { GeometryBVH } from '../core/GeometryBVH.js';
import { BYTES_PER_NODE, UINT32_PER_NODE } from '../core/Constants.js';
import { proxy, proxyFn } from './nodes/NodeProxy.js';
import {
bvhNodeStruct,
transformStruct,
} from './tsl/structs.js';
import { appendBVHData, appendBVHSubtree, appendIndexData, appendGeometryData, getSubtreeNodeCount, getMaxNodeDepth } from './utils/packBVHBufferUtils.js';
import { getShapecastFn } from './shapecastFns/getShapecastFn.js';
import { getRaycastFirstHitFn } from './shapecastFns/getRaycastFirstHitFn.js';
import { getSampleTrianglePointFn } from './shapecastFns/getSampleTrianglePointFn.js';
import { getClosestPointToPointFn } from './shapecastFns/getClosestPointToPointFn.js';
import { SAH } from '../core/Constants.js';
import { ClusteredBVH } from './ClusteredBVH.js';
import { BVH_STACK_DEPTH } from './tsl/constants.js';
// 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 support for other geometry types (tris, lines, custom BVHs etc)
// scratch
const _matrix = /* @__PURE__ */ new Matrix4();
const _inverseMatrix = /* @__PURE__ */ new Matrix4();
const _range = { start: 0, count: 0, vertexStart: 0, vertexCount: 0 };
// functions
function isObjectVisible( object ) {
let curr = object;
while ( curr ) {
if ( curr.visible === false ) {
return false;
}
curr = curr.parent;
}
return true;
}
function getTotalBVHByteLength( bvh ) {
return bvh._roots.reduce( ( v, root ) => v + root.byteLength, 0 );
}
function getTransformKey( compositeId, root ) {
return `${ compositeId }_${ root }`;
}
/**
* Packs one or more scene objects into GPU-accessible BVH buffers (TLAS + BLAS) for use
* in WebGPU compute shaders via the Three.js TSL node system. After construction, call
* {@link BVHComputeData#update} to populate the storage buffers, then reference
* `this.storage` and `this.fns` in your compute shader nodes.
*
* @note This API is unstable and subject to change in future releases.
*
* @note This class requires three.js r185 or higher.
*/
export class BVHComputeData {
/**
* @param {Object3D|BufferGeometry|GeometryBVH|Array} objects
* Scene objects to include. A single item or array of Object3D, BufferGeometry, or GeometryBVH instances are
* all accepted and wrapped automatically in a BVH.
* @param {Object} [options]
* @param {Record<string,string>} [options.attributes={ position: 'vec4f' }]
* WGSL type map for the interleaved per-vertex attribute buffer. Keys are geometry
* attribute names; values are WGSL type strings (e.g. `'vec3f'`, `'vec4f'`).
* @param {boolean} [options.autogenerateBvh=true]
* When true, a {@link MeshBVH} is automatically built for any object that does not
* already have `geometry.boundsTree` set.
*/
constructor( objects, options = {} ) {
const {
attributes = { position: 'vec4f' },
autogenerateBvh = true,
} = options;
// convert the arguments to a list of objects
if ( ! Array.isArray( objects ) ) {
objects = [ objects ];
}
objects = objects.map( item => {
if ( item.isObject3D ) {
return item;
} else if ( item.isBufferGeometry ) {
return new Mesh( item );
} else if ( item instanceof GeometryBVH ) {
const dummy = new Mesh();
dummy.geometry.boundsTree = item;
return dummy;
}
} );
this._bvhCache = new Map();
this.autogenerateBvh = autogenerateBvh;
this.attributes = attributes;
this.objects = objects;
this.bvh = null;
// storage buffers and structs are populated in "update"; their members are accessed through
// proxy nodes so the functions below can reference them up front and keep working across rebuilds
this.storage = new NodeProxyObject();
this.structs = new NodeProxyObject( { transform: transformStruct } );
this.fns = new NodeProxyObject( {
raycastFirstHit: getRaycastFirstHitFn( this ),
closestPointToPoint: getClosestPointToPointFn( this ),
sampleTrianglePoint: null,
}, proxyFn );
}
/**
* Returns the representative root object for the scene to be constructed.
* @returns {Object3D}
*/
getRootObject() {
// convert the arguments to a list of objects
let { objects } = this;
if ( objects.isObject3D ) {
return objects;
}
if ( ! Array.isArray( objects ) ) {
objects = [ objects ];
}
objects = objects.map( item => {
if ( item.isObject3D ) {
return item;
} else if ( item.isBufferGeometry ) {
return new Mesh( item );
} else if ( item instanceof GeometryBVH ) {
const dummy = new Mesh();
dummy.geometry.boundsTree = item;
return dummy;
}
} );
const result = new Group();
result.children = objects;
return result;
}
/**
* Builds a WGSL shapecast function that traverses the TLAS and per-cluster BLAS in a single
* merged stack/loop for a custom shape type. The returned function signature is:
* `fn name( shape: ShapeStruct[, result: ptr<function, ResultStruct>] ) -> bool`
*
* @param {Object} options
* @param {string} [options.name] - Function name. Defaults to a random identifier.
* @param {StructTypeNode} options.shapeStruct - TSL struct or definition describing the query shape.
* @param {StructTypeNode|null} [options.resultStruct] - TSL struct for the accumulated result, or null.
* @param {Function|null} [options.prefixFn] - function node that runs before the bvh traversal - useful for resetting or initializing necessary module variables.
* @param {Function|null} [options.boundsOrderFn] - function node controlling left/right child traversal order.
* @param {Function} options.intersectsBoundsFn - function node testing the shape against a BVH node's bounds.
* @param {Function} options.intersectRangeFn - function node testing the shape against a leaf triangle range.
* @param {Function|null} [options.transformShapeFn] - function node that transforms the shape into object local space.
* @param {Function|null} [options.transformResultFn] - function node that transforms a hit result back to world space.
* @param {Function|null} [options.resetShapeFn] - function node called after each BLAS traversal to reset any per-object state set by `transformShapeFn`.
* @returns {Function} TSL function node for the traversal.
*/
getShapecastFn( options ) {
return getShapecastFn( this, options );
}
/**
* Rebuilds all GPU storage buffers from the current scene state. Must be called at least
* once before using `this.storage` or `this.fns` in a shader, and again whenever the
* scene topology changes (objects added/removed, geometry modified).
*/
update() {
// TODO
// - check if the total object geometries have changed somehow. We should sort the objects to
// a deterministic order and then check the BVH. Anything different in this case will trigger
// a full refresh (detecting batched mesh / instance differences by geometry id + count).
// - If the geometries are the same then we check whether they've changed (attribute versions,
// skinned mesh bone tex versions, morph target versions). If a geometry _has_ changed then
// the BVH needs to be refit (get BVH needs to continue to return consistent objects that are
// of the same structure - how to confirm this? requires caching?), then it should be written
// to the bvh nodes while refitting the TLAS.
// - If only non-structural attributes have changed then we can just write those (eg normals)
// TODO: we should include some kind of heuristic here for using a clustered or non-clustered
// BVH. Something like number of leaf objects, etc?
// "objects" may be a single item rather than an array
const root = this.getRootObject();
let total = 0;
root.traverse( c => {
// TODO: this needs to be in-sync with how clustered bvh totals count
if ( c.isMesh ) {
total ++;
}
} );
this.bvh = new ClusteredBVH( root, {
strategy: SAH,
getBVH: ( object, instance ) => this.getBVH( object, instance, _range ),
primitiveLimit: total < 3 ? Infinity : 64,
} );
// free any buffers from a previous update before swapping in the new ones
this.dispose();
const { attributes, structs, bvh } = this;
// collect the BVHs
const bvhInfo = [];
// per referenced cluster subtree (deduped by bvh + root + node): { data, root, node, size, base }.
// only these subtrees are copied into the node buffer - the upper nodes above the cluster cuts,
// which no TLAS leaf enters, are never uploaded.
const subtreeInfo = [];
const subtreeMap = new Map();
// accumulate the sizes of the bvh nodes buffer, number of objects, and geometry buffers
let bvhNodesBufferLength = getTotalBVHByteLength( bvh );
let indexBufferLength = 0;
let attributesBufferLength = 0;
// per primitive (in final tree order): the { transformSlot, subtree } for its TLAS leaf. The
// leaf's node offset is resolved to the subtree's packed base once all subtrees are laid out.
const primitiveInfo = [];
// tracks the deepest packed cluster subtree for the traversal depth bound below
let maxSubtreeDepth = 0;
// the transform slots, derived from the same primitive buffer walk "updateTransforms" uses
const transformMap = this._getTransformMap( bvh );
const { primitiveBuffer, primitiveBufferStride } = bvh;
for ( let i = 0, l = primitiveBuffer.length; i < l; i += primitiveBufferStride ) {
const compositeId = primitiveBuffer[ i ];
const compositeNodeId = primitiveBuffer[ i + 1 ];
const object = bvh.objects[ bvh.getObjectId( compositeId ) ];
const instanceId = bvh.getInstanceId( compositeId );
const range = { start: 0, count: 0, vertexStart: 0, vertexCount: 0 };
const primBvh = this.getBVH( object, instanceId, range );
if ( ! primBvh ) {
throw new Error( 'BVHComputeData: BVH not found.' );
}
// dedupe the geometry ( index + attributes ) once per bvh
let data = bvhInfo.find( info => info.bvh === primBvh );
if ( ! data ) {
data = {
index: bvhInfo.length,
bvh: primBvh,
range: range,
geometryOffset: 0,
};
// the whole geometry is packed once per bvh; only referenced subtrees ( below ) contribute
// to the node buffer
indexBufferLength += data.range.count;
attributesBufferLength += data.range.vertexCount;
bvhInfo.push( data );
}
const root = bvh.getBVHRootIndex( compositeNodeId );
const node = bvh.getBVHNodeIndex( compositeNodeId ) / UINT32_PER_NODE;
// dedupe the referenced cluster subtree - the TLAS leaf only enters this subtree, so packing
// just its contiguous node range skips the unreferenced upper nodes above the cut
const subtreeKey = `${ data.index }_${ root }_${ node }`;
let subtree = subtreeMap.get( subtreeKey );
if ( subtree === undefined ) {
const size = getSubtreeNodeCount( primBvh._roots[ root ], node );
subtree = { data, root, node, size, base: 0 };
subtreeMap.set( subtreeKey, subtree );
subtreeInfo.push( subtree );
bvhNodesBufferLength += size * BYTES_PER_NODE;
maxSubtreeDepth = Math.max( maxSubtreeDepth, getMaxNodeDepth( primBvh._roots[ root ], node ) );
}
// nodeOffset is resolved to the subtree's packed base after the subtrees are laid out
primitiveInfo.push( {
transformSlot: transformMap.get( getTransformKey( compositeId, root ) ).slot,
subtree,
} );
}
// Get the max depth of the tlas traversal and infer the needed depth for the compute stack. Subtract one since
// the TLAS leaf pushes one node instead of two.
const tlasDepth = getMaxNodeDepth( bvh._roots[ 0 ] );
const maxTraversalDepth = tlasDepth + maxSubtreeDepth - 1;
if ( maxTraversalDepth > BVH_STACK_DEPTH.value ) {
throw new Error( 'BVHComputeData: BVH depth overruns the compute stack depth.' );
}
//
// @note These buffer lengths are increased to a minimum size of 2 to avoid TSL converting storage buffers
// with length 1 being converted to a scalar value.
// TODO: remove this when fixed in three
const transformBufferLength = Math.max( transformMap.size, 2 );
indexBufferLength = Math.max( indexBufferLength, 2 );
attributesBufferLength = Math.max( attributesBufferLength, 2 );
// construct the attribute struct
const attributeStruct = new StructTypeNode( attributes, 'bvh_GeometryStruct' );
// write the geometry buffer attributes & bvh data
let attributesOffset = 0;
let indexOffset = 0;
const indexBuffer = new Uint32Array( indexBufferLength );
const attributesBuffer = new ArrayBuffer( attributesBufferLength * attributeStruct.getLength() * 4 );
const bvhNodesBuffer = new ArrayBuffer( bvhNodesBufferLength );
// pack each unique geometry ( index + attributes ) once, recording its triangle base so the
// referenced subtrees' leaves can be rebased into it
bvhInfo.forEach( info => {
info.geometryOffset = indexOffset / 3;
// A non-indirect bvh built over a sub-range stores absolute leaf offsets, so shift the
// base back by the range's triangle start.
if ( ! info.bvh.indirect ) {
info.geometryOffset -= info.range.start / 3;
}
appendIndexData( info.bvh, info.range, attributesOffset, indexOffset, indexBuffer );
appendGeometryData( info.bvh, info.range, attributesOffset, attributesBuffer, attributeStruct, this );
indexOffset += info.range.count;
attributesOffset += info.range.vertexCount;
} );
// pack only the referenced cluster subtrees into the node buffer, after the TLAS region. Each
// subtree's write base becomes the node offset written into its TLAS leaves.
let nodeWriteOffset = getTotalBVHByteLength( bvh ) / BYTES_PER_NODE;
subtreeInfo.forEach( subtree => {
subtree.base = nodeWriteOffset;
appendBVHSubtree( subtree.data.bvh._roots[ subtree.root ], subtree.node, subtree.size, subtree.data.geometryOffset, nodeWriteOffset, bvhNodesBuffer );
nodeWriteOffset += subtree.size;
} );
// resolve each TLAS leaf's node offset now that the subtree bases are known, then pack the TLAS
primitiveInfo.forEach( info => info.nodeOffset = info.subtree.base );
appendBVHData( bvh, primitiveInfo, 0, bvhNodesBuffer );
//
const transformArrayBuffer = new ArrayBuffer( structs.transform.getLength() * transformBufferLength * 4 );
//
// set up the storage buffers
// if itemSize for StorageBufferAttribute == arraySize,
// then buffer is treated not as array of structs, but as a single struct
// And that breaks code. For now itemSize = 1 does not seem to break anything
const bvhNodesStorage = storage( new StorageBufferAttribute( new Uint32Array( bvhNodesBuffer ), 1 ), bvhNodeStruct ).toReadOnly().setName( 'bvh_nodes' );
const transformsBuffer = new StorageBufferAttribute( new Uint32Array( transformArrayBuffer ), 1 );
const transformsStorage = storage( transformsBuffer, structs.transform ).toReadOnly().setName( 'bvh_transforms' );
const indexStorage = storage( new StorageBufferAttribute( indexBuffer, 1 ), 'uint' ).toReadOnly().setName( 'bvh_index' );
const attributesStorage = storage( new StorageBufferAttribute( new Uint32Array( attributesBuffer ), attributeStruct.getLength() ), attributeStruct ).toReadOnly().setName( 'bvh_attributes' );
this.storage.transforms = transformsStorage;
this.storage.nodes = bvhNodesStorage;
this.storage.index = indexStorage;
this.storage.attributes = attributesStorage;
this.structs.attributes = attributeStruct;
// writes every transform
_inverseMatrix.copy( bvh.matrixWorld ).invert();
transformMap.forEach( info => {
this.writeTransformData( info, _inverseMatrix, info.slot, transformArrayBuffer );
} );
// depends on the resolved attribute struct, so it must be built here rather than up front
this.fns.sampleTrianglePoint = getSampleTrianglePointFn( this );
// clear our cache for now. In the future we will need to keep this around.
this._bvhCache.clear();
}
/**
* Refits the clustered BVH and rewrites every entry in the transform buffer from the objects'
* current world matrices. Call this when object transforms or visibility change but the scene
* topology does not. The transform slots are derived from the clustered BVH's primitive buffer,
* so they match those written by {@link BVHComputeData#update}.
*/
updateTransforms() {
const { bvh, storage } = this;
bvh.refit();
// the TLAS occupies the head of the node buffer - rewrite just those nodes' bounds. A null
// "primitiveInfo" leaves the leaf encodings, and the cluster subtrees that follow them, in place.
const nodesAttribute = storage.nodes.proxyNode.value;
appendBVHData( bvh, null, 0, nodesAttribute.array.buffer );
nodesAttribute.needsUpdate = true;
const transformsAttribute = storage.transforms.proxyNode.value;
const transformArrayBuffer = transformsAttribute.array.buffer;
_inverseMatrix.copy( bvh.matrixWorld ).invert();
this._getTransformMap( bvh ).forEach( info => {
this.writeTransformData( info, _inverseMatrix, info.slot, transformArrayBuffer );
} );
transformsAttribute.needsUpdate = true;
}
/**
* Writes the world/inverse-world matrices and visibility flag for one transform entry
* into a raw ArrayBuffer. Override this in a subclass to inject additional per-object
* data (e.g. material index).
*
* @private
* @param {Object} info - Transform entry from the internal transform map.
* @param {Matrix4} premultiplyMatrix - Matrix pre-multiplied onto the object's world matrix (usually the inverse TLAS root matrix).
* @param {number} writeOffset - Index of the transform slot to write into.
* @param {ArrayBuffer} targetBuffer - Destination buffer.
*/
writeTransformData( info, premultiplyMatrix, writeOffset, targetBuffer ) {
const { structs } = this;
const transformBufferF32 = new Float32Array( targetBuffer );
const transformBufferU32 = new Uint32Array( targetBuffer );
const { object, instanceId } = info;
if ( object.isInstancedMesh || object.isBatchedMesh ) {
object.getMatrixAt( instanceId, _matrix );
_matrix.premultiply( object.matrixWorld );
} else {
_matrix.copy( object.matrixWorld );
}
// write transform
_matrix.premultiply( premultiplyMatrix );
_matrix.toArray( transformBufferF32, writeOffset * structs.transform.getLength() );
// write inverse transform
_matrix.invert();
_matrix.toArray( transformBufferF32, writeOffset * structs.transform.getLength() + 16 );
let visible = isObjectVisible( object );
if ( object.isBatchedMesh ) {
visible = visible && object.getVisibleAt( instanceId );
}
transformBufferU32[ writeOffset * structs.transform.getLength() + 32 ] = visible ? 1 : 0;
}
/**
* Returns the BVH for a given object/instance, populating `rangeTarget` with the
* corresponding index and vertex ranges within the packed geometry buffers. Override
* to support custom BVH types or caching strategies.
*
* @private
* @param {Object3D} object - The object to generate a BVH for.
* @param {number} instanceId - Instance index (relevant for InstancedMesh / BatchedMesh).
* @param {{start:number,count:number,vertexStart:number,vertexCount:number}} rangeTarget - Populated with the object's geometry range.
* @returns {MeshBVH|SkinnedMeshBVH|null}
*/
getBVH( object, instanceId, rangeTarget ) {
const { autogenerateBvh, _bvhCache } = this;
let bvh = null;
if ( object.boundsTree || object.isSkinnedMesh ) {
// this is a case where a mesh has morph targets and skinned meshes
const geometry = object.geometry;
rangeTarget.count = geometry.index ? geometry.index.count : geometry.attributes.position.count;
rangeTarget.vertexCount = geometry.attributes.position.count;
bvh = object.boundsTree || null;
if ( bvh === null && autogenerateBvh ) {
const id = object.uuid;
bvh = _bvhCache.get( id ) || new SkinnedMeshBVH( object );
_bvhCache.set( id, bvh );
}
} else if ( object.isBatchedMesh ) {
const geometryId = object.getGeometryIdAt( instanceId );
const range = object.getGeometryRangeAt( geometryId );
Object.assign( rangeTarget, range );
bvh = object.boundsTrees[ geometryId ] || null;
if ( bvh === null && autogenerateBvh ) {
const id = `batched_${ object.geometry.uuid }_${ range.start }_${ range.count }`;
bvh = _bvhCache.get( id ) || new MeshBVH( object.geometry, { range: { ...rangeTarget } } );
_bvhCache.set( id, bvh );
}
} else {
const geometry = object.geometry;
rangeTarget.count = geometry.index ? geometry.index.count : geometry.attributes.position.count;
rangeTarget.vertexCount = geometry.attributes.position.count;
bvh = object.geometry.boundsTree || null;
if ( bvh === null && autogenerateBvh ) {
const id = geometry.uuid;
bvh = _bvhCache.get( id ) || new MeshBVH( geometry );
_bvhCache.set( id, bvh );
}
}
return bvh;
}
/**
* Returns the default vec4 value written to the attribute buffer for vertices that lack
* a given attribute. Override to change per-attribute defaults.
*
* @private
* @param {string} key - Attribute name (e.g. `'position'`, `'normal'`).
* @param {Vector4} target - Receives the default value.
* @returns {Vector4}
*/
getDefaultAttributeValue( key, target ) {
switch ( key ) {
case 'position':
case 'color':
target.set( 1, 1, 1, 1 );
break;
default:
target.set( 0, 0, 0, 0 );
}
return target;
}
/**
* Releases GPU resources held by this instance.
*/
dispose() {
const { storage } = this;
for ( const key in storage ) {
storage[ key ].value?.dispose();
}
}
// Provides a consistent, deduplicated list of the transforms from the clustered BVH
_getTransformMap( bvh ) {
const { primitiveBuffer, primitiveBufferStride } = bvh;
const transformMap = new Map();
for ( let i = 0, l = primitiveBuffer.length; i < l; i += primitiveBufferStride ) {
const compositeId = primitiveBuffer[ i ];
const root = bvh.getBVHRootIndex( primitiveBuffer[ i + 1 ] );
const key = getTransformKey( compositeId, root );
// each bvh root gets its own transform so per-group data like materials can
// be attached. The many cluster primitives of a root all share it, so matrices
// are not duplicated.
if ( transformMap.has( key ) ) {
continue;
}
const slot = transformMap.size;
const object = bvh.objects[ bvh.getObjectId( compositeId ) ];
const instanceId = bvh.getInstanceId( compositeId );
transformMap.set( key, { object, instanceId, compositeId, root, slot } );
}
return transformMap;
}
}
// A container whose string members are returned as stable proxy nodes. Assigning a member stores
// the underlying node.
// TODO: we should automatically infer a proxy node vs fn. Perhaps in r185 we won't need the difference?
class NodeProxyObject {
constructor( initialization = {}, createProxy = proxy ) {
const proxies = {};
// the raw backing object holds the underlying nodes and is the proxy target. "createProxy"
// selects the proxy variant - "proxy" for plain nodes, "proxyFn" for callable function nodes.
return new Proxy( { ...initialization }, {
get( target, property ) {
if ( ! proxies[ property ] ) {
proxies[ property ] = createProxy( property, target );
}
return proxies[ property ];
},
set( target, property, value ) {
target[ property ] = value;
return true;
},
} );
}
}