TSL Transpiler: Introduce Linker and improvements

[sunag]

Fixes #31274

Description

Many improvements have been made to the transpiler including the Linker which helps with flow control when encoding in TSL, and in the future in WGSL.

In the comparison images below, we can see the code much closer to that created by a programmer than by a string-based transpiler.

Summary

• Introduce Linker
  • It is now possible to optimize flow control with Linker, as it is possible to know which variable was assigned or accessed.
• Optimize variable creations and number conversions
• Add While support
• Add Comment support
• Fix unnecessary imports
• Refactorings in the code

Code

float FunctionTest( float alpha1, float alpha2 ) {

	// variable never assigned does not create .toVar()
	int tempA = 1 + 2;

	// variable is assigned then .toVar() was created
	float tempB = 1 + 1;
	tempB += tempA;

	// parameter is assigned, parameters are immutable in WGSL, so .toVar() was created
	// alpha2 was preserved
	alpha1 += 1;

	// alpha2 was preserved
	int i = 0;

	// +1 is incremented and returns the previous value
	int b = i++;

	// since the output is void, that is, it is not assigned to any variable,
	// it converts from i++ to ++i
	i++;

	return alpha1 + alpha2;

}

Old

// Three.js Transpiler r178dev

import { float, int, Fn } from 'three/tsl';

export const FunctionTest = /*#__PURE__*/ Fn( ( [ alpha1_immutable, alpha2_immutable ] ) => {

	const alpha2 = float( alpha2_immutable ).toVar();
	const alpha1 = float( alpha1_immutable ).toVar();
	const tempA = int( int( 1 ).add( int( 2 ) ) ).toVar();
	const tempB = float( int( 1 ).add( int( 1 ) ) ).toVar();
	tempB.addAssign( tempA );
	alpha1.addAssign( int( 1 ) );
	const i = int( int( 0 ) ).toVar();
	const b = int( i.increment() ).toVar();
	i.incrementBefore();

	return alpha1.add( alpha2 );

}, { alpha1: 'float', alpha2: 'float', return: 'float' } );

Now

// Three.js Transpiler r178dev

import { float, int, Fn } from 'three/tsl';

export const FunctionTest = /*#__PURE__*/ Fn( ( [ alpha1_immutable, alpha2 ] ) => {

	const alpha1 = alpha1_immutable.toVar();

	// variable never assigned does not create .toVar()

	const tempA = int( 1 + 2 );

	// variable is assigned then .toVar() was created

	const tempB = float( 1 + 1 ).toVar();
	tempB.addAssign( tempA );

	// parameter is assigned, parameters are immutable in WGSL, so .toVar() was created
	// alpha2 was preserved

	alpha1.addAssign( 1 );

	// alpha2 was preserved

	const i = int( 0 ).toVar();

	// +1 is incremented and returns the previous value

	const b = i.increment();

	// since the output is void, that is, it is not assigned to any variable,
	// it converts from i++ to ++i

	i.incrementBefore();

	return alpha1.add( alpha2 );

}, { alpha1: 'float', alpha2: 'float', return: 'float' } );

[linbingquan]

Test this PR in webgpu_tsl_transpiler.html, I found float is not used, but import it.

// Three.js Transpiler r178dev

import { pow2, float, sub, sqrt, EPSILON, max, div, Fn } from 'three/tsl';

// Put here your GLSL code to transpile to TSL:

export const V_GGX_SmithCorrelated = /*#__PURE__*/ Fn( ( [ alpha, dotNL, dotNV ] ) => {

	const a2 = pow2( alpha );
	const gv = dotNL.mul( sqrt( a2.add( sub( 1.0, a2 ).mul( pow2( dotNV ) ) ) ) );
	const gl = dotNV.mul( sqrt( a2.add( sub( 1.0, a2 ).mul( pow2( dotNL ) ) ) ) );

	return div( 0.5, max( gv.add( gl ), EPSILON ) );

}, { alpha: 'float', dotNL: 'float', dotNV: 'float', return: 'float' } );

[RenaudRohlinger]

Other test failing :

float mag = dot(p, p);
p = abs(p) / mag + vec3(-0.5, -0.4, -1.487);

const mag = dot( p, p );
p.assign( abs( p ).div( mag ).add( vec3( - 0.5.sub( 0.4 ).sub( 1.487 ) ) ) );

vec3(-0.5, -0.4, -1.487) => vec3( - 0.5.sub( 0.4 ).sub( 1.487 ) ) )

[cmhhelgeson]

Just going through some of the old issues:

#30858: TSL Transpiler not working with vec2[4] name
Still erroring, but produces a different error than before (generally arrays still need to be implemented)

vec2[4] marching_square(vec2 uv, float level)
{
    vec2 edge[4] = vec2[](vec2(-50.), vec2(-50.),vec2(-50.), vec2(-50.));
    return edge;
}

// Three.js Transpiler r178dev

import { property } from 'three/tsl';

const [ = property( 'vec2' ), float = property( 'vec2' );

#30885: TSL Transpiler not working with struct

struct Edges {
    vec2 a
    vec2 b
    vec2 c
    vec2 d
}

Edges marching_square(vec2 uv, float level) {
    Edges e
    e.a = vec2(-50.)
    e.b = vec2(-50.)
    e.c = vec2(-50.)
    e.d = vec2(-50.)
    return e
}

// Three.js Transpiler r178dev

// Put here your GLSL code to transpile to TSL:

- 50.;

#31094: Sampler 2D Conversion Trouble (seems like the sampler issue has been solved, but there's a dangling comma in the layout properties of the main function)

varying float nShiftPower;

uniform sampler2D map;

uniform float time;

varying vec2 vUv;

vec4 rgbShift( float n, vec2 uvs, sampler2D m ){

	vec4 final = vec4(0.0, 0.0, 0.0, 1.0);

    if( n > 0.0 ){

        float d1 = n * 0.001;
        float d2 = n * 0.002;

        vec2 q   = uvs;

        float x =  n * sin(0.3*time+q.y*21.0)*sin(0.7*time+q.y*29.0)*sin(0.3+0.33*time+q.y*31.0)*0.0017;
        
        float R = texture2D(map,vec2(x + uvs.x+ d1 ,uvs.y+ d1)).x+0.05;
        float G = texture2D(map,vec2(x + uvs.x  ,uvs.y- d2)).y+0.05;
        float B = texture2D(map,vec2(x + uvs.x- d2 ,uvs.y    )).z+0.05;

        float nshiftQuater = n * 0.75;
        
        R += 0.08*texture2D(map,nshiftQuater*vec2(x+0.025, -0.027)+ n * vec2(uvs.x+0.001,uvs.y+0.001)).x;
        G += 0.05*texture2D(map,nshiftQuater*vec2(x-0.022, -0.02) + n * vec2(uvs.x+0.000,uvs.y-0.002)).y;
        B += 0.08*texture2D(map,nshiftQuater*vec2(x-0.02, -0.018) + n * vec2(uvs.x-0.002,uvs.y+0.000)).z;

        final = vec4(R, G, B, 1.0);

    }
    else {

        final = texture2D(m, uvs);
    }

	return final;

}

void main(){

	gl_FragColor = getColor( nShiftPower,vUv,  map );
}

// Three.js Transpiler r178dev

import { varying, float, texture, uniform, vec2, vec4, mul, sin, add, If, Fn } from 'three/tsl';

const nShiftPower = varying( float(), 'nShiftPower' );
const map = texture( /* <THREE.Texture> */ );
const time = uniform( 'float' );
const vUv = varying( vec2(), 'vUv' );

export const rgbShift = /*#__PURE__*/ Fn( ( [ n, uvs, m ] ) => {

	const final = vec4( 0.0, 0.0, 0.0, 1.0 ).toVar();

	If( n.greaterThan( 0.0 ), () => {

		const d1 = n.mul( 0.001 );
		const d2 = n.mul( 0.002 );
		const q = uvs;
		const x = n.mul( sin( mul( 0.3, time ).add( q.y.mul( 21.0 ) ) ) ).mul( sin( mul( 0.7, time ).add( q.y.mul( 29.0 ) ) ) ).mul( sin( add( 0.3, mul( 0.33, time ) ).add( q.y.mul( 31.0 ) ) ) ).mul( 0.0017 );
		const R = map.sample( vec2( x.add( uvs.x ).add( d1 ), uvs.y.add( d1 ) ) ).x.add( 0.05 ).toVar();
		const G = map.sample( vec2( x.add( uvs.x ), uvs.y.sub( d2 ) ) ).y.add( 0.05 ).toVar();
		const B = map.sample( vec2( x.add( uvs.x.sub( d2 ) ), uvs.y ) ).z.add( 0.05 ).toVar();
		const nshiftQuater = n.mul( 0.75 );
		R.addAssign( mul( 0.08, map.sample( nshiftQuater.mul( vec2( x.add( 0.025 ), - 0.027 ) ).add( n.mul( vec2( uvs.x.add( 0.001 ), uvs.y.add( 0.001 ) ) ) ) ).x ) );
		G.addAssign( mul( 0.05, map.sample( nshiftQuater.mul( vec2( x.sub( 0.022 ), - 0.02 ) ).add( n.mul( vec2( uvs.x.add( 0.000 ), uvs.y.sub( 0.002 ) ) ) ) ).y ) );
		B.addAssign( mul( 0.08, map.sample( nshiftQuater.mul( vec2( x.sub( 0.02 ), - 0.018 ) ).add( n.mul( vec2( uvs.x.sub( 0.002 ), uvs.y.add( 0.000 ) ) ) ) ).z ) );
		final.assign( vec4( R, G, B, 1.0 ) );

	} ).Else( () => {

		final.assign( m.sample( uvs ) );

	} );

	return final;

}, { n: 'float', uvs: 'vec2', m: 'sampler2D', return: 'vec4' } );

export const main = /*#__PURE__*/ Fn( () => {

	gl_FragColor.assign( getColor( nShiftPower, vUv, map ) );

}, { , return: 'void' } );

[sunag]

Thanks for the tests.

Just going through some of the old issues:

Thanks for reminding, I'll leave them to other PRs due to complexity.