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

[gkjohnson]

Related to gkjohnson/three-mesh-bvh#832

This PR is aiming to add TLAS structure using three-mesh-bvh's "ObjectBVH" which will remove the need for merging the scene into a "super geometry" & simplify the handling of batched and instanced mesh instances. Longer term it should enable us to update specific transforms when they change. It will also fix gkjohnson/three-mesh-bvh#832. Once this is complete I'll move it back to three-mesh-bvh.

There are some quirks with the interaction between TSL nodes, structs, etc and wgslFn that I'll want to improve the ergonomics of at some point if I can. Also if we aim to support the "indirect" MeshBVH flag then we'll have to continue embedding the material index as a vertex attribute because otherwise "derefencing" the index will break the coherence of the group material indices.

Later

• add skinned mesh support
• make a callable version of the "proxy" return function so we can fire proxied functions

cc @TheBlek

[gkjohnson]

@TheBlek sorry this is taking longer (and gotten bigger) than expected. I'm running into a number of TSL ergonomics and feature issues getting in the way of writing this cleanly and generically. I've added this template tag function which should help us avoid needing to explicitly name or track our structs and functions for use in the wgsl functions.

Here's a before, for example:

wgslFn( /* wgsl */`
fn traverse( ray: Ray ) -> RayIntersection {
  // ...    
  let hit = interesectsBounds( ray );
  let vertex0 = attributes.value[ hit.indices.x ];
  // ...
  return hit;
}`, [
  rayStruct, intersectionStruct, attributesStruct,
  attributesStorage, intersectsBoundsFn, constants
] );

And after:

wgslTagFn`
// raw includes
${ [ constants ] }

// fn
fn traverse( ray: ${ rayStruct } ) -> ${ intersectionStruct } {
  // ...    
  let hit = ${ intersectsBoundsFn( 'ray' ) };
  let vertex0 = ${ attributesStorage }[ hit.indices.x ];
  // ...
  return hit;
}`;

All the dependencies, structs, etc used in the template arguments are implicitly added to the node graph dependencies, similar to TSL functions, so there's no need to declare them separately. I'm hoping this should make writing complex wgsl functions with interdependencies a bit more manageable but I'd like feedback if you have any.

Once this is resolved I'll take a look at #719. I expect we'll have with #717, though. Thanks for your work on those!

[gkjohnson]

@TheBlek This should be ready now - most of the changed lines are from copied code from three-mesh-bvh (wgsl functions, structs, ObjectBVH). Once we've been able to use this a bit I'll plan to move the logic back over to three-mesh-bvh so it can be reused.

primitive scene with 3 white instanced mesh spheres

A quick overview of some of the functions and additions:

BVHComputeData

This is the class responsible for packing geometry + bvh data into storage buffers and generating the necessary structs and functions for querying it. It only supports passing an "ObjectBVH" as a root BVH. It generates a "raycastFirstHit" function by default but the class also includes a "getShapecastFn" which can be used to generate a TSL node for performing custom spatial queries similar to the "shapecast" function in three-mesh-bvh.

There's a "PathtracerBVHComputeData" which includes material information in addition to per-object transforms. In the future maybe it will make sense to allow for passing a single mesh bvh as well as allowing for a more simple way to update matrices or material mapping / properties without having to incure the cost of regenerating the full BVH buffers.

NodeProxy

A "proxy" node is a custom node that's similar to the built-in TSL "reference" node so that a node in another object can be indirectly included in a graph. This will let us swap out our block of "BVHData" without having to rebuild the compute kernels every time the scene changes:

const indirectData = {
  bvhData: null,
};

const computeFn = wgslFn( /* wgsl */`
  fn compute() -> void {
    let result = raycastFirstHit( ray );
    // ...
  }
`, [ proxy( 'bvhData.fns.raycastFirstHit', indirectData ) ] );

// "indirectData" can be modified later to point to nodes so we can avoid rebuilding
// the we can avoid regenerating the computeFn.
indirectData.bvhData = new BVHComputeData( bvh );

wgslTagFn, wgslTagCode

Described in #716 (comment), these will let us write functions and code snippets using wgsl syntax without having to rely on fixed function or struct names, etc.

wgslTagFn/* wgsl */`
  fn traverse( ray: ${ rayStruct } ) -> ${ intersectionStruct } {
    // ...    
    let hit = ${ intersectsBoundsFn( 'ray' ) };
    let vertex0 = ${ attributesStorage }[ hit.indices.x ];
    // ...
    return hit;
  }
`;

If these look good and turn out to be functional for us I'll probably make a suggestion to three.js to include this kind of functionality.

--

Take a look and let me know what you think. Once we can merge this I'll take a look at #719 and can help in getting that PR updated with the necessary changes from this PR.

[TheBlek]

This is a very clever way to write wgsl nodes with everything already embedded. Didn't know about literal functions before - cool stuff.

However, what I worry here about is complexity: this adds another syntax to write shaders with. That is additional burden for the maintainer and to those who want to contribute. Getting into nodes is not easy as it is (at least it wasn't for me), especially when trying to figure where a bug comes from. Or do you think benefits of automatically figuring out the names and dependencies outweigh that?

What I want here is to learn your judgement better. Interested to hear your thoughts.

[gkjohnson]

Regarding complexity and maintenance - I'm, of course, sympathetic to this. It's always the case that there's a balance between adding layers of abstraction to improve dev experience and the added overhead of both communicating what the layers do and added code. In practice I think the code used to enable this is fairly small (a few hundred lines including padding) and is more similar than it is different to templating languages, like those used in React, Vue, lit-html, etc, which have already become ubiquitous in Javascript. So hopefully leaning on this existing knowledge base can help. If we can prove the utility of this out further I'd also like to push for this to be built in to TSL which would remove the burden from this project and ensure it's not broken by internal TSL development over time.

The other half of this is that I feel the complexity required by not enabling a system like this far outweighs any that's added. If we choose to go the route of using raw strings with wgslFn then all dependencies must have predetermined names associated with them - a requirement that is otherwise not necessary for use of any TSL node. And if we want to handle names generically (eg templated function definitions or user provided functions) then we have insert fn.name, anyway, in addition to adding fn to the dependency list.

That's just for the simple cases. For "storage" nodes you cannot just use the declared names directly. The storage accessors are nested into a "value" member of a buffer object struct so now, for storage buffers only, you have to read the data using myStorageBufferName.value[ i ]. This nesting and "value" variable naming is an arbitrary, undocumented structure so it may change on any release and isn't obvious on its own. By letting TSL resolve this naming we can let the TSL shader construction operate as it's intended and rely on it to give us the right accessor name however it's determined to be structured.

From a library perspective it's quite important to me to allow the API surface of the project to align with the mental model that's being promoted by the core three.js project & TSL. As mentioned above TSL does not require explicit names for anything in order to function. Aside from these native function nodes, human-readable names are more or less a debug symbol. If we want to expose any user-facing "function builder" behavior, as we have with the "getShapecastFn" interface, and we want to use wgslFn then a user must specify a name, which is inconsistent with what the broader TSL ecosystem mandates. The fact that a fixed name is required is really an artifact of our choice to use wgslFn internally. If we'd used TSL "Fn" and node chaining it wouldn't be required, so now the user has to be aware of what should be an implementation detail. And if they forget to define name, as they may have naturally become accustomed to not having to define, then things will not work and the library looks fragile (understandably so).

I think there's a lot of potential for some cool user-land modifications and function slotting with three-mesh-bvh, the pathtracer, and TSL in general so I'd like to look forward to how to prepare the project to handle that ergonomically. For this project, for instance, you can imagine user-defined functions for support custom ray-intersections on scenes constructed from SDFs (see tools like MagicaCSG), for example, or the ability for users to specify custom procedural BRDFs for materials (see #479).

--

A bit of a longwinded explanation but hopefully it makes some of my thinking more clear.

[TheBlek]

It all makes sense now, thanks.

[gkjohnson]

Happy to explain - let me know if there's anything you think should be addressed before merging.

[TheBlek]

Looks good to me