TSL - SetLayout Function Being Constructed Multiple Times in WGSL Output

[cmhhelgeson]

Description

Encountered this error when trying to encapsulate bitonic sort.

The function:

export const getBitonicFlipIndices = /*@__PURE__*/ Fn( ( [ index, swapSpan ] ) => {

	const blockOffset = ( index.mul( 2 ).div( swapSpan ) ).mul( swapSpan );
	const halfHeight = swapSpan.div( 2 );
	const idx = uvec2(
		index.mod( halfHeight ),
		swapSpan.sub( index.mod( halfHeight ) ).sub( 1 )
	);
	idx.x.addAssign( blockOffset );
	idx.y.addAssign( blockOffset );

	return idx;

}, { index: 'uint', swapSpan: 'uint', return: 'uvec2' } );

And it's invocation in the ComputeNodes:

class BitonicSort {
...

_globalCompareAndSwapTSL( idxBefore, idxAfter, dataBuffer, tempBuffer ) {
	// If the later element is less than the current element
	If( dataBuffer.element( idxAfter ).lessThan( dataBuffer.element( idxBefore ) ), () => {
		tempBuffer.element( idxBefore ).assign( dataBuffer.element( idxAfter ) );
		tempBuffer.element( idxAfter ).assign( dataBuffer.element( idxBefore ) );
	} ).Else( () => {
		// Otherwise apply the existing values to temporary storage.
		tempBuffer.element( idxBefore ).assign( dataBuffer.element( idxBefore ) );
		tempBuffer.element( idxAfter ).assign( dataBuffer.element( idxAfter ) );
	} );
}

_getFlipGlobal() {
	const { infoStorage, tempBuffer, dataBuffer } = this;
	const currentSwapSpan = infoStorage.element( 1 );
	const fnDef = Fn( () => {
		const idx = getBitonicFlipIndices( instanceIndex, currentSwapSpan );
		this._globalCompareAndSwapTSL( idx.x, idx.y, dataBuffer, tempBuffer );
	} )().compute( this.dispatchSize, [ this.workgroupSize ] );
	return fnDef;
}

}

Generates this WGSL output, where the function fn2 is being called two more times than necessary ( the relevant line to search for is "fn2( instanceIndex, BitonicSortInfo.value[ 1u ] );" )

// Three.js r180dev - Node System

// codes
fn fn2 ( index : u32, swapSpan : u32 ) -> vec2<u32> {

	var nodeVar0 : u32;
	var nodeVar1 : vec2<u32>;
	var nodeVar2 : u32;

	nodeVar0 = ( swapSpan / 2u );
	nodeVar1 = vec2<u32>( ( index % nodeVar0 ), ( ( swapSpan - ( index % nodeVar0 ) ) - 1u ) );
	nodeVar2 = ( ( ( index * 2u ) / swapSpan ) * swapSpan );
	nodeVar1.x = ( nodeVar1.x + nodeVar2 );
	nodeVar1.y = ( nodeVar1.y + nodeVar2 );

	return nodeVar1;

}

@compute @workgroup_size( 64, 1, 1 )
fn main( @builtin( global_invocation_id ) globalId : vec3<u32>,
	@builtin( workgroup_id ) workgroupId : vec3<u32>,
	@builtin( local_invocation_id ) localId : vec3<u32>,
	@builtin( num_workgroups ) numWorkgroups : vec3<u32>,
	@builtin( subgroup_size ) subgroupSize : u32 ) {

	// system
	instanceIndex = globalId.x
		+ globalId.y * ( 64 * numWorkgroups.x )
		+ globalId.z * ( 64 * numWorkgroups.x ) * ( 1 * numWorkgroups.y );
	
	var nodeVar0 : vec2<u32>;
	var nodeVar1 : vec2<u32>;
	var nodeVar2 : vec2<u32>;

        // fn2 being called 2 extra times

	nodeVar0 = fn2( instanceIndex, BitonicSortInfo.value[ 1u ] );
	if ( ( Elements.value[ nodeVar0.y ] < Elements.value[ nodeVar0.x ] ) ) {
		nodeVar1 = fn2( instanceIndex, BitonicSortInfo.value[ 1u ] );
		TempStorage.value[ nodeVar1.x ] = Elements.value[ nodeVar1.y ];
		TempStorage.value[ nodeVar1.y ] = Elements.value[ nodeVar1.x ];
	} else {
		nodeVar2 = fn2( instanceIndex, BitonicSortInfo.value[ 1u ] );
		TempStorage.value[ nodeVar2.x ] = Elements.value[ nodeVar2.x ];
		TempStorage.value[ nodeVar2.y ] = Elements.value[ nodeVar2.y ];
	}
}

Reproduction steps

1. Currently this is occurring on a non-functional version of an encapsulated bitonic sort that I am working on for both webgpu_compute_sort_bitonic and eventually birds, but the most up to date version of the branch is here. I've included a console.log of this.computeShader in the githack. The shader generating the display on the left is obviously non-functional.

2. Search for the relevant line of text 'fn2( instanceIndex, BitonicSortInfo.value[ 1u ] );' in the Chrome devTools. The global flip shader should at most be compiled around three seconds after the example begins.

Code

Seen above...can post full class (albeit incomplete) if requested.

Live example

https://raw.githack.com/cmhhelgeson/three.js/cmh/birds_hash/examples/webgpu_compute_sort_bitonic.html

...

Screenshots

No response

Version

r180

Device

Desktop

Browser

Chrome

OS

Windows

[kenjinp]

Out of curiosity, how are you able to see the compiled WGSL output of a compute node generated by TSL?

[cmhhelgeson]

Out of curiosity, how are you able to see the compiled WGSL output of a compute node generated by TSL?

I'm usually working within the Three.js library code itself, so for working on the Bitonic Sort, I just go into the WGSLNodeBuilder and log this.computeShader. There's also a DebugNode that could offer the same functionality, but I haven't had the chance to use it yet.

// Lines 1922 - 1935 of WGSLNodeBuilder
if ( this.material !== null ) {

	this.vertexShader = this._getWGSLVertexCode( shadersData.vertex );
	this.fragmentShader = this._getWGSLFragmentCode( shadersData.fragment );

} else {

	// Early strictly validated in computeNode

	const workgroupSize = this.object.workgroupSize;

	this.computeShader = this._getWGSLComputeCode( shadersData.compute, workgroupSize );
        console.log( this.computeShader );

}

three.js/src/nodes/utils/DebugNode.js

Line 5 in c604c09

class DebugNode extends TempNode {

[kenjinp]

I was able to see the compiled output by using the webgpu inspector and enabling the chrome://flags/#enable-webgpu-developer-features flag

very handy to see what the variables look like on the other side :)

[cmhhelgeson]

I was able to see the compiled output by using the webgpu inspector and enabling the chrome://flags/#enable-webgpu-developer-features flag

very handy to see what the variables look like on the other side :)

That makes sense. I should try to use the webgpu inspector more often, I'm sure it's probably much more robust than when I first checked it out. It would also keep me from having to make perfunctory commits removing my console.log( this.computeShader ) code for every PR 😅