TSL Struct Issues

[cmhhelgeson]

Description

Related PRs: #30394

A couple of things I noticed while experimenting with structs. I'll mostly be using code from the webgpu_compute_water sample as my test case here...

Ability To Label Structs:

I'm not totally familiar with an alternative solution with setLayout, but currently, I can only make the return type a struct by setting the function's output type to the auto-generated name of the struct within the code output. Ideally, the user should be able to name the struct and use that name as the output type for their function, rather than waiting to see the index the struct resolves to once the code is constructed by its respective NodeBuilder.

// CURRENT CASE: 
const HeightStruct = struct( {
	heightIndex: 'uint',
	prevHeightIndex: 'uint'
} );
const GetHeightIndices = Fn( ( [ index ] ) => {
     const hs = HeightStruct();
     hs.get( 'heightIndex' ).assign( index );
     hs.get( 'prevHeightIndex' ).assign( index );
     return hs;
} ).setLayout( {
	name: 'GetHeightVariables',
        // The name the NodeBuilder provides the struct
	type: 'StructType0', 
	inputs: [
		{ name: 'index', type: 'uint' }
	]
} );

// IDEAL CASE
const HeightStruct = struct( {
	heightIndex: 'uint',
	prevHeightIndex: 'uint'
} ).label('Height')
const GetHeightIndices = Fn( ( [ index ] ) => {
     ...
} ).setLayout( {
	type: 'HeightStructType', 
} );

Functions Matching Equivalent TSL Calls:

In WGSL, the following TSL code....

const hs = GetHeightIndices( instanceIndex );
const height = heightStorage.element( hs.get( 'heightIndex' ) ).toVar();
const prevHeight = prevHeightStorage.element( hs.get( 'prevHeightIndex' ) ).toVar();

resolves to this WGSL output

heightVariable = Height.value[ u32( GetHeightIndices( instanceIndex ).heightIndex ) ];
prevHeightVariable = PrevHeight.value[ u32( GetHeightIndices( instanceIndex ).prevHeightIndex ) ];

Despite being called only once, the function is executed inline each time its corresponding TSL variable is referenced, even though half of the work of each function call is being unused. Adding a .toVar() declaration to the GetHeightIndices(instanceIndex) call will fix this issue.

const hs = GetHeightVariables(instanceIndex)
heightVariable = Height.value[ u32( hs.heightIndex ) ];
prevHeightVariable = PrevHeight.value[ u32( hs.prevHeightIndex ) ];

Perhaps this is more of a problem with setLayout than the structs themselves, but I feel like the intention behind calling a function is somewhat different than the intention behind a generic mathematical operation in that the user would probably expect the function to only be called once in this case. However, given the brevity and efficiency of inlining these function calls when the outputs are base types, perhaps this issue can simply be resolved by calling out this behavior somewhere in the documentation.

Potential Mathematical Bugs
Thought I'd flag this even though I'm not sure whether my TSL is semantically correct, as to me this seems like a bug. When I attempt to use the element types of a struct in a mathematical calculation, the code output seems to break.

TSL FUNCTION

const NeighborStruct = struct( {
	north: 'float',
	east: 'float',
	south: 'float',
	west: 'float'
} );

const GetNeighborValues = ( index, store ) => {
	const neighborIndices = GetNeighborIndices( index ).toVar( 'neighborIndices' );
	const neighbors = NeighborStruct().toVar( 'neighborValues' );
	neighbors.get( 'north' ).assign( store.element( neighborIndices.get( 'northIndex' ) ) );
	neighbors.get( 'south' ).assign( store.element( neighborIndices.get( 'southIndex' ) ) );
	neighbors.get( 'east' ).assign( store.element( neighborIndices.get( 'eastIndex' ) ) );
	neighbors.get( 'west' ).assign( store.element( neighborIndices.get( 'westIndex' ) ) );
	return neighbors;
};

TSL INPUT 1

const GetNeighborValues = ( index, store ) => {
	const neighborIndices = GetNeighborIndices( index ).toVar( 'neighborIndices' );
	const neighbors = NeighborStruct().toVar( 'neighborValues' );
	neighbors.get( 'north' ).assign( store.element( neighborIndices.get( 'northIndex' ) ) );
	neighbors.get( 'south' ).assign( store.element( neighborIndices.get( 'southIndex' ) ) );
	neighbors.get( 'east' ).assign( store.element( neighborIndices.get( 'eastIndex' ) ) );
	neighbors.get( 'west' ).assign( store.element( neighborIndices.get( 'westIndex' ) ) );
	return neighbors;
};

const neighbors = GetNeighborValues( instanceIndex, heightStorage );
const neighborHeight = neighbors.get( 'north' ).add( neighbors.get( 'south' ) ).toVar( 'neighborHeight' );
neighborHeight.mulAssign( 0.5 );
neighborHeight.subAssign( prevHeight );
const newHeight = neighborHeight.mul( viscosity );

TSL INPUT 2

const neighbors = GetNeighborValues( instanceIndex, heightStorage );
const neighborHeight = neighbors.get( 'north' ).add( neighbors.get( 'south' ) );
neighborHeight.mulAssign( 0.5 );
neighborHeight.subAssign( prevHeight );
const newHeight = neighborHeight.mul( viscosity );

WGSL OUTPUT 1

struct StructType1 {
	north : f32,
	east : f32,
	south : f32,
	west : f32
};

var nodeVar0 : f32;
var nodeVar1 : f32;
var neighborIndices : StructType0;
var nodeVar2 : StructType1;
var neighborValues : StructType1;
// Type of neighborHeight is void
var neighborHeight : void;

// Initial setup and call to getNeighborIndices

nodeVar2 = StructType1( 0.0, 0.0, 0.0, 0.0 );
neighborValues = nodeVar2;
neighborValues.north = Height.value[ u32( neighborIndices.northIndex ) ];
neighborValues.south = Height.value[ u32( neighborIndices.southIndex ) ];
neighborValues.east = Height.value[ u32( neighborIndices.eastIndex ) ];
neighborValues.west = Height.value[ u32( neighborIndices.westIndex ) ];
neighborHeight = ;
neighborHeight = ;
neighborHeight = ;
( neighborHeight * object.viscosity ) = ;

WGSL OUTPUT 2

var nodeVar0 : f32;
var nodeVar1 : f32;
var neighborIndices : StructType0;
var nodeVar2 : StructType1;
var neighborValues : StructType1;

// flow

nodeVar0 = Height.value[ instanceIndex ];
nodeVar1 = PrevHeight.value[ instanceIndex ];
neighborIndices = GetNeighborIndices( instanceIndex );
nodeVar2 = StructType1( 0.0, 0.0, 0.0, 0.0 );
neighborValues = nodeVar2;
neighborValues.north = Height.value[ u32( neighborIndices.northIndex ) ];
neighborValues.south = Height.value[ u32( neighborIndices.southIndex ) ];
neighborValues.east = Height.value[ u32( neighborIndices.eastIndex ) ];
neighborValues.west = Height.value[ u32( neighborIndices.westIndex ) ];
( neighborValues.north + neighborValues.south ) = ;
( neighborValues.north + neighborValues.south ) = ;
(  * object.viscosity ) = ;

Reproduction steps

The section on labeling structs and and issues with setLayout are mostly self-contained, existing merely as minor modifications of the first part of the computeHeight function, with all the relevant code contained therein. For the section on using struct values in mathematical operations, I've simply replaced the functions with TSL suffixes with their counterparts as seen in the example code in the Mathematical bugs section.

Code

// NA

Live example

Live WebGPU Example

Screenshots

No response

Version

r173

Device

Desktop

Browser

Chrome

OS

Windows

[sunag]

I'm not totally familiar with an alternative solution with setLayout, but currently, I can only make the return type a struct by setting the function's output type to the auto-generated name of the struct within the code output. Ideally, the user should be able to name the struct and use that name as the output type for their function, rather than waiting to see the index the struct resolves to once the code is constructed by its respective NodeBuilder.

You could use the second parameter name of struct( members, name ) as done in webgpu_struct_drawindirect.

struct().toVar() should not be necessary.

I'll check out the other use cases soon.

[sunag]

Can you confirm if this solved the problem? It seems that the other problems related to members were also because of .toVar().

[cmhhelgeson]

Can you confirm if this solved the problem? It seems that the other problems related to members were also because of .toVar().

Yes that seemed to solve the issue. I removed the toVars() on struct() calls, but the main issue was the return value of the GetNeighborValues(). I tested with toVar() and without toVar(), and the sample now works.