tl;dr: jump to the Tag Template Literal Functions section for a look at a suggested new mechanism for writing TSL-compatible wgsl.
I’ve recently taken a first dive into three.js’ new Node Materials and TSL. And admittedly I’m jumping straight into the deep end by working on upgrading three-mesh-bvh and three-gpu-pathtracer to enable generic, compute-based spatial queries and rendering.
Like many I have opinions about TSL but as someone who has been trying to write fairly complex, reusable shaders using brittle string replacement and “onBeforeCompile” for nearly a decade it’s hard not to see TSL as bringing some huge improvements. Implicit dependency inclusion & deduplication, indirect value references for uniforms, and being able to slot functions into each other, to name a few, help make code feel significantly more generalize-able. For the work I do, at least, It’s a step that needed to happen.
Writing the code itself proves to be more difficult, though. As you might imagine declaring variables, performing math operations, and writing ray tracing logic with the function-chaining approach recommended by TSL was not ergonomic for me. This isn’t TSL’s fault - it’s the nature of trying to enable this kind of meta-programming in a language like Javascript. But I still wanted to see if I could find a better way. At least for the way I want to think about things.
As a baseline, here’s a notional example of what TSL might look like for calculating a box intersection for a ray and assigning it to a struct to return. It feels fairly verbose and uses some (for me) unintuitive syntax for things like if statements. For a full blown pathtracer I didn’t feel this would be approachable:
// tsl syntax
const getBoxIntersectionPoint = Fn( ( { ray, box } ) => {
const dist = intersectsBoundsFn( { ray, box } ).toVar();
const hit = intersectionStruct().toVar();
If( dist.lessThan( float( 0.0 ) ), () => {
hit.didHit.assign( false );
} ).Else( () => {
hit.didHit.assign( true );
hit.distance.assign( dist );
hit.point.assign( ray.origin.add( ray.direction.mul( dist ) ) );
} );
return hit;
} ).setLayout( {
name: 'getBoxIntersectionPoint',
type: intersectionStruct,
inputs: [
{ name: 'ray', type: rayStruct },
{ name: 'box', type: boxBoundsStruct },
],
} );
What about wgslFn?
But there are “native” function nodes! These nodes allow you to write string literals containing wgsl (or glsl) directly, meaning you can use the traditional shader syntax you’re familiar with while still integrating with TSL’s node syntax via dependencies and function call sites:
// wgsl
const getBoxIntersectionPoint = wgslFn( /* wgsl */`
fn getBoxIntersectionPoint(
ray: Ray,
box: BoxBounds
) -> IntersectionResult {
let dist = intersectsBounds( ray, box );
var hit: IntersectionResult;
if ( dist < 0.0 ) {
hit.didHit = false;
} else {
hit.didHit = true;
hit.distance = dist;
hit.point = ray.origin + ray.direction * dist;
}
return hit;
}
`, [ rayStruct, boxBoundsStruct, intersectionStruct ] );
With some comment-tag syntax highlighting in an IDE this can actually start to feel pretty good! You can write functions in a familiar format and still use nodes from the broader TSL ecosystem. But it pretty quickly becomes apparent that there are a lot of gotchas and limitations. All structs, functions, uniforms, and variables used must have names explicitly declared so you can use them in the function, which are otherwise not available until build time (notice that “Ray”, “IntersectionResult” etc, struct names must all match). This is a pretty heavy restriction especially when trying to write generic code for users to integrate or slotting user functions into your wrapper and can quickly can start to feel too limiting.
Tag Template Literal Functions
Tag Template Literal functions are functions that allow you to process template string segments and arguments by effectively passing the function to the template function without parentheses. The string tokens and args from the template string are then are then effectively passed to the function as arrays so they can be processed into whatever format is needed. These are an incredibly powerful Javascript feature that I don’t often see get much use. In a lot of ways they’re great for exactly what we’re talking about here.
For our use we can write a tag template literal function that can take TSL nodes as template arguments and return a node that can be used, then evaluated at build-time by the TSL node graph. This turns our function declaration into something like so:
// tag template wgsl
const getBoxIntersectionPoint = wgslTagFn/* wgsl */`
fn getBoxIntersectionPoint(
ray: ${ rayStruct },
box: ${ boxBoundsStruct }
) -> ${ intersectionStruct } {
let dist = ${ intersectsBounds( 'ray', 'box' ) };
var hit: ${ intersectionStruct };
if ( dist < 0.0 ) {
hit.didHit = false;
} else {
hit.didHit = true;
hit.distance = dist;
hit.point = ray.origin + ray.direction * dist;
}
return hit;
}
`;
Notice that the structs and functions can be used directly as template string arguments and the names no longer need to be declared! The node returned by wgslTagFn caches the referenced nodes and can evaluate the variable names to inline when the shader is being built. This can work for structs, function names, uniforms, literals, etc. With some other clever processing we can allow for strings to be passed to the functions being inlined in order to refer to local function variables, as you see above. And, of course, all these nodes are then implicitly added to the dependency graph and deduped.
Opinions will vary, I’m sure, but for me this is nearly exactly what I’d like to be writing in terms of shader definitions: all the benefits of dependencies and TSL nodes while still being able to write code in a familiar syntax with run-time composition and requirements for brittle dependencies on things like fixed naming.
What’s Next
If you’re interested in trying this out I’ve separated the code into this repository along with some more descriptions of the available features - please feel free to give feedback, as well! As development on three-mesh-bvh and the pathtracer continue I plan to keep exercising this approach and if it continues to prove useful perhaps there are ways it can be made more reusable or recommended as an addition to the main project.
