WebGPURenderer: Support computing texture in TSL

[RenaudRohlinger]

Description

Just like in the examples webgpu_compute_points and webgpu_compute_particles_rain would it be possible to use the TSL Syntax with tslFn to compute texture storage data? @sunag

For example I started to try to port the webgpu_compute_texture example to tsl but it seems that texture_storage_2d is not supported yet or I don't know how to use it:

const width = 512, height = 512;
const storageTexture = new StorageTexture( width, height );

const computeShaderFn = tslFn( ( [ storageTexture ] ) => {

	// https://www.shadertoy.com/view/Xst3zN

	const posX = instanceIndex.mod( width );
	const posY = instanceIndex.div( height );
	const indexUV = vec2( posX, posY );

	const x = float( posX ).div( 50. );
	const y = float( posY ).div( 50. );

	const v1 = sin( x );
	const v2 = sin( y );
	const v3 = sin( x.add( y ) );
	const v4 = sin( sqrt( x.mul( x ).add( y.mul( y ) ) ).add( 5. ) );
	const v = v1.add( v2 ).add( v3 ).add( v4 );

	const PI = 3.14159265359;

	const r = sin( v );
	const g = sin( v.add( PI ) );
	const b = sin( v.add( PI ).sub( 0.5 ) );


	return textureStore( storageTexture, indexUV, vec4( r, g, b, 1 ) );

} );
const computeNode = computeShaderFn().compute( width * height );

Something to keep in mind would be to be able to, just like in #27367, have a WebGL fallback where we would somehow ping pong basic textures I guess.

[Spiri0]

The textureStore node came in r155 or r156 at my request for the ability to store textures with WGSL shaders. I use the textureStore node a lot, but only like this:

const myStorageTexture = new StorageTexture(width, height);

const myWGSLShader = wgslFn(`
   fn computeMain(
      writeTexture: texture_storage_2d<rgba32float, write>,
      index: u32,
      ...,
   ) -> void {
      ...,
      //the textureStore here is a wgsl command and not the textureStore node!
      textureStore(writeTexture, indexUV, color);
   }
`);

const computeTexture = myWGSLShader({
   writeTexture: textureStore(myStorageTexture),
   index: instanceIndex,
   ...,
}).compute(width,height);

The textureStore node is not the same as the wgsl command textureStore(writeTexture, index, color);

I think this is perhaps an illustrative analogy:

texture(getTexture)
textureStore(setTexture)

Maybe textureStore recognizes whether only a storageTexture is passed or whether a storageTexture, index and vec4 data is passed and then knows that it is a tsl. For me, this extended range of functions was not relevant, as I have only worked with WGSL so far and quite massively. But I don't think that such a comprehensive range of functions was integrated from the start, because first of all it was about the simplest possible case and that was the WGSL ability to be able to store textures. TSL is then a more abstract level. But Sunag will know that better than me.

[sunag]

This is very interesting. Node system can generate code from a single Node, this would allow us to have multiple shaders from a single TSL function, generating a single optimized code for each textureStore, this could prevent possible conflicts with transform feedback in WebGL fallback for use as ping/pong.

I think that at first we could make a RTTNode too (Render to Texture), like the previous node system. This could be useful for simpler examples like this and would require less initial effort, maybe even a gateway to fallback.

const rtt = new RTTNode( width, height, ... );
rtt.outputNode = tslFn( () => {

	// https://www.shadertoy.com/view/Xst3zN

	const uv0 = uv();

	const x = float( uv0.x ).div( 50. );
	const y = float( uv0.y ).div( 50. );

	const v1 = sin( x );
	const v2 = sin( y );
	const v3 = sin( x.add( y ) );
	const v4 = sin( sqrt( x.mul( x ).add( y.mul( y ) ) ).add( 5. ) );
	const v = v1.add( v2 ).add( v3 ).add( v4 );

	const PI = 3.14159265359;

	const r = sin( v );
	const g = sin( v.add( PI ) );
	const b = sin( v.add( PI ).sub( 0.5 ) );

	return vec4( r, g, b, 1 );

} );

rtt.render( renderer );

[RenaudRohlinger]

This would be awesome! I would be very curious to see how nodes such as GaussianBlurNode could benefit from this. If I can be of any help on the WebGL part I would be happy to help as that's a feature I'm very looking forward to on the WebGPU Side.

[RenaudRohlinger]

Partial implementation has been done with #27582 with the capacity to write onto a storage texture.

Currently the storage access type is the default one (write-only).

What would be interesting with StorageTexture is to be able to perform multiple computation of a texture through read-write operations, for example for sorting algorithms or we could even replace the webgpu_compute_texture_pingpong example with a single compute program instead of ping-pong implementation.

But for that we will need to wait that browsers remove the flag --enable-dawn-features=allow_unsafe_apis.

It's possible to follow the progress of the read-write access for textureStorage here: https://github.com/gpuweb/gpuweb/issues/3838

As a reminder here how it would look like:

const width = 512, height = 512;
const storageTexture = new StorageTexture( width, height, 'read-write' ); // <-- we can specify the access mode here

const computeShaderFn = tslFn( ( { storageTexture } ) => {

	const posX = instanceIndex.mod( width );
	const posY = instanceIndex.div( height );
	const indexUV = vec2( posX, posY );


        const newStorageTexture = textureLoad( storageTexture, indexUV ).toVar(); // <-- we can fetch the previous pixel information now

        newStorageTexture.rgb.addAssign(0.01);

	return textureStore( storageTexture, indexUV, newStorageTexture );

} );
const computeNode = computeShaderFn().compute( width * height );

I will let this Issue open as a reminder so when there is some progress on that we can start implementing this new feature (read-write). /cc @sunag

[Spiri0]

The read/write topic is the same as mine in #27502. We have a very general problem between tslFn / wgslFn and RenderTargets. That's why I had the idea of ​​preparing a specially adapted RenderTarget that uses storageTextures for webgpu tslFn / wgslFn. To do this, a renderTarget would have to be in jsm/renderers/common/ next to the storageTexture. That's not great art. In my opinion, this renderTarget would only have to use storageTexture instead of Texture with the write flag. This also applies to the depthTexture (so basically a depthStorageTexture, although I don't like the name). In jsm/renderers/common/Bindings.js I see that there is a function that checks whether new bindings are necessary. I can't get through the whole topic, I need a lot more time and I only have a limited amount of it alongside my projects and my job. Maybe we can make a little progress together. In any case, we need different bindings for read and write. The depthTexture inherits from the Texture class just like the storageTexture. Then it would be just as easy to create a depthTexture in jsm/renderers/common, which then works in the same way as the storageTexture, only for depthTextures. Both require different bindings for read/write in wgsl. Maybe in the future it would also be conceivable to simply rename the storageTexture to Texture. Since it is outside of three.module.js there are no problems and the nodes (texture/textureStore) decide whether you use it for read or write. So the names Texture/depthTexture would be appropriate again because storageTexture/storageDepthTexture seems somehow inappropriate to me.

Feel free to correct my thoughts! I'm not nearly deep enough into the three.js code yet. /cc. @RenaudRohlinger @sunag