TSL: textureStore() and example (#26648)

* TextureStoreNode

* Add exampe: webgpu_compute_texture

* cleanup

* cleanup

* cleanup

examples/files.json

@@ -312,6 +312,7 @@
 		"webgpu_backdrop_area",
 		"webgpu_clearcoat",
 		"webgpu_compute",
+		"webgpu_compute_texture",
 		"webgpu_cubemap_adjustments",
 		"webgpu_cubemap_dynamic",
 		"webgpu_cubemap_mix",

examples/jsm/nodes/Nodes.js

@@ -86,6 +86,7 @@ export { default as SceneNode, backgroundBlurriness, backgroundIntensity } from
 export { default as StorageBufferNode, storage } from './accessors/StorageBufferNode.js';
 export { default as TangentNode, tangentGeometry, tangentLocal, tangentView, tangentWorld, transformedTangentView, transformedTangentWorld } from './accessors/TangentNode.js';
 export { default as TextureNode, texture, /*textureLevel,*/ sampler } from './accessors/TextureNode.js';
+export { default as TextureStoreNode, textureStore } from './accessors/TextureStoreNode.js';
 export { default as UVNode, uv } from './accessors/UVNode.js';
 export { default as UserDataNode, userData } from './accessors/UserDataNode.js';
 

examples/jsm/nodes/accessors/TextureStoreNode.js

@@ -0,0 +1,29 @@
+import { addNodeClass } from '../core/Node.js';
+import TextureNode from './TextureNode.js';
+import { nodeProxy } from '../shadernode/ShaderNode.js';
+
+class TextureStoreNode extends TextureNode {
+
+	constructor( value, uvNode, storeNode = null ) {
+
+		super( value, uvNode );
+
+		this.storeNode = storeNode;
+
+		this.isStoreTextureNode = true;
+
+	}
+
+	getNodeType( /*builder*/ ) {
+
+		return 'void';
+
+	}
+
+}
+
+export default TextureStoreNode;
+
+export const textureStore = nodeProxy( TextureStoreNode );
+
+addNodeClass( TextureStoreNode );

examples/jsm/renderers/common/Bindings.js

@@ -80,9 +80,11 @@ class Bindings extends DataMap {
 
 		for ( const binding of bindings ) {
 
-			if ( binding.isSampler || binding.isSampledTexture ) {
+			if ( binding.isSampledTexture ) {
 
-				this.textures.updateTexture( binding.texture );
+				const store = binding.store === true;
+
+				this.textures.updateTexture( binding.texture, { store } );
 
 			} else if ( binding.isStorageBuffer ) {
 

examples/jsm/renderers/common/SampledTexture.js

@@ -12,6 +12,7 @@ class SampledTexture extends Binding {
 
 		this.texture = texture;
 		this.version = texture ? texture.version : 0;
+		this.store = false;
 
 		this.isSampledTexture = true;
 

examples/jsm/renderers/common/Textures.js

@@ -112,11 +112,13 @@ class Textures extends DataMap {
 		options.height = height;
 		options.depth = depth;
 		options.needsMipmaps = this.needsMipmaps( texture );
-		options.levels = this.getMipLevels( texture, width, height, options.needsMipmaps );
+		options.levels = options.needsMipmaps ? this.getMipLevels( texture, width, height ) : 1;
 
 		//
 
-		if ( isRenderTarget ) {
+		if ( isRenderTarget || options.store === true ) {
+
+			//if ( options.store === true ) options.levels = 1; /* no mipmaps? */
 
 			backend.createSampler( texture );
 			backend.createTexture( texture, options );
@@ -239,7 +241,7 @@ class Textures extends DataMap {
 
 	}
 
-	getMipLevels( texture, width, height, needsMipmaps ) {
+	getMipLevels( texture, width, height ) {
 
 		let mipLevelCount;
 
@@ -247,13 +249,9 @@ class Textures extends DataMap {
 
 			mipLevelCount = texture.mipmaps.length;
 
-		} else if ( needsMipmaps ) {
-
-			mipLevelCount = Math.floor( Math.log2( Math.max( width, height ) ) ) + 1;
-
 		} else {
 
-			mipLevelCount = 1; // a texture without mipmaps has a base mip (mipLevel 0)
+			mipLevelCount = Math.floor( Math.log2( Math.max( width, height ) ) ) + 1;
 
 		}
 

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

@@ -308,13 +308,14 @@ class WGSLNodeBuilder extends NodeBuilder {
 
 				}
 
+				texture.store = node.isStoreTextureNode === true;
 				texture.setVisibility( gpuShaderStageLib[ shaderStage ] );
 
 				// add first textures in sequence and group for last
 				const lastBinding = bindings[ bindings.length - 1 ];
 				const index = lastBinding && lastBinding.isUniformsGroup ? bindings.length - 1 : bindings.length;
 
-				if ( shaderStage === 'fragment' && this.isUnfilterable( node.value ) === false ) {
+				if ( shaderStage === 'fragment' && this.isUnfilterable( node.value ) === false && texture.store === false ) {
 
 					const sampler = new NodeSampler( `${uniformNode.name}_sampler`, uniformNode.node );
 					sampler.setVisibility( gpuShaderStageLib[ shaderStage ] );
@@ -404,7 +405,7 @@ class WGSLNodeBuilder extends NodeBuilder {
 
 	isReference( type ) {
 
-		return super.isReference( type ) || type === 'texture_2d' || type === 'texture_cube';
+		return super.isReference( type ) || type === 'texture_2d' || type === 'texture_cube' || type === 'texture_storage_2d';
 
 	}
 
@@ -594,7 +595,7 @@ class WGSLNodeBuilder extends NodeBuilder {
 
 				const texture = uniform.node.value;
 
-				if ( shaderStage === 'fragment' && this.isUnfilterable( texture ) === false ) {
+				if ( shaderStage === 'fragment' && this.isUnfilterable( texture ) === false && uniform.node.isStoreTextureNode !== true ) {
 
 					if ( texture.isDepthTexture === true && texture.compareFunction !== null ) {
 
@@ -622,6 +623,11 @@ class WGSLNodeBuilder extends NodeBuilder {
 
 					textureType = 'texture_external';
 
+				} else if ( uniform.node.isStoreTextureNode === true ) {
+
+					// @TODO: Add support for other formats
+					textureType = 'texture_storage_2d<rgba8unorm, write>';
+
 				} else {
 
 					textureType = 'texture_2d<f32>';

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

@@ -2,7 +2,7 @@ import NodeFunction from '../../../nodes/core/NodeFunction.js';
 import NodeFunctionInput from '../../../nodes/core/NodeFunctionInput.js';
 
 const declarationRegexp = /^[fn]*\s*([a-z_0-9]+)?\s*\(([\s\S]*?)\)\s*[\-\>]*\s*([a-z_0-9]+)?/i;
-const propertiesRegexp = /[a-z_0-9]+/ig;
+const propertiesRegexp = /[a-z_0-9]+|<(.*?)>+/ig;
 
 const wgslTypeLib = {
 	f32: 'float'
@@ -46,7 +46,7 @@ const parse = ( source ) => {
 
 			// precision
 
-			if ( i < propsMatches.length && /^[fui]\d{2}$/.test( propsMatches[ i ][ 0 ] ) === true )
+			if ( i < propsMatches.length && propsMatches[ i ][ 0 ].startsWith( '<' ) === true )
 				i ++;
 
 			// add input

examples/jsm/renderers/webgpu/utils/WebGPUBindingUtils.js

@@ -59,6 +59,12 @@ class WebGPUBindingUtils {
 
 				bindingGPU.externalTexture = {}; // GPUExternalTextureBindingLayout
 
+			} else if ( binding.isSampledTexture && binding.store ) {
+
+				const format = this.backend.get( binding.texture ).texture.format;
+
+				bindingGPU.storageTexture = { format }; // GPUStorageTextureBindingLayout
+
 			} else if ( binding.isSampledTexture ) {
 
 				const texture = {}; // GPUTextureBindingLayout

examples/jsm/renderers/webgpu/utils/WebGPUTextureUtils.js

@@ -111,6 +111,12 @@ class WebGPUTextureUtils {
 
 		let usage = GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST | GPUTextureUsage.COPY_SRC;
 
+		if ( options.store === true ) {
+
+			usage |= GPUTextureUsage.STORAGE_BINDING;
+
+		}
+
 		if ( texture.isCompressedTexture !== true ) {
 
 			usage |= GPUTextureUsage.RENDER_ATTACHMENT;

examples/webgpu_compute_texture.html

@@ -0,0 +1,133 @@
+<html lang="en">
+	<head>
+		<title>three.js - WebGPU - Compute Texture</title>
+		<meta charset="utf-8">
+		<meta name="viewport" content="width=device-width, user-scalable=no, minimum-scale=1.0, maximum-scale=1.0">
+		<link type="text/css" rel="stylesheet" href="main.css">
+	</head>
+	<body>
+
+		<div id="info">
+			<a href="https://threejs.org" target="_blank" rel="noopener">three.js</a> WebGPU - Compute Texture
+			<br>Texture generated using GPU Compute.
+		</div>
+
+		<script async src="https://unpkg.com/[email protected]/dist/es-module-shims.js"></script>
+
+		<script type="importmap">
+			{
+				"imports": {
+					"three": "../build/three.module.js",
+					"three/addons/": "./jsm/",
+					"three/nodes": "./jsm/nodes/Nodes.js"
+				}
+			}
+		</script>
+
+		<script type="module">
+
+			import * as THREE from 'three';
+			import { ShaderNode, texture, textureStore, wgslFn, instanceIndex, MeshBasicNodeMaterial } from 'three/nodes';
+
+			import WebGPU from 'three/addons/capabilities/WebGPU.js';
+			import WebGPURenderer from 'three/addons/renderers/webgpu/WebGPURenderer.js';
+
+			let camera, scene, renderer;
+
+			init();
+
+			function init() {
+
+				if ( WebGPU.isAvailable() === false ) {
+
+					document.body.appendChild( WebGPU.getErrorMessage() );
+
+					throw new Error( 'No WebGPU support' );
+
+				}
+
+				camera = new THREE.OrthographicCamera( - 1.0, 1.0, 1.0, - 1.0, 0, 1 );
+				camera.position.z = 1;
+
+				scene = new THREE.Scene();
+
+				// texture
+
+				const width = 512, height = 512;
+
+				const storageTexture = new THREE.Texture();
+				storageTexture.image = { width, height };
+				storageTexture.magFilter = THREE.LinearFilter;
+				storageTexture.minFilter = THREE.NearestFilter;
+
+				// create function
+
+				const computeShaderNode = new ShaderNode( ( stack ) => {
+
+					// the first function will be the main one
+
+					const computeWGSL = wgslFn( `
+						fn computeWGSL( storageTex: texture_storage_2d<rgba8unorm, write>, index:u32 ) -> void {
+
+							let posX = index % ${ width };
+							let posY = index / ${ width };
+							let indexUV = vec2<u32>( posX, posY );
+							let uv = getUV( posX, posY );
+
+							textureStore( storageTex, indexUV, vec4f( uv, 0, 1 ) );
+
+						}
+
+						fn getUV( posX:u32, posY:u32 ) -> vec2<f32> {
+
+							let uv = vec2<f32>( f32( posX ) / ${ width }.0, f32( posY ) / ${ height }.0 );
+
+							return uv;
+
+						}
+					` );
+
+					stack.add( computeWGSL( { storageTex: textureStore( storageTexture ), index: instanceIndex } ) );
+
+				} );
+
+				// compute
+
+				const computeNode = computeShaderNode.compute( width * height );
+
+				const material = new MeshBasicNodeMaterial( { color: 0x00ff00 } );
+				material.colorNode = texture( storageTexture );
+
+				const plane = new THREE.Mesh( new THREE.PlaneGeometry( 1, 1 ), material );
+				scene.add( plane );
+
+				renderer = new WebGPURenderer();
+				renderer.setPixelRatio( window.devicePixelRatio );
+				renderer.setSize( window.innerWidth, window.innerHeight );
+				renderer.setAnimationLoop( animate );
+				document.body.appendChild( renderer.domElement );
+
+				// compute texture
+				renderer.compute( computeNode );
+
+				window.addEventListener( 'resize', onWindowResize );
+
+			}
+
+			function onWindowResize() {
+
+				camera.updateProjectionMatrix();
+
+				renderer.setSize( window.innerWidth, window.innerHeight );
+
+			}
+
+			function animate() {
+
+				renderer.render( scene, camera );
+
+			}
+
+		</script>
+	</body>
+</html>

test/e2e/puppeteer.js

@@ -111,6 +111,7 @@ const exceptionList = [
 	'webgpu_backdrop_area',
 	'webgpu_clearcoat',
 	'webgpu_compute',
+	'webgpu_compute_texture',
 	'webgpu_cubemap_adjustments',
 	'webgpu_cubemap_dynamic',
 	'webgpu_cubemap_mix',