WebGPURenderer: Compute Texture using TSL (#27582)

examples/jsm/nodes/accessors/TextureStoreNode.js

@@ -14,16 +14,69 @@ class TextureStoreNode extends TextureNode {
 
 	}
 
-	getNodeType( /*builder*/ ) {
+	getInputType( /*builder*/ ) {
 
-		return 'void';
+		return 'storageTexture';
+
+	}
+
+	setup( builder ) {
+
+		super.setup( builder );
+
+		const properties = builder.getNodeProperties( this );
+		properties.storeNode = this.storeNode;
+
+	}
+
+	generate( builder, output ) {
+
+		let snippet;
+
+		if ( this.storeNode !== null ) {
+
+			snippet = this.generateStore( builder );
+
+		} else {
+
+			snippet = super.generate( builder, output );
+
+		}
+
+		return snippet;
+
+	}
+
+	generateStore( builder ) {
+
+		const properties = builder.getNodeProperties( this );
+
+		const { uvNode, storeNode } = properties;
+
+		const textureProperty = super.generate( builder, 'property' );
+		const uvSnippet = uvNode.build( builder, 'uvec2' );
+		const storeSnippet = storeNode.build( builder, 'vec4' );
+
+		const snippet = builder.generateTextureStore( builder, textureProperty, uvSnippet, storeSnippet );
+
+		builder.addLineFlowCode( snippet );
 
 	}
 
 }
 
 export default TextureStoreNode;
 
-export const textureStore = nodeProxy( TextureStoreNode );
+const textureStoreBase = nodeProxy( TextureStoreNode );
+
+export const textureStore = ( value, uvNode, storeNode ) => {
+
+	const node = textureStoreBase( value, uvNode, storeNode );
+
+	if ( storeNode !== null ) node.append();
+
+	return node;
+
+};
 
 addNodeClass( 'TextureStoreNode', TextureStoreNode );

examples/jsm/nodes/core/NodeBuilder.js

@@ -462,7 +462,7 @@ class NodeBuilder {
 
 	isReference( type ) {
 
-		return type === 'void' || type === 'property' || type === 'sampler' || type === 'texture' || type === 'cubeTexture';
+		return type === 'void' || type === 'property' || type === 'sampler' || type === 'texture' || type === 'cubeTexture' || type === 'storageTexture';
 
 	}
 
@@ -523,7 +523,7 @@ class NodeBuilder {
 	getVectorType( type ) {
 
 		if ( type === 'color' ) return 'vec3';
-		if ( type === 'texture' ) return 'vec4';
+		if ( type === 'texture' || type === 'cubeTexture' || type === 'storageTexture' ) return 'vec4';
 
 		return type;
 

examples/jsm/renderers/webgpu/nodes/WGSLNodeBuilder.js

@@ -208,6 +208,12 @@ class WGSLNodeBuilder extends NodeBuilder {
 
 	}
 
+	generateTextureStore( texture, textureProperty, uvIndexSnippet, valueSnippet ) {
+
+		return `textureStore( ${ textureProperty }, ${ uvIndexSnippet }, ${ valueSnippet } )`;
+
+	}
+
 	isUnfilterable( texture ) {
 
 		return texture.isDataTexture === true && texture.type === FloatType;
@@ -283,7 +289,7 @@ class WGSLNodeBuilder extends NodeBuilder {
 			const name = node.name;
 			const type = node.type;
 
-			if ( type === 'texture' || type === 'cubeTexture' ) {
+			if ( type === 'texture' || type === 'cubeTexture' || type === 'storageTexture' ) {
 
 				return name;
 
@@ -336,11 +342,11 @@ class WGSLNodeBuilder extends NodeBuilder {
 
 			const bindings = this.bindings[ shaderStage ];
 
-			if ( type === 'texture' || type === 'cubeTexture' ) {
+			if ( type === 'texture' || type === 'cubeTexture' || type === 'storageTexture' ) {
 
 				let texture = null;
 
-				if ( type === 'texture' ) {
+				if ( type === 'texture' || type === 'storageTexture' ) {
 
 					texture = new NodeSampledTexture( uniformNode.name, uniformNode.node );
 
@@ -723,7 +729,7 @@ ${ flowData.code }
 
 		for ( const uniform of uniforms ) {
 
-			if ( uniform.type === 'texture' || uniform.type === 'cubeTexture' ) {
+			if ( uniform.type === 'texture' || uniform.type === 'cubeTexture' || uniform.type === 'storageTexture' ) {
 
 				const texture = uniform.node.value;
 

examples/webgpu_compute_texture.html

@@ -24,7 +24,7 @@
 		<script type="module">
 
 			import * as THREE from 'three';
-			import { texture, textureStore, wgslFn, instanceIndex, MeshBasicNodeMaterial } from 'three/nodes';
+			import { texture, textureStore, tslFn, instanceIndex, float, uvec2, vec4, MeshBasicNodeMaterial } from 'three/nodes';
 
 			import WebGPU from 'three/addons/capabilities/WebGPU.js';
 			import WebGPURenderer from 'three/addons/renderers/webgpu/WebGPURenderer.js';
@@ -60,39 +60,34 @@
 
 				// create function
 
-				const computeWGSL = wgslFn( `
-					fn computeWGSL( storageTex: texture_storage_2d<rgba8unorm, write>, index: u32 ) -> void {
+				const computeTexture = tslFn( ( { storageTexture } ) => {
 
-						let posX = index % ${ width };
-						let posY = index / ${ width };
-						let indexUV = vec2u( posX, posY );
+					const posX = instanceIndex.remainder( width );
+					const posY = instanceIndex.div( width );
+					const indexUV = uvec2( posX, posY );
 
-						// https://www.shadertoy.com/view/Xst3zN
+					// https://www.shadertoy.com/view/Xst3zN
 
-						let x = f32( posX ) / 50.0;
-						let y = f32( posY ) / 50.0;
+					const x = float( posX ).div( 50.0 );
+					const y = float( posY ).div( 50.0 );
 
-						let v1 = sin( x );
-						let v2 = sin( y );
-						let v3 = sin( x + y );
-						let v4 = sin( sqrt( x * x + y * y ) + 5.0 );
-						let v = v1 + v2 + v3 + v4;
+					const v1 = x.sin();
+					const v2 = y.sin();
+					const v3 = x.add( y ).sin();
+					const v4 = x.mul( x ).add( y.mul( y ) ).sqrt().add( 5.0 ).sin();
+					const v = v1.add( v2, v3, v4 );
 
-						let PI = 3.14159265359;
+					const r = v.sin();
+					const g = v.add( Math.PI ).sin();
+					const b = v.add( Math.PI ).sub( 0.5 ).sin();
 
-						let r = sin( v );
-						let g = sin( v + PI );
-						let b = sin( v + PI - 0.5 );
+					textureStore( storageTexture, indexUV, vec4( r, g, b, 1 ) );
 
-						textureStore( storageTex, indexUV, vec4f( r, g, b, 1 ) );
-
-					}
-				` );
+				} );
 
 				// compute
 
-				const computeWGSLCall = computeWGSL( { storageTex: textureStore( storageTexture ), index: instanceIndex } );
-				const computeNode = computeWGSLCall.compute( width * height );
+				const computeNode = computeTexture( { storageTexture } ).compute( width * height );
 
 				const material = new MeshBasicNodeMaterial( { color: 0x00ff00 } );
 				material.colorNode = texture( storageTexture );