TSL: Function Layout

[sunag]

Related issue: Closes #23350

Description

Node System can generate function code in WGSL/GLSL using tslFn, and create a cache so that the same generated code can be reused in other materials.

The implementation of tslFn.setLayout() will guide the construction at the head of the interface, while the entire body of the code is generated using the current NodeBuilder in an independent flow.

We cache the code based on the NodeBuilder type, for example WGSLNodeBuilder with the TSL function, to be able to reuse it in other shaders.

D_GGX

TSL

const D_GGX = tslFn( ( inputs ) => {

	const { alpha, dotNH } = inputs;

	const a2 = alpha.pow2();

	const denom = dotNH.pow2().mul( a2.oneMinus() ).oneMinus(); // avoid alpha = 0 with dotNH = 1

	return a2.div( denom.pow2() ).mul( 1 / Math.PI );

} ).setLayout( {
	name: 'D_GGX',
	type: 'float',
	inputs: [
		{ name: 'alpha', type: 'float' },
		{ name: 'dotNH', type: 'float' }
	]
} );

WGSL

Note that only one variable was created, as this Node was used more than once in the code. Node System avoids creating variables where only one use is made.

fn D_GGX ( alpha : f32, dotNH : f32 ) -> f32 {

	var nodeVar0 : f32;

	nodeVar0 = pow( alpha, 2.0 );

	return ( ( nodeVar0 / pow( ( 1.0 - ( pow( dotNH, 2.0 ) * ( 1.0 - nodeVar0 ) ) ), 2.0 ) ) * 0.3183098861837907 );

}

GLSL

float D_GGX ( float alpha, float dotNH ) {

	float nodeVar0;

	nodeVar0 = pow( alpha, 2.0 );

	return ( ( nodeVar0 / pow( ( 1.0 - ( pow( dotNH, 2.0 ) * ( 1.0 - nodeVar0 ) ) ), 2.0 ) ) * 0.3183098861837907 );

}

Call tree

TSL

const desaturateSubTest = tslFn( ( { value } ) => {

	return vec3( 0.299, 0.587, 0.114 ).dot( value.xyz );

} ).setLayout( {
	name: 'desaturateSubTest',
	type: 'vec3',
	inputs: [
		{ name: 'value', type: 'vec3' }
	]
} );

const desaturate = tslFn( ( { value, mix } ) => {

	return desaturateSubTest( { value: value.xy } )

} ).setLayout( {
	name: 'desaturate',
	type: 'vec3',
	inputs: [
		{ name: 'value', type: 'vec3' }
	]
} );

WGSL

fn desaturateSubTest ( value : vec3<f32> ) -> vec3<f32> {

	return vec3<f32>( dot( vec3<f32>( 0.299, 0.587, 0.114 ), value ) );

}

fn desaturate ( value : vec3<f32> ) -> vec3<f32> {

	return desaturateSubTest( vec3<f32>( value.xy, 0.0 ) );

}

WGSL Inline sub-calls ( without desaturateSubTest.setLayout() )

fn desaturate ( value : vec3<f32> ) -> vec3<f32> {

	return vec3<f32>( dot( vec3<f32>( 0.299, 0.587, 0.114 ), vec3<f32>( value.xy, 0.0 ).xyz ) );

}

Thanks to @fyoudine who took the first step in this direction in this PR.

[LeviPesin]

Maybe we can create a separate node for this like ReferenceNode (current ReferenceNode we can merge with MaterialReferenceNode maybe)?

[sunag]

ReferenceNode has other uses as used in userData(). The function signature looks good for me, mainly because the declaration comes after the name, which would be a fundamental parameter.

[LeviPesin]

So, if I understand correctly, using layout allows compiling TSL functions directly to GLSL/WGSL (unlike other TSL functions that are every time run in JS, generating nodes, and then building the resulting node)? I love such approach!

[sunag]

So, if I understand correctly, using layout allows compiling TSL functions directly to GLSL/WGSL (unlike other TSL functions that are every time run in JS, generating nodes, and then building the resulting node)?

@LeviPesin That's right, JS would only be executed once per language.

[sunag]

One detail to note, tslFn.setInterface() does not yet work with global variables, for example normalWorld, Varyings and Uniforms, these values must be received as function inputs.